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)