diff --git a/PhiML b/PhiML index ab1ffbc45..b44051b9f 160000 --- a/PhiML +++ b/PhiML @@ -1 +1 @@ -Subproject commit ab1ffbc4523b9d4febe5d60bb7d0451992badb26 +Subproject commit b44051b9f93edec4dc46b7f8d3bc7e2b522ce700 diff --git a/phi/field/_grid.py b/phi/field/_grid.py index 1e178191d..34e63ea03 100644 --- a/phi/field/_grid.py +++ b/phi/field/_grid.py @@ -66,10 +66,10 @@ def with_bounds(self, bounds: Box): return type(self)(self.values, extrapolation=self.extrapolation, bounds=bounds) def __value_attrs__(self): - return '_values', '_extrapolation' + return '_values', def __variable_attrs__(self): - return '_values', + return '_values', '_extrapolation' def __expand__(self, dims: Shape, **kwargs) -> 'Grid': return self.with_values(math.expand(self.values, dims, **kwargs)) diff --git a/phi/field/_point_cloud.py b/phi/field/_point_cloud.py index 2b720ce85..8b123e95e 100644 --- a/phi/field/_point_cloud.py +++ b/phi/field/_point_cloud.py @@ -87,10 +87,10 @@ def with_bounds(self, bounds: Box): return PointCloud(elements=self.elements, values=self.values, extrapolation=self.extrapolation, add_overlapping=self._add_overlapping, bounds=bounds) def __value_attrs__(self): - return '_values', '_extrapolation' + return '_values', '_extrapolation', '_elements' def __variable_attrs__(self): - return '_values', '_elements' + return '_values', '_elements', '_extrapolation' def __expand__(self, dims: Shape, **kwargs) -> 'PointCloud': return self.with_values(expand(self.values, dims, **kwargs)) diff --git a/phi/field/_scene.py b/phi/field/_scene.py index b5807c3db..c267a2447 100644 --- a/phi/field/_scene.py +++ b/phi/field/_scene.py @@ -214,7 +214,7 @@ def at(directory: Union[str, tuple, typing_list, math.Tensor, 'Scene'], id: Unio id = math.wrap(id) paths = math.map(lambda d, i: join(d, f"sim_{i:06d}"), directory, id) # test all exist - for path in math.flatten(paths, flatten_batch=True): + for path in math.flatten(wrap(paths), flatten_batch=True): if not isdir(path): raise IOError(f"There is no scene at '{path}'") return Scene(paths) @@ -241,10 +241,7 @@ def single_subpath(path): return path result = math.map(single_subpath, self._paths) - if result.rank == 0: - return result.native() - else: - return result + return result def _init_properties(self): if self._properties is not None: diff --git a/phi/geom/_sphere.py b/phi/geom/_sphere.py index b7976a0d1..3a559838e 100644 --- a/phi/geom/_sphere.py +++ b/phi/geom/_sphere.py @@ -107,6 +107,9 @@ def scaled(self, factor: Union[float, Tensor]) -> 'Geometry': def __variable_attrs__(self): return '_center', '_radius' + def __value_attrs__(self): + return '_center', '_radius' + def __getitem__(self, item): item = slicing_dict(self, item) return Sphere(self._center[_keep_vector(item)], self._radius[item]) diff --git a/phi/geom/_stack.py b/phi/geom/_stack.py index 4b3012e54..e359e172c 100644 --- a/phi/geom/_stack.py +++ b/phi/geom/_stack.py @@ -13,10 +13,12 @@ class GeometryStack(Geometry): def __init__(self, geometries: Tensor): self.geometries = geometries - self._shape = shape_stack(geometries.shape, *[g.shape for g in geometries]) + inner_dims = math.merge_shapes(*[g.shape for g in geometries], allow_varying_sizes=True) + self._stack_dim = geometries.shape.without(inner_dims) + self._shape = geometries.shape def unstack(self, dimension) -> tuple: - if dimension == self.geometries.shape.name: + if dimension == self._stack_dim.name: return tuple(self.geometries) else: # return GeometryStack([g.unstack(dimension) for g in self.geometries], self.geometries.shape) @@ -25,7 +27,7 @@ def unstack(self, dimension) -> tuple: @property def center(self): centers = [g.center for g in self.geometries] - return math.stack(centers, self.geometries.shape) + return math.stack(centers, self._stack_dim) @property def spatial_rank(self) -> int: @@ -37,41 +39,41 @@ def shape(self) -> Shape: @property def volume(self) -> math.Tensor: - if self.geometries.shape.type == INSTANCE_DIM: + if self._stack_dim.type == INSTANCE_DIM: raise NotImplementedError("instance dimensions not yet supported") - return math.stack([g.volume for g in self.geometries], self.geometries.shape) + return math.stack([g.volume for g in self.geometries], self._stack_dim) @property def shape_type(self) -> Tensor: types = [g.shape_type for g in self.geometries] - return math.stack(types, self.geometries.shape) + return math.stack(types, self._stack_dim) def lies_inside(self, location: math.Tensor): - if self.geometries.shape in location.shape: - location = location.unstack(self.geometries.shape.name) + if self._stack_dim in location.shape: + location = location.unstack(self._stack_dim.name) else: location = [location] * len(self.geometries) inside = [g.lies_inside(loc) for g, loc in zip(self.geometries, location)] - return math.stack(inside, self.geometries.shape) + return math.stack(inside, self._stack_dim) def approximate_signed_distance(self, location: math.Tensor): raise NotImplementedError() def bounding_radius(self): radii = [expand(g.bounding_radius(), non_channel(g)) for g in self.geometries] - return math.stack(radii, self.geometries.shape) + return math.stack(radii, self._stack_dim) def bounding_half_extent(self): values = [expand(g.bounding_half_extent(), non_channel(g)) for g in self.geometries] - return math.stack(values, self.geometries.shape) + return math.stack(values, self._stack_dim) def at(self, center: Tensor) -> 'Geometry': - geometries = [self.geometries[idx].native().at(center[idx]) for idx in self.geometries.shape.meshgrid()] - return GeometryStack(math.layout(geometries, self.geometries.shape)) + geometries = [self.geometries[idx].at(center[idx]) for idx in self._stack_dim.meshgrid()] + return GeometryStack(math.layout(geometries, self._stack_dim)) def rotated(self, angle): geometries = [g.rotated(angle) for g in self.geometries] - return GeometryStack(math.layout(geometries, self.geometries.shape)) + return GeometryStack(math.layout(geometries, self._stack_dim)) def push(self, positions: Tensor, outward: bool = True, shift_amount: float = 0) -> Tensor: raise NotImplementedError('GeometryStack.push() is not yet implemented.') @@ -79,7 +81,7 @@ def push(self, positions: Tensor, outward: bool = True, shift_amount: float = 0) def __eq__(self, other): return isinstance(other, GeometryStack) \ and self._shape == other.shape \ - and self.geometries.shape == other.stack_dim \ + and self._stack_dim == other.stack_dim \ and self.geometries == other.geometries def shallow_equals(self, other): @@ -87,7 +89,7 @@ def shallow_equals(self, other): return True if not isinstance(other, GeometryStack) or self._shape != other.shape: return False - if self.geometries.shape != other.geometries.shape: + if self._stack_dim != other.geometries.shape: return False return all(g1.shallow_equals(g2) for g1, g2 in zip(self.geometries, other.geometries)) @@ -96,7 +98,6 @@ def __hash__(self): def __getitem__(self, item): selected = self.geometries[slicing_dict(self, item)] - if selected.shape.volume > 1: - return GeometryStack(selected) - else: - return next(iter(selected)) + if isinstance(selected, Geometry): + return selected + return GeometryStack(selected) diff --git a/phi/vis/_vis.py b/phi/vis/_vis.py index 843f90955..3307fb002 100644 --- a/phi/vis/_vis.py +++ b/phi/vis/_vis.py @@ -437,28 +437,28 @@ def layout_sub_figures(data: Union[Tensor, SampledField], dim0 = reduced = data.shape[0] if dim0.only(overlay): for overlay_index in dim0.only(overlay).meshgrid(names=True): # overlay these fields - e_rows, e_cols, d_non_reduced, d_reduced = layout_sub_figures(data[overlay_index].native(), row_dims, col_dims, animate, overlay, offset_row, offset_col, positioning, indices, {**base_index, **overlay_index}) + e_rows, e_cols, d_non_reduced, d_reduced = layout_sub_figures(data[overlay_index], row_dims, col_dims, animate, overlay, offset_row, offset_col, positioning, indices, {**base_index, **overlay_index}) rows = max(rows, e_rows) cols = max(cols, e_cols) non_reduced &= d_non_reduced reduced = merge_shapes(reduced, d_reduced, allow_varying_sizes=True) elif dim0.only(animate): - data = math.stack(data.native(), dim0) + data = math.stack(data, dim0) return layout_sub_figures(data, row_dims, col_dims, animate, overlay, offset_row, offset_col, positioning, indices, base_index) else: elements = unstack(data, dim0) for item_name, e in zip(dim0.get_item_names(dim0.name) or range(dim0.size), elements): index = dict(base_index, **{dim0.name: item_name}) if dim0.only(row_dims): - e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e.native(), row_dims, col_dims, animate, overlay, offset_row + rows, offset_col, positioning, indices, index) + e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e, row_dims, col_dims, animate, overlay, offset_row + rows, offset_col, positioning, indices, index) rows += e_rows cols = max(cols, e_cols) elif dim0.only(col_dims): - e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e.native(), row_dims, col_dims, animate, overlay, offset_row, offset_col + cols, positioning, indices, index) + e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e, row_dims, col_dims, animate, overlay, offset_row, offset_col + cols, positioning, indices, index) cols += e_cols rows = max(rows, e_rows) else: - e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e.native(), row_dims, col_dims, animate, overlay, offset_row, offset_col, positioning, indices, index) + e_rows, e_cols, e_non_reduced, e_reduced = layout_sub_figures(e, row_dims, col_dims, animate, overlay, offset_row, offset_col, positioning, indices, index) cols = max(cols, e_cols) rows = max(rows, e_rows) non_reduced &= e_non_reduced diff --git a/setup.py b/setup.py index 909845b42..7d203e08e 100644 --- a/setup.py +++ b/setup.py @@ -40,7 +40,7 @@ url='https://github.com/tum-pbs/PhiFlow', include_package_data=True, install_requires=[ - 'phiml==1.2.1', + 'phiml==1.5.1', 'matplotlib>=3.5.0', # also required by dash for color maps 'packaging', ], diff --git a/tests/commit/geom/test__sphere.py b/tests/commit/geom/test__sphere.py index fa6e9a678..f7ce75d6f 100644 --- a/tests/commit/geom/test__sphere.py +++ b/tests/commit/geom/test__sphere.py @@ -58,10 +58,8 @@ def test_reshaping_const_radius(self): s = expand(s, batch(b=100)) s = rename_dims(s, 'b', 'bat') s = unpack_dim(s, 'points', spatial(x=10, y=5)) - assert not s.radius.shape assert batch(bat=100) & spatial(x=10, y=5) & channel(vector='x,y') == s.shape s = pack_dims(s, 'x,y', instance('particles')) - assert not s.radius.shape assert batch(bat=100) & instance(particles=50) & channel(vector='x,y') == s.shape s = flatten(s) assert batch(bat=100) & instance(flat=50) & channel(vector='x,y') == s.shape