Write a filter¶

A filter is a Python function that takes a dataset and returns a field or a new dataset, running kernels to do the work. The built-in filters are written that way, with the same public API you have, so a filter of your own can do anything theirs do. This chapter rewrites five of them. Each rewrite gives the built-in's result in every part, including topology, geometry, fields and their units, sets, ghost marks, global ids, origin, spatial unit and time. The example program checks this, and so does the test suite (test_public_filters.py).

A filter has to compute the right values, and it also has to keep what the dataset means:

the positions are relative to an origin;
fields have units;
a piece's ghost cells and points are copies that other pieces own;
hidden cells hold no valid data;
the arrays live on a device.

Each section below shows where a filter of that kind deals with one of these. Write Mesh Algorithms covers the kernels themselves (views, for_each, ids); read it first.

The data¶

    w = td.wavelet(6)
    data = td.DataSet(w.topology, w.geometry, fields=dict(w.fields),
                      origin=(1.0e6, 2.0e6, 0.0), unit="m", time=3.0, time_unit="s")
    data.fields["cell_value"] = td.values_at_centers(data, "RTData")
    data.sets["some"] = td.Set(data, "cells", [0, 5, 9, 100])
    pieces = td.generate_global_ids(td.split(data, 4, ghost_layers=1))

The origin is far from the cells, so a filter that confuses world and relative positions gives visibly wrong answers. The dataset has a unit and a time, a cell set, and pieces with a layer of ghost cells and global ids. The example runs every filter on the image, on its tetrahedra and on a ghosted piece.

A field from positions: `point_elevation`¶

@tack.kernel
def _elevations(positions, lx, ly, lz, dx, dy, dz, length2, low, span, out):
    for i in range(out.shape[0]):
        p = positions[i]
        s = ((p[0] - lx) * dx + (p[1] - ly) * dy + (p[2] - lz) * dz) / length2
        out[i] = low + min(max(s, 0.0), 1.0) * span


def elevation(data, low_point, high_point, range):
    """Each point's position along the line from low_point to high_point,
    scaled to range: td.point_elevation."""
    # The two points are world positions; the kernel sees positions relative
    # to the dataset's origin. Convert on the host, in f64.
    low = np.asarray(low_point, float) - np.asarray(data.origin)
    direction = np.asarray(high_point, float) - np.asarray(data.origin) - low
    n = size_of(data.geometry.values)               # an implicit array has no .shape
    out = tack.empty((n,), data.dtype, device=data.device)
    _elevations(materialize(data.geometry.values), *map(float, low), *map(float, direction),
                float(direction @ direction), float(range[0]), float(range[1] - range[0]),
                out)
    return td.Field(data.geometry.space, out)

Positions are relative. A dataset's positions are origin + value, and kernels only see the values, which stay precise in f32 however far the origin is (see Understand a Dataset). The two end points come in world positions, so the filter moves them to the origin on the host, in f64, before any kernel sees them.
Allocate on the dataset's device. Use device=data.device, and data.dtype for values in the geometry's precision. Every public algorithm runs on its input's device. A filter that allocates on the default device instead fails when the dataset lives on another one: a kernel refuses arrays from two devices.
Arrays may be implicit. An image's geometry is a UniformCoordinates array, which computes positions instead of storing them, and has no .shape. size_of gives its length, and materialize gives a stored array for a kernel that takes plain arrays. A kernel can also index it directly, since implicit arrays are indexable in kernels.
The result is a Field on the geometry's space, one value per point.

A subset of cells: `threshold`¶

@tack.kernel
def _all_corners_inside(cells, u, lower, upper, keep):
    for c in cells:
        inside = 0
        for j in range(cells.NUM_POINTS):
            v = u.corner_value(c, j)
            if lower <= v and v <= upper:
                inside += 1
        keep[cells.entity_id(c)] = 1 if inside == cells.NUM_POINTS else 0


def threshold(data, name, lower, upper):
    """The cells whose corners all lie in [lower, upper]: td.threshold."""
    field = data.fields[name]
    lower = units.magnitude(lower, field.unit, "threshold's bounds")   # "300 K" or 300.0
    upper = units.magnitude(upper, field.unit, "threshold's bounds")
    keep = tack.empty((data.num_cells,), tack.i32, device=data.device)
    td.for_each(_all_corners_inside, data, "cells", field, lower, upper, keep)
    return td.extract_cells(data, td.Set.from_flags(data, "cells", keep))

Parameters take quantities. units.magnitude turns "300 K", (26.85, "degC") or a plain 300.0 into a number in the field's unit, so threshold(data, "T", "-123.15 degC", (250, "K")) works as the built-in's does. A plain number passes through untouched. For a length, use units.magnitude(value, data.unit, ...).
A kernel's choice becomes a set. The kernel writes one flag per cell at cells.entity_id(c), and Set.from_flags turns the flags into a Set. extract_cells then keeps those cells with all their fields, sets, ghost marks, global ids and frame, and drops hidden cells.

New geometry on the same cells: `shrink`¶

@tack.kernel
def _shrink_corners(positions, offsets, factor, n_cells):
    for c in range(n_cells):
        first = offsets[c]
        end = offsets[c + 1]
        center = positions[first]
        for k in range(first + 1, end):
            center += positions[k]
        center = center / (end - first)
        for k in range(first, end):
            positions[k] = center + factor * (positions[k] - center)


def shrink(data, factor):
    """Each cell shrunk toward its corners' mean: td.shrink."""
    l2 = td.L2(data)                                # each cell its own copy of its corners
    corners = td.project(data, data.geometry, l2)
    if data.num_cells:
        _shrink_corners(corners.values, corners.space.offsets, float(factor), data.num_cells)
    fields = {}
    for name, f in data.fields.items():
        if isinstance(f.space, td.Values) and f.space.on == "points":
            f = td.Field(td.H1(data), f.values, f.unit)
        fields[name] = td.project(data, f, l2) if isinstance(f.space, td.H1) else f
    # The same cells and points as the input: keep its marks, ids, origin,
    # unit and time.
    return td.keep_marks(data, td.DataSet(data.topology, corners, fields=fields,
                                          sets=dict(data.sets)))

shrink keeps the input's cells and points but builds a new DataSet around them. A dataset built by hand starts with nothing, so the filter hands it the input's ghost marks, global ids, origin, unit and time with keep_marks.

Forgetting them is a silent error. Contours of pieces without ghost cells cannot place their points' ghost marks, so they carry a note that the marks were lost, and reductions over such pieces refuse to run rather than count shared points twice. A filter that drops the note turns that refusal into a wrong answer. When a filter's output cells and points don't correspond to the input's at all, use drop_marks(data, out, reason) to record the loss instead.

Fields on faces, and units of results: `jump`¶

@tack.kernel
def _cell_jump(faces, values, out):
    for f in faces:
        inside = values[faces.side_cell(f, 0)]
        out[faces.entity_id(f)] = (values[faces.side_cell(f, 1)] - inside
                                   if faces.num_sides(f) == 2 else inside * 0.0)


def jump(data, name):
    """A cell field's difference across each face: td.jump."""
    field = data.fields[name]
    values = materialize(field.values)
    out = tack.empty((data.topology.faces().num_faces,), values.dtype, device=data.device)
    td.for_each(_cell_jump, data, "faces", values, out)
    # A difference of degrees Celsius is in kelvin; unknown stays unknown.
    return td.Field(td.Values(data, "faces"), out, units.delta(field.unit))

A field a filter computes gets the unit of what it is, so a difference gets units.delta: a difference of degrees Celsius is in kelvin. For other results:

units.derived((unit, power), ...) gives the unit of a product of powers (a gradient is the field's unit over the spatial unit);
units.named folds a unit to its SI name.

An unknown unit (None) gives None, never a guess, and passing no unit at all also gives None. The kernel here is a face kernel: side_cell(f, 0) is the cell on the face's inner side, and a boundary face has one side.

New cells: `triangulate`¶

_TRIANGLES = {int(td.shapes.TRIANGLE): [0, 1, 2], int(td.shapes.QUAD): [0, 1, 2, 0, 2, 3],
              int(td.shapes.PIXEL): [0, 1, 3, 0, 3, 2]}


@tack.kernel
def _piece_counts(cells, pieces, counts):
    for c in cells:
        counts[cells.entity_id(c)] = pieces


@tack.kernel
def _split(cells, table, pieces, starts, connectivity, sources):
    for c in cells:
        e = cells.entity_id(c)
        at = starts[e]
        for t in range(pieces):
            for v in range(3):
                connectivity[(at + t) * 3 + v] = cells.point_id(c, table[t * 3 + v])
            sources[at + t] = e


def triangulate(data):
    """Each polygon of a surface as triangles: td.triangulate."""
    data = td.visible(data)                         # hidden cells are never processed
    n, idt, device = data.num_cells, data.id_dtype, data.device
    groups = [g for g in data.topology.groups() if g.count and int(g.shape.ID) in _TRIANGLES]
    # Count each cell's triangles, scan the counts into starts, fill.
    counts = tack.zeros((n,), tack.i32, device=device)
    for g in groups:
        _piece_counts(g.view(), len(_TRIANGLES[int(g.shape.ID)]) // 3, counts)
    starts = tack.empty((n,), idt, device=device)
    count = exclusive_scan(counts, starts, n) if n else 0
    connectivity = tack.empty((3 * count,), idt, device=device)
    sources = tack.empty((count,), idt, device=device)
    for g in groups:
        table = _TRIANGLES[int(g.shape.ID)]
        _split(g.view(), tack.array(np.array(table, np.int32), device=device),
               len(table) // 3, starts, connectivity, sources)
    out = td.UnstructuredTopology(ConstantArray(int(td.shapes.TRIANGLE), count, dtype=tack.u8,
                                                device=device),
                                  CountingArray(count + 1, start=0, step=3, dtype=idt),
                                  connectivity, num_points=data.num_points)
    # The input's points; each triangle a piece of the cell `sources` names.
    return td.carry(data, out, points=td.Same(), cells=td.Pieces(sources))

Start from visible(data). It leaves out hidden cells (HIDDENCELL, data not valid there) and AMR cells a finer level covers. Every filter that makes cells starts there.
Count, scan, fill. One kernel counts each cell's output, an exclusive_scan turns the counts into each cell's first slot and the total, and a second kernel fills the slots. No atomics are needed, and the output is in input order on every backend.
A topology's arrays may be implicit. Every output cell is a triangle, so the types are a ConstantArray and the offsets a CountingArray. The ConstantArray is made with device=device, like every array the filter allocates: an implicit array lives on the device it is made for.
carry brings everything else. It takes the input dataset, the new topology, and a map for each kind of entity, saying where the output's entities came from. It then carries fields, sets, ghost marks, global ids and the frame by those maps:

Map	Output entity `i` is…	Used for
`Same()`	input entity `i`	points or cells kept as they are
`Take(ids)`	input entity `ids[i]`	a subset, a reordering
`Pieces(ids)`	a part of input cell `ids[i]`	cells cut from cells: triangles, tetrahedra
`Interpolate(ends, weights)`	a point between two input points	points made on edges: contours, clips

A kind of entity without a map carries no fields. faces_to_cells= says the output's cells are the input's faces (an extracted surface), and geometry= gives a geometry the filter computed itself.

A reduction over pieces: `entropy`¶

def entropy(pieces, name, bins=10):
    """The entropy of a point field's histogram over every piece: td.entropy."""
    blocks = tc.Blocks([p.device for p in pieces], pieces.world)
    # Owned values only: a point a ghost layer copies is counted once.
    values = {g: td.owned_values(p, name).numpy() for g, p in zip(blocks.local, pieces)}
    lo = blocks.all_reduce({g: v.min() for g, v in values.items()}, op="min")
    hi = blocks.all_reduce({g: v.max() for g, v in values.items()}, op="max")
    counts = {}
    for (g, v), piece in zip(values.items(), pieces):
        c, _ = histogram(tack.array(v, device=piece.device), bins=bins, range=(lo, hi))
        counts[g] = np.asarray(c.numpy() if hasattr(c, "numpy") else c, float)
    p = blocks.all_reduce(counts)                    # summed in piece order, on every rank
    p = p[p > 0] / p.sum()
    return 0.0 - float((p * np.log2(p)).sum())

Count each value once. owned_values gives the values of the cells and points a piece owns, leaving out its ghost copies and hidden cells, so a value shared by two pieces is counted once.
Reduce over the blocks. tack.communication.Blocks holds this rank's pieces, and all_reduce combines one value per piece in global piece order, so a float sum has the same bits however the pieces are spread over ranks. The same function runs unchanged under MPI: pass a World(comm) to the PartitionedDataSet.
Algorithms that need neighbours' data rather than a reduction (exchanges, ghost plans, routing points to pieces) are in Communicate Between Blocks.

What to keep, and what forgetting it gives¶

Concern	Do	Forgetting it gives
Origin	Convert world parameters with `data.origin` on the host; build datasets through `carry`, `keep_marks` or with `origin=`	Wrong positions, visible only far from zero
Device	Allocate with `device=data.device`	An error on a dataset on another device
Implicit arrays	`size_of`, `materialize`, or index in a kernel	`AttributeError` on `.shape`
Hidden cells	Start from `visible(data)`	Cells with invalid data processed
Ghost marks	`carry` maps, `keep_marks`, or `drop_marks`	Shared cells and points counted once per piece
Units	`units.magnitude` for parameters; `units.delta`/`derived` for results	A number in the wrong unit; a wrong unit on the result
Pieces	`owned_values`, `Blocks.all_reduce`	Ghost copies counted, or results that depend on the rank count

The Data Model reference documents each of these names.

Run and check¶

uv run python docs/examples/data_filters.py --arch cpu --check

The check holds each rewrite to its built-in, on the image, its tetrahedra and a ghosted piece, and a threshold with bounds in degrees Celsius and kelvin to the built-in's. Download the complete program.

Continue with Sample and Connect Simulations to connect filters into a workflow, or Communicate Between Blocks to write one that exchanges data between pieces.

Write a filter¶

The data¶

A field from positions: point_elevation¶

A subset of cells: threshold¶

New geometry on the same cells: shrink¶

Fields on faces, and units of results: jump¶

New cells: triangulate¶

A reduction over pieces: entropy¶