Staggered grids¶

Google Collab Book

Staggered grids are a key component of the marker-and-cell (MAC) method [Harlow and Welch 1965]. They sample the velocity components not at the cell centers but in staggered form at the corresponding face centers. Their main advantage is that the divergence of a cell can be computed exactly.

ΦFlow only stores valid velocity values in memory. This may require non-uniform tensors for the values since the numbers of horizontal and vertical faces are generally not equal. Depending on the boundary conditions, the outer-most values may also be redundant and, thus, not stored.

image

ΦFlow represents staggered grids as instances of StaggeredGrid. They have the same properties as CenteredGrid but the values field may reference a non-uniform tensor to reflect the varying number of x, y and z sample points.

In [1]:
# !pip install --quiet phiflow
from phi.flow import *

grid = StaggeredGrid(0, extrapolation.BOUNDARY, x=10, y=10)
grid.values
Out[1]:
(~vectorᵈ=x,y, xˢ=~(x=11, y=10) int64, yˢ=~(x=10, y=11) int64) const 0.0

Here, each component of the values tensor has one more sample point in the direction it is facing. If the extrapolation was extrapolation.ZERO, it would be one less (see above image).

Creating Staggered Grids¶

The StaggeredGrid constructor supports two modes:

  • Direct construction StaggeredGrid(values: Tensor, extrapolation, bounds). All required fields are passed as arguments and stored as-is. The values tensor must have the correct shape considering the extrapolation.
  • Construction by resampling StaggeredGrid(values: Any, extrapolation, bounds, resolution, **resolution). When specifying the resolution as a Shape or via keyword arguments, non-Tensor values can be passed for values, such as geometries, other fields, constants or functions (see the documentation).

Examples:

In [2]:
domain = dict(x=10, y=10, bounds=Box(x=1, y=1), extrapolation=extrapolation.ZERO)

grid = StaggeredGrid((1, -1), **domain)  # from constant vector
grid = StaggeredGrid(Noise(), **domain)  # sample analytic field
grid = StaggeredGrid(grid, **domain)  # resample existing field
grid = StaggeredGrid(lambda x: math.exp(-x), **domain)  # function value(location)
grid = resample(Sphere(x=0, y=0, radius=1), StaggeredGrid(0, **domain))  # no anti-aliasing
grid = resample(Sphere(x=0, y=0, radius=1), StaggeredGrid(0, **domain), soft=True)  # with anti-aliasing
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 3
      1 domain = dict(x=10, y=10, bounds=Box(x=1, y=1), extrapolation=extrapolation.ZERO)
----> 3 grid = StaggeredGrid((1, -1), **domain)  # from constant vector
      4 grid = StaggeredGrid(Noise(), **domain)  # sample analytic field
      5 grid = StaggeredGrid(grid, **domain)  # resample existing field

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:176, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
    174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
    175 assert values.shape.instance_rank == 0, f"Instance dimensions not supported for grids. Got values with shape {values.shape}"
--> 176 return Field(elements, values, extrapolation)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:46, in _FieldType.__call__(cls, geometry, values, boundary, variable_attrs, value_attrs, **sampling_kwargs)
     44         values = expand(wrap(values), non_batch(geometry) - 'vector')
     45 result = cls.__new__(cls, geometry, values, boundary, variable_attrs, value_attrs)
---> 46 result.__init__(geometry, values, boundary, variable_attrs, value_attrs)  # also calls __post_init__()
     47 return result

File <string>:8, in __init__(self, geometry, values, boundary, variable_attrs, value_attrs)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:85, in Field.__post_init__(self)
     84 def __post_init__(self):
---> 85     at = self.sampled_at
     86     if at in {'center', 'face'}:
     87         math.merge_shapes(self.values, non_batch(self.sampled_elements).non_channel)  # shape check

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/functools.py:998, in cached_property.__get__(self, instance, owner)
    996 val = cache.get(self.attrname, _NOT_FOUND)
    997 if val is _NOT_FOUND:
--> 998     val = self.func(instance)
    999     try:
   1000         cache[self.attrname] = val

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:399, in Field.sampled_at(self)
    397         if all(_size_equal(v_size, set_size) for v_size, set_size in zip(v_sizes, set_shape.sizes)):
    398             return name
--> 399 raise ValueError(f"Could not determine where the values of this Field are sampled. Geometry sets: {self.geometry.sets}, Field values shape: {v_shape}")

ValueError: Could not determine where the values of this Field are sampled. Geometry sets: {'center': (xˢ=10, yˢ=10), 'face': (~vectorᵈ=x,y, xˢ=~(x=11, y=10) int64, yˢ=~(x=10, y=11) int64), 'node': (xˢ=11, yˢ=11)}, Field values shape: (~vectorᵈ=x,y, xˢ=~(x=9, y=10) int64, yˢ=~(x=10, y=9) int64)

To construct a StaggeredGrid from NumPy arrays (or TensorFlow/PyTorch/Jax tensors), the tensors first need to be converted to ΦFlow tensors using tensor() or wrap().

In [3]:
vx = tensor(np.zeros([33, 32]), spatial('x,y'))
vy = tensor(np.zeros([32, 33]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.BOUNDARY)

vx = tensor(np.zeros([32, 32]), spatial('x,y'))
vy = tensor(np.zeros([32, 32]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.PERIODIC)

vx = tensor(np.zeros([31, 32]), spatial('x,y'))
vy = tensor(np.zeros([32, 31]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), 0)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 11
      9 vx = tensor(np.zeros([31, 32]), spatial('x,y'))
     10 vy = tensor(np.zeros([32, 31]), spatial('x,y'))
---> 11 StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), 0)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:176, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
    174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
    175 assert values.shape.instance_rank == 0, f"Instance dimensions not supported for grids. Got values with shape {values.shape}"
--> 176 return Field(elements, values, extrapolation)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:46, in _FieldType.__call__(cls, geometry, values, boundary, variable_attrs, value_attrs, **sampling_kwargs)
     44         values = expand(wrap(values), non_batch(geometry) - 'vector')
     45 result = cls.__new__(cls, geometry, values, boundary, variable_attrs, value_attrs)
---> 46 result.__init__(geometry, values, boundary, variable_attrs, value_attrs)  # also calls __post_init__()
     47 return result

File <string>:8, in __init__(self, geometry, values, boundary, variable_attrs, value_attrs)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:85, in Field.__post_init__(self)
     84 def __post_init__(self):
---> 85     at = self.sampled_at
     86     if at in {'center', 'face'}:
     87         math.merge_shapes(self.values, non_batch(self.sampled_elements).non_channel)  # shape check

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/functools.py:998, in cached_property.__get__(self, instance, owner)
    996 val = cache.get(self.attrname, _NOT_FOUND)
    997 if val is _NOT_FOUND:
--> 998     val = self.func(instance)
    999     try:
   1000         cache[self.attrname] = val

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:399, in Field.sampled_at(self)
    397         if all(_size_equal(v_size, set_size) for v_size, set_size in zip(v_sizes, set_shape.sizes)):
    398             return name
--> 399 raise ValueError(f"Could not determine where the values of this Field are sampled. Geometry sets: {self.geometry.sets}, Field values shape: {v_shape}")

ValueError: Could not determine where the values of this Field are sampled. Geometry sets: {'center': (xˢ=32, yˢ=32), 'face': (~vectorᵈ=x,y, xˢ=~(x=33, y=32) int64, yˢ=~(x=32, y=33) int64), 'node': (xˢ=33, yˢ=33)}, Field values shape: (~vectorᵈ=x,y, xˢ=~(x=31, y=32) int64, yˢ=~(x=32, y=31) int64)

Staggered grids can also be created from other fields using field.at() or @ by passing an existing StaggeredGrid.

Some field functions also return StaggeredGrids:

  • spatial_gradient() with type=StaggeredGrid
  • stagger()

Values Tensor¶

For non-periodic staggered grids, the values tensor is non-uniform to reflect the different number of sample points for each component.

In [4]:
grid.values
Out[4]:
(~vectorᵈ=x,y, xˢ=~(x=11, y=10) int64, yˢ=~(x=10, y=11) int64) const 0.0

Functions to get a uniform tensor:

  • uniform_values() interpolates the staggered values to the cell centers and returns a CenteredGrid
  • at_centers() interpolates the staggered values to the cell centers and returns a CenteredGrid
  • staggered_tensor() pads the internal tensor to an invariant shape with n+1 entries along all dimensions.
In [5]:
grid.uniform_values()
Out[5]:
(xˢ=11, yˢ=11, vectorᶜ=x,y) const 0.0

Slicing¶

Like tensors, grids can be sliced using the standard syntax. When selecting a vector component, such as x or y, the result is represented as a CenteredGrid with shifted locations.

In [6]:
grid.vector['x']  # select component
Out[6]:
(xˢ=11, yˢ=10) const 0.0 @ UniformGrid[(vectorᶜ=x,y, xˢ=11, yˢ=10)]:center, BC=zero-gradient

Grids do not support slicing along spatial dimensions because the result would be ambiguous with StaggeredGrids. Instead, slice the values directly.

In [7]:
grid.values.x[3:4]  # spatial slice
Out[7]:
(~vectorᵈ=x,y, xˢ=1, yˢ=~(x=10, y=11) int64) const 0.0
In [8]:
grid.values.x[0]  # spatial slice
Out[8]:
(~vectorᵈ=x,y, yˢ=~(x=10, y=11) int64) const 0.0

Slicing along batch dimensions has no special effect, this just slices the values.

In [9]:
grid.batch[0]  # batch slice
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[9], line 1
----> 1 grid.batch[0]  # batch slice

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/dataclasses/_dataclasses.py:63, in sliceable.<locals>.wrap.<locals>.__dataclass_getattr__(obj, name)
     61     raise RuntimeError(f"Evaluation of property '{type(obj).__name__}.{name}' failed.")
     62 else:
---> 63     raise AttributeError(f"'{type(obj)}' instance has no attribute '{name}'")

AttributeError: '<class 'phi.field._field.Field'>' instance has no attribute 'batch'

Fields can also be sliced using unstack(). This returns a tuple of all slices along a dimension.