From 0d3a0bb06786e06f8f80b45fdec11aad5a6ca417 Mon Sep 17 00:00:00 2001 From: Philipp Holl Date: Wed, 31 Jan 2024 15:47:11 +0100 Subject: [PATCH] [field,geom] Support changing grid dim order --- phi/field/_field.py | 7 ++++++- phi/geom/_grid.py | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/phi/field/_field.py b/phi/field/_field.py index 3953b760c..53b3c93d0 100644 --- a/phi/field/_field.py +++ b/phi/field/_field.py @@ -348,7 +348,12 @@ def with_boundary(self, boundary): def with_bounds(self, bounds: Box): """ Returns a copy of this field with `bounds` replaced. """ - return Field(self._geometry, self._values, self._boundary) + order = list(bounds.vector.item_names) + geometry = self._geometry.vector[order] + geometry[order] + new_shape = self._values.shape.without(order) & self._values.shape.only(order, reorder=True) + values = math.transpose(self._values, new_shape) + return Field(geometry, values, self._boundary) def with_geometry(self, elements: Geometry): """ Returns a copy of this field with `elements` replaced. """ diff --git a/phi/geom/_grid.py b/phi/geom/_grid.py index 9115a7e40..7c44de711 100644 --- a/phi/geom/_grid.py +++ b/phi/geom/_grid.py @@ -147,7 +147,10 @@ def __getitem__(self, item): bounds = Box(lower, upper) gather_dict[dim] = slice(start, stop) resolution = self._resolution.after_gather(gather_dict) - return UniformGrid(resolution, bounds[{d: s for d, s in item.items() if d != 'vector'}]) + bounds = bounds[{d: s for d, s in item.items() if d != 'vector'}] + if 'vector' in item: + bounds = bounds[item['vector']] # resolution[item['vector']] will be done automatically + return UniformGrid(resolution, bounds) def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: Optional[int], **kwargs) -> 'Cuboid': return math.pack_dims(self.center_representation(), dims, packed_dim, pos, **kwargs)