def staggered_elements(resolution: Shape, bounds: Box, extrapolation: math.Extrapolation): cells = GridCell(resolution, bounds) grids = [] for dim in resolution.names: lower, upper = extrapolation.valid_outer_faces(dim) grids.append(cells.stagger(dim, lower, upper)) return geom.stack(grids, channel(staggered_direction=resolution.names))
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 sample(value: Geometry or Field or int or float or callable, resolution: Shape, box: Box, extrapolation=math.extrapolation.ZERO): if isinstance(value, Geometry): value = SoftGeometryMask(value) if isinstance(value, Field): elements = GridCell(resolution, box) data = value.sample_in(elements) else: if callable(value): x = GridCell(resolution, box).center value = value(x) value = wrap(value) data = math.zeros(resolution) + value return CenteredGrid(data, box, extrapolation)
def __getitem__(self, item: dict): values = self._values[{ dim: sel for dim, sel in item.items() if dim not in self.shape.spatial }] for dim, sel in item.items(): if dim in self.shape.spatial: raise AssertionError( "Cannot slice StaggeredGrid along spatial dimensions.") # sel = slice(sel, sel + 1) if isinstance(sel, int) else sel # values = [] # for vdim, val in zip(self.shape.spatial.names, self.values.unstack('vector')): # if vdim == dim: # values.append(val[{dim: slice(sel.start, sel.stop + 1)}]) # else: # values.append(val[{dim: sel}]) # values = math.stack(values, channel('vector')) extrapolation = self._extrapolation[item] bounds = GridCell(self._resolution, self._bounds)[item].bounds if 'vector' in item: selection = item['vector'] if isinstance(selection, str) and ',' in selection: selection = parse_dim_order(selection) if isinstance(selection, str): # single item name item_names = self.shape.get_item_names('vector', fallback_spatial=True) assert selection in item_names, f"Accessing field.vector['{selection}'] failed. Item names are {item_names}." selection = item_names.index(selection) if isinstance(selection, int): dim = self.shape.spatial.names[selection] comp_cells = GridCell(self.resolution, bounds).stagger( dim, *self.extrapolation.valid_outer_faces(dim)) return CenteredGrid(values, bounds=comp_cells.bounds, extrapolation=extrapolation) else: assert isinstance( selection, slice) and not selection.start and not selection.stop return StaggeredGrid(values, bounds=bounds, extrapolation=extrapolation)
def __init__(self, values: Tensor, resolution: Shape, bounds: Box, extrapolation=math.extrapolation.ZERO): SampledField.__init__(self, GridCell(resolution, bounds), values, extrapolation) self._bounds = bounds assert_same_rank( self.values.shape, bounds, 'data dimensions %s do not match box %s' % (self.values.shape, bounds))
def unstack(self, dimension='vector'): if dimension == 'vector': result = [] for dim, data in zip(self.resolution.spatial.names, self.values.vector.unstack()): comp_cells = GridCell(self.resolution, self._bounds).extend_symmetric(dim, 1) result.append( CenteredGrid(data, comp_cells.bounds, self.extrapolation)) return tuple(result) else: values = self.values.unstack(dimension) return tuple(self.with_(values=v) for v in values)
def cells(self): return GridCell(self.resolution, self.bounds)
def sample(value: Field or Geometry or callable or Tensor or float or int, resolution: Shape, bounds: Box, extrapolation=math.extrapolation.ZERO) -> 'StaggeredGrid': """ Creates a StaggeredGrid from `value`. `value` has to be one of the following: * Geometry: sets inside values to 1, outside to 0 * Field: resamples the Field to the staggered sample points * float, int: uses the value for all sample points * tuple, list: interprets the sequence as vector, used for all sample points * Tensor compatible with grid dims: uses tensor values as grid values Args: value: values to use for the grid resolution: grid resolution bounds: physical grid bounds extrapolation: return: Sampled values in staggered grid form matching domain resolution (Default value = math.extrapolation.ZERO) value: Field or Geometry or callable or Tensor or float or int: resolution: Shape: bounds: Box: Returns: Sampled values in staggered grid form matching domain resolution """ if isinstance(value, Geometry): value = HardGeometryMask(value) if isinstance(value, Field): assert_same_rank( value.spatial_rank, bounds.spatial_rank, 'rank of value (%s) does not match domain (%s)' % (value.spatial_rank, bounds.spatial_rank)) if isinstance(value, StaggeredGrid) and value.bounds == bounds and np.all( value.resolution == resolution): return value else: components = value.vector.unstack(bounds.spatial_rank) tensors = [] for dim, comp in zip(resolution.spatial.names, components): comp_cells = GridCell(resolution, bounds).extend_symmetric(dim, 1) comp_grid = CenteredGrid.sample(comp, comp_cells.resolution, comp_cells.bounds, extrapolation) tensors.append(comp_grid.values) return StaggeredGrid(math.channel_stack(tensors, 'vector'), bounds, extrapolation) else: # value is function or constant if callable(value): points = GridCell(resolution, bounds).face_centers() value = value(points) value = wrap(value) components = (value.staggered if 'staggered' in value.shape else value.vector).unstack(resolution.spatial_rank) tensors = [] for dim, component in zip(resolution.spatial.names, components): comp_cells = GridCell(resolution, bounds).extend_symmetric(dim, 1) tensors.append(math.zeros(comp_cells.resolution) + component) return StaggeredGrid(math.channel_stack(tensors, 'vector'), bounds, extrapolation)
def cells(self) -> GridCell: """ Returns the geometry of all cells as a `Box` object. The box will have spatial dimensions matching the resolution of the Domain, i.e. `domain.cells.shape == domain.resolution`. """ return GridCell(self.resolution, self.bounds)
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)