Finite element layer

Spaces, batched element geometry, the AD kernels, assembly and boundary conditions.

hippymfem.fem.spaces

Finite element space wrappers.

A thin layer over mfem.ParFiniteElementSpace whose main job is lifetime management: MFEM keeps raw pointers to the mesh and the finite element collection, so a space built from a temporary collection crashes the moment it is used. FunctionSpace owns its pieces.

class hippymfem.fem.spaces.FunctionSpace(mesh, fec, vdim=1, ordering=None)[source]

Bases: KeepAlive

A parallel finite element space together with everything it points at.

Parameters:
  • mesh (mfem.ParMesh)

  • fec (mfem.FiniteElementCollection)

  • vdim (int) – Number of solution components.

  • ordering (int) – mfem.Ordering.byNODES (default) or byVDIM. The element vdof array is component-major in both cases, which is what the element kernels assume.

classmethod H1(mesh, order, vdim=1, ordering=None)[source]

Continuous Lagrange space of degree order.

classmethod L2(mesh, order, vdim=1, ordering=None)[source]

Discontinuous Lagrange space of degree order.

classmethod ND(mesh, order)[source]

Nedelec H(curl) space of degree order.

Fields from this space carry val (a vector) and curl, not grad; see hippymfem.fem.vectorfe.

classmethod RT(mesh, order)[source]

Raviart-Thomas H(div) space of degree order.

Fields from this space carry val (a vector) and div. Note MFEM’s numbering: RT(0) is the lowest-order space.

classmethod wrap(fes, mesh=None, fec=None)[source]

Adopt an existing ParFiniteElementSpace without rebuilding it.

property dim

Reference (topological) dimension of the mesh.

property sdim

Space (physical) dimension of the mesh.

property order

Polynomial degree (of element 0; of the collection on a rank without an element, where MFEM stops the process at the question for one).

GetTrueVSize()[source]
GlobalTrueVSize()[source]
GetVSize()[source]
vector()[source]

A zero true-dof ParVector on this space.

The layout is built once and shared, so repeated calls inside a loop do not communicate.

gridfunction(name=None)[source]

A zero ParGridFunction on this space.

The function keeps the space alive, not the other way round: a space keeping every grid function it made would leak one state-sized array per call.

to_gridfunction(v, gf=None)[source]

Scatter a true-dof vector into a ParGridFunction (with ghosts).

from_gridfunction(gf, v=None)[source]

Restrict a ParGridFunction to a true-dof vector.

local_values(v, out=None, device=False)[source]

Local dof values of a true-dof vector, including ghosts, as numpy.

This is P v, the gather element kernels need. A numpy array is returned rather than an mfem.Vector so that callers never have to reach for GetDataArray(), whose views do not keep their owner alive (see hippymfem.common.parvector.to_numpy()).

With device=True, for a caller that hands the values to the element kernels, a JAX array on the kernels’ GPU is returned instead when MFEM is on that GPU as well (hippymfem.common.devicebridge): P v is formed on the device and copied there, where the numpy route would bring it to the host only for the kernels to upload it again.

assemble_dual(local, v=None)[source]

Assemble a local dual (residual) vector to true dofs: P^T local.

local may be a numpy array of local dof values or an mfem.Vector.

project(fn, v=None)[source]

Project a python callable fn(x) -> float or array onto the space.

On a nodal (Lagrange) space the projection is interpolation at the dof nodes, and fn is evaluated there directly: once on the whole (n, sdim) array of node coordinates when it accepts one, point by point otherwise, with no MFEM callback per point either way. Other spaces (H(curl), H(div)) go through MFEM’s ProjectCoefficient. The point-by-point evaluation reproduces the callback route bit for bit; a vectorized fn runs numpy’s array kernels, whose transcendental functions can differ from the scalar ones in the last bit.

property is_nodal

Whether the dofs are point values at nodes (H1 and L2 Lagrange spaces).

coordinates()[source]

Physical coordinates of the scalar dofs as an (ntdof, sdim) array.

Only meaningful for Lagrange (nodal) spaces with vdim == 1.

hippymfem.fem.spaces.as_space(v)[source]

Coerce v (FunctionSpace or ParFiniteElementSpace) to a FunctionSpace.

hippymfem.fem.elementbatch

Batched element data: quadrature geometry, basis tables, and dof maps.

The element kernels are vectorized over elements, so everything they need is precomputed once and stored as contiguous arrays:

  • per quadrature point, the reference basis values N[q, i] and reference gradients G[q, i, k], identical for every element of a given geometry and polynomial order, so computed once;

  • per element and quadrature point, Jinv[e, q, k, d] = dxi_k / dx_d and the weight wdet[e, q] = w_q |det J|, from MFEM’s batched GeometricFactors;

  • per element, the (signed) dof indices into the local dof vector.

Physical gradients are not stored: they are formed inside the JIT’d kernel as G @ Jinv, which costs one fused contraction and saves a factor of ndof/sdim in memory.

MFEM’s GeometricFactors uses a single integration rule for the whole mesh, so meshes with more than one element geometry are split into groups and handled one group at a time.

hippymfem.fem.elementbatch.KEEP_GEOMETRIC_FACTORS = False

Leave MFEM’s cached GeometricFactors in place after the element geometry has been copied out. Off by default, since nothing in hIPPyMFEM reads them again; set HIPPYMFEM_KEEP_GEOMETRIC_FACTORS=1 if your own code holds a pointer from mesh.GetGeometricFactors on a mesh hIPPyMFEM also batches.

hippymfem.fem.elementbatch.AFFINE_GEOMETRY = True

Keep the inverse Jacobians of a group once an element instead of once a quadrature point where every element of the group is mapped affinely (straight simplices, parallelepipeds: any uniformly refined box). Jinv then has the shape (ne, 1, dim, sdim) and the kernels apply the one matrix at every point (hippymfem.fem.kernel._eval_field()). With 64 points that is 72 bytes an element instead of 4.6 kB: 0.12 GB instead of 8 GB for 1.7 million hexahedra, on the host and in every pass of an element kernel. A group of that size streams its geometry to the card with every chunk (kernel.GEOMETRY_STREAM_FRACTION), and the stream was four fifths of a residual; of an affine group the inverse Jacobians and the weights then stay on the card and only the coordinates are left on the host, for the kernels that read them (hippymfem.fem.kernel.GroupKernel._kept_on_device()). On eight MI250X with 1.7 million quadratic hexahedra each, in double precision, a residual takes 0.10 s instead of 0.49 s and two Newton-CG steps of the README’s benchmark 59.8 s instead of 72.2 s; a rank holds 34.5 GiB of host memory instead of 41.9 GiB. The inverse Jacobian kept is the mean over the points, which differ in their last digits where the map is affine, so the element arrays change at round-off. On by default since November 2026; HIPPYMFEM_AFFINE_GEOMETRY=0 keeps one inverse Jacobian a quadrature point everywhere, as before (for the batches built afterwards). A group with any element that is not affine keeps them all anyway.

hippymfem.fem.elementbatch.AFFINE_TOL = 1e-12

How far the inverse Jacobians of an element may differ between its quadrature points, relative to their largest entry, for its map to count as affine.

hippymfem.fem.elementbatch.GEOMETRY_SLICE = 0

Elements per slice when the geometry is built here. 0 sizes a slice at about 128 MB of Jacobians; a positive value pins it and also forces this path, which is how the two routes are compared. Set HIPPYMFEM_GEOMETRY_SLICE.

hippymfem.fem.elementbatch.VERTEX_RULE = -1

A quadrature degree that stands for the rule at the vertices of an element, with equal weights (integration_rule()): the rule that makes the mass matrix of first-order elements diagonal, the sum of each row of the consistent one on an affine element (a prior’s lumped_mass).

hippymfem.fem.elementbatch.integration_rule(geom, degree)[source]

MFEM’s rule of degree on the geometry geom, or for VERTEX_RULE the rule at its vertices (one object a geometry, kept for the life of the process: MFEM keys what it caches for a rule on the rule’s address).

hippymfem.fem.elementbatch.default_quadrature_degree(spaces, extra=2)[source]

A deliberately generous default: twice the highest degree present, plus extra.

The residual density may be non-polynomial (exp(m), 1/m, …) so no exact rule exists. What matters for correctness of the inverse problem is that every derivative block uses the same rule, which is automatic because they are all differentiated from one quadrature sum.

class hippymfem.fem.elementbatch.ElementGroup(mesh, geom, elems, quadrature_degree, batched=None, sliced=None)[source]

Bases: KeepAlive

Quadrature geometry for the elements of one geometry type.

geom

mfem.Geometry type.

Type:

int

elems

Mesh element indices in this group, ascending.

Type:

ndarray of int

ir
Type:

mfem.IntegrationRule

w

Reference quadrature weights.

Type:

ndarray, shape (nq,)

wdet

w_q * |det J|.

Type:

ndarray, shape (ne, nq)

Jinv

(ne, 1, dim, sdim) where affine is set (AFFINE_GEOMETRY).

Type:

ndarray, shape (ne, nq, dim, sdim)

X

Physical coordinates of the quadrature points.

Type:

ndarray, shape (ne, nq, sdim)

affine

Whether Jinv is kept once an element: every element’s map is affine.

Type:

bool

full_Jinv()[source]

Jinv at every quadrature point, (ne, nq, dim, sdim): the array itself, or a read-only view of the compact one that repeats it.

make_tables(space)[source]

Basis tables of space on this group (a boundary group overrides this).

hippymfem.fem.elementbatch.apply_trial_transformation(mats, trial_tables, elements=None)[source]

A T_trial: the column side of MFEM’s TransformDual, for a row space without a transformation.

A block with H1 or L2 rows and H(curl) columns of order two and higher on simplices needs it as much as one with H(curl) on both sides; dropping it gives the matrix in the reference basis of the columns. elements, a (start, stop) pair, says that mats holds that stretch of the group’s elements only (a chunk), as for SpaceTables.transform_dual_matrix().

class hippymfem.fem.elementbatch.SpaceTables(space, group)[source]

Bases: KeepAlive

Basis tables and element dof maps of one space on one element group.

N
Type:

ndarray, shape (nq, nd)

G
Type:

ndarray, shape (nq, nd, dim)

edofs

Indices into the space’s local dof vector, component-major.

Type:

ndarray of int, shape (ne, vdim*nd)

signs

+-1 multipliers decoded from MFEM’s signed vdof encoding.

Type:

ndarray, shape (ne, vdim*nd)

gather(local_array)[source]

Element dof values, shape (ne, vdim*nd), from a local dof array.

gather_device(local_array)[source]

Element dof values as a device array, gathered where the kernels run.

local_array may be numpy (it is copied to the device once) or already a device array. The dof map and the signs are cached on the device, so repeated assemblies upload only the dof values.

jax_tables()[source]

Arrays the field evaluator needs: the reference basis and its gradient.

jax_axes()[source]

in_axes for jax_tables(); both are shared by the whole group.

evaluator()[source]

A function (dofs, tables, Jinv) -> Field for this slot.

transform_dual_matrix(mats, trial_tables, elements=None)[source]

Map element matrices from the reference dof basis to the global one.

Nothing to do on the rows for H1/L2, where the two coincide; the columns still take the trial space’s DofTransformation when it has one (H(curl) of order two and higher on simplices). elements, a (start, stop) pair, says that mats holds that stretch of the group’s elements only (a chunk).

transform_dual_vector(vecs)[source]

Map element vectors from the reference dof basis to the global one.

hippymfem.fem.elementbatch.group_tables(space, group)[source]

Cached SpaceTables for one space on one element group.

Memoized on the identity of the space’s ParFiniteElementSpace and the group. Building tables is purely rank-local (shape evaluation and the element-to-dof table), so a hit on one rank and a miss on another is harmless.

hippymfem.fem.elementbatch.clear_table_cache()[source]

Drop every cached SpaceTables (for tests, and after refinement).

hippymfem.fem.elementbatch.clear_batch_cache()[source]

Drop every cached MeshBatches and their tables (after refinement).

The batches are cached on the mesh object, which a refinement in place keeps: their element lists and geometry would describe the mesh before it.

class hippymfem.fem.elementbatch.MeshBatches(mesh, quadrature_degree, comm=None)[source]

Bases: KeepAlive

All element groups of a mesh for one quadrature degree, with table caching.

Construction is collective: whether the batched GeometricFactors path can be used is decided from the global number of element geometries, because that call is collective on a ParMesh and the local geometry count is not the same on every rank.

tables(space)[source]

Cached SpaceTables for space, one per group.

property nelem
hippymfem.fem.elementbatch.get_batches(mesh, quadrature_degree, comm=None)[source]

Cached MeshBatches for (mesh, quadrature_degree).

The cache key is the mesh identity and the degree, both of which every rank agrees on, so the cached construction is hit or missed uniformly and its internal collective stays synchronized.

hippymfem.fem.kernel

Element kernels generated by automatic differentiation.

The user supplies a pointwise residual density

def pde_varf(u, m, p, x):
    return jnp.exp(m.val) * inner(u.grad, p.grad) - f(x) * p.val

where u, m, p are Field records carrying val and grad at one quadrature point, and x is that point’s physical position. It plays the role a UFL form plays in a form-language library: it is the weak residual of the forward problem, read as the PDE part of the Lagrangian.

QuadratureKernel turns it into element-level arrays:

grad_block(i)

element vectors dR/d(dofs_i)

hess_block(i, j)

element matrices d2R/d(dofs_i)d(dofs_j)

third_block(i, j, k)

element vectors, d3R contracted with two given directions

with i a field slot. For the stationary problem those are exactly the blocks a form-language library obtains from ufl.derivative: hess_block(ADJOINT, STATE) is the forward Jacobian A, hess_block(ADJOINT, PARAMETER) is C, hess_block(STATE, STATE) is W_uu, and so on.

Because every block is differentiated from one quadrature sum, all of them are consistent with each other and with the same integration rule, which is what makes the adjoint gradient exact for the discretized problem rather than merely close.

hippymfem.fem.kernel.matrix_precision()[source]

"fp32" when the element matrices are computed in single precision.

hippymfem.fem.kernel.vector_precision()[source]

"fp32" when the element vectors and values are computed in single precision.

hippymfem.fem.kernel.FP32_EPS = 1.2e-07

relative accuracy of a quantity computed by a single-precision kernel

hippymfem.fem.kernel.residual_floor()[source]

The relative size below which an assembled residual cannot be driven with the present precision of the element vectors: a tolerance below it is unreachable.

hippymfem.fem.kernel.gradient_floor()[source]

The relative size below which the norm of a reduced gradient cannot be driven when the element vectors are in single precision: the state, the adjoint and the gradient are then computed to about that accuracy (more on a fine mesh, where the errors of the two solves grow with the condition number), and an optimizer that asks for less ends in a line search that cannot find a decrease. Zero in double precision and in the mixed mode.

hippymfem.fem.kernel.matrix_tolerance(tol, single=1e-05)[source]

A tolerance for a comparison that involves assembled element matrices: tol when they are computed in double precision. In single precision an entry carries a relative error near FP32_EPS and a product with the matrix somewhat more, by cancellation, so that no such comparison holds tighter than single.

hippymfem.fem.kernel.set_precision(mode)[source]

Choose the element-kernel precision ("fp64", "mixed" or "fp32", see PRECISION); returns the old value.

Clears the compiled-kernel cache of every kernel built so far, since the dtype is baked into the jitted code.

hippymfem.fem.kernel.HESSIAN_MODE = 'element'

How the element Hessian blocks are differentiated. element (the default) pushes one forward tangent per element dof of the column slot through the whole element. quadrature differentiates the density twice at each point, with respect to the fields’ values and gradients (one tangent per such feature), skips the parts of that pointwise Hessian that are identically zero, and contracts the rest with the basis values and physical gradients; it is exact too, since the fields are linear in the dofs. Which route is the faster on a GPU depends on the element, the size of the mesh and the card: for the Jacobian of a scalar diffusion density on meshes that fill the card, an assembly by this route took 0.9 to 2.7 times as long as by the element route for degree two and higher (L40S, H100, MI250X), where on small meshes it had been 1.6 to 4.4 times faster on an L40S. On P1 the element route is the faster, and so it is on the host, except for vector-valued Jacobians. auto times both once per kernel program, column slot and device and keeps the faster (_ROUTES), so every kernel built alike takes the same route in a run, while which one a run takes, and its round-off, can differ from run to run. Spaces with a Piola map and interior facets always take the element route. Set HIPPYMFEM_HESSIAN or set_hessian_mode().

hippymfem.fem.kernel.set_hessian_mode(mode)[source]

Choose how element Hessians are differentiated (HESSIAN_MODE); returns the old mode. Clears every kernel’s compiled Hessians.

hippymfem.fem.kernel.set_device(kind='cpu', index=None, comm=None)[source]

Choose where element kernels run.

Parameters:
  • kind ({"cpu", "gpu", "cuda", None}) – None resets to the default. Asking for a GPU requires HIPPYMFEM_DEVICE=gpu or auto (or JAX_PLATFORMS=cuda) in the environment before hippymfem is imported, because JAX reads the platform list once.

  • index (int, optional) – Which device. The default assigns one GPU per MPI rank by the rank’s index within its node, so the ranks of a job do not all land on device 0.

  • comm (mpi4py communicator, optional) – Used for the node-local rank; defaults to COMM_WORLD.

  • device. (Returns the chosen)

hippymfem.fem.kernel.device()[source]

The device element kernels run on.

Honors HIPPYMFEM_DEVICE=gpu (and auto, a GPU when the process can see one) on first use, so a script only has to set the environment variable; set_device overrides it.

hippymfem.fem.kernel.on_gpu()[source]

True when element kernels are running on a GPU.

hippymfem.fem.kernel.gather_rows(values, idx, signs=None)

values[idx] (times signs), compiled.

The element dof gather of every assembly. Eager indexing normalizes and checks the index array with separate array operations at every call, and at 1.6e6 tetrahedra it took 0.10 s a slot against 0.0036 s compiled; a gather and a sign flip are exact, so the values are the same.

hippymfem.fem.kernel.to_host(x)[source]

A numpy view/copy of a kernel result, wherever it lives.

hippymfem.fem.kernel.any_nonzero(arrays)[source]

Whether any entry of any array is nonzero, checked where the data lives.

hippymfem.fem.kernel.all_finite(arrays)[source]

Whether every entry of every array is finite, checked where the data lives.

Reducing on the device first means one scalar comes back instead of the whole element array, which matters because this is called on every assembly.

class hippymfem.fem.kernel.Field(val, grad)[source]

Bases: NamedTuple

A field at one quadrature point.

val is a scalar for vdim == 1 and shape (vdim,) otherwise; grad is (sdim,) or (vdim, sdim) correspondingly. It is a NamedTuple so jax.vmap can map over quadrature points transparently.

Parameters:
  • val (jax.numpy.ndarray)

  • grad (jax.numpy.ndarray)

val: jax.numpy.ndarray

Alias for field number 0

grad: jax.numpy.ndarray

Alias for field number 1

hippymfem.fem.kernel.ELEMENT_CHUNK = 0

Largest number of elements handed to one kernel launch. 0, the default, plans the size from the device’s free memory (the whole group on the host) and shrinks it on an out-of-memory error; a fixed size would cost extra launches on every problem to protect the few that need it. Set HIPPYMFEM_ELEMENT_CHUNK to pin a size.

hippymfem.fem.kernel.AD_DOUBLES_PER_TANGENT_QP = 16.0

Doubles of device working set per (tangent, quadrature point) pair, used to size a chunk before the first launch. Forward-mode AD carries nd_total tangents through nq points; measured between 7 and 13 across element types, so 16 leaves margin. Raising it makes chunks smaller and safer.

hippymfem.fem.kernel.AD_VECTOR_TANGENTS = 8.0

Tangents that a chunk of an element vector (a residual, a gradient) is planned for. The plan above is for the element matrices, which carry one tangent for every dof of every variable; a vector is one reverse sweep, whose working set does not grow with the dofs. Measured for the residual of quadratic hexahedra with 64 quadrature points on an MI250X: 16 kB an element (31 doubles a point, two tangents’ worth), where the matrices’ plan is 508 kB, so eight leave a margin of four. With a quarter of that card for JAX and 1.7 million such elements a residual then runs in chunks of 40 000 to 80 000 elements instead of 5000 to 10 000 and takes 0.47 s instead of 0.70 s, a gradient 0.56 s instead of 1.02 s (a launch costs 1.5 ms; what is left is mostly the geometry on its way to the card), and the operations of two Newton-CG steps 5 to 7 % less. An estimate that is too small for some form lands in the retry of GroupKernel._chunked(), as one for a matrix does. 0 plans a vector as a matrix. Set HIPPYMFEM_AD_VECTOR_TANGENTS.

hippymfem.fem.kernel.GEOMETRY_STREAM_FRACTION = 0.25

Share of JAX’s device budget above which a split batch keeps its per-element geometry (Jinv, wdet, X) on the host and transfers it a chunk at a time. Below it the geometry is cached on the device: streaming a small geometry costs transfers for nothing, while caching one that is most of the budget is what runs a chunked Jacobian out of memory. 0 never streams. Set HIPPYMFEM_GEOMETRY_STREAM.

hippymfem.fem.kernel.PINNED_STREAM = True

How a chunk loop moves the geometry it streams. On (the default, where hippymfem.common.devicebridge can be used), the group’s geometry is registered with the GPU runtime once (locked in RAM), and each chunk’s slice of it goes to the device by one copy of the runtime at the rate of the bus, and only the arrays the kernel reads. Off, the slices are handed to the kernel as numpy arrays and JAX moves them through a staging buffer of its own: 5 to 10 GB/s on an H100 against 55, which made an element pass over two million quadratic hexahedra on one card take a second. Set HIPPYMFEM_PINNED_STREAM=0 for that route.

hippymfem.fem.kernel.GPU_MEM_RESERVE = 0.0

Device memory, in GiB, kept aside for whatever else shares the card (above all hypre’s matrices and multigrid hierarchy). The chunk planner subtracts it from its budget: an adaptive counterpart of the import-time HIPPYMFEM_GPU_MEM_FRACTION. Set HIPPYMFEM_GPU_MEM_RESERVE.

hippymfem.fem.kernel.HOST_MEM_FRACTION = 0.5

Share of the node’s available memory one rank may plan a chunk against on the host, the whole node’s figure divided by the ranks sharing it. Without this the host planner has no budget at all (it asks the card, and there is none), so a rank takes its whole element group in one launch; that is right until the group is large, and at 524288 P2 hexahedra a rank it is what the OOM killer ends. Half is left to hypre, the matrices and everything outside the kernel, which is the same split the device path makes. 0 restores the unplanned host behaviour. Set HIPPYMFEM_HOST_MEM_FRACTION.

hippymfem.fem.kernel.RETRIES = 0

How many times an element pass of this process ran out of memory and was tried again in smaller chunks. The group’s passes of that weight keep the smaller chunk for the rest of the run (GroupKernel._chunked()), and a warning on that rank is all that says so. For a benchmark’s record: 0 where every pass ran as planned.

hippymfem.fem.kernel.HOST_BATCH = 2048

Elements per step of an element kernel on the host. XLA runs a vmap over the whole batch as one program whose intermediates have the batch’s size, and on a CPU core those leave the cache from about 1e4 elements on: a P1 tetrahedron’s Jacobian costs 0.5 us per element in a batch of 3072 and 1.8 us in one of 196608, a third derivative 0.2 and 0.7 us. On the host the batch is therefore mapped in steps of this many elements inside the compiled program (a loop over windows of the arguments), which keeps the working set in cache. Steps of 512 elements and more gave the whole-batch arrays bit for bit in every case measured; steps of a few elements round differently, by an ulp, as a chunked batch does (benchmarks/DESIGN_NOTES.md, section 1). A GPU maps the whole batch, which is what it is fast at, and so does the host when the batch fits in one step. 0 maps the whole batch on the host too. Set HIPPYMFEM_HOST_BATCH or hm.config.host_batch; it is read at every call, each size compiling its own program.

class hippymfem.fem.kernel.GroupKernel(density, tables, group, nparams=0)[source]

Bases: object

Compiled element kernels for one ElementGroup.

Each field slot contributes its own basis tables and its own rule for turning element dofs into a field record at the quadrature points. For an H1/L2 slot the tables are shared by every element of the group and the record is Field; for an H(curl) or H(div) slot the Piola map makes them per-element and the record is a HCurlField / HDivField. The two cases differ only in the in_axes handed to jax.vmap, so nothing is broadcast and a scalar problem pays nothing for the vector-element support.

A BoundaryGroup is served the same way: its tables are per boundary element (the face-to-volume map differs from face to face), which the slot axes already express, and it carries the outward unit normal, which the density then takes after the coordinates. Chunking, fp32 and the single Hessian primitive are the same for both.

hessian_route

{(device, column slot): "element" | "quadrature"}, what auto chose

single_matrices

Whether the "mixed" precision computes this kernel’s matrices in single precision. Off unless the owner of the kernel turns it on: single-precision matrices need the consistent Jacobian and the refined solves that go with them, which PDEVariationalProblem has for the element matrices of its domain term. Every other kernel computes in double under "mixed".

will_chunk(weight=1)[source]

Whether this group’s batch will be split on the current device.

Asked before assembling, so a caller can pick the scatter that suits: splitting already gives up bit-identical results against an unsplit run, so where it happens the fused scatter costs no further exactness and saves the largest allocation an assembly makes. Where it does not happen, the glued path stays, and stays exact.

in_chunks(fn, dofs, mapped, pr)[source]

Run a batched kernel chunk by chunk, yielding (start, stop, result).

The same splitting _chunked() does, without the concatenation at the end. fn is the wrapped kernel, so calling it on a slice would chunk the slice again; the unwrapped one is taken from _raw.

mapped(streaming=False)[source]

(tables, geometry): the mapped arguments, in the order the kernels expect.

The per-element geometry is cached on the device unless _stream_geometry() finds the batch split and the geometry too large a share of the budget; it is then left on the host, and a chunk loop slices it, so that only a chunk’s rows are transferred instead of the whole (ne, nq, dim, sdim) staying resident (gigabytes on a large mesh, against megabytes per chunk). streaming is the request of a chunk loop. It used to get the host arrays whatever was cached, which transferred the whole geometry, chunk by chunk, at every assembly of a split batch although the device held a copy for the other kernels: 1.7 GB for 262 144 second-order hexahedra, more than half of what such an assembly took outside its kernel. Now the device copy serves the chunk loop too. Tables are per group and stay on the device either way.

grad_block(i)[source]

Batched dR/d(dofs_i); returns (ne, nd_i).

consistent_block(i, j, symmetric=False)[source]

Block (i, j) in single precision, exact on the constants.

An element matrix computed in single precision is off by about 1e-7 of its entries, and on a mesh of like elements every element makes the same error. The assembled matrix then errs coherently: applied to a smooth function, whose values on an element are nearly equal, the errors of a row add up to the defect of its row sum, the same on every element, where the exact entries cancel. A solve amplifies that by the condition number. Measured on the Jacobian of the model problem at 16^3 second-order hexahedra, with a constant coefficient: an incremental solve was off by 2.2e-4, and by 2.7e-7 with the row sums put right.

So the block is computed in single precision and corrected by a matrix of rank one per component: its product with the constants of slot j, and that of its transpose with the constants of slot i, are replaced by the double-precision ones, which are two derivative passes with one tangent each (a matrix pass has one per element dof). symmetric says that the block is a symmetric matrix: it is then symmetrized first, one pass serves both sides, and the result is symmetric to the last bit.

hess_block(i, j)[source]

Batched d2R/d(dofs_i) d(dofs_j); returns (ne, nd_i, nd_j).

Row index is i (the test space), column index is j (the trial space), which is MFEM’s elmat convention. Rows i of the slot pass _hess_slot() (j), selected after the compiled function.

hess_cols(cols)[source]

The Hessian blocks whose column slot is in cols: H[i][k] of shape (ne, nd_i, nd_cols[k]) for every row slot i.

A Newton linearization point never asks for the adjoint columns (its blocks are C, W_uu, W_um, W_mm: columns u and m), and a residual linear in the state needs the m columns alone. One slot pass per column (_hess_slot()) does that much work and no more, holds (ne, nd, nd_j) arrays instead of (ne, nd, nd), so a batch splits into fewer chunks, and is bit-identical to hess_block() by construction.

hess_all()[source]

Every Hessian block at once: H[i][j] of shape (ne, nd_i, nd_j), hess_cols() over every slot.

hvp_block(i, j)[source]

Batched d2R/d_i d_j applied to a direction in slot j.

Returns a callable (dofs, jdir, tabs, geo, params) -> (ne, nd_i): one forward tangent through the slot-i gradient, where the assembled block takes nd_j tangents and a scatter into a matrix. For a block applied once or twice at a point (the sample loop of a sample-average Hessian) this is the cheaper way to its products.

third_block(i, j, k)[source]

Batched d3R/d_i d_j d_k contracted with directions in slots j, k.

Returns a callable (dofs, jdir, kdir, tabs, geo, params) -> (ne, nd_i), where jdir and kdir are element dof arrays for arguments j and k.

third_dir_block(i, nmodes, present=None)[source]

Batched \(\sum_m w_m\, D^2(\partial_i R)[t_m, t_m]\): the second directional derivative of the slot-i gradient along directions t_m that may span every slot, weighted and summed over nmodes directions.

Returns a callable (dofs, T, w, tabs, geo, params) -> (ne, nd_i), where T is a tuple over the slots of element dof arrays (ne, nmodes, nd_s) and w the nmodes weights. With t = (t_j, t_k) over two slots it is third_block() summed over both orders and both diagonals, \(R_{ijj}[t_j, t_j] + 2 R_{ijk}[t_j, t_k] + R_{ikk}[t_k, t_k]\), in one pass instead of four; the directions share the gather of the point and the scatter. present names the slots the directions have components in (every slot by default), and T then holds those slots’ arrays alone: the others are not differentiated at all (see hvp_block()).

residual_value()[source]

Batched R itself (one scalar per element); used for form integrals.

geometry()[source]

Deprecated alias for the geometry arrays alone; prefer mapped().

gather(local_arrays)[source]

Element dof arrays for this group from local (ghosted) dof vectors.

The indexing runs where the kernels run: the local dof vector (one value per local dof) is copied to the device and gathered there, rather than gathered on the host and copied element by element. For a P1 quadrilateral space that is a quarter of the transfer, and it keeps the element dof array off the host entirely.

class hippymfem.fem.kernel.QuadratureKernel(density, spaces, batches, nparams=0)[source]

Bases: object

Element kernels for a residual density over a whole (possibly mixed) mesh.

Parameters:
  • density (callable) – density(u, m, p, x, *aux, *params) -> scalar, JAX-traceable, with Field arguments for the fields.

  • spaces (sequence of FunctionSpace) – One per field slot, in the order the density expects them. The stationary problem uses three, [STATE, PARAMETER, ADJOINT]; a one-step time-dependent problem uses four, with the previous time level in slot 1, so that the adjoint’s time coupling is an exact derivative rather than a difference.

  • batches (MeshBatches)

  • nparams (int) – Number of trailing scalar arguments (time, step size, …). They are passed as traced values, so changing them between calls reuses the compiled kernel instead of triggering a recompile.

property groups
tables(gi)[source]
element_matrices(i, j, local_arrays, params=(), consistent=None)[source]

Element matrices of block (i, j), one array per element group.

consistent ("symmetric" or "two-sided") asks, where the matrices are computed in single precision, for the ones that are exact on the constants (GroupKernel.consistent_block()): for a matrix that will be inverted.

will_chunk(weight=1)[source]
element_matrices_or_chunks(i, j, local_arrays, params=(), consistent=None, comm=None)[source]

The arrays when the batch fits whole, a chunk thunk when it does not.

Bit-identical to element_matrices() in the first case, and in the second the splitting has already cost that, so the fused scatter is free. With comm the choice is made once for all ranks: the two results take different paths through the assembly, with their collectives in different places, and ranks with different element counts may judge differently.

element_matrix_chunks(i, j, local_arrays, params=(), consistent=None)[source]

Yield (group, start, stop, array) for block (i, j).

The same work element_matrices() does, handed over a chunk at a time instead of glued into one (ne, nd, nd) array per group, the largest resident allocation an assembly would make; the scatter consumes each chunk (data_fused()).

element_matrix_chunks_many(pairs, local_arrays, params=())[source]

Yield (group, start, stop, {(i, j): array}) from one differentiation.

element_matrices_many() glues every block’s full (ne, nd, nd) array on the device and concatenates the chunk parts to get there; that glue, resident plus a transient copy, is a device peak no chunk size changes. This yields the same blocks a chunk at a time, from the same single pass, so the scatter can consume each chunk and nothing larger than a chunk is ever resident.

element_matrices_many(pairs, local_arrays, params=())[source]

Several blocks from one differentiation pass.

pairs is a sequence of (i, j). Wherever a linearization point needs more than one block, which is every Newton step, this is the cheaper way to get them: one differentiation pass over the column slots the pairs use (GroupKernel.hess_cols()) instead of one pass per block. Returns {(i, j): [array per group]}, holding exactly the arrays element_matrices() returns for each pair.

element_vectors(i, local_arrays, params=())[source]

Element residual vectors of block i, one array per element group.

element_hvp(i, j, local_arrays, jdir_local, params=())[source]

Element vectors of the second-derivative block (i, j) applied to a direction in slot j, without forming the block (GroupKernel.hvp_block()).

element_third(i, j, k, local_arrays, jdir_local, kdir_local, params=())[source]

Element vectors of the third-derivative block contracted with two directions.

element_third_dir(i, local_arrays, dir_locals, weights, params=())[source]

Element vectors of \(\sum_m w_m\, D^2(\partial_i R)[t_m, t_m]\) (see GroupKernel.third_dir_block()).

dir_locals holds one direction per weight, each a sequence over the slots of local dof arrays, None where the direction has no component.

element_values(local_arrays, params=())[source]

Per-element integrals of the density (used for functional evaluation).

class hippymfem.fem.kernel.BoundaryKernel(density, spaces, batches, nparams=0)[source]

Bases: QuadratureKernel

AD kernels for a boundary density: QuadratureKernel over BoundaryBatches.

density(u, m, p, x, n, *params) -> scalar with the same Field records as the domain density, plus the outward unit normal n, which the boundary groups carry and the group kernel hands to the density after the coordinates. Everything else is the domain kernel’s.

hippymfem.fem.vectorfe

H(curl) and H(div) spaces: Nedelec and Raviart-Thomas.

Two things set these families apart from the Lagrange ones in the AD route, and both are handled here.

The basis is vector valued and mapped. A Nedelec basis function is pushed forward covariantly (J^{-T} w_ref) and a Raviart-Thomas one contravariantly (J w_ref / det J), so unlike Lagrange shape functions they cannot be tabulated once for the whole element group, since the map depends on the element. MFEM’s CalcVShape(ElementTransformation&, ...) applies the right map, and it is used here rather than a hand-rolled Piola transform, so the convention is MFEM’s by construction. The derivative a field of this kind carries is not its gradient (which is not in the space) but its curl for H(curl) and its divergence for H(div), so the densities are written with u.curl / u.div instead of u.grad:

def varf(u, m, p, x):                 # curl-curl with a varying coefficient
    return jnp.exp(m.val) * u.curl * p.curl + u.val @ p.val

Element dofs may need a transformation, not just a sign. MFEM encodes edge and face orientation partly as a sign on the dof index (which the batched gather already decodes) and, for H(curl) of order above one on a geometry that is not a tensor product (triangles, tetrahedra, wedges), partly as a small dense DofTransformation T per element. The convention, from MFEM’s own assembly loop, is

  • gather: u_ref = T^{-1} (signs * u_global)

  • matrix: A_global = T^T A_ref T, then the signs

  • vector: b_global = T^T b_ref, then the signs

and this module applies all three. T is built only when MFEM reports it is not the identity, which is never for H(div), for H(curl) on quadrilaterals and hexahedra, or for lowest-order H(curl) anywhere.

class hippymfem.fem.vectorfe.HCurlField(val, curl)[source]

Bases: NamedTuple

An H(curl) field at one quadrature point: a vector value and its curl.

curl is a scalar in 2D and a (3,) vector in 3D, as in MFEM.

Parameters:
val: ndarray

Alias for field number 0

curl: ndarray

Alias for field number 1

class hippymfem.fem.vectorfe.HDivField(val, div)[source]

Bases: NamedTuple

An H(div) field at one quadrature point: a vector value and its divergence.

Parameters:
val: ndarray

Alias for field number 0

div: ndarray

Alias for field number 1

class hippymfem.fem.vectorfe.VectorSpaceTables(space, group)[source]

Bases: KeepAlive

Mapped vector basis tables and dof maps of one vector space on one group.

V

Physical (Piola-mapped) vector shape functions.

Type:

ndarray, shape (ne, nq, nd, sdim)

D

Physical curl, (ne, nq, nd) in 2D and (ne, nq, nd, 3) in 3D, or physical divergence, (ne, nq, nd).

Type:

ndarray

edofs, signs
Type:

ndarray, shape (ne, nd)

T, Tinv

MFEM’s per-element DofTransformation, present only when it is not the identity.

Type:

ndarray or None, shape (ne, nd, nd)

gather(local_array)[source]

Element dof values in the reference dof basis.

The signs come first and the inverse transformation second, which is the order MFEM’s GridFunction::GetElementDofValues uses.

gather_device(local_array)[source]

The same gather, performed where the kernels run.

jax_tables()[source]
jax_axes()[source]

Per element: the Piola map depends on the element Jacobian.

evaluator()[source]

A function (dofs, tables, Jinv) -> HCurlField or HDivField.

transform_dual_matrix(mats, trial_tables, elements=None)[source]

A_global = T_test^T A_ref T_trial, MFEM’s TransformDual.

Applied on whichever side the element matrices already live, so a device result is not brought to the host just to be multiplied. elements, a (start, stop) pair, says that mats holds that stretch of the group’s elements only, a chunk of an assembly that does not fit whole, and the transformations are taken for it alone.

transform_dual_vector(vecs)[source]

b_global = T^T b_ref.

hippymfem.fem.vectorfe.is_vector_space(space)[source]

True for an H(curl) or H(div) space, whose basis functions are vectors.

hippymfem.fem.vectorfe.vector_kind(space)[source]

"hcurl", "hdiv" or None for the element family of space.

hippymfem.fem.boundary

Boundary integrals in the AD route.

The domain kernels in kernel cover \int_\Omega only, which leaves out Neumann and Robin conditions, boundary sources, and any misfit or quantity of interest that lives on the boundary. UFL gets those from ds; this module is the equivalent here.

A boundary density is a function of the same fields plus the outward unit normal:

def bdr_varf(u, m, p, x, n):
    return kappa * (u.val - g(x)) * p.val          # Robin
    # or jnp.exp(m.val) * jnp.dot(u.grad, n) * p.val   (flux)

and it is differentiated exactly like the domain density, so every block (A, C, W_uu, W_um, W_mm, third derivatives) picks up its boundary part automatically.

The fields carry their full gradient, not just their trace. A boundary element’s own finite element can only give tangential derivatives, so instead each face quadrature point is mapped back into the adjacent volume element with MFEM’s FaceElementTransformations; the shape functions and gradients are evaluated there, and the element dofs are the volume element’s. That is what makes u.grad, and therefore the normal flux dot(u.grad, n), available. The price is that the basis tables vary per boundary element instead of being shared, which costs nbe * nq * nd doubles: boundary elements are a lower-dimensional set, so this is small next to the domain tables.

Geometry is built with a Python loop over boundary elements for the same reason (nbe grows only as N^((d-1)/d)), and the loop avoids MFEM’s face-geometry batching, whose availability depends on rank-local face counts.

class hippymfem.fem.boundary.BoundaryGroup(mesh, bdr_elems, quadrature_degree, geom=None)[source]

Bases: KeepAlive

Face quadrature data for boundary elements sharing one geometry pair.

elems

Boundary element indices, ascending. Named elems so that the assembly pattern, which orders entries by group.elems to match MFEM’s own loop, needs no special case.

Type:

ndarray

vol_elems

Index of the volume element adjacent to each boundary element; the dof maps come from these.

Type:

ndarray

wdet

w_q times the face measure.

Type:

ndarray, shape (ne, nq)

Jinv

Inverse Jacobian of the volume element at the mapped point.

Type:

ndarray, shape (ne, nq, dim, sdim)

X
Type:

ndarray, shape (ne, nq, sdim)

nor

Outward unit normal.

Type:

ndarray, shape (ne, nq, sdim)

eip

Volume reference coordinates of the face quadrature points.

Type:

ndarray, shape (ne, nq, dim)

h

The adjacent element’s measure divided by the face measure, which is what a Nitsche or DG penalty is scaled by, and the same quantity MFEM’s DGDiffusionIntegrator uses on a boundary face. A density that takes it after the normal is given it; one that does not is left alone.

Type:

ndarray, shape (ne, nq)

make_tables(space)[source]
class hippymfem.fem.boundary.BoundarySpaceTables(space, group)[source]

Bases: KeepAlive

Shape values and gradients of the adjacent volume element, per boundary element.

The layout matches SpaceTables except that N and G carry a leading element axis, because the face-to-volume map differs from face to face.

gather(local_array)[source]

Element dof values, shape (ne, vdim*nd), from a local dof array.

gather_device(local_array)[source]

Element dof values as a device array, gathered where the kernels run.

local_array may be numpy (it is copied to the device once) or already a device array. The dof map and the signs are cached on the device, so repeated assemblies upload only the dof values.

jax_tables()[source]
jax_axes()[source]

Per boundary element: the face-to-volume map differs from face to face.

evaluator()[source]
transform_dual_matrix(mats, trial_tables, elements=None)[source]
transform_dual_vector(vecs)[source]
class hippymfem.fem.boundary.BoundaryBatches(mesh, quadrature_degree, bdr_attributes='all', comm=None, space=None)[source]

Bases: KeepAlive

Boundary element groups of a mesh for one quadrature degree and marker.

Grouped by boundary geometry and by the number of dofs of the adjacent volume element, so that every group has uniform array shapes even on a mixed mesh where a triangle and a quadrilateral share an edge of the boundary.

tables(space)[source]
property nelem
hippymfem.fem.boundary.assemble_boundary_matrix(test_space, trial_space, groups, element_matrices, test_ess=None, trial_ess=None, diag_policy='one')[source]

Assemble a boundary block.

The element arrays are indexed by boundary element but carry the adjacent volume element’s dofs.

hippymfem.fem.boundary.assemble_boundary_vector(space, groups, element_vectors, ess=None, out=None)[source]

Assemble a boundary dual vector through the direct-CSR path.

hippymfem.fem.boundary.boundary_marker(mesh, bdr_attributes='all')[source]

An MFEM attribute marker array from "all", None, or a list of attributes.

hippymfem.fem.facets

Interior facet terms: the dS of a discontinuous Galerkin form.

boundary covers \(\int_{\partial\Omega}\), where a face has one adjacent element. An interior face has two, and a DG form is written in their traces: the jump \([\![u]\!] = u^- - u^+\) and the average \(\{u\} = (u^- + u^+)/2\), with \(-\) the side the face normal points out of (MFEM’s Elem1). A facet density is a function of those:

def facet_varf(u, m, p, x, n, h):
    return (-jnp.dot(avg_grad(u), n) * jump(p)
            - jnp.dot(avg_grad(p), n) * jump(u)
            + kappa * 0.5 * (1.0 / h[0] + 1.0 / h[1]) * jump(u) * jump(p))

and it is differentiated exactly like a domain density, so every block picks up its facet part. Each field arrives as a FacetField carrying the two traces, and h holds the two elements’ measures, each divided by the face measure (see InteriorFacetGroup).

One slot, two elements. The dof vector of a slot on a facet is the two elements’ dofs concatenated, the layout of MFEM’s AssembleFaceMatrix, so a second derivative of the density is a face matrix, with no rearrangement. The geometry carries both elements’ inverse Jacobians stacked, which is the only thing the group kernels of kernel need to serve two sides.

Faces shared between ranks. The second element of a shared face lives on a neighbour, so a facet block couples dofs this rank does not have. It is assembled the way MFEM assembles its own DG forms: over local dofs, with the neighbour’s columns carrying their global numbers, and folded to true dofs by the prolongation afterwards. Each rank takes the rows of the element on its side of a shared face and the other rank takes the rest, so the face is counted once. csrassemble’s pattern cannot be reused for this, since it holds the couplings within one element and a facet couples two.

class hippymfem.fem.facets.FacetField(minus, plus)[source]

Bases: NamedTuple

A field at an interior facet point, from both sides.

minus is the element the face normal points out of, plus the other; each is a Field with val and grad. Use jump() and avg() on it rather than reaching for the sides, so a density reads like the form it comes from.

Parameters:
minus: object

Alias for field number 0

plus: object

Alias for field number 1

hippymfem.fem.facets.jump(f)[source]

\([\![f]\!] = f^- - f^+\), for a field or either of its parts.

hippymfem.fem.facets.avg(f)[source]

\(\{f\} = (f^- + f^+)/2\).

hippymfem.fem.facets.check_facet_symmetry(density, spaces, sdim, nparams=0, rtol=1e-09, seed=20260927)[source]

Warn if an interior-facet density changes when the two sides of a face trade places; returns the relative change it found (None if the density could not be evaluated on the random traces it tries).

A face shared by two ranks is assembled by each from its own side (the element it owns is the one the normal points out of there), and a face interior to a rank from the side MFEM’s orientation picks; so the density must be the same with the traces swapped, the normal reversed and the two element measures exchanged, as a DG form written in jumps, averages and the normal is. One that is not gives results that depend on the partition and on the face orientation (a matrix-free product of apply_ij_at(), for one, then differs from the assembled block’s on more than one rank). The check evaluates the density once each way on random traces.

hippymfem.fem.facets.jump_grad(f)[source]

\([\![\nabla f]\!]\).

hippymfem.fem.facets.avg_grad(f)[source]

\(\{\nabla f\}\).

class hippymfem.fem.facets.InteriorFacetGroup(mesh, quadrature_degree, geom=None)[source]

Bases: KeepAlive

Quadrature data of interior faces sharing one geometry and one dof count.

elems

Mesh face indices of the local interior faces, then -1 for each face shared with another rank.

Type:

ndarray

e1, e2

The two adjacent elements. For a shared face e2 is mesh.GetNE() + k, the face-neighbour numbering MFEM uses.

Type:

ndarray

wdet

w_q times the face measure.

Type:

ndarray, shape (nf, nq)

Jinv

Inverse Jacobians of the two elements at the mapped points, minus first.

Type:

ndarray, shape (nf, nq, 2, dim, sdim)

X
Type:

ndarray, shape (nf, nq, sdim)

nor

Unit normal, pointing out of e1.

Type:

ndarray, shape (nf, nq, sdim)

h

Both elements’ measures divided by the face measure, the quantity MFEM’s DG integrators penalise with: their {h^{-1}} is the average (1/h[0] + 1/h[1])/2. A density receives it as a pair, so any other convention (the smaller of the two, say) is the user’s to write.

Type:

ndarray, shape (nf, nq, 2)

eip1, eip2

Reference coordinates of the face points inside each element.

Type:

ndarray, shape (nf, nq, dim)

h_always = True

the face measure reaches every facet density, which cannot penalise without it

make_tables(space)[source]

Cached FacetSpaceTables for space on this group.

class hippymfem.fem.facets.FacetSpaceTables(space, group)[source]

Bases: KeepAlive

Shape values and gradients of both adjacent elements, per interior face.

The layout matches BoundarySpaceTables twice over: N1, G1 for the minus side and N2, G2 for plus, each with a leading face axis, because the face-to-element map differs from face to face. A slot’s dof vector is the two elements’ dofs concatenated, so nd_total is their sum.

jax_tables()[source]

The four per-face tables the two-sided evaluator needs.

jax_axes()[source]

Every table carries the face axis: the map differs from face to face.

evaluator()[source]

(dofs, tables, Jinv) -> FacetField; dofs is both sides concatenated.

gather(values)[source]

Both elements’ dof values, (nf, 2*vdim*nd).

values is what facet_values() returns: this rank’s local dof values followed by the face-neighbour ones, which is where the far side of a shared face lives. A purely local batch may pass the local array alone.

gather_device(local_array)[source]
transform_dual_matrix(mats, trial_tables, elements=None)[source]

No reference-to-global map on the rows: the facet kernels cover H1 and L2 spaces; a column space with a DofTransformation still takes it.

transform_dual_vector(vecs)[source]
class hippymfem.fem.facets.FacetBatches(mesh, quadrature_degree, comm=None)[source]

Bases: KeepAlive

The interior faces of a mesh, in the shape the group kernels take.

One group, with one face geometry: its quadrature rule and reference points serve every face. A mesh whose interior faces have more than one geometry (wedges and pyramids, which have triangular and quadrilateral faces, or tetrahedra next to hexahedra) is refused, since the one rule would integrate the other faces wrongly. Built like BoundaryBatches so that QuadratureKernel takes either without noticing.

tables(space)[source]

Cached FacetSpaceTables for space, one per group.

property nfaces
hippymfem.fem.facets.get_facet_batches(mesh, quadrature_degree=None, space=None, comm=None)[source]

Cached FacetBatches for a mesh and a quadrature degree.

hippymfem.fem.facets.clear_facet_cache()[source]

Drop every cached facet batch, pattern and dof map (after refinement, and in tests).

The assembly skeleton and the global dof numbers are cached on the spaces and the mesh they were built from, so a refinement that keeps those objects alive must say so; clear_pattern_cache() is the domain counterpart.

hippymfem.fem.facets.global_ldof_map(space)[source]

Global ldof numbers of the extended dofs: this rank’s own, then its neighbours’.

FacetSpaceTables numbers the far side of a shared face after this rank’s own dofs, and this map turns that numbering into the global one a parallel matrix takes as its columns. They are ldof numbers, one per rank per shared dof, not true dofs; the facet block is folded onto true dofs by the prolongation, exactly as MFEM folds its own DG blocks. MFEM keeps the same map in face_nbr_glob_dof_map and hands it back as a bare pointer, so it is rebuilt here by exchanging the numbers themselves. Collective; cached per space.

hippymfem.fem.facets.assemble_facet_matrix(test_space, groups, element_matrices, trial_space=None, test_tables=None, trial_tables=None)[source]

Assemble interior facet terms into a parallel true-dof matrix.

element_matrices is one (nf, nd_test1 + nd_test2, nd_trial1 + nd_trial2) array per group, as element_matrices gives for a facet batch: rows the test dofs of both elements, columns the trial dofs of both, minus side first. With trial_space the block may be rectangular, as the parameter blocks of an inverse problem are; square and rectangular blocks go the same way.

hippymfem.fem.facets.assemble_facet_vector(space, groups, element_vectors, tables=None, out=None)[source]

Assemble a facet residual into a dual (true-dof) vector.

element_vectors is one (nf, 2*nd) array per group, the two elements’ entries concatenated, as a first derivative of a facet density gives them.

No communication. Each rank scatters the entries of the element the normal leaves, which it owns, for every face it sees; a face shared with another rank is seen by both, once from each side, so between them the face’s whole contribution lands. A face interior to this rank is seen once and both halves are scattered here. That is the same division of labour MFEM’s own DG assembly makes, and it is why the matrix needs MFEM’s face-neighbour dofs while the vector does not.

hippymfem.fem.facets.facet_values(space, v, gf=None)[source]

Local dof values followed by MFEM’s face-neighbour values, for a facet gather.

The far side of a face shared with another rank is not in this rank’s dof vector; MFEM carries those values in ParGridFunction.FaceNbrData once ExchangeFaceNbrData has run, and FacetSpaceTables indexes the two as one array. Collective.

hippymfem.fem.csrassemble

Direct CSR assembly: scatter element arrays without MFEM’s per-element callback.

Handing each element array to MFEM through a PyBilinearFormIntegrator lets MFEM own the dof combination, the essential-bc elimination and the parallel reduction, but its Python call per element is a large share of an assembly on a fine mesh (benchmarks/DESIGN_NOTES.md, section 2). The test suite keeps that route as its reference (hippymfem.test.reference_assembly). This module, the route of hippymfem.fem.assemble.assemble_matrix(), removes that call and does what ParBilinearForm::ParallelAssemble does, one level up:

  1. scatter the element arrays into the local ldof x ldof CSR graph with one vectorized reduction (no Python loop over elements);

  2. turn that local CSR into the true-dof HypreParMatrix. When both prolongations are boolean (every conforming space without a DofTransformation, decided collectively so all ranks take one branch), the ldof entries are mapped straight onto true-dof slots (tdofassemble): the rank’s own rows land in the diagonal and off-diagonal blocks directly, and the rows of shared dofs it does not own go to their owner in one Alltoallv, so P^T A P is never formed, which is a few times cheaper. Otherwise the CSR is wrapped as a block-diagonal HypreParMatrix over the ldof partition and hypre forms P^T A_local P (_triple()), which is right for non-conforming interfaces and for faces that carry a dof transformation alike. HIPPYMFEM_PARMAT forces either (set_parmat_mode()). With hypre on a device every live matrix gets its own copy of the pattern’s index arrays, for the reason in _as_sparse();

  3. eliminate essential dofs with the same HypreParMatrix calls MFEM uses (EliminateRowsCols for square blocks, EliminateCols then EliminateRows for rectangular ones) or, when the prolongation is boolean, by masking the ldof rows and columns before step 2, which gives the same matrix without a second pass.

One deliberate difference from the callback route. MFEM inserts element entries with skip_zeros=1, so a slot that is exactly zero in the element matrices never enters the sparsity pattern. This module keeps it, because the pattern is built once and reused: an entry that vanishes for one parameter can be nonzero for the next (P1 stiffness matrices on right triangles have exact zeros from orthogonal gradients; an anisotropic coefficient fills them in). The cost is a few percent more nonzeros and, because hypre then sums matvec rows in a different order, results that differ from the callback route at round-off (which a long BFGS run can amplify into visibly different iterates). The assembled matrices themselves are identical: the test suite compares them entry by entry on 1, 2 and 4 ranks.

The sparsity pattern and the scatter map depend only on the mesh and the spaces, so they are built once per pair of spaces and reused by every later assembly. That is what pays off in an inverse problem, where a Newton-CG run reassembles A, C, W_uu and the other blocks hundreds of times on a fixed mesh.

class hippymfem.fem.csrassemble.ScatterPattern(test_tables, trial_tables, test_space=None, trial_space=None)[source]

Bases: KeepAlive

CSR graph of the local ldof x ldof element-assembly map.

Built once per (test space, trial space, element groups) triple. Holds

  • indptr, indices: the CSR structure over local dof indices. Within a row the columns ascend, except that on a square pattern the diagonal entry comes first, the order hypre keeps a square matrix in, so hypre has nothing to permute and the matrix can alias these arrays without a copy;

  • slot: for every entry of every element array, the position of its target in the CSR data array. Assembly is then one reduction over slot, with no Python-level loop;

  • sign: the product of the test and trial dof signs, or None when every sign is +1 (always, for H1/L2; not for H(curl)/H(div)).

Parameters:
  • test_tables (sequence of SpaceTables) – One per element group, in the same order as the element arrays will be.

  • trial_tables (sequence of SpaceTables) – One per element group, in the same order as the element arrays will be.

release_slot()[source]

Give the slot map up where it can be made again and is large (LEAN): four bytes an element-matrix entry of host memory. Returns whether it was given up.

nent

the element-matrix entries of the pattern (the length of slot)

property slot

The CSR slot of every element-matrix entry. A map that was given up (release_slot()) is made again here, by whatever made it.

slots_of(rows, cols)[source]

CSR slots of the local (row, col) pairs, which must all be in the graph.

For entries that are not element-array entries of this pattern’s groups: a boundary residual’s element matrices, whose dofs are those of the adjacent volume element, land in the domain block’s own slots this way, so one matrix is built with the boundary terms already in it (instead of two matrices added and eliminated afterwards). One sort of the graph’s keys per pattern, then one searchsorted per call.

masked_slots(row_mask, col_mask)[source]

Indices of the CSR slots in the masked rows and columns.

Zeroing these in the data array before the triple product is what essential-dof elimination does afterwards, and costs an indexed write of the eliminated entries rather than a second parallel matrix. Cached on the identity of the two marker arrays, which are themselves cached per (space, dof set).

hypre_indices(roff, coff, same)[source]

(indptr, shifted indices, row starts, col starts) for hypre.

Cached on (roff, coff, same): the pattern and a rank’s ldof offsets are both fixed for the life of the spaces, so these arrays serve every one of an inverse problem’s hundreds of assemblies. None is returned for the column starts of a square block, where hypre must be handed one pointer for both partitions.

data(element_matrices, zero_slots=None, host=True, dtype=None)[source]

CSR data array for one set of element matrices.

The reduction runs where the element matrices are: when the kernels ran on a GPU the scatter happens there and only the nnz assembled values come back, so the (ne, nd, nd) array never touches the host. With host=False they do not come back either: the device array is returned as it is, for a consumer that builds the matrix on the device (hippymfem.fem.tdofassemble.TrueDofPattern.finish()).

On a device the sums are made in the precision of the element matrices unless dtype names another (on the host in double precision, whatever theirs is).

fusable()[source]

Whether chunks can be scattered through slices of slot as it is.

A chunk of elements must map onto a slice of the map, which slot gives unless the entries were permuted into mesh-element order across several geometry groups, the only case perm is set for. The fused scatter slices a copy in the order of the element arrays there (_array_order()).

data_fused(chunks, zero_slots=None, host=True, dtype=None)[source]

Scatter each chunk of element matrices as it arrives.

The chunks are never glued into one (ne, nd, nd) array, which only the scatter would read and which would be the largest resident allocation of the assembly.

chunks yields (group, start, stop, array). Written as fused_begin() / fused_add() / fused_end() so that a loop over chunks that carry several blocks can feed several patterns at once. dtype is the accumulator’s (fused_begin()).

fused_begin(dtype=None)[source]

A zeroed accumulator on the current device, and the maps to feed it.

The accumulator is the largest allocation an assembly makes, so it is borrowed from the buffers earlier assemblies kept and reset in place rather than allocated afresh (FUSED_KEEP). It is of double precision unless dtype says otherwise: a matrix that is to live in single precision is accumulated in single precision (numpy.float32), at half the memory.

fused_add(acc, maps, g, a, bnd, arr)[source]

Add one chunk’s element matrices into the accumulator; returns it.

The accumulator passed in is donated to the update (see _fused_add()), so only the returned one may be used afterwards.

fused_end(acc, zero_slots=None, host=True)[source]

Bring the accumulator to the host and zero the eliminated slots there.

The host copy comes back writeable (host_writable()), as data()’s does. With host=False nothing is copied: the eliminated slots are zeroed on the device, by a jitted update that is given the accumulator and writes into it, and the device array is returned. Its consumer hands the buffer back with give_back() once it has taken the values.

give_back(acc)[source]

Return a device accumulator whose values have been taken, so that the next assembly of this size reuses the buffer (FUSED_KEEP).

hippymfem.fem.csrassemble.set_deterministic(flag=True)[source]

Turn the order-independent device scatter on or off; returns the old value.

hippymfem.fem.csrassemble.set_fold_elimination(flag=True)[source]

Turn the folded essential-dof elimination on or off; returns the old value.

class hippymfem.fem.csrassemble.VectorPattern(tables, space=None)[source]

Bases: KeepAlive

Local dof indices of every element-vector entry, in mesh-element order.

scatter(element_vectors, host=True)[source]

Local dof array holding the summed element vectors.

Reduced where the element vectors are, as in ScatterPattern.data(); with host=False a device result is returned as the device array it is.

hippymfem.fem.csrassemble.get_vector_pattern(space, groups)[source]

Cached VectorPattern for a space on a set of element groups.

hippymfem.fem.csrassemble.get_pattern(test_space, trial_space, groups)[source]

Cached ScatterPattern for a space pair on a set of element groups.

The cache is keyed on the identity of the two spaces and the groups, and building a pattern involves no communication, so a hit on one rank and a miss on another is harmless. (A cache that wraps a collective cannot be keyed this way; see Layout.)

One pattern must not serve two sets of groups, though: the true-dof pattern built on it is found or built per pattern, and building it is collective (get_tdof_pattern()). A rank that holds no element of the groups has none to key on, so there the list itself stands for them: pass the batches’ own groups, one object for one set on every rank. Keyed on the spaces alone, the second boundary term on a pair of spaces found the first one’s pattern on such a rank, which then skipped the construction the other ranks were making.

hippymfem.fem.csrassemble.assemble_matrix_csr(test_space, trial_space, groups, element_matrices, nelem, test_ess=None, trial_ess=None, diag_policy='one', boundary=None, symmetric=None)[source]

Assemble a block from per-group element matrices without the callback.

Signature-compatible with hippymfem.fem.assemble.assemble_matrix(); nelem is accepted and unused (the pattern already knows which elements each group holds). Runs plan_block(), the scatter and finish_block(). With boundary = (test_tables, trial_tables, element_matrices) over the boundary groups, those entries are added into the same matrix (add_boundary_entries()).

hippymfem.fem.csrassemble.add_boundary_entries(p, acc, boundary)[source]

Add a boundary residual’s element matrices into a domain block’s accumulator.

boundary is (test_tables, trial_tables, element_matrices) over the boundary groups. The dofs are the adjacent volume elements’, so every entry has a slot in the domain pattern (ScatterPattern.slots_of()), mapped on to the true-dof target where that route is taken. The slots the folded elimination zeroes are zeroed again afterwards, so the elimination holds for the sum and one matrix comes out, instead of two to add and eliminate. The streamed pass of a linearization point works the same way (PDEVariationalProblem._blocks_streamed).

hippymfem.fem.csrassemble.assemble_vector_csr(space, groups, element_vectors, nelem, ess=None, out=None)[source]

Assemble a dual (residual) vector by direct scatter plus P^T.

The element vectors are summed into the local dof array with one numpy.bincount(), then reduced to true dofs by the space’s prolongation transpose, which is the same P^T that ParLinearForm::ParallelAssemble applies.

hippymfem.fem.csrassemble.clear_pattern_cache()[source]

Drop every cached pattern (for tests, and after a mesh is refined).

hippymfem.fem.pattern

Sparsity patterns and the scatter of element arrays into them.

ScatterPattern is the value-independent part of a matrix assembly (the CSR graph, the slot of every element entry, the masks and the device copies) and VectorPattern its counterpart for dual vectors; both are cached per (space pair, element groups) in an IdentityCache.

hippymfem.fem.pattern.DETERMINISTIC = False

When true, the device-side scatter gathers into each CSR slot in a fixed order instead of using atomics, so repeated GPU assemblies are bit-identical. Costs nnz * maxc doubles of working memory. Set HIPPYMFEM_GPU_DETERMINISTIC=1.

hippymfem.fem.pattern.FUSED_KEEP = True

Keep the nnz-sized accumulator of a fused assembly and reset it in place instead of allocating one per assembly. JAX’s arena never returns memory, so the kept buffer costs nothing that freeing it would give back, while asking for a fresh one costs a contiguous nnz block: at two million P2 hexahedra on a rank that is 7.7 GiB, more than half the arena at a quarter-card cap, and once the arena has grown to its cap and small allocations have landed in the freed block, the next assembly fails with RESOURCE_EXHAUSTED inside the scatter. Buffers are kept per (device, nnz) and borrowed, so what is held is the peak an assembly needs at once, never more. Set HIPPYMFEM_FUSED_KEEP=0 to allocate per assembly.

hippymfem.fem.pattern.FUSED_KEEP_SHARE = 0.25

Share of JAX’s arena above which an accumulator is kept. Only the large ones are worth it: a small buffer is one the allocator can fit in the holes it already has, while a buffer of a third of the arena or more cannot be found again once the arena is full and its old block has been split. Keeping the small ones as well raises the high-water mark for nothing, since A, C and the W blocks are each assembled in a pass of their own: at 128^3 on two L40S, where the accumulator is a fifth of the arena, keeping it takes JAX’s peak from 7.6 to 11.7 GiB. Set HIPPYMFEM_FUSED_KEEP_SHARE.

hippymfem.fem.pattern.DEVICE_PATTERN = False

the slot of every element-matrix entry (4 bytes per entry, which the fused scatter otherwise uploads chunk by chunk at every assembly) and the column indices of the two blocks of its true-dof matrix (4 bytes per nonzero, which a matrix built on the device otherwise uploads once per matrix). For the Jacobian of 2.1 million quadratic-hexahedron dofs that is 0.76 GB and 0.54 GB of JAX’s arena, and it saves 0.11 s of a 0.55 s forward solve on an H100. Off by default, since device memory is the scarcer resource. Set HIPPYMFEM_DEVICE_PATTERN=1.

Type:

Keep a pattern’s index arrays on the device

hippymfem.fem.pattern.COMPACT_PATTERN = True

for every element and test dof (a row of the element matrix) two base slots, and for every entry one byte that picks one of them and adds an offset below 128 (two bytes and offsets below 32768 where that does not reach). A row of an element matrix lands in at most two runs of its matrix row (the diagonal and the off-diagonal block of the true-dof layout, or the send buffer), each shorter than the row, so the map takes 1 + 8 / nd_trial bytes an entry instead of 4: for quadratic hexahedra 1.3. The slot is rebuilt on the device, in the scatter, from the two, and the sum is the same. It is a third of the upload of an assembly from the host, and a third of the device memory of DEVICE_PATTERN. A map that does not have the form (a row in three runs) stays as it is. Set HIPPYMFEM_COMPACT_PATTERN=0 for the full map always.

Type:

The scatter map of a fused assembly in a compact form

hippymfem.fem.pattern.LEAN = True

Build a large pattern without arrays of the size of the element matrices beside what it keeps, and keep no map twice: the graph of one group of elements comes from its two dof tables (build_tables()), the true-dof graph is laid out row by row (tdof_graph()), the compact map of COMPACT_PATTERN is made in slabs, and a map with one entry per element-matrix entry (slot) is given up once what is made from it exists: the ldof pattern’s when its true-dof target has its own, the target’s when the compact map has been made. A map given up is made again if something asks for it. With a million quadratic hexahedra on a rank the index tables of a Newton step then take 13 GiB of host memory instead of 30, and their builds 3 GiB more instead of 12. The arrays are the same. HIPPYMFEM_PATTERN_LEAN=0 builds and keeps them as before.

hippymfem.fem.pattern.LEAN_MIN_ENTRIES = 16777216

from this many entries.

Type:

A map is given up only where it is worth it

hippymfem.fem.pattern.set_device_pattern(flag)[source]

Turn the device copies of the index arrays on or off; returns the old setting.

hippymfem.fem.pattern.set_deterministic(flag=True)[source]

Turn the order-independent device scatter on or off; returns the old value.

hippymfem.fem.pattern.host_writable(acc)[source]

A writeable host array from an accumulator that may still be a device array.

np.asarray of a JAX array is the one device-to-host copy; JAX hands it back read-only but owning its buffer. Every accumulator here is consumed by its caller, which zeroes the eliminated slots, adds the boundary entries and builds the matrix on top of it, so the flag is lifted rather than paying a second nnz-sized copy. Where it cannot be lifted (a view of a host JAX buffer), the copy is made. np.add.at does not check the flag, so a read-only accumulator would fail only at its first indexed assignment.

hippymfem.fem.pattern.clear_accumulators()[source]

Drop every kept accumulator (for tests, and after a mesh is refined).

class hippymfem.fem.pattern.ScatterPattern(test_tables, trial_tables, test_space=None, trial_space=None)[source]

Bases: KeepAlive

CSR graph of the local ldof x ldof element-assembly map.

Built once per (test space, trial space, element groups) triple. Holds

  • indptr, indices: the CSR structure over local dof indices. Within a row the columns ascend, except that on a square pattern the diagonal entry comes first, the order hypre keeps a square matrix in, so hypre has nothing to permute and the matrix can alias these arrays without a copy;

  • slot: for every entry of every element array, the position of its target in the CSR data array. Assembly is then one reduction over slot, with no Python-level loop;

  • sign: the product of the test and trial dof signs, or None when every sign is +1 (always, for H1/L2; not for H(curl)/H(div)).

Parameters:
  • test_tables (sequence of SpaceTables) – One per element group, in the same order as the element arrays will be.

  • trial_tables (sequence of SpaceTables) – One per element group, in the same order as the element arrays will be.

release_slot()[source]

Give the slot map up where it can be made again and is large (LEAN): four bytes an element-matrix entry of host memory. Returns whether it was given up.

nent

the element-matrix entries of the pattern (the length of slot)

property slot

The CSR slot of every element-matrix entry. A map that was given up (release_slot()) is made again here, by whatever made it.

slots_of(rows, cols)[source]

CSR slots of the local (row, col) pairs, which must all be in the graph.

For entries that are not element-array entries of this pattern’s groups: a boundary residual’s element matrices, whose dofs are those of the adjacent volume element, land in the domain block’s own slots this way, so one matrix is built with the boundary terms already in it (instead of two matrices added and eliminated afterwards). One sort of the graph’s keys per pattern, then one searchsorted per call.

masked_slots(row_mask, col_mask)[source]

Indices of the CSR slots in the masked rows and columns.

Zeroing these in the data array before the triple product is what essential-dof elimination does afterwards, and costs an indexed write of the eliminated entries rather than a second parallel matrix. Cached on the identity of the two marker arrays, which are themselves cached per (space, dof set).

hypre_indices(roff, coff, same)[source]

(indptr, shifted indices, row starts, col starts) for hypre.

Cached on (roff, coff, same): the pattern and a rank’s ldof offsets are both fixed for the life of the spaces, so these arrays serve every one of an inverse problem’s hundreds of assemblies. None is returned for the column starts of a square block, where hypre must be handed one pointer for both partitions.

data(element_matrices, zero_slots=None, host=True, dtype=None)[source]

CSR data array for one set of element matrices.

The reduction runs where the element matrices are: when the kernels ran on a GPU the scatter happens there and only the nnz assembled values come back, so the (ne, nd, nd) array never touches the host. With host=False they do not come back either: the device array is returned as it is, for a consumer that builds the matrix on the device (hippymfem.fem.tdofassemble.TrueDofPattern.finish()).

On a device the sums are made in the precision of the element matrices unless dtype names another (on the host in double precision, whatever theirs is).

fusable()[source]

Whether chunks can be scattered through slices of slot as it is.

A chunk of elements must map onto a slice of the map, which slot gives unless the entries were permuted into mesh-element order across several geometry groups, the only case perm is set for. The fused scatter slices a copy in the order of the element arrays there (_array_order()).

data_fused(chunks, zero_slots=None, host=True, dtype=None)[source]

Scatter each chunk of element matrices as it arrives.

The chunks are never glued into one (ne, nd, nd) array, which only the scatter would read and which would be the largest resident allocation of the assembly.

chunks yields (group, start, stop, array). Written as fused_begin() / fused_add() / fused_end() so that a loop over chunks that carry several blocks can feed several patterns at once. dtype is the accumulator’s (fused_begin()).

fused_begin(dtype=None)[source]

A zeroed accumulator on the current device, and the maps to feed it.

The accumulator is the largest allocation an assembly makes, so it is borrowed from the buffers earlier assemblies kept and reset in place rather than allocated afresh (FUSED_KEEP). It is of double precision unless dtype says otherwise: a matrix that is to live in single precision is accumulated in single precision (numpy.float32), at half the memory.

fused_add(acc, maps, g, a, bnd, arr)[source]

Add one chunk’s element matrices into the accumulator; returns it.

The accumulator passed in is donated to the update (see _fused_add()), so only the returned one may be used afterwards.

fused_end(acc, zero_slots=None, host=True)[source]

Bring the accumulator to the host and zero the eliminated slots there.

The host copy comes back writeable (host_writable()), as data()’s does. With host=False nothing is copied: the eliminated slots are zeroed on the device, by a jitted update that is given the accumulator and writes into it, and the device array is returned. Its consumer hands the buffer back with give_back() once it has taken the values.

give_back(acc)[source]

Return a device accumulator whose values have been taken, so that the next assembly of this size reuses the buffer (FUSED_KEEP).

hippymfem.fem.pattern.PATTERN_TIMING = False

Print how long each phase of a pattern build took, on rank 0. The build is the one-time cost a large run pays before its first assembly (115 s of the 121 s a 256^3 linearization point took on 8 H100), and which phase carries it depends on where the sort runs, so it is worth being able to ask. Set HIPPYMFEM_PATTERN_TIMING=1.

hippymfem.fem.pattern.get_pattern(test_space, trial_space, groups)[source]

Cached ScatterPattern for a space pair on a set of element groups.

The cache is keyed on the identity of the two spaces and the groups, and building a pattern involves no communication, so a hit on one rank and a miss on another is harmless. (A cache that wraps a collective cannot be keyed this way; see Layout.)

One pattern must not serve two sets of groups, though: the true-dof pattern built on it is found or built per pattern, and building it is collective (get_tdof_pattern()). A rank that holds no element of the groups has none to key on, so there the list itself stands for them: pass the batches’ own groups, one object for one set on every rank. Keyed on the spaces alone, the second boundary term on a pair of spaces found the first one’s pattern on such a rank, which then skipped the construction the other ranks were making.

hippymfem.fem.pattern.clear_pattern_cache()[source]

Drop every cached pattern (for tests, and after a mesh is refined).

class hippymfem.fem.pattern.VectorPattern(tables, space=None)[source]

Bases: KeepAlive

Local dof indices of every element-vector entry, in mesh-element order.

scatter(element_vectors, host=True)[source]

Local dof array holding the summed element vectors.

Reduced where the element vectors are, as in ScatterPattern.data(); with host=False a device result is returned as the device array it is.

hippymfem.fem.pattern.get_vector_pattern(space, groups)[source]

Cached VectorPattern for a space on a set of element groups.

hippymfem.fem.prolongation

The prolongation P of a space: is it the identity, is it boolean, and the essential-dof mask it induces on local dofs.

Every answer here selects between code paths with different collectives, so the ones that matter are decided collectively and memoized on the space (IdentityCache).

hippymfem.fem.parmat

From a local CSR to the parallel HypreParMatrix: the constructors, the triple product P^T A P, and the route selection (HIPPYMFEM_PARMAT).

hippymfem.fem.parmat.PARMAT_MODE = 'auto'

How the local CSR becomes a HypreParMatrix. "tdof" scatters straight into true-dof rows and exchanges the few entries in ghost rows (hippymfem.fem.tdofassemble), which needs boolean prolongations and is two to three times faster than a triple product on a few ranks. "mfem" builds the block-diagonal matrix over the local dofs and forms the parallel triple product (_triple()), which is right for every space. "auto", the default, is "tdof" wherever the prolongations are boolean and "mfem" otherwise. Set HIPPYMFEM_PARMAT.

hippymfem.fem.parmat.COPYING_CONSTRUCTOR = False

Build the local matrix on the host with the constructor that copies the arrays into hypre’s memory, as a run with hypre on a device always does, instead of the one that aliases them. For the test suite, which holds the two to exact equality.

hippymfem.fem.parmat.set_parmat_mode(mode)[source]

Choose how the parallel matrix is built; returns the old mode.

hippymfem.fem.parmat.local_par_matrix(pattern, data, test_space, trial_space)[source]

Block-diagonal HypreParMatrix over the ldof partition.

The local CSR carries local column indices; they are shifted into the global ldof numbering here. Nothing is communicated: every column of a rank’s element matrices is one of that rank’s own local dofs.

Two constructors can do this, with bit-identical results. On the host the arrays are wrapped as an mfem.SparseMatrix (_as_sparse()) that MFEM’s constructor aliases, copying nothing. With hypre on a device the raw arrays go to the constructor that copies them into hypre’s memory: the aliasing one registers the arrays with MFEM’s memory manager, and a Python proxy drops the arrays it keeps before its C++ destructor runs, so hypre’s destroy would find the device mirrors already freed (see TrueDofPattern.finish).

hippymfem.fem.elimination

Essential-dof elimination: folded into the scatter where the prolongation is boolean (HIPPYMFEM_FOLD_ELIMINATION), through MFEM’s own calls otherwise.

hippymfem.fem.elimination.FOLD_ELIMINATION = True

When true (the default), essential-dof elimination is folded into the scatter as a mask on the local CSR slots, instead of a second parallel matrix from EliminateRowsCols, which is a large share of a GPU assembly. HIPPYMFEM_FOLD_ELIMINATION=0 uses MFEM’s calls instead; the test suite compares the two entry by entry.

hippymfem.fem.elimination.set_fold_elimination(flag=True)[source]

Turn the folded essential-dof elimination on or off; returns the old value.

hippymfem.fem.assemble

Assembly of AD-generated element arrays into parallel matrices and vectors.

The element arrays are scattered directly (csrassemble). Essential boundary conditions are applied block by block as the table in bcs states, so that the matrices behave exactly as hIPPYlib’s do.

hippymfem.fem.assemble.empty_ess()[source]

A shared empty essential-dof list.

hippymfem.fem.assemble.assemble_matrix(test_space, trial_space, groups, element_matrices, nelem, test_ess=None, trial_ess=None, diag_policy='one', symmetric=None)[source]

Assemble a block from per-group element matrices.

Parameters:
  • test_space (FunctionSpace) – Row and column spaces.

  • trial_space (FunctionSpace) – Row and column spaces.

  • groups (sequence of ElementGroup)

  • element_matrices (sequence of ndarray, or callable) – One (ne, nd_test, nd_trial) array per group. A callable returning an iterator of (group, start, stop, array) chunks lets each chunk be scattered as it is produced, so the full array is never formed.

  • nelem (int) – Total number of mesh elements.

  • test_ess (mfem.intArray, optional) – Essential true dofs to eliminate on each side.

  • trial_ess (mfem.intArray, optional) – Essential true dofs to eliminate on each side.

  • diag_policy ({"one", "zero", "keep"}) – Diagonal entry left on eliminated rows of a square block. "one" for the forward Jacobian (so that the essential block is the identity and solves work on a right-hand side with zeroed essential entries); "zero" for second-order blocks, which must annihilate the essential subspace.

  • symmetric (bool, optional) – Whether a block with the same space on both sides is eliminated on both (the default); False for a parameter’s block, whose columns stay (see plan_block()).

Returns:

mfem.HypreParMatrix

hippymfem.fem.assemble.assemble_vector(space, groups, element_vectors, nelem, ess=None, out=None)[source]

Assemble a dual (residual) vector from per-group element vectors.

The result lives on true dofs: the local-dof sums are reduced with P^T, which is the correct reduction for a linear functional.

hippymfem.fem.assemble.assemble_scalar(comm, element_values)[source]

Sum per-element integrals into one global scalar.

The per-group sum is taken where the values are, so a device result returns one scalar rather than one value per element.

hippymfem.fem.assemble.mass_functional(space, w=None, coeff=None, out=None)[source]

M w as a linear form, without building M.

\(\ell_i = \int c\,w_h\,\phi_i\), with \(w_h\) the finite element field w (a true-dof vector on space; None means the constant 1) and c an optional mfem.Coefficient multiplying it. This is the mass matrix’s action on w, to round-off: the rule is MassIntegrator’s own, 2p + Trans.OrderW() (MFEM adds nothing for a coefficient), so the two agree on curved elements as well as affine ones, where the linear form’s default rule of 2p + 1 would not. Building M for a single product costs a matrix and hypre’s assembly temporaries; this costs one element loop and one vector. Scalar H1/L2 spaces only.

hippymfem.fem.assemble.assemble_native_matrix(space, integrators, bdr_integrators=(), ess=None, diag_policy='one', trial_space=None, trial_ess=None)[source]

Assemble a block from hand-written MFEM integrators instead of AD kernels.

Used by the priors (whose forms are fixed bilinear forms that MFEM already has integrators for) and available to users through LinearPDEProblem.

hippymfem.fem.coefficients

Coefficients built from mesh attributes.

hippymfem.fem.coefficients.attribute_indicator(mesh, attributes)[source]

A piecewise-constant coefficient: 1 on the elements whose attribute is in attributes, 0 elsewhere.

Returns (coefficient, keep). MFEM’s coefficient holds a raw pointer to the marker vector, so the caller keeps everything in keep alive with it.

hippymfem.fem.bcs

Essential (Dirichlet) boundary conditions.

hIPPYlib takes two boundary condition lists: bc with the real data, used for the forward solve, and bc0, its homogeneous counterpart, used for the adjoint and incremental solves. The same pair appears here.

The discrete convention, which matches what hIPPYlibx does through dolfinx, is a single uniform rule:

for any operator block, zero the essential entries of the input when the trial space carries essential conditions, and of the output when the test space does.

Matrices follow suit, block by block. This table is the one statement of the convention; assemble implements it and PDEVariationalProblem relies on it.

block

essential-dof treatment

A, A^T

rows and columns eliminated, unit diagonal, so a solve against a right-hand side that is zero on the essential dofs returns an increment that is zero there

W_uu

rows and columns eliminated, zero diagonal

C, W_um

rows of the test space eliminated

W_mm

untouched

C^T, W_mu

the transposes, which carry the condition to the correct side

With those eliminations baked into the matrices, apply_ij is a plain matrix product and needs no masking, and the result agrees with hIPPYlib’s matrix-free set_bc bookkeeping block for block.

hippymfem.fem.bcs.zero_essential(v, ess)[source]

Set the true-dof entries ess of v to zero, where v is.

For a vector MFEM holds on a GPU the write is done there, from a device copy of ess that is built once and kept as long as the array ess lives (the BC objects pass the same array every time). Otherwise it is an indexed write through the host array, which for a device vector would mean a copy of the whole vector down and another up.

class hippymfem.fem.bcs.DirichletBC(space, value=None, bdr_attributes='all', marker=None, component=-1)[source]

Bases: KeepAlive

Essential boundary condition on part of the boundary of a space.

Parameters:
  • space (FunctionSpace or ParFiniteElementSpace)

  • value (float, callable, or None) – Boundary data. A callable receives physical coordinates as a numpy array and returns a float (or a length-vdim array). On a vector space a single number applies to every component. None means homogeneous.

  • bdr_attributes (sequence of int, "all", or None) – Mesh boundary attributes (MFEM’s 1-based numbering) where the condition applies. "all" selects the whole boundary.

  • marker (callable, optional) – Alternative to bdr_attributes: marker(x) -> bool applied to the coordinates of boundary element vertices; a boundary element is selected when all of its vertices satisfy it.

  • component (int, optional) – For vector spaces, restrict the condition to one component.

homogeneous()[source]

The same condition with zero data (hIPPYlib’s bc0).

apply(v)[source]

Set the essential entries of the true-dof vector v to the data.

The boundary values depend only on the data, the marker and the space, so they are projected once and reused; projecting into a fresh ParGridFunction on every call and keeping it would leak one state-sized array (with its device mirror) per forward solve. Reassigning value invalidates the cache.

zero(v)[source]

Set the essential entries of v to zero.

class hippymfem.fem.bcs.BCSet(bcs=None, space=None)[source]

Bases: KeepAlive

A collection of DirichletBC on one space, acting as one.

homogeneous()[source]
apply(v)[source]
zero(v)[source]
zero_copy(v, out=None)[source]

A copy of v with essential entries zeroed (leaves v alone).

hippymfem.fem.bcs.as_bcset(bcs, space=None)[source]

Coerce bcs (None / DirichletBC / list / BCSet) to a BCSet.

hippymfem.fem.jaxops

UFL-flavoured helpers for writing residual densities.

These are one-line wrappers over jax.numpy whose only purpose is to let a residual density read like the UFL form it replaces:

def pde_varf(u, m, p, x):
    return jnp.exp(m.val) * inner(u.grad, p.grad) - f(x) * p.val
hippymfem.fem.jaxops.inner(a, b)[source]

Full contraction of two arrays of equal shape (scalars included).

hippymfem.fem.jaxops.dot(a, b)[source]

Single contraction over the last axis of a and the first of b.

hippymfem.fem.jaxops.div(grad_v)[source]

Divergence of a vector field from its gradient (vdim, sdim).

hippymfem.fem.jaxops.tr(A)[source]

Trace.

hippymfem.fem.jaxops.sym(A)[source]

Symmetric part.

hippymfem.fem.jaxops.skew(A)[source]

Antisymmetric part.

hippymfem.fem.jaxops.Identity(n)[source]

n x n identity.

hippymfem.fem.jaxops.outer(a, b)[source]

Outer product.

hippymfem.fem.jaxops.cross(a, b)[source]

3-vector cross product.

hippymfem.fem.jaxops.nabla_grad(grad_v)[source]

grad_v transposed, i.e. the other gradient convention.

hippymfem.fem.jaxops.sqrt(x)[source]
hippymfem.fem.jaxops.exp(x)[source]
hippymfem.fem.jaxops.ln(x)[source]
hippymfem.fem.jaxops.log(x)
hippymfem.fem.jaxops.norm2(a)[source]

Squared Euclidean norm.

hippymfem.fem.io

Writing fields out for visualization, and reading them back.

ParaView is the default because MFEM’s ParaViewDataCollection handles high-order fields natively, so a P2 state is written as P2 rather than silently sampled onto vertices.

hippymfem.fem.io.write_paraview(basename, mesh, fields, cycle=0, time=0.0, order=None, data_format='binary', high_order=True)[source]

Write named fields to a ParaView collection.

Parameters:
  • basename (str) – Output path; the directory is created and the collection named after the final path component.

  • mesh (mfem.ParMesh)

  • fields (dict) – {name: (space, ParVector)} or {name: ParGridFunction}.

  • order (int, optional) – Levels of detail; defaults to the highest polynomial degree present, so high-order fields are not under-resolved in the output.

class hippymfem.fem.io.ParaViewWriter(basename, mesh, spaces, order=None, data_format='binary', high_order=True)[source]

Bases: object

Incremental ParaView output, for time-dependent or iteration histories.

Register the fields once, then call save() per time level; the grid functions are reused so nothing is reallocated per step.

save(values, time=None, cycle=None)[source]

Write one time level; values maps field name to a ParVector.

hippymfem.fem.io.write_glvis(basename, mesh, fields)[source]

Write MFEM’s native mesh/solution files, readable by GLVis.

hippymfem.fem.io.save_vector(path, v, comm=None)[source]

Save a distributed vector as a single .npy on rank 0.

hippymfem.fem.io.load_vector(path, v, comm=None)[source]

Load a vector saved by save_vector() back into v.

hippymfem.fem.io.write_point_csv(path, points, values, comm, header=None)[source]

Write point coordinates and values as CSV on rank 0.