Halos: stencils on patches

Download notebookView on GitHubRun when the site was built.

A field allocated with halo=h has h cells (or nodes) around every patch. Fill them with exchange_halos, and a finite-difference stencil can be applied to each patch on its own, right up to its edges. restrict goes the other way, from fine patches to the coarse points under them.

import numpy as np

from patchwork import AMRHierarchy, plot_patches

Filling the halos

A halo point is filled from a patch of the same level where there is one, across periodic faces too; elsewhere by bilinear interpolation of the level below. Halo points outside the domain are left as they are, for the boundary condition. Here a linear function is set on the patches’ own nodes only, and the halos start at NaN.

mesh = AMRHierarchy(base_shape=(16, 16))
mesh.add_patch(level=1, box=((8, 20), (0, 32)))  # from wall to wall in y
mesh.add_patch(level=1, box=((20, 28), (8, 20)))  # its neighbour
mesh.add_patch(level=2, box=((24, 40), (16, 32)))
mesh.allocate_field("u", halo=1, centering="node")


def u(x, y):
    return 1.0 + 2.0 * x + 3.0 * y


for patch in mesh:
    storage = patch.field("u").local_with_halos
    storage[...] = np.nan
    x, y = patch.local_node_coordinates(mesh, "u")
    patch.field("u").local[...] = u(x, y)

mesh.exchange_halos("u")

fine = mesh.patches[-1].field("u").local_with_halos
print("NaN halo points left on the level-2 patch:", int(np.isnan(fine).sum()))
print(
    "NaN halo points on the base, outside the domain:",
    int(np.isnan(mesh.patches[0].field("u").local_with_halos).sum()),
)
fig, ax = plot_patches(mesh, field="u")

NaN halo points left on the level-2 patch: 0
NaN halo points on the base, outside the domain: 72

A stencil on every patch

With the halos filled, the centred difference $\partial u / \partial x$ can be taken on all of a patch’s own nodes. For $u = 1 + 2x + 3y$ it is 2 everywhere, also on the nodes at a patch’s edge, whose neighbour comes from another patch or from the level below.

Bilinear interpolation is exact for a linear function only. For a curved one, a halo point filled from the level below carries the coarse level’s interpolation error, and a stencil that reaches into it sees that error: place patch edges where the solution is smooth on the coarser level.

for patch in mesh.patches[1:]:
    dx, _ = mesh.spacing(patch.level)
    storage = patch.field("u").local_with_halos
    derivative = (storage[2:, 1:-1] - storage[:-2, 1:-1]) / (2 * dx)
    print(
        f"level {patch.level} {patch.box}: du/dx between "
        f"{derivative.min():.12f} and {derivative.max():.12f}"
    )

level 1 ((8, 20), (0, 32)): du/dx between 2.000000000000 and 2.000000000000
level 1 ((20, 28), (8, 20)): du/dx between 2.000000000000 and 2.000000000000
level 2 ((24, 40), (16, 32)): du/dx between 2.000000000000 and 2.000000000000

Restriction

restrict gives coarse cells the mean of the fine cells over them, and coarse nodes the fine node at the same place, finest level first. Afterwards the coarse level agrees with the fine one wherever they overlap; call exchange_halos after it if the halos are needed.

for patch in mesh:
    if patch.level == 0:
        patch.field("u").local[...] = 0.0  # forget the coarse values

mesh.restrict("u")

base = mesh.patches[0]
x, y = base.local_node_coordinates(mesh, "u")
restored = ~np.isclose(base.field("u").local, 0.0)
print("base nodes under a finer patch, now with the fine value:", int(restored.sum()))
print(
    "largest difference to u there:",
    float(np.abs(base.field("u").local - u(x, y))[restored].max()),
)

base nodes under a finer patch, now with the fine value: 147
largest difference to u there: 0.0

Both calls are collective under MPI and work wherever the patches are: each patch can live on its own ranks (add_patch(..., ranks=...)), and the data travel in one exchange over the hierarchy’s communicator. The schedules are built at the first call for a field and kept until a patch is added.