def test_expand_copy_item_names(self): a = math.zeros(channel(vector=2)) try: math.expand(a, channel(vector=3)) self.fail() except IncompatibleShapes: pass b = math.expand(a, channel(vector='x,y')) self.assertEqual(('x', 'y'), b.vector.item_names) try: math.expand(b, channel(vector='a,b')) self.fail() except IncompatibleShapes: pass
def vector_grid(self, value: Field or Tensor or Number or Geometry or callable = 0., type: type = CenteredGrid, extrapolation: math.Extrapolation or str = 'vector') -> CenteredGrid or StaggeredGrid: """ Creates a vector grid matching the resolution and bounds of the domain. The grid is created from the given `value` which must be one of the following: * Number (int, float, complex or zero-dimensional tensor): all grid values will be equal to `value`. This has a near-zero memory footprint. * Field: the given value is resampled to the grid cells of this Domain. * Tensor with spatial dimensions matcing the domain resolution: grid values will equal `value`. * Geometry: grid values are determined from the volume overlap between grid cells and geometry. Non-overlapping = 0, fully enclosed grid cell = 1. * function(location: Tensor) returning one of the above. The returned grid will have a vector dimension with size equal to the rank of the domain. Args: value: constant, Field, Tensor or function specifying the grid values type: class of Grid to create, must be either CenteredGrid or StaggeredGrid extrapolation: (optional) grid extrapolation, defaults to Domain.boundaries['vector'] Returns: Grid of specified type """ extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation] result = type(value, resolution=self.resolution, bounds=self.bounds, extrapolation=extrapolation) if result.shape.channel_rank == 0: result = result.with_values(math.expand(result.values, channel(vector=self.rank))) else: assert result.shape.get_size('vector') == self.rank return result
def concat(geometries: tuple or list, dim: str, sizes: tuple or list or None = None): """ Concatenates multiple geometries of the same type. Args: geometries: sequence of `phi.geom.Geometry` objects of the same type sizes: implicit dim: dimension to concatenate Returns: New `phi.geom.Geometry` object """ if all(isinstance(g, type(geometries[0])) for g in geometries): characteristics = [g.__characteristics__() for g in geometries] if sizes is not None: characteristics = [{ key: math.expand(val, dim, size) for key, val in c.items() } for c, size in zip(characteristics, sizes)] new_attributes = {} for key in characteristics[0].keys(): concatenated = math.concat([c[key] for c in characteristics], dim) new_attributes[key] = concatenated return geometries[0].__with__(**new_attributes) else: raise NotImplementedError()
def __init__(self, lower: Tensor or float or int = None, upper: Tensor or float or int = None, **size: int or Tensor): """ Args: lower: physical location of lower corner upper: physical location of upper corner **size: Upper l """ if lower is not None: self._lower = wrap(lower) if upper is not None: self._upper = wrap(upper) else: lower = [] upper = [] for item in size.values(): if isinstance(item, (tuple, list)): assert len( item ) == 2, f"Box kwargs must be either dim=upper or dim=(lower,upper) but got {item}" lo, up = item lower.append(lo) upper.append(up) elif item is None: lower.append(-INF) upper.append(INF) else: lower.append(0) upper.append(item) lower = [-INF if l is None else l for l in lower] upper = [INF if u is None else u for u in upper] self._upper = math.wrap(upper, math.channel(vector=tuple(size.keys()))) self._lower = math.wrap(lower, math.channel(vector=tuple(size.keys()))) vector_shape = self._lower.shape & self._upper.shape self._lower = math.expand(self._lower, vector_shape) self._upper = math.expand(self._upper, vector_shape) if self.size.vector.item_names is None: warnings.warn( "Creating a Box without item names prevents certain operations like project()", DeprecationWarning, stacklevel=2)
def concat(*fields: SampledField, dim: str): assert all(isinstance(f, SampledField) for f in fields) assert all(isinstance(f, type(fields[0])) for f in fields) if any(f.extrapolation != fields[0].extrapolation for f in fields): raise NotImplementedError("Concatenating extrapolations not supported") if isinstance(fields[0], Grid): values = math.concat([f.values for f in fields], dim=dim) return fields[0].with_(values=values) elif isinstance(fields[0], PointCloud): elements = geom.concat([f.elements for f in fields], dim, sizes=[f.shape.get_size(dim) for f in fields]) values = math.concat([ math.expand(f.values, dim, f.shape.get_size(dim)) for f in fields ], dim) colors = math.concat( [math.expand(f.color, dim, f.shape.get_size(dim)) for f in fields], dim) return fields[0].with_(elements=elements, values=values, color=colors) raise NotImplementedError(type(fields[0]))
def test_boolean_mask_batched(self): for backend in BACKENDS: with backend: x = math.expand(math.range(spatial('x'), 4), batch(batch=2)) * math.tensor([1, -1]) mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x')) selected = math.boolean_mask(x, 'x', mask) expected_0 = math.tensor([(0, -0), (2, -2)], spatial('x'), channel('vector')) expected_1 = math.tensor([(1, -1)], spatial('x'), channel('vector')) math.assert_close(selected.batch[0], expected_0, msg=backend.name) math.assert_close(selected.batch[1], expected_1, msg=backend.name) math.assert_close(selected, x.x[mask], msg=backend.name)
def test_scatter_1d(self): for backend in BACKENDS: with backend: base = math.ones(spatial(x=4)) indices = math.wrap([1, 2], instance('points')) values = math.wrap([11, 12], instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, [1, 11, 12, 1]) updated = math.scatter(base, indices, values, mode='add', outside_handling='undefined') math.assert_close(updated, [1, 12, 13, 1]) # with vector dim indices = math.expand(indices, channel(vector=1)) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, [1, 11, 12, 1])
def concat(fields: List[SampledFieldType], dim: Shape) -> SampledFieldType: """ Concatenates the given `SampledField`s along `dim`. See Also: `stack()`. Args: fields: List of matching `SampledField` instances. dim: Concatenation dimension as `Shape`. Size is ignored. Returns: `SampledField` matching concatenated fields. """ assert all(isinstance(f, SampledField) for f in fields) assert all(isinstance(f, type(fields[0])) for f in fields) if any(f.extrapolation != fields[0].extrapolation for f in fields): raise NotImplementedError("Concatenating extrapolations not supported") if isinstance(fields[0], Grid): values = math.concat([f.values for f in fields], dim) return fields[0].with_values(values) elif isinstance(fields[0], PointCloud): elements = geom.concat([f.elements for f in fields], dim, sizes=[f.shape.get_size(dim) for f in fields]) values = math.concat( [math.expand(f.values, f.shape.only(dim)) for f in fields], dim) colors = math.concat( [math.expand(f.color, f.shape.only(dim)) for f in fields], dim) return PointCloud(elements=elements, values=values, color=colors, extrapolation=fields[0].extrapolation, add_overlapping=fields[0]._add_overlapping, bounds=fields[0]._bounds) raise NotImplementedError(type(fields[0]))
def expand_staggered(values: Tensor, resolution: Shape, extrapolation: math.Extrapolation): """ Add missing spatial dimensions to `values` """ cells = GridCell( resolution, Box( 0, math.wrap((1, ) * resolution.rank, channel(vector=resolution.names)))) components = values.vector.unstack(resolution.spatial_rank) tensors = [] for dim, component in zip(resolution.spatial.names, components): comp_cells = cells.stagger(dim, *extrapolation.valid_outer_faces(dim)) tensors.append(math.expand(component, comp_cells.resolution)) return math.stack(tensors, channel(vector=resolution.names))
def test_collapsed_non_uniform_tensor(self): non_uniform = math.stack( [math.zeros(spatial(a=2)), math.ones(spatial(a=3))], batch('b')) e = math.expand(non_uniform, channel('vector')) assert e.shape.without('vector') == non_uniform.shape
def grad(_x, _y, df): return math.flatten(math.expand(df * 0, batch(tmp=2))),
def __init__(self, values: Any, extrapolation: Any = 0., bounds: Box = None, resolution: int or Shape = None, **resolution_: int or Tensor): """ Args: values: Values to use for the grid. Has to be one of the following: * `phi.geom.Geometry`: sets inside values to 1, outside to 0 * `Field`: resamples the Field to the staggered sample points * `Number`: uses the value for all sample points * `tuple` or `list`: interprets the sequence as vector, used for all sample points * `phi.math.Tensor` compatible with grid dims: uses tensor values as grid values * Function `values(x)` where `x` is a `phi.math.Tensor` representing the physical location. The spatial dimensions of the grid will be passed as batch dimensions to the function. extrapolation: The grid extrapolation determines the value outside the `values` tensor. Allowed types: `float`, `phi.math.Tensor`, `phi.math.extrapolation.Extrapolation`. bounds: Physical size and location of the grid as `phi.geom.Box`. resolution: Grid resolution as purely spatial `phi.math.Shape`. **resolution_: Spatial dimensions as keyword arguments. Typically either `resolution` or `spatial_dims` are specified. """ if resolution is None and not resolution_: assert isinstance( values, math.Tensor ), "Grid resolution must be specified when 'values' is not a Tensor." resolution = values.shape.spatial bounds = bounds or Box(0, math.wrap(resolution, channel('vector'))) elements = GridCell(resolution, bounds) else: if isinstance(resolution, int): assert not resolution_, "Cannot specify keyword resolution and integer resolution at the same time." resolution = spatial( **{ dim: resolution for dim in bounds.size.shape.get_item_names('vector') }) resolution = (resolution or math.EMPTY_SHAPE) & spatial(**resolution_) bounds = bounds or Box(0, math.wrap(resolution, channel('vector'))) elements = GridCell(resolution, bounds) if isinstance(values, math.Tensor): values = math.expand(values, resolution) elif isinstance(values, Geometry): values = reduce_sample(HardGeometryMask(values), elements) elif isinstance(values, Field): values = reduce_sample(values, elements) elif callable(values): values = math.map_s2b(values)(elements.center) assert isinstance( values, math.Tensor ), f"values function must return a Tensor but returned {type(values)}" else: if isinstance( values, (tuple, list)) and len(values) == resolution.rank: values = math.tensor(values, channel(vector=resolution.names)) values = math.expand(math.tensor(values), resolution) if values.dtype.kind not in (float, complex): values = math.to_float(values) assert resolution.spatial_rank == bounds.spatial_rank, f"Resolution {resolution} does not match bounds {bounds}" Grid.__init__(self, elements, values, extrapolation, values.shape.spatial, bounds)