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)