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:
lo_valid, up_valid = extrapolation.valid_outer_faces(dim)
width_lower = {dim: (int(lo_valid), int(up_valid) - 1)}
width_upper = {
dim: (int(lo_valid or up_valid) - 1, int(lo_valid
and up_valid))
}
all_lower.append(
math.pad(field.values, width_lower, field.extrapolation))
all_upper.append(
math.pad(field.values, width_upper, field.extrapolation))
all_upper = math.stack(all_upper, channel('vector'))
all_lower = math.stack(all_lower, channel('vector'))
values = face_function(all_lower, all_upper)
result = StaggeredGrid(values,
bounds=field.bounds,
extrapolation=extrapolation)
assert result.shape.spatial == field.shape.spatial
return result
elif type == CenteredGrid:
left, right = math.shift(field.values, (-1, 1),
padding=field.extrapolation,
stack_dim=channel('vector'))
values = face_function(left, right)
return CenteredGrid(values,
bounds=field.bounds,
extrapolation=extrapolation)
else:
raise ValueError(type)
def shift(grid: CenteredGrid, offsets: tuple, stack_dim='shift'):
"""
Wraps :func:`math.shift` for CenteredGrid.
Args:
grid: CenteredGrid:
offsets: tuple:
stack_dim: (Default value = 'shift')
Returns:
"""
data = math.shift(grid.values,
offsets,
padding=grid.extrapolation,
stack_dim=stack_dim)
return [
CenteredGrid(data[i], grid.box, grid.extrapolation)
for i in range(len(offsets))
]