def stagger(field: CenteredGrid,
face_function: Callable,
extrapolation: math.extrapolation.Extrapolation,
type: type = StaggeredGrid):
"""
Creates a new grid by evaluating `face_function` given two neighbouring cells.
One layer of missing cells is inferred from the extrapolation.
This method returns a Field of type `type` which must be either StaggeredGrid or CenteredGrid.
When returning a StaggeredGrid, the new values are sampled at the faces of neighbouring cells.
When returning a CenteredGrid, the new grid has the same resolution as `field`.
Args:
field: centered grid
face_function: function mapping (value1: Tensor, value2: Tensor) -> center_value: Tensor
extrapolation: extrapolation mode of the returned grid. Has no effect on the values.
type: one of (StaggeredGrid, CenteredGrid)
field: CenteredGrid:
face_function: Callable:
extrapolation: math.extrapolation.Extrapolation:
type: type: (Default value = StaggeredGrid)
Returns:
grid of type matching the `type` argument
"""
all_lower = []
all_upper = []
if type == StaggeredGrid:
for dim in field.shape.spatial.names:
all_upper.append(
math.pad(field.values, {dim: (0, 1)}, field.extrapolation))
all_lower.append(
math.pad(field.values, {dim: (1, 0)}, field.extrapolation))
all_upper = math.channel_stack(all_upper, 'vector')
all_lower = math.channel_stack(all_lower, 'vector')
values = face_function(all_lower, all_upper)
return StaggeredGrid(values, field.bounds, extrapolation)
elif type == CenteredGrid:
left, right = math.shift(field.values, (-1, 1),
padding=field.extrapolation,
stack_dim='vector')
values = face_function(left, right)
return CenteredGrid(values, field.bounds, extrapolation)
else:
raise ValueError(type)
def unstack_staggered_tensor(data: Tensor) -> TensorStack:
sliced = []
for dim, component in zip(data.shape.spatial.names,
data.unstack('vector')):
sliced.append(component[{
d: slice(None, -1)
for d in data.shape.spatial.without(dim).names
}])
return math.channel_stack(sliced, 'vector')
def _rotate(self, location):
sin = math.sin(self.angle)
cos = math.cos(self.angle)
y, x = location.vector.unstack()
if GLOBAL_AXIS_ORDER.is_x_first:
x, y = y, x
rot_x = cos * x - sin * y
rot_y = sin * x + cos * y
return math.channel_stack([rot_y, rot_x], 'vector')
def stack_staggered_components(data: Tensor) -> Tensor:
padded = []
for dim, component in zip(data.shape.spatial.names,
data.unstack('vector')):
padded.append(
math.pad(
component,
{d: (0, 1)
for d in data.shape.spatial.without(dim).names},
mode=math.extrapolation.ZERO))
return math.channel_stack(padded, 'vector')
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if reduce_channels:
shape = self._shape.non_channel.without(reduce_channels)
assert len(reduce_channels) == 1
geoms = geometry.unstack(reduce_channels[0])
assert all(isinstance(g, GridCell) for g in geoms)
components = [self.grid_sample(g.resolution, g.grid_size, shape) for g in geoms]
return math.channel_stack(components, 'vector')
if isinstance(geometry, GridCell):
return self.grid_sample(geometry.resolution, geometry.grid_size)
raise NotImplementedError(f"{type(geometry)} not supported. Only GridCell allowed.")
def closest_values(self, points: Tensor, reduce_channels=()):
if not reduce_channels:
channels = [
component.sample_at(points) for component in self.unstack()
]
else:
assert len(reduce_channels) == 1
points = points.unstack(reduce_channels[0])
channels = [
component.closest_values(p)
for p, component in zip(points, self.unstack())
]
return math.channel_stack(channels, 'vector')
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if geometry == self.elements and reduce_channels:
return self.values
if not reduce_channels:
channels = [
component.sample_in(geometry) for component in self.unstack()
]
else:
assert len(reduce_channels) == 1
geometries = geometry.unstack(reduce_channels[0])
channels = [
component.sample_in(g)
for g, component in zip(geometries, self.unstack())
]
return math.channel_stack(channels, 'vector')
def extrapolate_valid(grid: GridType,
valid: GridType,
distance_cells=1) -> tuple:
"""
Extrapolates values of `grid` which are marked by nonzero values in `valid` using `phi.math.extrapolate_valid_values().
If `values` is a StaggeredGrid, its components get extrapolated independently.
Args:
grid: Grid holding the values for extrapolation
valid: Grid (same type as `values`) marking the positions for extrapolation with nonzero values
distance_cells: Number of extrapolation steps
Returns:
grid: Grid with extrapolated values.
valid: binary Grid marking all valid values after extrapolation.
"""
assert isinstance(
valid, type(grid)), 'Type of valid Grid must match type of grid.'
if isinstance(grid, CenteredGrid):
new_values, new_valid = extrapolate_valid_values(
grid.values, valid.values, distance_cells)
return grid.with_(values=new_values), valid.with_(values=new_valid)
elif isinstance(grid, StaggeredGrid):
new_values = []
new_valid = []
for cgrid, cvalid in zip(grid.unstack('vector'),
valid.unstack('vector')):
new_tensor, new_mask = extrapolate_valid(
cgrid, valid=cvalid, distance_cells=distance_cells)
new_values.append(new_tensor.values)
new_valid.append(new_mask.values)
return grid.with_(
values=math.channel_stack(new_values, 'vector')), valid.with_(
values=math.channel_stack(new_valid, 'vector'))
else:
raise NotImplementedError()
def sample_at(self, points, reduce_channels=()) -> math.Tensor:
distances = points - self.location
strength = self.strength if self.falloff is None else self.strength * self.falloff(
distances)
if reduce_channels:
assert len(reduce_channels) == 1
velocities = [
math.cross_product(strength, dist).vector[i]
for i, dist in enumerate(distances.unstack(reduce_channels[0]))
] # TODO this is inefficient, computes components that are discarded
velocity = math.channel_stack(velocities, 'vector')
else:
velocity = math.cross_product(strength, distances)
velocity = math.sum(velocity,
self.location.shape.batch.without(points.shape))
return velocity
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if not reduce_channels:
if geometry == self.elements:
return self.values
elif isinstance(geometry, GridCell):
return self._grid_scatter(geometry.bounds, geometry.resolution)
elif isinstance(geometry, GeometryStack):
sampled = [self.sample_at(g) for g in geometry.geometries]
return math.batch_stack(sampled, geometry.stack_dim_name)
else:
raise NotImplementedError()
else:
assert len(reduce_channels) == 1
components = self.unstack('vector') if 'vector' in self.shape else (self,) * geometry.shape.get_size(reduce_channels[0])
sampled = [c.sample_in(p) for c, p in zip(components, geometry.unstack(reduce_channels[0]))]
return math.channel_stack(sampled, 'vector')
def downsample2x(grid: Grid) -> GridType:
if isinstance(grid, CenteredGrid):
values = math.downsample2x(grid.values, grid.extrapolation)
return CenteredGrid(values, grid.bounds, grid.extrapolation)
elif isinstance(grid, StaggeredGrid):
values = []
for dim, centered_grid in zip(grid.shape.spatial.names,
grid.unstack()):
odd_discarded = centered_grid.values[{dim: slice(None, None, 2)}]
others_interpolated = math.downsample2x(
odd_discarded,
grid.extrapolation,
dims=grid.shape.spatial.without(dim))
values.append(others_interpolated)
return StaggeredGrid(math.channel_stack(values, 'vector'), grid.bounds,
grid.extrapolation)
else:
raise ValueError(type(grid))
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor:
if reduce_channels:
assert len(reduce_channels) == 1
geometries = geometry.unstack(reduce_channels[0])
components = self.vector.unstack(len(geometries))
sampled = [c.sample_in(g) for c, g in zip(components, geometries)]
return math.channel_stack(sampled, 'vector')
if isinstance(geometry, GeometryStack):
sampled = [self.sample_in(g) for g in geometry.geometries]
return math.batch_stack(sampled, geometry.stack_dim_name)
if isinstance(geometry, GridCell):
if self.elements == geometry:
return self.values
elif math.close(self.dx, geometry.size):
fast_resampled = self._shift_resample(geometry.resolution,
geometry.bounds)
if fast_resampled is not NotImplemented:
return fast_resampled
return self.sample_at(geometry.center, reduce_channels)
def sample_at(self, points, reduce_channels=()) -> Tensor:
local_points = self.box.global_to_local(points) * self.resolution - 0.5
if len(reduce_channels) == 0:
return math.grid_sample(self.values, local_points,
self.extrapolation)
else:
assert self.shape.channel.sizes == points.shape.get_size(
reduce_channels)
if len(reduce_channels) > 1:
raise NotImplementedError(
f"{len(reduce_channels)} > 1. Only 1 reduced channel allowed."
)
channels = []
for i, channel in enumerate(self.values.vector.unstack()):
channels.append(
math.grid_sample(channel, local_points[{
reduce_channels[0]:
i
}], self.extrapolation))
return math.channel_stack(channels, 'vector')
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 face_centers(self, staggered_name='staggered'):
face_centers = [
self.extend_symmetric(dim, 1).center
for dim in self.shape.spatial.names
]
return math.channel_stack(face_centers, staggered_name)