Mixed Assembly

Many PDE systems couple several unknown fields in one weak form: velocity and pressure in Stokes / Navier-Stokes, displacement and pressure in poroelasticity, velocity, pressure and temperature in Boussinesq convection. Discretizing them stably usually requires different function spaces per field — the classic example is the Taylor-Hood pair (quadratic velocity, linear pressure), which satisfies the inf-sup (LBB) condition that equal-order pairs violate.

MixedElementAssembler assembles such multi-field bilinear forms into one block sparse matrix. You declare the fields once, write the scalar integrand of the coupled weak form as a single forward, and everything downstream — boundary conditions via Condenser, .solve(), autograd — works on the block system exactly as it does for a single field:

from tensormesh import Condenser, Field, Mesh, MixedElementAssembler

class StokesAssembler(MixedElementAssembler):
    fields = [
        Field(trial="u", test="v", order=2, components=2),  # P2 velocity
        Field(trial="p", test="q", order=1),                # P1 pressure
    ]

    def __post_init__(self, mu=1.0):
        self.mu = mu

    def forward(self, gradu, p, gradv, q):
        return self.mu * (gradu * gradv).sum() \
            - p * gradv.diagonal().sum() \
            - q * gradu.diagonal().sum()

mesh = Mesh.gen_rectangle(chara_length=0.05, order=2).double()
assembler = StokesAssembler.from_mesh(mesh, mu=1.0)
K = assembler()          # SparseMatrix over all velocity + pressure DOFs

The rest of this page walks through the pieces: field declarations, the integrand contract, the block DOF layout, load vectors, the data channels, and generalized order pairs — fields whose polynomial order differs from the mesh order.

Declaring fields

fields is a class attribute: a list of Field declarations, one per unknown. Each field names its trial and test function (the argument names your forward will use), its polynomial order, and its number of vector components:

fields = [
    Field(trial="u", test="v", order=2, components=2),
    Field(trial="p", test="q", order=1),
]

Three conventions to remember:

  • Trial indexes columns, test indexes rows. The block \(K_{\beta\alpha}\) couples trial field \(\alpha\) (columns) to test field \(\beta\) (rows). For a symmetric weak form you never notice; for non-symmetric couplings (convection, divergence constraints) the declaration makes the orientation explicit.

  • Names are global. All trial and test names must be distinct, valid identifiers, must not be "x" (reserved for coordinates) and must not start with "grad" (gradients are derived as "grad" + name).

  • Order is per-field and independent of the mesh order — see Generalized order pairs.

The scalar integrand

forward returns the scalar integrand of the coupled bilinear form at one quadrature point — the same by-name dispatch as ElementAssembler (see The weak-form contract), but the field arguments are now tensor-valued:

Argument

Per-point shape

Meaning

scalar field value (p, q)

0-d []

the field at the quadrature point

scalar field gradient (gradp)

[D]

physical-space gradient

vector field value (u, v)

[c]

all \(c\) components together

vector field gradient (gradu)

[c, D]

the Jacobian: row \(i\) is \(\nabla u_i\)

so gradu.diagonal().sum() is \(\nabla\cdot u\), (gradu * gradv).sum() is \(\nabla u : \nabla v\), and (gradu @ w).dot(v) is the convection term \((w\cdot\nabla)u\cdot v\). Only the arguments you name are computed: a pass that touches neither field of a block is skipped, so absent couplings (e.g. the zero pressure-pressure block of Stokes) cost nothing beyond their sparsity-pattern slot.

Important

The integrand must be bilinear: every term must contain exactly one trial factor and one test factor. Internally each \((\alpha, \beta)\) block is extracted by evaluating the integrand with the other fields zeroed — an identity for bilinear forms, silently wrong otherwise. The assembler guards this with a cheap zero-evaluation check and raises if your integrand has a constant or non-bilinear part. Nonlinear weak forms are handled the usual FEM way: linearize (Picard / Newton) and feed the previous iterate in as data, e.g. the lagged convection velocity w below.

A Picard-linearized Navier-Stokes step, for example, adds one data argument to the Stokes integrand (see Lid-Driven Cavity for the full script):

def forward(self, gradu, p, v, gradv, q, w):
    convection = self.rho * (gradu @ w).dot(v)
    diffusion = self.mu * (gradu * gradv).sum()
    return convection + diffusion \
        - p * gradv.diagonal().sum() \
        - q * gradu.diagonal().sum()

Building and calling

from_mesh() builds the assembler; extra positional/keyword arguments go to your __post_init__ (physical parameters), exactly like the single-field assemblers:

assembler = StokesAssembler.from_mesh(mesh, mu=1.0)
K = assembler()                       # -> SparseMatrix  [N, N]

The default quadrature degree is 2 * max(field orders) — exact for the products of two highest-order basis functions on affine elements. Pass quadrature_order= to override (e.g. for data-weighted or convective terms); the shipped tables go up to degree 7.

The result is a square SparseMatrix of size \(N = \sum_f n_f\, c_f\). Repeated calls (Picard iterations, time steps) reuse the same cached index tensors, so a Condenser built once keeps its condensation layout across calls.

Note

Mixed assembly is currently single-device: the distributed() decorator (Distributed FEM) covers the single-field assemblers only. Distributing the block layout across ranks needs a global DOF numbering layer that does not exist yet — it is a planned ROADMAP item.

The block DOF layout

The global DOF vector stacks the fields in declaration order, node-major within each field:

\[\mathrm{dof} \;=\; \mathrm{offset}_f + n_{\mathrm{local}}\cdot c_f + \mathrm{comp}.\]

You never index this by hand — assembler.layout (a BlockLayout) provides the vocabulary:

layout = assembler.layout

layout.n_dofs                      # total N
layout.names                       # ["u", "p"]
layout.n_nodes("u")                # nodes of the P2 velocity space
layout.points("u")                 # their coordinates  [n_u, D]

# masks and indices (all in global block-DOF numbering)
layout.dof_mask("u")                          # every u DOF
layout.dof_mask("u", node_mask, component=0)  # u_x on selected nodes
layout.dof_mask("p", where=lambda x: x[:, 0] < 1e-6)   # coordinate predicate
layout.boundary_mask("u")                     # topological boundary DOFs
layout.dof_index("p", int(layout.node_ids("p")[0]))    # a single DOF (pressure pin)

# moving between the block vector and per-field tensors
fields = layout.split(sol)         # {"u": [n_u, 2], "p": [n_p]}
sol = layout.cat(u=vel, p=0.0)     # assemble a block vector
p_on_mesh = layout.prolong("p", fields["p"])   # P1 -> all mesh points
f_on_p = layout.restrict("p", point_values)    # mesh points -> P1 nodes

Boundary conditions then go through the unchanged Condenser — a mixed system is just a bigger DOF vector with a boolean mask:

bc_mask = layout.dof_mask("u", mesh.boundary_mask)     # no-slip walls
bc_mask[layout.dof_index("p", int(layout.node_ids("p")[0]))] = True

bc_val = torch.zeros(layout.n_dofs, dtype=torch.float64)
bc_val[layout.dof_mask("u", is_lid, component=0)] = 1.0

condenser = Condenser(bc_mask, bc_val[bc_mask])
K_, f_ = condenser(K, f)
sol = condenser.recover(K_.solve(f_))

boundary_mask detects the boundary topologically (a facet referenced by exactly one cell), so it also works for meshes without is_boundary point data — e.g. gmsh imports — and for generalized-order fields whose DOFs are not mesh nodes.

Load vectors

assemble_vector() is the linear form \(l(v)\) of the mixed space: same dispatch, test arguments only, output a dense [N] vector in the same block layout. Override forward_vector for the “natural” RHS of the problem, or pass a one-off func=:

class TransientNS(MixedElementAssembler):
    fields = [...]

    def forward(self, u, gradu, p, v, gradv, q, w):
        ...  # backward-Euler matrix side

    def forward_vector(self, v, uprev):
        return self.rho / self.dt * uprev.dot(v)

f = assembler.assemble_vector(field_data={"uprev": ("u", u_prev)})

# one-off linear forms: e.g. L2-project the vorticity of a P2
# velocity onto the P1 pressure space
omega_rhs = assembler.assemble_vector(
    func=lambda q, gradw: (gradw[1, 0] - gradw[0, 1]) * q,
    field_data={"w": ("u", velocity)},
)

Because the load is integrated in each field’s own space, sources on a generalized-order field are consistent by construction — no manual interpolation between spaces.

Data channels

All the single-field channels work in forward / forward_vector — x, point_data (interpolated with the mesh-order basis, grad{key} available), element_data, scalar_data — plus one new channel:

field_data={"name": (field, values)}

Data living on a field’s own DOFs, interpolated with that field’s basis (grad{name} available too). This is the channel for Picard / time stepping on generalized-order fields: the previous iterate of a topological P2 velocity has no mesh-point representation, so it must ride on its own DOF vector:

w = layout.split(sol)["u"]                    # [n_u, c] on the P2 DOFs
K = assembler(field_data={"w": ("u", w)})

On a mesh whose order equals the field order the two channels agree — point_data is then the simpler spelling.

Generalized order pairs

Field orders are not tied to the mesh order. Besides order 1 and the mesh order (whose DOFs sit on mesh nodes), any other order builds a topological Lagrange DOF map: one DOF per vertex, \(k-1\) per unique (oriented) edge, plus cell interiors. In practice this means:

  • Taylor-Hood on a linear gmsh mesh. A Field(order=2) on an imported P1 mesh creates the quadratic space topologically — no re-meshing at order 2. The cylinder flow and flow past obstacles examples run exactly this way.

  • Higher pairs. P3-P2 and beyond work the same; edge DOFs are stored on oriented edges so adjacent elements agree and the spaces stay \(C^0\)-conforming at any order.

  • Geometry stays isoparametric at the mesh order. A P2 field on a curved (order-2) mesh integrates on the curved geometry even when paired with a P1 pressure; sub- and super-parametric fields are both correct.

Current limits, checked with explicit errors: in 3D the topological map covers vertex + edge DOFs (tetra P2 works; orders needing face DOFs raise NotImplementedError), and the quadrature tables stop at degree 7 — about P3-P2 in practice.

For fields with topological DOFs, layout.node_ids (a mesh-node concept) is unavailable; use layout.points(name), layout.boundary_mask(name) and layout.dof_mask(name, where=...), which work for every field kind, and move data with layout.prolong / layout.restrict (isoparametric interpolation).

Worked examples

The fluid example gallery is built entirely on mixed assembly, in increasing order of complexity: