MRX API Reference

Module contents

class mrx.Any(*args, **kwargs)

Bases: object

Special type indicating an unconstrained type.

Any is compatible with every type.
Any assumed to have all methods.
All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class mrx.Array

Bases: object

Array base class for JAX

jax.Array is the public interface for instance checks and type annotation of JAX arrays and tracers. Its main applications are in instance checks and type annotations; for example:

x = jnp.arange(5)
isinstance(x, jax.Array)  # returns True both inside and outside traced functions.

def f(x: Array) -> Array:  # type annotations are valid for traced and non-traced types.
  return x

jax.Array should not be used directly for creation of arrays; instead you should use array creation routines offered in jax.numpy, such as jax.numpy.array(), jax.numpy.zeros(), jax.numpy.ones(), jax.numpy.full(), jax.numpy.arange(), etc.

abstract property T

Compute the all-axis array transpose.

Refer to jax.numpy.transpose() for details.

addressable_data(index: int) → Array: Return an array of the addressable data at a particular index.

property addressable_shards: Sequence[Any]: List of addressable shards.

Test whether all array elements along a given axis evaluate to True.

Refer to jax.numpy.all() for the full documentation.

Test whether any array elements along a given axis evaluate to True.

Refer to jax.numpy.any() for the full documentation.

abstract argmax(axis: int | None = None, out: None = None, keepdims: bool | None = None) → Array

Return the index of the maximum value.

Refer to jax.numpy.argmax() for the full documentation.

abstract argmin(axis: int | None = None, out: None = None, keepdims: bool | None = None) → Array

Return the index of the minimum value.

Refer to jax.numpy.argmin() for the full documentation.

abstract argpartition(kth: int, axis: int = -1) → Array

Return the indices that partially sort the array.

Refer to jax.numpy.argpartition() for the full documentation.

abstract argsort(axis: int | None = -1, *, kind: None = None, order: None = None, stable: bool = True, descending: bool = False) → Array

Return the indices that sort the array.

Refer to jax.numpy.argsort() for the full documentation.

Copy the array and cast to a specified dtype.

This is implemented via jax.lax.convert_element_type(), which may have slightly different behavior than numpy.ndarray.astype() in some cases. In particular, the details of float-to-int and int-to-float casts are implementation dependent.

abstract property at

Helper property for index update functionality.

The at property provides a functionally pure equivalent of in-place array modifications.

In particular:

Alternate syntax	Equivalent In-place expression
`x = x.at[idx].set(y)`	`x[idx] = y`
`x = x.at[idx].add(y)`	`x[idx] += y`
`x = x.at[idx].subtract(y)`	`x[idx] -= y`
`x = x.at[idx].multiply(y)`	`x[idx] *= y`
`x = x.at[idx].divide(y)`	`x[idx] /= y`
`x = x.at[idx].power(y)`	`x[idx] **= y`
`x = x.at[idx].min(y)`	`x[idx] = minimum(x[idx], y)`
`x = x.at[idx].max(y)`	`x[idx] = maximum(x[idx], y)`
`x = x.at[idx].apply(ufunc)`	`ufunc.at(x, idx)`
`x = x.at[idx].get()`	`x = x[idx]`

None of the x.at expressions modify the original x; instead they return a modified copy of x. However, inside a jit() compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place.

Unlike NumPy in-place operations such as x[idx] += y, if multiple indices refer to the same location, all updates will be applied (NumPy would only apply the last update, rather than applying all updates.) The order in which conflicting updates are applied is implementation-defined and may be nondeterministic (e.g., due to concurrency on some hardware platforms).

By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound index semantics can be specified via the mode parameter (see below).

Parameters:

mode –
string specifying out-of-bound indexing mode. Options are:
- "promise_in_bounds": (default) The user promises that indices are in bounds. No additional checking will be performed. In practice, this means that out-of-bounds indices in get() will be clipped, and out-of-bounds indices in set(), add(), etc. will be dropped.
- "clip": clamp out of bounds indices into valid range.
- "drop": ignore out-of-bound indices.
- "fill": alias for "drop". For get(), the optional fill_value argument specifies the value that will be returned.
See jax.lax.GatherScatterMode for more details.
wrap_negative_indices – If True (default) then negative indices indicate position from the end of the array, similar to Python and NumPy indexing. If False, then negative indices are considered out-of-bounds and behave according to the mode parameter.
fill_value – Only applies to the get() method: the fill value to return for out-of-bounds slices when mode is 'fill'. Ignored otherwise. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.
indices_are_sorted – If True, the implementation will assume that the (normalized) indices passed to at[] are sorted in ascending order, which can lead to more efficient execution on some backends. If True but the indices are not actually sorted, the output is undefined.
unique_indices – If True, the implementation will assume that the (normalized) indices passed to at[] are unique, which can result in more efficient execution on some backends. If True but the indices are not actually unique, the output is undefined.

Examples

>>> x = jnp.arange(5.0)
>>> x
Array([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[2].get()
Array(2., dtype=float32)
>>> x.at[2].add(10)
Array([ 0.,  1., 12.,  3.,  4.], dtype=float32)

By default, out-of-bound indices are ignored in updates, but this behavior can be controlled with the mode parameter:

>>> x.at[10].add(10)  # dropped
Array([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[20].add(10, mode='clip')  # clipped
Array([ 0.,  1.,  2.,  3., 14.], dtype=float32)

For get(), out-of-bound indices are clipped by default:

>>> x.at[20].get()  # out-of-bounds indices clipped
Array(4., dtype=float32)
>>> x.at[20].get(mode='fill')  # out-of-bounds indices filled with NaN
Array(nan, dtype=float32)
>>> x.at[20].get(mode='fill', fill_value=-1)  # custom fill value
Array(-1., dtype=float32)

Negative indices count from the end of the array, but this behavior can be disabled by setting wrap_negative_indices = False:

>>> x.at[-1].set(99)
Array([ 0.,  1.,  2.,  3., 99.], dtype=float32)
>>> x.at[-1].set(99, wrap_negative_indices=False, mode='drop')  # dropped!
Array([0., 1., 2., 3., 4.], dtype=float32)

abstract byteswap() → Array

Swap the bytes of the array elements.

This switches between a little-endian and big-endian data representation.

Returns:: An array with the same dtype as self, with underlying bytes of each entry reversed.

Examples

>>> import jax.numpy as jnp
>>> x = jnp.arange(5, dtype='int32')
>>> x
Array([0, 1, 2, 3, 4], dtype=int32)
>>> x.byteswap()
Array([       0, 16777216, 33554432, 50331648, 67108864], dtype=int32)

When the resulting bytes are viewed as a big-endian dtype (possible in NumPy, but not in JAX) they represent the original values:

>>> import numpy as np
>>> np.array(x.byteswap()).view('>i4')  # view as big-endian
array([0, 1, 2, 3, 4], dtype='>i4')

Calling byteswap twice will return the original array:

>>> x.byteswap().byteswap()
Array([0, 1, 2, 3, 4], dtype=int32)

Construct an array choosing from elements of multiple arrays.

Refer to jax.numpy.choose() for the full documentation.

Return an array whose values are limited to a specified range.

Refer to jax.numpy.clip() for full documentation.

property committed: bool

Whether the array is committed or not.

An array is committed when it is explicitly placed on device(s) via JAX APIs. For example, jax.device_put(np.arange(8), jax.devices()[0]) is committed to device 0. While jax.device_put(np.arange(8)) is uncommitted and will be placed on the default device.

Computations involving some committed inputs will happen on the committed device(s) and the result will be committed on the same device(s). Invoking an operation on arguments that are committed to different device(s) will raise an error.

Examples

>>> a = jax.device_put(np.arange(8), jax.devices()[0])
>>> b = jax.device_put(np.arange(8), jax.devices()[1])
>>> a + b  
Traceback (most recent call last):
  ...
ValueError: Received incompatible devices for jitted computation.

Return selected slices of this array along given axis.

Refer to jax.numpy.compress() for full documentation.

abstract conj() → Array

Return the complex conjugate of the array.

Refer to jax.numpy.conj() for the full documentation.

abstract conjugate() → Array

Return the complex conjugate of the array.

Refer to jax.numpy.conjugate() for the full documentation.

abstract copy() → Array

Return a copy of the array.

Refer to jax.numpy.copy() for the full documentation.

copy_to_host_async()

Copies an Array to the host asynchronously.

For arrays that live an an accelerator, such as a GPU or a TPU, JAX may cache the value of the array on the host. Normally this happens behind the scenes when the value of an on-device array is requested by the user, but waiting to initiate a device-to-host copy until the value is requested requires that JAX block the caller while waiting for the copy to complete.

copy_to_host_async requests that JAX populate its on-host cache of an array, but does not wait for the copy to complete. This may speed up a future on-host access to the array’s contents.

Return the cumulative product of the array.

Refer to jax.numpy.cumprod() for the full documentation.

Return the cumulative sum of the array.

Refer to jax.numpy.cumsum() for the full documentation.

property device: Any

Array API-compatible device attribute.

For single-device arrays, this returns a Device. For sharded arrays, this returns a Sharding.

abstract diagonal(offset: int = 0, axis1: int = 0, axis2: int = 1) → Array

Return the specified diagonal from the array.

Refer to jax.numpy.diagonal() for the full documentation.

Compute the dot product of two arrays.

Refer to jax.numpy.dot() for the full documentation.

property dtype: dtype: The data type (numpy.dtype) of the array.

abstract property flat

Use flatten() instead.

Type:: Not implemented

abstract flatten(order: str = 'C', *, out_sharding=None) → Array

Flatten array into a 1-dimensional shape.

Refer to jax.numpy.ravel() for the full documentation.

property global_shards: Sequence[Any]: List of global shards.

abstract property imag: Array: Return the imaginary part of the array.

property is_fully_addressable: bool

Is this Array fully addressable?

A jax.Array is fully addressable if the current process can address all of the devices named in the Sharding. is_fully_addressable is equivalent to “is_local” in multi-process JAX.

Note that fully replicated is not equal to fully addressable i.e. a jax.Array which is fully replicated can span across multiple hosts and is not fully addressable.

property is_fully_replicated: bool: Is this Array fully replicated?

abstract item(*args: int) → bool | int | float | complex: Copy an element of an array to a standard Python scalar and return it.

abstract property itemsize: int: Length of one array element in bytes.

abstract property mT

Compute the (batched) matrix transpose.

Refer to jax.numpy.matrix_transpose() for details.

Return the maximum of array elements along a given axis.

Refer to jax.numpy.max() for the full documentation.

Return the mean of array elements along a given axis.

Refer to jax.numpy.mean() for the full documentation.

Return the minimum of array elements along a given axis.

Refer to jax.numpy.min() for the full documentation.

abstract property nbytes: int: Total bytes consumed by the elements of the array.

property ndim: int: The number of dimensions in the array.

Return indices of nonzero elements of an array.

Refer to jax.numpy.nonzero() for the full documentation.

Return product of the array elements over a given axis.

Refer to jax.numpy.prod() for the full documentation.

abstract ptp(axis: int | Sequence[int] | None = None, out: None = None, keepdims: bool = False) → Array

Return the peak-to-peak range along a given axis.

Refer to jax.numpy.ptp() for the full documentation.

abstract ravel(order: str = 'C', *, out_sharding=None) → Array

Flatten array into a 1-dimensional shape.

Refer to jax.numpy.ravel() for the full documentation.

abstract property real: Array: Return the real part of the array.

Construct an array from repeated elements.

Refer to jax.numpy.repeat() for the full documentation.

abstract reshape(*args: Any, order: str = 'C', out_sharding=None) → Array

Returns an array containing the same data with a new shape.

Refer to jax.numpy.reshape() for full documentation.

abstract round(decimals: int = 0, out: None = None) → Array

Round array elements to a given decimal.

Refer to jax.numpy.round() for full documentation.

Perform a binary search within a sorted array.

Refer to jax.numpy.searchsorted() for full documentation.

property shape: tuple[int, ...]: The shape of the array.

property sharding: Any: The sharding for the array.

property size: int: The total number of elements in the array.

abstract sort(axis: int | None = -1, *, kind: None = None, order: None = None, stable: bool = True, descending: bool = False) → Array

Return a sorted copy of an array.

Refer to jax.numpy.sort() for full documentation.

abstract squeeze(axis: int | Sequence[int] | None = None) → Array

Remove one or more length-1 axes from array.

Refer to jax.numpy.squeeze() for full documentation.

Compute the standard deviation along a given axis.

Refer to jax.numpy.std() for full documentation.

Sum of the elements of the array over a given axis.

Refer to jax.numpy.sum() for full documentation.

abstract swapaxes(axis1: int, axis2: int) → Array

Swap two axes of an array.

Refer to jax.numpy.swapaxes() for full documentation.

Take elements from an array.

Refer to jax.numpy.take() for full documentation.

abstract to_device(device: Device | Sharding, *, stream: int | Any | None = None)

Return a copy of the array on the specified device

Parameters:

device – Device or Sharding to which the created array will be committed.
stream – not implemented, passing a non-None value will lead to an error.

Returns:

copy of array placed on the specified device or devices.

Return the sum along the diagonal.

Refer to jax.numpy.trace() for full documentation.

abstract transpose(*args: Any) → Array

Returns a copy of the array with axes transposed.

Refer to jax.numpy.transpose() for full documentation.

Compute the variance along a given axis.

Refer to jax.numpy.var() for full documentation.

abstract view(dtype: str | type[Any] | dtype | SupportsDType | None = None, type: None = None) → Array

Return a bitwise copy of the array, viewed as a new dtype.

This is fuller-featured wrapper around jax.lax.bitcast_convert_type().

If the source and target dtype have the same bitwidth, the result has the same shape as the input array. If the bitwidth of the target dtype is different from the source, the size of the last axis of the result is adjusted accordingly.

>>> jnp.zeros([1,2,3], dtype=jnp.int16).view(jnp.int8).shape
(1, 2, 6)
>>> jnp.zeros([1,2,4], dtype=jnp.int8).view(jnp.int16).shape
(1, 2, 2)

Conversions involving booleans are not well-defined in all situations. With regards to the shape of result as explained above, booleans are treated as having a bitwidth of 8. However, when converting to a boolean array, the input should only contain 0 or 1 bytes. Otherwise, results may be unpredictable or may change depending on how the result is used.

This conversion is guaranteed and safe:

>>> jnp.array([1, 0, 1], dtype=jnp.int8).view(jnp.bool_)
Array([ True, False,  True], dtype=bool)

However, there are no guarantees about the results of any expression involving a view such as this: jnp.array([1, 2, 3], dtype=jnp.int8).view(jnp.bool_). In particular, the results may change between JAX releases and depending on the platform. To safely convert such an array to a boolean array, compare it with 0:

>>> jnp.array([1, 2, 0], dtype=jnp.int8) != 0
Array([ True,  True, False], dtype=bool)

Parameters:

dtype – An optional output dtype. If not specified, the output dtype is the same as the input dtype.
type – Not implemented; accepted for NumPy compatibility.

Returns:

The array, viewed as the new dtype. Unlike NumPy, the array may or may not be a copy of the input array.

class mrx.BoundaryConditionPair(free: 'Optional[object]' = None, dbc: 'Optional[object]' = None)

Bases: Module

__init__(free: object | None = None, dbc: object | None = None) → None

dbc: object | None = None

free: object | None = None

class mrx.BoundaryIterativeRuntimeTuning(free: 'IterativeRuntimeTuning' = <factory>, dbc: 'IterativeRuntimeTuning' = <factory>)

Bases: Module

__init__(free: ~mrx.operators.IterativeRuntimeTuning = <factory>, dbc: ~mrx.operators.IterativeRuntimeTuning = <factory>) → None

dbc: IterativeRuntimeTuning

free: IterativeRuntimeTuning

class mrx.BoundaryOperator(Λ, types)

Bases: object

A lazy boundary operator for handling boundary conditions in differential forms.

This class implements boundary condition operators for differential forms on cube-like domains. It supports different types of boundary conditions and form degrees.

k

Degree of the differential form (0, 1, 2, or 3)

Type:: int

Lambda_0

Type:: DifferentialForm

types

Tuple of boundary condition types for each direction.

Type:: tuple

nr

Number of points in r-direction after boundary conditions

Type:: int

nt

Number of points in θ-direction after boundary conditions

Type:: int

nz

Number of points in ζ-direction after boundary conditions

Type:: int

dr

Number of points in r-direction

Type:: int

dt

Number of points in θ-direction

Type:: int

dz

Number of points in ζ-direction

Type:: int

n1

Size of first component

Type:: int

n2

Size of second component

Type:: int

n3

Size of third component

Type:: int

n

Total size of the operator

Type:: int

M: Assembled operator matrix

__init__(Λ, types)

Initialize the boundary operator.

Parameters:

Λ (DifferentialForm)
types (tuple) – Tuple of boundary condition types for each direction. Can be ‘dirichlet’ (zero at boundaries), ‘half’ (zero only at x=1) or other types (no boundary conditions).

_element(row_idx, col_idx)

Compute the operator element at specified indices.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index

Returns:

The operator element value

Return type:

jnp.ndarray

_unravel_index(idx)

Convert linear index to multi-dimensional coordinates.

Parameters:

idx (int) – Linear index

Returns:

(category, i, j, k) where category indicates the vector: component and (i,j,k) are the spatial coordinates

Return type:

tuple

_vector_index(idx)

Convert linear index to vector component and local index.

Parameters:

idx (int) – Linear index

Returns:

(category, local_index) where category indicates the vector: component and local_index is the index within that component

Return type:

tuple

build_extraction()

Build the MatrixFreeExtraction by probing each row against all columns.

Maps over rows sequentially, computing one row at a time with batched map over columns. Non-zero indices and values are collected into a MatrixFreeExtraction (gather/scatter apply, no matrix stored).

class mrx.BoundaryProjector(seq: DeRhamSequence, k: Literal[0, 1, 2, 3])

Bases: object

Project a k-form onto the Dirichlet boundary DOFs via a surface integral.

Computes the boundary load vector

b_i = ∫_{r=1} g(ξ) · trace(φ_i)(ξ) dS,

then selects the BC DOF values via the e_k_bc extraction operator.

g follows the same convention as Projector: for k = 0 and 3, a scalar function ξ → (1,); for k = 1 and 2, a vector function ξ → (3,) in the physical (x, y, z) frame.

All quadrature-dependent quantities (surface Jacobian, boundary quad points, r-spline values at r = 1) are computed once in __init__ and reused across calls.

__call__(g: Callable[[Array], Array] | Array) → Array

Compute the boundary load vector for prescribed boundary data g.

Parameters:

g (callable or array) –

If callable: ξ → (1,) for k = 0 or 3; ξ → (3,) in physical frame for k = 1 or 2. Evaluated at the boundary quad points.

If array of shape (ny*nx*nz, d): precomputed values at the full 3D quad grid (e.g. from oneform_projection). The θ,ζ quad points are the same as for the boundary; the r-dimension is irrelevant for boundary data, so slice [:, 0, :, :] is used.

Return type:

Array of shape (n_k_bc,)

_eval_trace_1form(u: Array) → Array

Return logical components E_log at r = 1, shape (nt, nz, 3).

No DF is applied here; physical E_phys = DF^{-T} E_log via einsum('tzji,tzj->tzi', self._DF_inv_bdy, E_log) if needed.

_eval_trace_2form(u: Array) → Array

Return B_phys · surf_normal = B_log_r at r = 1, shape (nt, nz).

J cancels: B_phys · surf_normal = (1/J)(DF B_log) · surf_normal: = (1/J) J B_log_r = B_log_r.

This is the unscaled normal flux (integrated against the surface element). To get the pointwise normal component B_phys · n̂ divide by the surface Jacobian ‖∂_θF × ∂_ζF‖, accessible as bp.surf_jac().

_project_1form(g_jk: Array) → Array: Transform physical → logical covariant (DF^{-1}) then integrate.

_project_2form(g_jk: Array) → Array: Pull back g to logical covariant 2-form (DF^T g / J) and integrate against each reference basis group weighted by surf_jac.

evaluate_trace(u: Array) → Array

Evaluate the trace of a discrete k-form at the boundary quad points.

Given the full (unreduced) DOF vector u of shape (n_k,), reconstruct the field values at the (nt, nz) boundary quad points.

No coordinate map evaluation is needed:

k = 0: scalar f(1, θ, ζ), shape (nt, nz).
k = 1: logical components E_log = (E_r, E_θ, E_ζ) at r = 1, shape (nt, nz, 3). The physical tangential vector is DF^{-T} E_log using the precomputed self._DF_inv_bdy.
k = 2: normal flux B_log_r = B_phys · (∂_θF × ∂_ζF) at r = 1, shape (nt, nz). The Jacobian J cancels exactly because B_phys = (1/J) DF B_log, so no DF evaluation is needed.

Parameters:: u (Array, shape (n_k,)) – Full DOF vector in the unreduced space (i.e. not BC-extracted).
Return type:: Array of shape (nt, nz) for k = 0 or 2, (nt, nz, 3) for k = 1.

class mrx.BoundaryShiftedIterativeRuntimeTuning(free: 'ShiftedIterativeRuntimeTuning' = <factory>, dbc: 'ShiftedIterativeRuntimeTuning' = <factory>)

Bases: Module

__init__(free: ~mrx.operators.ShiftedIterativeRuntimeTuning = <factory>, dbc: ~mrx.operators.ShiftedIterativeRuntimeTuning = <factory>) → None

dbc: ShiftedIterativeRuntimeTuning

free: ShiftedIterativeRuntimeTuning

class mrx.DeRhamSequence(ns, ps, q, types, *legacy_args, polar, tol=1e-12, maxiter=10000, r_scale=1.0, n_inner=5, betti_numbers=(1, 1, 0, 0))

Bases: object

Discrete de Rham sequence on a mapped 3-D domain.

Holds four DifferentialForm objects (basis_0 … basis_3), a QuadratureRule, a SequenceGeometry, and extraction/boundary operators for each form degree. After calling assemble_all_sparse() (or the individual assemble_* methods), operator application methods become available.

ns

Number of basis functions in each direction (n_r, n_θ, n_ζ).

Type:: tuple of int

ps

Polynomial degree in each direction (p_r, p_θ, p_ζ).

Type:: tuple of int

basis_0, basis_1, basis_2, basis_3

Spline bases for 0-, 1-, 2-, and 3-forms respectively.

Type:: DifferentialForm

quad

Tensor-product Gauss quadrature rule used for assembly.

Type:: QuadratureRule

geometry

Metric and Jacobian data derived from the logical-to-physical map.

Type:: SequenceGeometry

e0, e1, e2, e3

Extraction operators mapping constrained DOF vectors to the full spline basis for each form degree (no Dirichlet BCs).

Type:: MatrixFreeExtraction

e0_dbc, e1_dbc, e2_dbc, e3_dbc

Extraction operators with homogeneous Dirichlet BCs applied at the radial boundary (or axis in polar coordinates).

Type:: MatrixFreeExtraction

basis_r_jk

Radial 0-form basis splines evaluated at radial quadrature points. Shape (n_qr, n_r). Populated by evaluate_1d().

Type:: jnp.ndarray

basis_t_jk

Poloidal 0-form basis splines evaluated at poloidal quadrature points. Shape (n_qθ, n_θ). Populated by evaluate_1d().

Type:: jnp.ndarray

basis_z_jk

Toroidal 0-form basis splines evaluated at toroidal quadrature points. Shape (n_qζ, n_ζ). Populated by evaluate_1d().

Type:: jnp.ndarray

d_basis_r_jk

Radial derivative splines evaluated at radial quadrature points. Shape (n_qr, n_r). Populated by evaluate_1d().

Type:: jnp.ndarray

d_basis_t_jk

Poloidal derivative splines evaluated at poloidal quadrature points. Shape (n_qθ, n_θ). Populated by evaluate_1d().

Type:: jnp.ndarray

d_basis_z_jk

Toroidal derivative splines evaluated at toroidal quadrature points. Shape (n_qζ, n_ζ). Populated by evaluate_1d().

Type:: jnp.ndarray

__init__(ns, ps, q, types, *legacy_args, polar, tol=1e-12, maxiter=10000, r_scale=1.0, n_inner=5, betti_numbers=(1, 1, 0, 0))

Construct a de Rham sequence.

Parameters:

ns (list of int) – Number of basis functions [n_r, n_θ, n_ζ] for each direction.
ps (list of int) – Polynomial degree [p_r, p_θ, p_ζ] of the spline basis.
q (int) – Number of quadrature points per direction.
types (list of str) – Boundary-condition type string per direction, e.g. ['periodic', 'periodic', 'periodic'].
polar (bool) – If True, apply polar extraction operators that enforce regularity at the magnetic axis.
tol (float, optional) – Convergence tolerance for iterative solvers.
maxiter (int, optional) – Maximum iteration count for iterative solvers.
r_scale (float, optional) – Exponent used to cluster radial knots toward the axis (knot spacing proportional to r**r_scale).
n_inner (int, optional) – Number of inner CG iterations used by block preconditioners.
betti_numbers (tuple of 4 ints, optional) – (b0, b1, b2, b3) for the physical domain. Determines how many harmonic k-forms each Hodge Laplacian has, and hence the shapes of the nullspace arrays stored on SequenceOperators. Defaults to (1, 1, 0, 0) which matches a solid torus.

Notes

Geometry is no longer installed during construction. Call set_map() or set_spline_map() explicitly after building the sequence.

_add_boundary_dual(dv_dual, boundary_dual, operator_name): Add a prescribed boundary functional in the operator’s dual target space.

_apply_reference_mass_matrix(v, dirichlet=True): Apply the reference-domain 0-form mass matrix to v.

_apply_reference_mass_matrix_preconditioner(v, dirichlet=True): Apply the diagonal (Jacobi) preconditioner for the reference mass matrix.

_compute_nullspaces(betti_numbers=None, eps=1e-06)

Iteratively compute harmonic forms and store them on self.operators.

betti_numbers defaults to self.betti_numbers. Returns the info dict from compute_nullspaces_iterative().

_find_nullspace_vectors(k, n_vectors, eps, dirichlet=True): Find n_vectors nullspace vectors of the k-form Laplacian via inverse iteration.

_form_comp_info(k)

Return component metadata for tensor-product evaluation of the k-th form.

Returns:

comp_info (list of tuple) – Each entry (output_dim, R_jk, T_jk, Z_jk) describes one component: the physical vector index and the three 1-D basis arrays (one differentiated per form degree).
comp_shapes (list of int) – Number of DOFs for each component block.

_get_nullspace(k, dirichlet): Return the nullspace basis for the k-form Laplacian.

_get_saddle_point_nullspaces(k, dirichlet): Return the pair of nullspace bases for the k-th saddle-point system.

_grad_1d(d_basis, boundary_type): Return the 1-D gradient matrix for the given derivative basis and BC type.

_require_geometry(): Return the attached geometry or raise when none is installed.

_require_operators(operators=None): Return an explicit operator bundle or raise when none is available.

_require_reference_mass_matrix(): Raise if the reference-domain mass matrix has not been assembled.

_resolve_operators(operators=None): Use an explicit operator bundle when provided, else fall back to the cache.

_sync_operators(): Mirror bundled operators onto legacy fields during the transition.

apply_bc_mass_correction(g: Array, k: int) → Array

Compute the DBC-space RHS correction for a non-zero Dirichlet BC.

For a k-form mass-matrix system M_dbc @ u = rhs where the boundary DOFs are prescribed as g, the corrected right-hand side is:

rhs_corrected = rhs - seq.apply_bc_mass_correction(g, k)

The correction is E_dbc @ M_full @ E_bc^T @ g, i.e. the DBC-space projection of the mass matrix applied to the BC lift.

Requires assemble_all_sparse() (or the relevant assemble_M{k} call) to have been called first.

Parameters:

g (array of shape (n_k_bc,))
k (int)

Return type:

array of shape (n_k_dbc,)

apply_derivative_matrix(v, k, dirichlet_in=True, dirichlet_out=True, transpose=False, operators=None)

Apply the derivative matrix Dk (mapping k-forms to (k+1)-forms) to a vector v:: k=0: D0_ij = ∫ Λ1_i · G⁻¹ grad Λ0_j det DF dx (grad) k=1: D1_ij = ∫ Λ2_i · G curl Λ1_j (det DF)⁻¹ dx (curl) k=2: D2_ij = ∫ Λ3_i div Λ2_j (det DF)⁻¹ dx (div)

If transpose=True, apply Dk.T instead (mapping (k+1)-forms to k-forms).

apply_hodge_laplacian(v, k, dirichlet=True, operators=None): Backward-compatible alias for apply_laplacian.

apply_hodge_laplacian_approx(v, k, dirichlet=True, operators=None): Backward-compatible alias for apply_laplacian_approx.

apply_hodge_laplacian_preconditioner(v, k, dirichlet=True, operators=None, kind='auto'): Backward-compatible alias for apply_laplacian_preconditioner.

apply_incidence_matrix(v, k, dirichlet_in=True, dirichlet_out=True, transpose=False, operators=None)

Apply the topological exterior-derivative incidence Gk to v.

Gk has entries in {-1, 0, +1} and is geometry-independent. On DoF spaces where the extraction operators are “unitary” (e @ e^T = I), this equals M_{k+1}^{-1} @ apply_derivative_matrix. For non-unitary extractions (e.g. polar axis gluing) the two differ; in that regime apply_strong_grad() / curl / div remain the mass-projected form and should be preferred when exact d∘d = 0 on extracted DoFs is required.

apply_inverse_hodge_laplacian(rhs, k, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', return_info=False): Backward-compatible alias for apply_inverse_laplacian.

apply_inverse_laplacian(rhs, k, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', return_info=False): Apply the inverse of the k-form Laplacian to a right-hand side.

apply_inverse_mass_matrix(rhs, k, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', return_info=False): Apply the inverse of the sparse mass matrix Mk⁻¹ for k-forms to a right-hand side, solved via CG with a structured mass preconditioner. An optional initial guess can be provided to warm-start the solver.

apply_inverse_mass_plus_eps_laplace_matrix(rhs, k, eps, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', return_info=False)

Solve (M_k + eps * L_k) x = rhs for the k-form x.

For k=0: (M_0 + eps * S_0) is SPD, solved with CG. For k>=1: uses MINRES on the symmetric saddle-point system:

M_k + eps*S_k eps*D_{k-1} | | u | | rhs |

eps*D_{k-1}^T -eps*M_{k-1} | | σ | = | 0 |

The system is nonsingular (no nullspace) since M_k + eps*L_k is SPD. Out-of-the-box diffusion preconditioners currently use the same mass-side defaults as the other inverse paths: Jacobi, tensor, and Chebyshev.

apply_inverse_reference_mass_matrix(rhs, dirichlet=True, guess=None, tol=None, maxiter=None): Apply the inverse of the cached reference-domain 0-form mass matrix.

apply_inverse_shifted_hodge_laplacian(rhs, k, eps, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', use_harmonic_coarse=None, return_info=False): Backward-compatible alias for apply_inverse_shifted_laplacian.

apply_inverse_shifted_laplacian(rhs, k, eps, dirichlet=True, guess=None, operators=None, tol=None, maxiter=None, preconditioner='auto', use_harmonic_coarse=None, return_info=False)

Solve (L_k + eps * M_k) x = rhs for the k-form x.

For eps=0 this reduces to the Hodge Laplacian solve; the system may be singular and nullspace deflation is applied automatically. For eps > 0 the system is nonsingular (shift-invert for L_k u = λ M_k u). The shifted solve itself does not require precomputed nullspace data; any harmonic coarse correction is optional and should stay disabled while inverse iteration is still constructing those vectors.

For k=0: solved with CG on (S_0 + eps M_0) u = rhs. For k>=1: MINRES on the symmetric saddle-point form of L_k + eps M_k:

S_k + eps*M_k D_{k-1} | | u | | rhs |

D_{k-1}^T -M_{k-1} | | σ | = | 0 |

apply_laplacian(v, k, dirichlet=True, operators=None)

Apply the k-form Laplacian L_k to a vector v.

Naming and structure used throughout MRX:

S_k is the k-form stiffness block,
L_k = S_k + D_{k-1} M_{k-1}^{-1} D_{k-1}^T,
equivalently L_k = G_k^T M_{k+1} G_k + M_k G_{k-1} M_{k-1}^{-1} G_{k-1}^T M_k.

For k >= 1 this is applied through the Schur form above. For k = 0, L_0 = S_0.

apply_laplacian_approx(v, k, dirichlet=True, operators=None)

Linear approximate Laplacian apply.

Replaces M_{k-1}^{-1} in the Schur term with a single configured mass-preconditioner apply. Linear, SPD, safe to nest inside Krylov solvers and to use as a preconditioner. Not exact unless the metric is tensor-separable on the reference domain.

apply_laplacian_preconditioner(v, k, dirichlet=True, operators=None, kind='auto')

Apply a preconditioner for the k-form Laplacian to a vector v.

kind selects between 'none' (identity), 'jacobi' (per-DoF diagonal) and 'tensor' (tensorized Hodge/Laplacian preconditioner; available for k = 0 when the tensor Hodge data are assembled, and for k = 3 via the tensor round-trip path). 'auto' (the default) uses 'tensor' when available and falls back to 'jacobi' otherwise.

apply_leray_projection(v, k=2, p_guess=None)

Apply the Leray projection to a 1 or 2-form v.

When k = 2:: Solves the system (k=3 Hodge Laplacian): div v = div σ (σ, ω) = -(p, div ω) ∀ω 2-forms -> div(v - σ) = 0 and σ.n = 0 on the boundary.
When k = 1:: Solves the k=0 Hodge Laplacian: (grad p, grad ω) = (v, grad ω) ∀ω 0-forms -> div(v - grad p) = 0 and p = 0 on the boundary.

Parameters:

v (jnp.ndarray) – The vector form DoFs
k (int) – The degree of the vector form
p_guess (jnp.ndarray) – Guess for pressure form DoFs

Returns:

v_out (jnp.ndarray) – divergence-cleaned v
p (jnp.ndarray) – The pressure form DoFs

apply_mass_matrix(v, k, dirichlet=True, operators=None)

Apply the sparse mass matrix Mk for k-forms to a vector v:: k=0: M0_ij = ∫ Λ0_i Λ0_j det DF dx k=1: M1_ij = ∫ Λ1_i · G⁻¹ Λ1_j det DF dx k=2: M2_ij = ∫ Λ2_i · G Λ2_j (det DF)⁻¹ dx k=3: M3_ij = ∫ Λ3_i Λ3_j (det DF)⁻¹ dx

apply_mass_matrix_preconditioner(v, k, dirichlet=True, operators=None, kind='auto'): Apply a configured mass-matrix preconditioner for Mk to a vector v.

apply_mass_plus_eps_laplace_matrix(v, k, eps, dirichlet=True, operators=None): Apply (M_k + eps * L_k) to a k-form vector.

apply_projection_matrix(v, k_in, k_out, dirichlet_in=True, dirichlet_out=True, operators=None): Apply the sparse projection matrix Pk_in_k_out to a vector v.

apply_stiffness(v, k, dirichlet=True, operators=None): Apply the stiffness matrix S_k to a k-form vector v.

k=0: grad_grad k=1: curl_curl k=2: div_div k=3: 0 (no stiffness)

apply_strong_curl(v, dirichlet_in=True, dirichlet_out=True): Apply the strong curl M2⁻¹ D1 to a 1-form DOF vector v.

apply_strong_div(v, dirichlet_in=True, dirichlet_out=True): Apply the strong divergence M3⁻¹ D2 to a 2-form DOF vector v.

apply_strong_grad(v, dirichlet_in=True, dirichlet_out=True): Apply the strong gradient M1⁻¹ D0 to a 0-form DOF vector v.

apply_weak_curl(v, dirichlet_in=True, dirichlet_out=True, boundary_dual=None)

Apply the weak curl operator to a vector v.

This returns M1^{-1} (D1.T v + boundary_dual) where boundary_dual is an optional prescribed boundary functional in the dual 1-form space.

apply_weak_div(v, dirichlet_in=True, dirichlet_out=True, boundary_dual=None)

Apply the weak divergence operator to a vector v.

This returns M0^{-1} (-D0.T v + boundary_dual) where boundary_dual is an optional prescribed boundary functional in the dual 0-form space.

apply_weak_grad(v, dirichlet_in=True, dirichlet_out=True, boundary_dual=None)

Apply the weak gradient operator to a vector v.

This returns M2^{-1} (-D2.T v + boundary_dual) where boundary_dual is an optional prescribed boundary functional in the dual 2-form space.

assemble_all_dense()

Assemble sparse operators without preconditioners, then cache dense matrices.

This is a courtesy path for debugging, densification, and direct-solve workflows. Dense operators are stored under self.get_operators().dense.

assemble_all_sparse(include_preconditioners: bool = True)

Assemble and cache all sparse operator matrices.

Builds mass matrices, derivative matrices, stiffness matrices, and Hodge-Laplacian operators for all form degrees, storing the result in self.operators and mirroring legacy fields. When include_preconditioners is true, also assemble the eager preconditioner payloads used by solver-facing convenience methods. Returns the operator bundle.

assemble_derivative_matrix(k)

Assemble and cache the weak derivative matrix mapping k-forms to (k+1)-forms.

Parameters:: k (int) – Form degree of the input form (0, 1, or 2).

assemble_hodge_laplacian(k): Backward-compatible alias for assemble_laplacian.

assemble_incidence_matrix(k)

Assemble and cache the topological incidence matrix Gk.

Parameters:: k (int) – Form degree of the input form (0, 1, or 2).

assemble_laplacian(k)

Assemble and cache the Laplacian stiffness data for k-forms.

Parameters:: k (int) – Form degree (0, 1, 2, or 3).

assemble_leray_projection(): Assemble the auxiliary operators required by apply_leray_projection().

assemble_mass_matrix(k)

Assemble and cache the mass matrix for k-forms.

Parameters:: k (int) – Form degree (0, 1, 2, or 3).

assemble_projection_matrix(k_from, k_to)

Assemble and cache the L²-projection matrix from k_from-forms to k_to-forms.

Parameters:

k_from (int) – Source form degree.
k_to (int) – Target form degree.

assemble_reference_mass_matrix(): Assemble and cache the 0-form mass matrix on the reference domain.

basis_0: DifferentialForm

basis_1: DifferentialForm

basis_2: DifferentialForm

basis_3: DifferentialForm

basis_r_jk: Array

basis_t_jk: Array

basis_z_jk: Array

bc_lift(g: Array, k: int) → Array

Embed boundary DOF values into the full spline basis space.

Parameters:

g (array of shape (n_k_bc,)) – DOF values at the Dirichlet boundary nodes.
k (int) – Form degree (0, 1, 2, 3).

Returns:

Full spline vector with g placed at the BC positions, zeros everywhere else. Multiply any full-spline-space operator by this vector to compute the BC contribution.

Return type:

array of shape (basis_k.n,)

build_spline_map(coefficients, extraction=None): Build a spline map using the sequence’s scalar spline basis.

compute_nullspaces(): Compute and cache the harmonic forms for all form degrees (closed-form).

cross_product_load(w, u, n, m, k, dirichlet_n=True, dirichlet_m=True, dirichlet_k=True)

Project a cross product of two differential forms onto an n-form.

Computes the n-form dual DOF vector

v_i = ∫ Λⁿ_i · (w × u) dx

with appropriate metric contractions depending on the form degrees n, m, k. Uses the tensor-product structure for efficient evaluation and integration.

Parameters:

w (array) – DOF vector of the m-form.
u (array) – DOF vector of the k-form.
n (int) – Form degree of the output (1 or 2).
m (int) – Form degree of the first input (1 or 2).
k (int) – Form degree of the second input (1 or 2).
dirichlet_n (bool, optional) – Use Dirichlet-constrained extraction for the output n-form.
dirichlet_m (bool, optional) – Use Dirichlet-constrained extraction for the input m-form.
dirichlet_k (bool, optional) – Use Dirichlet-constrained extraction for the input k-form.

Returns:

n-form dual DOF vector (apply M_n⁻¹ to obtain primal DOFs).

Return type:

array

d_basis_r_jk: Array

d_basis_t_jk: Array

d_basis_z_jk: Array

property e0

property e0_T

property e0_bc

property e0_bc_T

property e0_dbc

property e0_dbc_T

property e1

property e1_T

property e1_bc

property e1_bc_T

property e1_dbc

property e1_dbc_T

property e2

property e2_T

property e2_bc

property e2_bc_T

property e2_dbc

property e2_dbc_T

property e3

property e3_T

property e3_bc

property e3_bc_T

property e3_dbc

property e3_dbc_T

eval_basis_0_ijk(i, j, k): Evaluate the (i, j, k)-th 0-form basis function at all quadrature points.

eval_basis_1_ijk(i, j, k): Evaluate the (i, j, k)-th 1-form basis function at all quadrature points.

eval_basis_2_ijk(i, j, k): Evaluate the (i, j, k)-th 2-form basis function at all quadrature points.

eval_basis_3_ijk(i, j, k): Evaluate the (i, j, k)-th 3-form basis function at all quadrature points.

eval_d_basis_0_ijk(i, j, k): Evaluate the gradient of the (i, j, k)-th 0-form basis at all quadrature points.

eval_d_basis_1_ijk(i, j, k): Evaluate the curl of the (i, j, k)-th 1-form basis at all quadrature points.

eval_d_basis_2_ijk(i, j, k): Evaluate the divergence of the (i, j, k)-th 2-form basis at all quadrature points.

evaluate_1d()

Precompute 1-D spline and derivative values at quadrature points.

Populates basis_{r,t,z}_jk and d_basis_{r,t,z}_jk on self. These arrays drive the sum-factorized assembly and evaluation routines, and are required by geometry_from_spline_map() when using the fast spline-geometry path.

geometry: SequenceGeometry

geometry_from_spline_map(coefficients, extraction=None)

Construct geometry data from spline map coefficients.

Uses the sum-factorized path when the extraction operator is the sequence’s own e0 (so we have a precomputed transpose and 1D basis evaluations); otherwise falls back to the generic SequenceGeometry.from_map.

get_operators(): Return the cached operator bundle, if one is attached.

init_nullspaces(betti_numbers=None)

Initialise zero-valued nullspace arrays on self.operators.

Shapes are derived from betti_numbers (or self.betti_numbers).

interpolate(f, k: int, dirichlet: bool = False)

Compute primal DOFs by Greville interpolation (k=0) or histopolation (k=1,2,3).

Parameters:

f (callable)
k (int Form degree (0, 1, 2, 3).)
dirichlet (bool Use Dirichlet-constrained DOFs.)

property jacobian_j

l2_norm(v, k, dirichlet=True): Return the L² norm of a k-form DOF vector v.

l2_norm_sq(v, k, dirichlet=True): Return the squared L² norm of a k-form DOF vector v.

load(f, k: int, dirichlet: bool = False, bc: bool = False, frame: str = 'phys')

Assemble the dual k-form load vector v_i = ∫ Λ^k_i · f(ξ) w(ξ) dξ.

Parameters:

f (callable)
k (int Form degree (0, 1, 2, 3).)
dirichlet (bool Use Dirichlet-constrained DOFs.)
bc (bool Use boundary-trace DOFs (takes precedence over dirichlet).)
frame ({‘phys’, ‘ref’} Passed to mrx.projectors.load().)

property map

property metric_inv_jkl

property metric_jkl

ns: tuple[int, int, int]

property null_0

property null_0_dbc

property null_1

property null_1_dbc

property null_2

property null_2_dbc

property null_3

property null_3_dbc

pressure_load(p, u, gamma, dirichlet_p=True, dirichlet_u=True)

Evaluate the pressure projection -(grad p · u + γ p div u).

Computes the 0-form dual DOF vector:

q_i = ∫ Λ⁰_i (−∇p · u − γ p div u) w dx

The 0-form mass matrix weight J cancels with the 1/J from the wedge product (1-form · 2-form) and from div = (1/J) div_logical, so the integrand has no metric or Jacobian — only quad weights.

Parameters:

p (array – 0-form DOFs)
u (array – 2-form DOFs)
gamma (float – adiabatic exponent)
dirichlet_p (bool – Dirichlet BCs on p)
dirichlet_u (bool – Dirichlet BCs on u)

Returns:

q_dual

Return type:

array – 0-form dual DOFs (apply M0⁻¹ to get primal DOFs)

ps: tuple[int, int, int]

quad: QuadratureRule

set_geometry(geometry: SequenceGeometry): Replace the geometry attached to this sequence.

set_geometry_terms(metric_jkl, metric_inv_jkl, jacobian_j): Replace the geometry tensors used by mapped assembly and operators.

set_map(map): Update the active logical-to-physical map and derived geometry terms.

set_operators(operators, sync_legacy=True)

Attach an operator bundle to the sequence and optionally mirror legacy fields.

If operators has no nullspace arrays yet, they are initialised to zeros with shapes derived from self.betti_numbers.

set_spline_map(coefficients, extraction=None): Update the sequence geometry from spline map coefficients.

update_diffusion_runtime_tuning(k, eps=0.0, dirichlet=True, operators=None, preconditioner='auto'): Estimate and store runtime tuning for a polynomial diffusion preconditioner.

update_mass_runtime_tuning(k, dirichlet=True, operators=None, preconditioner='auto'): Estimate and store runtime tuning for a polynomial mass preconditioner.

update_scalar_hodge_runtime_tuning(k, eps=0.0, dirichlet=True, operators=None, preconditioner='auto'): Backward-compatible alias for update_scalar_laplacian_runtime_tuning.

update_scalar_laplacian_runtime_tuning(k, eps=0.0, dirichlet=True, operators=None, preconditioner='auto'): Estimate and store runtime tuning for a scalar Laplacian preconditioner.

update_schur_runtime_tuning(k, eps=0.0, dirichlet=True, operators=None, preconditioner='auto'): Estimate and store runtime tuning for a polynomial Schur-outer preconditioner.

class mrx.DegreeIterativeRuntimeTuning(k0: 'BoundaryIterativeRuntimeTuning' = <factory>, k1: 'BoundaryIterativeRuntimeTuning' = <factory>, k2: 'BoundaryIterativeRuntimeTuning' = <factory>, k3: 'BoundaryIterativeRuntimeTuning' = <factory>)

Bases: Module

__init__(k0: ~mrx.operators.BoundaryIterativeRuntimeTuning = <factory>, k1: ~mrx.operators.BoundaryIterativeRuntimeTuning = <factory>, k2: ~mrx.operators.BoundaryIterativeRuntimeTuning = <factory>, k3: ~mrx.operators.BoundaryIterativeRuntimeTuning = <factory>) → None

k0: BoundaryIterativeRuntimeTuning

k1: BoundaryIterativeRuntimeTuning

k2: BoundaryIterativeRuntimeTuning

k3: BoundaryIterativeRuntimeTuning

class mrx.DegreeShiftedIterativeRuntimeTuning(k0: 'BoundaryShiftedIterativeRuntimeTuning' = <factory>, k1: 'BoundaryShiftedIterativeRuntimeTuning' = <factory>, k2: 'BoundaryShiftedIterativeRuntimeTuning' = <factory>, k3: 'BoundaryShiftedIterativeRuntimeTuning' = <factory>)

Bases: Module

__init__(k0: ~mrx.operators.BoundaryShiftedIterativeRuntimeTuning = <factory>, k1: ~mrx.operators.BoundaryShiftedIterativeRuntimeTuning = <factory>, k2: ~mrx.operators.BoundaryShiftedIterativeRuntimeTuning = <factory>, k3: ~mrx.operators.BoundaryShiftedIterativeRuntimeTuning = <factory>) → None

k0: BoundaryShiftedIterativeRuntimeTuning

k1: BoundaryShiftedIterativeRuntimeTuning

k2: BoundaryShiftedIterativeRuntimeTuning

k3: BoundaryShiftedIterativeRuntimeTuning

class mrx.DenseSequenceOperators(m0: BoundaryConditionPair = <factory>, m1: BoundaryConditionPair = <factory>, m2: BoundaryConditionPair = <factory>, m3: BoundaryConditionPair = <factory>, d0: BoundaryConditionPair = <factory>, d1: BoundaryConditionPair = <factory>, d2: BoundaryConditionPair = <factory>, s0: BoundaryConditionPair = <factory>, s1: BoundaryConditionPair = <factory>, s2: BoundaryConditionPair = <factory>, s3: BoundaryConditionPair = <factory>, l0: BoundaryConditionPair = <factory>, l1: BoundaryConditionPair = <factory>, l2: BoundaryConditionPair = <factory>, l3: BoundaryConditionPair = <factory>, p21: BoundaryConditionPair = <factory>, p12: BoundaryConditionPair = <factory>, p03: BoundaryConditionPair = <factory>, p30: BoundaryConditionPair = <factory>)

Bases: Module

Optional dense cache for extracted operator matrices.

__init__(m0: ~mrx.preconditioners.BoundaryConditionPair = <factory>, m1: ~mrx.preconditioners.BoundaryConditionPair = <factory>, m2: ~mrx.preconditioners.BoundaryConditionPair = <factory>, m3: ~mrx.preconditioners.BoundaryConditionPair = <factory>, d0: ~mrx.preconditioners.BoundaryConditionPair = <factory>, d1: ~mrx.preconditioners.BoundaryConditionPair = <factory>, d2: ~mrx.preconditioners.BoundaryConditionPair = <factory>, s0: ~mrx.preconditioners.BoundaryConditionPair = <factory>, s1: ~mrx.preconditioners.BoundaryConditionPair = <factory>, s2: ~mrx.preconditioners.BoundaryConditionPair = <factory>, s3: ~mrx.preconditioners.BoundaryConditionPair = <factory>, l0: ~mrx.preconditioners.BoundaryConditionPair = <factory>, l1: ~mrx.preconditioners.BoundaryConditionPair = <factory>, l2: ~mrx.preconditioners.BoundaryConditionPair = <factory>, l3: ~mrx.preconditioners.BoundaryConditionPair = <factory>, p21: ~mrx.preconditioners.BoundaryConditionPair = <factory>, p12: ~mrx.preconditioners.BoundaryConditionPair = <factory>, p03: ~mrx.preconditioners.BoundaryConditionPair = <factory>, p30: ~mrx.preconditioners.BoundaryConditionPair = <factory>) → None

d0: BoundaryConditionPair

d1: BoundaryConditionPair

d2: BoundaryConditionPair

l0: BoundaryConditionPair

l1: BoundaryConditionPair

l2: BoundaryConditionPair

l3: BoundaryConditionPair

m0: BoundaryConditionPair

m1: BoundaryConditionPair

m2: BoundaryConditionPair

m3: BoundaryConditionPair

p03: BoundaryConditionPair

p12: BoundaryConditionPair

p21: BoundaryConditionPair

p30: BoundaryConditionPair

s0: BoundaryConditionPair

s1: BoundaryConditionPair

s2: BoundaryConditionPair

s3: BoundaryConditionPair

class mrx.DerivativeSpline(s: SplineBasis)

Bases: object

A class representing the derivative of a spline basis.

This class implements the derivative of a spline basis, supporting various types of splines (clamped, periodic, constant). It computes the derivative by adjusting the degree and number of basis functions based on the original spline type.

n

Number of derivative spline basis functions

Type:: int

p

Degree of the derivative spline

Type:: int

type

Type of spline (‘clamped’, ‘periodic’, or ‘constant’)

Type:: str

T

Knot vector for the derivative spline

Type:: jnp.ndarray

s

The underlying spline basis used for derivative computation

Type:: SplineBasis

__call__(x: float, i: int) → Array: Alias for evaluate().

__getitem__(i: int) → Callable[[float], Array]: Return lambda x: self(x, i).

__init__(s: SplineBasis) → None

Initialize a derivative spline basis.

Parameters:: s – The original SplineBasis object to compute derivatives from

evaluate(x: float, i: int) → Array

Evaluate the derivative of the ith spline at point x.

Computes the derivative based on the spline type: - For clamped splines: Uses a forward difference formula with appropriate scaling - For periodic splines: Handles wrapping of indices for periodic continuity - For constant splines: Returns 1.0 (derivative of constant function)

Derivative splines cannot be evaluated at a clamped boundary.

Parameters:

x – The point at which to evaluate the derivative
i – The index of the spline derivative to evaluate

Returns:

The value of the derivative at x

greville_spans() → Array

Return the consecutive Greville intervals of the parent spline basis.

For clamped splines, the parent Greville points include the endpoints, so the consecutive intervals form the natural histopolation cells.

Returns:: Array of shape (n, 2) where each row is [a, b] defining an integration interval.

histopolation_matrix(spans: Array | None = None, quadrature_order: int | None = None) → Array

Assemble the Greville-span histopolation matrix for this basis.

Parameters:

spans – Integration intervals of shape (n, 2). If omitted, the Greville spans from greville_spans() are used.
quadrature_order – Number of Gauss-Legendre quadrature points per span. Defaults to max(2, p + 2).

Returns:

Array of shape (n, n) where entry [k, i] is the integral of the i-th derivative basis function over spans[k].

class mrx.DifferentialForm(k, ns, ps, types, Ts=None)

Bases: object

Discrete k-form on a 3-D tensor-product spline space.

k=0 — scalar; k=1 — 1-form (edge); k=2 — 2-form (face); k=3 — volume form; k=-1 — vector field (3 copies of the 0-form space). Note: k=-1 is incompatible with polar setups because the polar extraction operator reduces the 0-form DOF count asymmetrically across the three components.

__call__(x, i): Alias for evaluate().

__getitem__(i): Return lambda x: self(x, i).

__init__(k, ns, ps, types, Ts=None)

Args: k: Form degree (0, 1, 2, 3, or -1 for a vector field). ns: Number of DOFs in each direction. ps: Polynomial degrees in each direction. types: Boundary condition types ('clamped', 'periodic',

'constant') for each direction.

Ts: Knot vectors; None uses uniform knots.

_ravel_index(c, i, j, k): Return the global DOF index for component c and grid indices (i,j,k).

_unravel_index(idx): Return (component, i, j, k) for a global DOF index.

_vector_index(idx): Return (component, local_index) for a global DOF index.

d: int

evaluate(x, i): Evaluate basis function i at logical point x.

k: int

n: int

nr: int

ns: Array

nt: int

nz: int

pr: int

pt: int

pz: int

class mrx.DiscreteFunction(dof, Λ, E=None)

Bases: object

A discrete function as a linear combination of k-form basis functions.

__call__(x): Evaluate at logical point x.

__init__(dof, Λ, E=None): Args: dof: Coefficient vector (DOFs). Λ: Underlying DifferentialForm. E: Extraction matrix; defaults to the identity.

class mrx.IterativeRuntimeTuning(lambda_max: jnp.ndarray | None = None, lambda_min: jnp.ndarray | None = None)

Bases: Module

Dynamic spectral data for an iterative polynomial preconditioner.

__init__(lambda_max: Array | None = None, lambda_min: Array | None = None) → None

lambda_max: Array | None = None

lambda_min: Array | None = None

class mrx.K0TensorHodgePreconditionerFactors(core_size: 'int', bulk_shape: 'tuple[int, int, int]', schur_inv: 'jnp.ndarray', schur_projector: 'Optional[jnp.ndarray]' = None, bulk_alpha: 'Optional[jnp.ndarray]' = None, bulk_V_r: 'Optional[jnp.ndarray]' = None, bulk_V_t: 'Optional[jnp.ndarray]' = None, bulk_V_z: 'Optional[jnp.ndarray]' = None, bulk_lam_r: 'Optional[jnp.ndarray]' = None, bulk_lam_t: 'Optional[jnp.ndarray]' = None, bulk_lam_z: 'Optional[jnp.ndarray]' = None, bulk_mass_r: 'Optional[jnp.ndarray]' = None, bulk_mass_t: 'Optional[jnp.ndarray]' = None, bulk_mass_z: 'Optional[jnp.ndarray]' = None, bulk_stiff_r: 'Optional[jnp.ndarray]' = None, bulk_stiff_t: 'Optional[jnp.ndarray]' = None, bulk_stiff_z: 'Optional[jnp.ndarray]' = None, bulk_term_mass_r: 'tuple[jnp.ndarray, ...]' = (), bulk_term_mass_t: 'tuple[jnp.ndarray, ...]' = (), bulk_term_mass_z: 'tuple[jnp.ndarray, ...]' = (), bulk_term_stiff_r: 'tuple[jnp.ndarray, ...]' = (), bulk_term_stiff_t: 'tuple[jnp.ndarray, ...]' = (), bulk_term_stiff_z: 'tuple[jnp.ndarray, ...]' = (), bulk_term_op_r: 'tuple[jnp.ndarray, ...]' = (), bulk_term_op_t: 'tuple[jnp.ndarray, ...]' = (), bulk_term_op_z: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_basis_r: 'Optional[jnp.ndarray]' = None, bulk_modal_basis_t: 'Optional[jnp.ndarray]' = None, bulk_modal_basis_z: 'Optional[jnp.ndarray]' = None, bulk_modal_mass_r: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_mass_t: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_mass_z: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_stiff_r: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_stiff_t: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_stiff_z: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_op_r: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_op_t: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_op_z: 'tuple[jnp.ndarray, ...]' = (), bulk_modal_denom: 'Optional[jnp.ndarray]' = None, bulk_modal_inv_denom: 'Optional[jnp.ndarray]' = None, cp_relative_error: 'Optional[float]' = None, cp_final_delta: 'Optional[float]' = None, precompute_coupling: 'bool' = True, core_coupling: 'Optional[jnp.ndarray]' = None, bulk_radial_block_inv: 'Optional[jnp.ndarray]' = None)

Bases: Module

__init__(core_size: int, bulk_shape: tuple[int, int, int], schur_inv: Array, schur_projector: Array | None = None, bulk_alpha: Array | None = None, bulk_V_r: Array | None = None, bulk_V_t: Array | None = None, bulk_V_z: Array | None = None, bulk_lam_r: Array | None = None, bulk_lam_t: Array | None = None, bulk_lam_z: Array | None = None, bulk_mass_r: Array | None = None, bulk_mass_t: Array | None = None, bulk_mass_z: Array | None = None, bulk_stiff_r: Array | None = None, bulk_stiff_t: Array | None = None, bulk_stiff_z: Array | None = None, bulk_term_mass_r: tuple[Array, ...] = (), bulk_term_mass_t: tuple[Array, ...] = (), bulk_term_mass_z: tuple[Array, ...] = (), bulk_term_stiff_r: tuple[Array, ...] = (), bulk_term_stiff_t: tuple[Array, ...] = (), bulk_term_stiff_z: tuple[Array, ...] = (), bulk_term_op_r: tuple[Array, ...] = (), bulk_term_op_t: tuple[Array, ...] = (), bulk_term_op_z: tuple[Array, ...] = (), bulk_modal_basis_r: Array | None = None, bulk_modal_basis_t: Array | None = None, bulk_modal_basis_z: Array | None = None, bulk_modal_mass_r: tuple[Array, ...] = (), bulk_modal_mass_t: tuple[Array, ...] = (), bulk_modal_mass_z: tuple[Array, ...] = (), bulk_modal_stiff_r: tuple[Array, ...] = (), bulk_modal_stiff_t: tuple[Array, ...] = (), bulk_modal_stiff_z: tuple[Array, ...] = (), bulk_modal_op_r: tuple[Array, ...] = (), bulk_modal_op_t: tuple[Array, ...] = (), bulk_modal_op_z: tuple[Array, ...] = (), bulk_modal_denom: Array | None = None, bulk_modal_inv_denom: Array | None = None, cp_relative_error: float | None = None, cp_final_delta: float | None = None, precompute_coupling: bool = True, core_coupling: Array | None = None, bulk_radial_block_inv: Array | None = None) → None

bulk_V_r: Array | None = None

bulk_V_t: Array | None = None

bulk_V_z: Array | None = None

bulk_alpha: Array | None = None

bulk_lam_r: Array | None = None

bulk_lam_t: Array | None = None

bulk_lam_z: Array | None = None

bulk_mass_r: Array | None = None

bulk_mass_t: Array | None = None

bulk_mass_z: Array | None = None

bulk_modal_basis_r: Array | None = None

bulk_modal_basis_t: Array | None = None

bulk_modal_basis_z: Array | None = None

bulk_modal_denom: Array | None = None

bulk_modal_inv_denom: Array | None = None

bulk_modal_mass_r: tuple[Array, ...] = ()

bulk_modal_mass_t: tuple[Array, ...] = ()

bulk_modal_mass_z: tuple[Array, ...] = ()

bulk_modal_op_r: tuple[Array, ...] = ()

bulk_modal_op_t: tuple[Array, ...] = ()

bulk_modal_op_z: tuple[Array, ...] = ()

bulk_modal_stiff_r: tuple[Array, ...] = ()

bulk_modal_stiff_t: tuple[Array, ...] = ()

bulk_modal_stiff_z: tuple[Array, ...] = ()

bulk_radial_block_inv: Array | None = None

bulk_shape: tuple[int, int, int]

bulk_stiff_r: Array | None = None

bulk_stiff_t: Array | None = None

bulk_stiff_z: Array | None = None

bulk_term_mass_r: tuple[Array, ...] = ()

bulk_term_mass_t: tuple[Array, ...] = ()

bulk_term_mass_z: tuple[Array, ...] = ()

bulk_term_op_r: tuple[Array, ...] = ()

bulk_term_op_t: tuple[Array, ...] = ()

bulk_term_op_z: tuple[Array, ...] = ()

bulk_term_stiff_r: tuple[Array, ...] = ()

bulk_term_stiff_t: tuple[Array, ...] = ()

bulk_term_stiff_z: tuple[Array, ...] = ()

core_coupling: Array | None = None

core_size: int

cp_final_delta: float | None = None

cp_relative_error: float | None = None

precompute_coupling: bool = True

schur_inv: Array

schur_projector: Array | None = None

class mrx.K1MassSurgeryPreconditionerFactors(surgery_indices: 'jnp.ndarray', bulk_indices: 'jnp.ndarray', r_indices: 'jnp.ndarray', theta_bulk_indices: 'jnp.ndarray', zeta_bulk_indices: 'jnp.ndarray', rt_indices: 'jnp.ndarray', surgery_size: 'int', rt_r_size: 'int', rt_theta_size: 'int', bulk_rt_size: 'int', bulk_zeta_size: 'int', apply_data: 'ExtractedMassApplyData', surgery_diaginv: 'jnp.ndarray', ass: 'jnp.ndarray', surgery_to_bulk_data: 'Optional[RestrictedExtractedMassApplyData]' = None, bulk_to_surgery_data: 'Optional[RestrictedExtractedMassApplyData]' = None, rt_atr_data: 'Optional[RestrictedExtractedMassApplyData]' = None, rt_art_data: 'Optional[RestrictedExtractedMassApplyData]' = None, rt_to_zeta_data: 'Optional[RestrictedExtractedMassApplyData]' = None, zeta_to_rt_data: 'Optional[RestrictedExtractedMassApplyData]' = None, coupling_sb: 'Optional[jnp.ndarray]' = None)

Bases: Module

__init__(surgery_indices: Array, bulk_indices: Array, r_indices: Array, theta_bulk_indices: Array, zeta_bulk_indices: Array, rt_indices: Array, surgery_size: int, rt_r_size: int, rt_theta_size: int, bulk_rt_size: int, bulk_zeta_size: int, apply_data: ExtractedMassApplyData, surgery_diaginv: Array, ass: Array, surgery_to_bulk_data: RestrictedExtractedMassApplyData | None = None, bulk_to_surgery_data: RestrictedExtractedMassApplyData | None = None, rt_atr_data: RestrictedExtractedMassApplyData | None = None, rt_art_data: RestrictedExtractedMassApplyData | None = None, rt_to_zeta_data: RestrictedExtractedMassApplyData | None = None, zeta_to_rt_data: RestrictedExtractedMassApplyData | None = None, coupling_sb: Array | None = None) → None

apply_data: ExtractedMassApplyData

ass: Array

bulk_indices: Array

bulk_rt_size: int

bulk_to_surgery_data: RestrictedExtractedMassApplyData | None = None

bulk_zeta_size: int

coupling_sb: Array | None = None

r_indices: Array

rt_art_data: RestrictedExtractedMassApplyData | None = None

rt_atr_data: RestrictedExtractedMassApplyData | None = None

rt_indices: Array

rt_r_size: int

rt_theta_size: int

rt_to_zeta_data: RestrictedExtractedMassApplyData | None = None

surgery_diaginv: Array

surgery_indices: Array

surgery_size: int

surgery_to_bulk_data: RestrictedExtractedMassApplyData | None = None

theta_bulk_indices: Array

zeta_bulk_indices: Array

zeta_to_rt_data: RestrictedExtractedMassApplyData | None = None

class mrx.K1TensorCurlCurlForwardModel(r_shape: 'tuple[int, int, int]', theta_shape: 'tuple[int, int, int]', zeta_shape: 'tuple[int, int, int]', curl_r_shape: 'tuple[int, int, int]', curl_theta_shape: 'tuple[int, int, int]', curl_zeta_shape: 'tuple[int, int, int]', rank: 'int', g_r: 'jnp.ndarray', g_t: 'jnp.ndarray', g_z: 'jnp.ndarray', rr_mass_r_terms: 'tuple[jnp.ndarray, ...]' = (), rr_mass_t_terms: 'tuple[jnp.ndarray, ...]' = (), rr_mass_z_terms: 'tuple[jnp.ndarray, ...]' = (), tt_mass_r_terms: 'tuple[jnp.ndarray, ...]' = (), tt_mass_t_terms: 'tuple[jnp.ndarray, ...]' = (), tt_mass_z_terms: 'tuple[jnp.ndarray, ...]' = (), zz_mass_r_terms: 'tuple[jnp.ndarray, ...]' = (), zz_mass_t_terms: 'tuple[jnp.ndarray, ...]' = (), zz_mass_z_terms: 'tuple[jnp.ndarray, ...]' = (), cp_relative_error: 'Optional[float]' = None, cp_final_delta: 'Optional[float]' = None)

Bases: Module

__init__(r_shape: tuple[int, int, int], theta_shape: tuple[int, int, int], zeta_shape: tuple[int, int, int], curl_r_shape: tuple[int, int, int], curl_theta_shape: tuple[int, int, int], curl_zeta_shape: tuple[int, int, int], rank: int, g_r: Array, g_t: Array, g_z: Array, rr_mass_r_terms: tuple[Array, ...] = (), rr_mass_t_terms: tuple[Array, ...] = (), rr_mass_z_terms: tuple[Array, ...] = (), tt_mass_r_terms: tuple[Array, ...] = (), tt_mass_t_terms: tuple[Array, ...] = (), tt_mass_z_terms: tuple[Array, ...] = (), zz_mass_r_terms: tuple[Array, ...] = (), zz_mass_t_terms: tuple[Array, ...] = (), zz_mass_z_terms: tuple[Array, ...] = (), cp_relative_error: float | None = None, cp_final_delta: float | None = None) → None

cp_final_delta: float | None = None

cp_relative_error: float | None = None

curl_r_shape: tuple[int, int, int]

curl_theta_shape: tuple[int, int, int]

curl_zeta_shape: tuple[int, int, int]

g_r: Array

g_t: Array

g_z: Array

r_shape: tuple[int, int, int]

rank: int

rr_mass_r_terms: tuple[Array, ...] = ()

rr_mass_t_terms: tuple[Array, ...] = ()

rr_mass_z_terms: tuple[Array, ...] = ()

theta_shape: tuple[int, int, int]

tt_mass_r_terms: tuple[Array, ...] = ()

tt_mass_t_terms: tuple[Array, ...] = ()

tt_mass_z_terms: tuple[Array, ...] = ()

zeta_shape: tuple[int, int, int]

zz_mass_r_terms: tuple[Array, ...] = ()

zz_mass_t_terms: tuple[Array, ...] = ()

zz_mass_z_terms: tuple[Array, ...] = ()

class mrx.K1TensorMassPreconditionerFactors(r_indices: 'jnp.ndarray', theta_bulk_indices: 'jnp.ndarray', zeta_bulk_indices: 'jnp.ndarray', rt_r_size: 'int', rt_theta_size: 'int', arr: 'TensorDiagonalBlockInverseFactors', theta: 'TensorDiagonalBlockInverseFactors', zeta: 'TensorDiagonalBlockInverseFactors', use_inner_schur: 'bool' = False, schur_inv: 'Optional[jnp.ndarray]' = None)

Bases: Module

__init__(r_indices: Array, theta_bulk_indices: Array, zeta_bulk_indices: Array, rt_r_size: int, rt_theta_size: int, arr: TensorDiagonalBlockInverseFactors, theta: TensorDiagonalBlockInverseFactors, zeta: TensorDiagonalBlockInverseFactors, use_inner_schur: bool = False, schur_inv: Array | None = None) → None

arr: TensorDiagonalBlockInverseFactors

r_indices: Array

rt_r_size: int

rt_theta_size: int

schur_inv: Array | None = None

theta: TensorDiagonalBlockInverseFactors

theta_bulk_indices: Array

use_inner_schur: bool = False

zeta: TensorDiagonalBlockInverseFactors

zeta_bulk_indices: Array

class mrx.K1TensorStiffnessPreconditioner(surgery: 'K1MassSurgeryPreconditionerFactors', factors: 'K1TensorMassPreconditionerFactors')

Bases: Module

__init__(surgery: K1MassSurgeryPreconditionerFactors, factors: K1TensorMassPreconditionerFactors) → None

factors: K1TensorMassPreconditionerFactors

surgery: K1MassSurgeryPreconditionerFactors

class mrx.K2MassSurgeryPreconditionerFactors(surgery_indices: 'jnp.ndarray', bulk_indices: 'jnp.ndarray', r_bulk_indices: 'jnp.ndarray', theta_indices: 'jnp.ndarray', zeta_indices: 'jnp.ndarray', surgery_size: 'int', r_bulk_size: 'int', theta_size: 'int', zeta_size: 'int', apply_data: 'ExtractedMassApplyData', surgery_diaginv: 'jnp.ndarray', ass: 'jnp.ndarray', surgery_to_bulk_data: 'Optional[RestrictedExtractedMassApplyData]' = None, bulk_to_surgery_data: 'Optional[RestrictedExtractedMassApplyData]' = None, r_to_theta_data: 'Optional[RestrictedExtractedMassApplyData]' = None, theta_to_r_data: 'Optional[RestrictedExtractedMassApplyData]' = None, rt_to_zeta_data: 'Optional[RestrictedExtractedMassApplyData]' = None, zeta_to_rt_data: 'Optional[RestrictedExtractedMassApplyData]' = None, coupling_sb: 'Optional[jnp.ndarray]' = None)

Bases: Module

__init__(surgery_indices: Array, bulk_indices: Array, r_bulk_indices: Array, theta_indices: Array, zeta_indices: Array, surgery_size: int, r_bulk_size: int, theta_size: int, zeta_size: int, apply_data: ExtractedMassApplyData, surgery_diaginv: Array, ass: Array, surgery_to_bulk_data: RestrictedExtractedMassApplyData | None = None, bulk_to_surgery_data: RestrictedExtractedMassApplyData | None = None, r_to_theta_data: RestrictedExtractedMassApplyData | None = None, theta_to_r_data: RestrictedExtractedMassApplyData | None = None, rt_to_zeta_data: RestrictedExtractedMassApplyData | None = None, zeta_to_rt_data: RestrictedExtractedMassApplyData | None = None, coupling_sb: Array | None = None) → None

apply_data: ExtractedMassApplyData

ass: Array

bulk_indices: Array

bulk_to_surgery_data: RestrictedExtractedMassApplyData | None = None

coupling_sb: Array | None = None

r_bulk_indices: Array

r_bulk_size: int

r_to_theta_data: RestrictedExtractedMassApplyData | None = None

rt_to_zeta_data: RestrictedExtractedMassApplyData | None = None

surgery_diaginv: Array

surgery_indices: Array

surgery_size: int

surgery_to_bulk_data: RestrictedExtractedMassApplyData | None = None

theta_indices: Array

theta_size: int

theta_to_r_data: RestrictedExtractedMassApplyData | None = None

zeta_indices: Array

zeta_size: int

zeta_to_rt_data: RestrictedExtractedMassApplyData | None = None

class mrx.K2TensorDivDivForwardModel(r_shape: 'tuple[int, int, int]', theta_shape: 'tuple[int, int, int]', zeta_shape: 'tuple[int, int, int]', scalar_shape: 'tuple[int, int, int]', rank: 'int', g_r: 'jnp.ndarray', g_t: 'jnp.ndarray', g_z: 'jnp.ndarray', mass_r_terms: 'tuple[jnp.ndarray, ...]' = (), mass_t_terms: 'tuple[jnp.ndarray, ...]' = (), mass_z_terms: 'tuple[jnp.ndarray, ...]' = (), component_mass_r_terms: 'tuple[jnp.ndarray, ...]' = (), component_mass_t_terms: 'tuple[jnp.ndarray, ...]' = (), component_mass_z_terms: 'tuple[jnp.ndarray, ...]' = (), cp_relative_error: 'Optional[float]' = None, cp_final_delta: 'Optional[float]' = None)

Bases: Module

__init__(r_shape: tuple[int, int, int], theta_shape: tuple[int, int, int], zeta_shape: tuple[int, int, int], scalar_shape: tuple[int, int, int], rank: int, g_r: Array, g_t: Array, g_z: Array, mass_r_terms: tuple[Array, ...] = (), mass_t_terms: tuple[Array, ...] = (), mass_z_terms: tuple[Array, ...] = (), component_mass_r_terms: tuple[Array, ...] = (), component_mass_t_terms: tuple[Array, ...] = (), component_mass_z_terms: tuple[Array, ...] = (), cp_relative_error: float | None = None, cp_final_delta: float | None = None) → None

component_mass_r_terms: tuple[Array, ...] = ()

component_mass_t_terms: tuple[Array, ...] = ()

component_mass_z_terms: tuple[Array, ...] = ()

cp_final_delta: float | None = None

cp_relative_error: float | None = None

g_r: Array

g_t: Array

g_z: Array

mass_r_terms: tuple[Array, ...] = ()

mass_t_terms: tuple[Array, ...] = ()

mass_z_terms: tuple[Array, ...] = ()

r_shape: tuple[int, int, int]

rank: int

scalar_shape: tuple[int, int, int]

theta_shape: tuple[int, int, int]

zeta_shape: tuple[int, int, int]

class mrx.K2TensorMassPreconditionerFactors(r_bulk_indices: 'jnp.ndarray', theta_indices: 'jnp.ndarray', zeta_indices: 'jnp.ndarray', r_bulk_size: 'int', theta_size: 'int', zeta_size: 'int', r_bulk: 'TensorDiagonalBlockInverseFactors', theta: 'TensorDiagonalBlockInverseFactors', zeta: 'TensorDiagonalBlockInverseFactors', use_inner_schur: 'bool' = False, schur_inv: 'Optional[jnp.ndarray]' = None)

Bases: Module

__init__(r_bulk_indices: Array, theta_indices: Array, zeta_indices: Array, r_bulk_size: int, theta_size: int, zeta_size: int, r_bulk: TensorDiagonalBlockInverseFactors, theta: TensorDiagonalBlockInverseFactors, zeta: TensorDiagonalBlockInverseFactors, use_inner_schur: bool = False, schur_inv: Array | None = None) → None

r_bulk: TensorDiagonalBlockInverseFactors

r_bulk_indices: Array

r_bulk_size: int

schur_inv: Array | None = None

theta: TensorDiagonalBlockInverseFactors

theta_indices: Array

theta_size: int

use_inner_schur: bool = False

zeta: TensorDiagonalBlockInverseFactors

zeta_indices: Array

zeta_size: int

class mrx.K2TensorStiffnessPreconditioner(surgery: 'K2MassSurgeryPreconditionerFactors', factors: 'K2TensorMassPreconditionerFactors')

Bases: Module

__init__(surgery: K2MassSurgeryPreconditionerFactors, factors: K2TensorMassPreconditionerFactors) → None

factors: K2TensorMassPreconditionerFactors

surgery: K2MassSurgeryPreconditionerFactors

class mrx.MassPreconditionerSpec(kind: 'str' = 'tensor', surgery_schur: 'bool' = False, steps: 'int' = 4, power_iterations: 'int' = 30, damping_safety: 'float' = 0.8, min_eig_fraction: 'float' = 0.001, lanczos_iterations: 'int' = 16, lanczos_max_eig_inflation: 'float' = 1.1, lanczos_min_eig_deflation: 'float' = 0.85, lanczos_min_eig_floor_fraction: 'float' = 0.001, schur_diag_mode: 'str' = 'tensor_probe', smoother: 'Optional[MassPreconditionerSpec]' = None)

Bases: object

damping_safety: float = 0.8

kind: str = 'tensor'

lanczos_iterations: int = 16

lanczos_max_eig_inflation: float = 1.1

lanczos_min_eig_deflation: float = 0.85

lanczos_min_eig_floor_fraction: float = 0.001

min_eig_fraction: float = 0.001

power_iterations: int = 30

schur_diag_mode: str = 'tensor_probe'

smoother: MassPreconditionerSpec | None = None

steps: int = 4

surgery_schur: bool = False

class mrx.MassPreconditioners(jacobi: 'Optional[JacobiMassPreconditioner]' = None, surgery: 'Optional[MassSurgeryPreconditioner]' = None, tensor: 'Optional[TensorMassPreconditioner]' = None)

Bases: Module

__init__(jacobi: JacobiMassPreconditioner | None = None, surgery: MassSurgeryPreconditioner | None = None, tensor: TensorMassPreconditioner | None = None) → None

jacobi: JacobiMassPreconditioner | None = None

surgery: MassSurgeryPreconditioner | None = None

tensor: TensorMassPreconditioner | None = None

class mrx.MatrixFreeExtraction(rows: Array, cols: Array, vals: Array, forward_shape: tuple, transposed: bool)

Bases: Module

Matrix-free polar/boundary extraction operator.

Applies E (forward) and E^T (transpose) as a cached gather/scatter using a static sparsity pattern instead of a stored BCSR matmul. The forward operator maps a full pre-extraction DoF vector (size forward_shape[1]) to the extracted/constrained vector (size forward_shape[0]); the transpose maps back.

The index pattern (rows, cols) and weights (vals) are computed once from the assembled sparse operator. The same pattern is reused by the surgery preconditioner through to_bcoo(), so no BCSR needs to be materialised or stored for the matvec path.

rows/cols/vals are always stored in the forward orientation; the transposed flag selects how they are consumed.

property T

__init__(rows: Array, cols: Array, vals: Array, forward_shape: tuple, transposed: bool) → None

cols: Array

property data: Nonzero values in the current orientation (BCOO-compatible).

property dtype

forward_shape: tuple

classmethod from_bcoo(bcoo, transposed: bool = False): Build a matrix-free extraction from an assembled BCOO matrix.

property indices: (nnz, 2) COO indices in the current orientation.

restrict_cols(col_indices)

Return a copy with the column dimension restricted to col_indices.

Works in the current orientation (respects transposed). The result keeps only nonzeros whose column (in current orientation) falls in col_indices, with columns remapped to a contiguous 0-based range. Returns a new MatrixFreeExtraction — no BCOO materialised.

restrict_rows(row_indices)

Return a copy with the row dimension restricted to row_indices.

Works in the current orientation (respects transposed). The result keeps only nonzeros whose row (in current orientation) falls in row_indices, with rows remapped to a contiguous 0-based range. Returns a new MatrixFreeExtraction — no BCOO materialised.

rows: Array

property shape

to_bcoo(): Materialise the (orientation-aware) sparse pattern as a BCOO.

todense()

transposed: bool

vals: Array

mrx.NamedTuple(typename, fields=None, /, **kwargs)

Typed version of namedtuple.

Usage:

class Employee(NamedTuple):
    name: str
    id: int

This is equivalent to:

Employee = collections.namedtuple('Employee', ['name', 'id'])

The resulting class has an extra __annotations__ attribute, giving a dict that maps field names to types. (The field names are also in the _fields attribute, which is part of the namedtuple API.) An alternative equivalent functional syntax is also accepted:

Employee = NamedTuple('Employee', [('name', str), ('id', int)])

class mrx.PolarExtractionOperator(Lambda, xi, zero_bc)

Bases: object

A class for extracting boundary conditions and handling polar mappings.

This class implements operators for handling boundary conditions and polar coordinate transformations.

k

Degree of the differential form

Type:: int

Λ

xi: Polar mapping coefficients

nr

Number of points in r-direction

Type:: int

nt

Number of points in θ-direction

Type:: int

nz

Number of points in ζ-direction

Type:: int

dr

Number of points in r-direction after boundary conditions

Type:: int

dt

Number of points in θ-direction after boundary conditions

Type:: int

dz

Number of points in ζ-direction after boundary conditions

Type:: int

o

Offset for boundary conditions (1 for zero BC, 0 otherwise)

Type:: int

n1

Size of first component

Type:: int

n2

Size of second component

Type:: int

n3

Size of third component

Type:: int

n

Total size of the operator

Type:: int

__init__(Lambda, xi, zero_bc)

Initialize the extraction operator.

Parameters:

Λ – Domain operator
ξ – Polar mapping coefficients
zero_bc (bool) – Whether to apply zero boundary conditions

_element(row_idx, col_idx)

Compute the operator element at specified indices.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index

Returns:

The operator element value

Return type:

jnp.ndarray

_inner_zeroform(row_idx, col_idx, nr, nt, nz)

Compute inner zero-form basis function.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index
nr (int) – Number of points in r-direction
nt (int) – Number of points in θ-direction
nz (int) – Number of points in ζ-direction

Returns:

The basis function value

Return type:

jnp.ndarray

_outer_zeroform(row_idx, col_idx, nr, nt, nz)

Compute outer zero-form basis function.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index
nr (int) – Number of points in r-direction
nt (int) – Number of points in θ-direction
nz (int) – Number of points in ζ-direction

Returns:

The basis function value

Return type:

jnp.ndarray

_threeform(row_idx, col_idx, nr, nt, nz)

Compute three-form basis function.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index
nr (int) – Number of points in r-direction
nt (int) – Number of points in θ-direction
nz (int) – Number of points in ζ-direction

Returns:

The basis function value

Return type:

jnp.ndarray

build_extraction(): Build the MatrixFreeExtraction from the explicit tensor-product sparsity pattern.

inner_oneform_r(row_idx, col_idx, nr, nt, nz)

Compute inner one-form basis function in r-direction.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index
nr (int) – Number of points in r-direction
nt (int) – Number of points in θ-direction
nz (int) – Number of points in ζ-direction

Returns:

The basis function value

Return type:

jnp.ndarray

inner_oneform_θ(row_idx, col_idx, nr, nt, nz)

Compute inner one-form basis function in θ-direction.

Parameters:

row_idx (int) – Row index
col_idx (int) – Column index
nr (int) – Number of points in r-direction
nt (int) – Number of points in θ-direction
nz (int) – Number of points in ζ-direction

Returns:

The basis function value

Return type:

jnp.ndarray

class mrx.Pullback(f, F, k)

Bases: object

Pullback of a k-form under the logical-to-physical map F.

Let J = det(DF). Transformation rules (ω evaluated at F(x)):

k= 0 F* ω = ω∘F k= 1 F* ω = DFᵀ · (ω∘F) k= 2 F* ω = J · DF⁻¹ · (ω∘F) (Piola) k= 3 F* ω = J · (ω∘F) k=−1 F* v = DF⁻¹ · (v∘F) (vector field)

__call__(x): Evaluate the pulled-back form at logical point x.

__init__(f, F, k): Args: f: The form to pull back. F: Logical-to-physical map. k: Form degree.

class mrx.Pushforward(f, F, k)

Bases: object

Pushforward of a k-form under the logical-to-physical map F.

Let J = det(DF). Transformation rules (ω evaluated at x):

k= 0 F_* ω = ω k= 1 F_* ω = (DFᵀ)⁻¹ · ω k= 2 F_* ω = DF · ω / J (Piola) k= 3 F_* ω = ω / J k=−1 F_* v = DF · v (vector field)

__call__(x): Evaluate the pushed-forward form at logical point x.

__init__(f, F, k): Args: f: The form to push forward. F: Logical-to-physical map. k: Form degree.

class mrx.QuadratureRule(form, p)

Bases: object

A class for handling quadrature rules in finite element analysis.

This class implements various quadrature rules for numerical integration in three-dimensional space. It supports different types of basis functions and provides efficient computation of quadrature points and weights.

x_x

Quadrature points in x-direction

Type:: array

x_y

Quadrature points in y-direction

Type:: array

x_z

Quadrature points in z-direction

Type:: array

w_x

Quadrature weights in x-direction

Type:: array

w_y

Quadrature weights in y-direction

Type:: array

w_z

Quadrature weights in z-direction

Type:: array

x

Combined quadrature points in 3D space

Type:: array

w

Combined quadrature weights

Type:: array

__init__(form, p)

Initialize the quadrature rule.

Parameters:

form – The differential form defining the basis functions
p (int) – Number of quadrature points per direction

class mrx.SaddlePointPreconditionerSpec(mass: 'MassPreconditionerSpec' = <factory>, schur: 'SchurPreconditionerSpec' = <factory>, coupled: 'bool' = False)

Bases: object

coupled: bool = False

mass: MassPreconditionerSpec

schur: SchurPreconditionerSpec

class mrx.SchurPreconditionerSpec(inner: 'MassPreconditionerSpec' = <factory>, outer: 'MassPreconditionerSpec' = <factory>)

Bases: object

inner: MassPreconditionerSpec

outer: MassPreconditionerSpec

class mrx.SequenceGeometry(map: Any, metric_jkl: jnp.ndarray = None, metric_inv_jkl: jnp.ndarray = None, jacobian_j: jnp.ndarray = None)

Bases: Module

Geometry data attached to a de Rham sequence.

An eqx.Module so that the three quadrature-grid arrays (metric_jkl, metric_inv_jkl, jacobian_j) are dynamic pytree leaves and can flow through jit / grad. map is kept as a normal field so that if it is itself a pytree (e.g. a SplineMap), its coefficient leaves are tracked; plain Callable maps are treated as opaque leaves.

__init__(map: Any, metric_jkl: Array = None, metric_inv_jkl: Array = None, jacobian_j: Array = None) → None

classmethod from_map(map: Callable, quad_x: Array) → SequenceGeometry

Build geometry by evaluating a map on the quadrature grid.

Parameters:

map – Differentiable logical-to-physical map F: R^3 -> R^3.
quad_x – Quadrature points, shape (N_q, 3).

Returns:

A fully populated SequenceGeometry.

classmethod from_spline_map(spline_map, seq) → SequenceGeometry

Sum-factorized geometry builder for tensor-product spline maps.

Requires that seq.evaluate_1d() has already been called (so seq.basis_{r,t,z}_jk / seq.d_basis_{r,t,z}_jk are populated) and that spline_map.extraction_T is set.

Parameters:

spline_map – A SplineMap with coefficients and extraction_T populated.
seq – A DeRhamSequence with evaluate_1d() already called.

Returns:

A fully populated SequenceGeometry.

jacobian_j: Array = None

map: Any

metric_inv_jkl: Array = None

metric_jkl: Array = None

class mrx.SequenceOperators(m0: Optional[jsparse.BCSR] = None, m1: Optional[jsparse.BCSR] = None, m2: Optional[jsparse.BCSR] = None, m3: Optional[jsparse.BCSR] = None, k0_tensor_hodge_precond: Optional[BoundaryConditionPair] = None, k1_tensor_stiff_model: Optional[K1TensorCurlCurlForwardModel] = None, k2_tensor_stiff_model: Optional[K2TensorDivDivForwardModel] = None, k1_tensor_stiff_precond: Optional[BoundaryConditionPair] = None, k2_tensor_stiff_precond: Optional[BoundaryConditionPair] = None, e0: Optional[MatrixFreeExtraction] = None, e0_T: Optional[MatrixFreeExtraction] = None, e0_dbc: Optional[MatrixFreeExtraction] = None, e0_dbc_T: Optional[MatrixFreeExtraction] = None, e0_bc: Optional[MatrixFreeExtraction] = None, e0_bc_T: Optional[MatrixFreeExtraction] = None, e1: Optional[MatrixFreeExtraction] = None, e1_T: Optional[MatrixFreeExtraction] = None, e1_dbc: Optional[MatrixFreeExtraction] = None, e1_dbc_T: Optional[MatrixFreeExtraction] = None, e1_bc: Optional[MatrixFreeExtraction] = None, e1_bc_T: Optional[MatrixFreeExtraction] = None, e2: Optional[MatrixFreeExtraction] = None, e2_T: Optional[MatrixFreeExtraction] = None, e2_dbc: Optional[MatrixFreeExtraction] = None, e2_dbc_T: Optional[MatrixFreeExtraction] = None, e2_bc: Optional[MatrixFreeExtraction] = None, e2_bc_T: Optional[MatrixFreeExtraction] = None, e3: Optional[MatrixFreeExtraction] = None, e3_T: Optional[MatrixFreeExtraction] = None, e3_dbc: Optional[MatrixFreeExtraction] = None, e3_dbc_T: Optional[MatrixFreeExtraction] = None, e3_bc: Optional[MatrixFreeExtraction] = None, e3_bc_T: Optional[MatrixFreeExtraction] = None, mass_preconds: Optional[MassPreconditioners] = None, runtime_tuning: SequenceRuntimeTuning = <factory>, d0: Optional[jsparse.BCSR] = None, d0_T: Optional[jsparse.BCSR] = None, d1: Optional[jsparse.BCSR] = None, d1_T: Optional[jsparse.BCSR] = None, d2: Optional[jsparse.BCSR] = None, d2_T: Optional[jsparse.BCSR] = None, g0: Optional[_MatrixFreeIncidence] = None, g0_T: Optional[_MatrixFreeIncidence] = None, g1: Optional[_MatrixFreeIncidence] = None, g1_T: Optional[_MatrixFreeIncidence] = None, g2: Optional[_MatrixFreeIncidence] = None, g2_T: Optional[_MatrixFreeIncidence] = None, inc_gram_inv_1: Optional[jsparse.BCSR] = None, inc_gram_inv_1_dbc: Optional[jsparse.BCSR] = None, inc_gram_inv_2: Optional[jsparse.BCSR] = None, inc_gram_inv_2_dbc: Optional[jsparse.BCSR] = None, inc_gram_inv_3: Optional[jsparse.BCSR] = None, inc_gram_inv_3_dbc: Optional[jsparse.BCSR] = None, g0_grad_00: Optional[jsparse.BCSR] = None, g0_grad_00_T: Optional[jsparse.BCSR] = None, g0_grad_01: Optional[jsparse.BCSR] = None, g0_grad_01_T: Optional[jsparse.BCSR] = None, g0_grad_10: Optional[jsparse.BCSR] = None, g0_grad_10_T: Optional[jsparse.BCSR] = None, g0_grad_11: Optional[jsparse.BCSR] = None, g0_grad_11_T: Optional[jsparse.BCSR] = None, g1_curl_00: Optional[jsparse.BCSR] = None, g1_curl_00_T: Optional[jsparse.BCSR] = None, g1_curl_01: Optional[jsparse.BCSR] = None, g1_curl_01_T: Optional[jsparse.BCSR] = None, g1_curl_10: Optional[jsparse.BCSR] = None, g1_curl_10_T: Optional[jsparse.BCSR] = None, g1_curl_11: Optional[jsparse.BCSR] = None, g1_curl_11_T: Optional[jsparse.BCSR] = None, grad_grad: Optional[jsparse.BCSR] = None, curl_curl: Optional[jsparse.BCSR] = None, div_div: Optional[jsparse.BCSR] = None, dd0_diaginv: Optional[object] = None, dd1_diaginv: Optional[object] = None, dd2_diaginv: Optional[object] = None, dd3_diaginv: Optional[object] = None, dd0_diaginv_dbc: Optional[object] = None, dd1_diaginv_dbc: Optional[object] = None, dd2_diaginv_dbc: Optional[object] = None, dd3_diaginv_dbc: Optional[object] = None, p21: Optional[jsparse.BCSR] = None, p12: Optional[jsparse.BCSR] = None, p03: Optional[jsparse.BCSR] = None, p30: Optional[jsparse.BCSR] = None, schur_diaginv_k1: Optional[jnp.ndarray] = None, schur_diaginv_k1_dbc: Optional[jnp.ndarray] = None, schur_diaginv_k2: Optional[jnp.ndarray] = None, schur_diaginv_k2_dbc: Optional[jnp.ndarray] = None, schur_diaginv_k3: Optional[jnp.ndarray] = None, schur_diaginv_k3_dbc: Optional[jnp.ndarray] = None, schur_diaginv_mode_k1: Optional[str] = None, schur_diaginv_mode_k1_dbc: Optional[str] = None, schur_diaginv_mode_k2: Optional[str] = None, schur_diaginv_mode_k2_dbc: Optional[str] = None, schur_diaginv_mode_k3: Optional[str] = None, schur_diaginv_mode_k3_dbc: Optional[str] = None, null_0: Optional[jnp.ndarray] = None, null_1: Optional[jnp.ndarray] = None, null_2: Optional[jnp.ndarray] = None, null_3: Optional[jnp.ndarray] = None, null_0_dbc: Optional[jnp.ndarray] = None, null_1_dbc: Optional[jnp.ndarray] = None, null_2_dbc: Optional[jnp.ndarray] = None, null_3_dbc: Optional[jnp.ndarray] = None, dense: Optional[DenseSequenceOperators] = None)

Bases: Module

Dynamic operator bundle for a de Rham sequence.

Stores geometry-dependent operator data explicitly so it can be carried through JAX transforms while the sequence object remains a static topology shell.

__init__(m0: ~jax.experimental.sparse.bcsr.BCSR | None = None, m1: ~jax.experimental.sparse.bcsr.BCSR | None = None, m2: ~jax.experimental.sparse.bcsr.BCSR | None = None, m3: ~jax.experimental.sparse.bcsr.BCSR | None = None, k0_tensor_hodge_precond: ~mrx.preconditioners.BoundaryConditionPair | None = None, k1_tensor_stiff_model: ~mrx.operators.K1TensorCurlCurlForwardModel | None = None, k2_tensor_stiff_model: ~mrx.operators.K2TensorDivDivForwardModel | None = None, k1_tensor_stiff_precond: ~mrx.preconditioners.BoundaryConditionPair | None = None, k2_tensor_stiff_precond: ~mrx.preconditioners.BoundaryConditionPair | None = None, e0: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e0_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e0_dbc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e0_dbc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e0_bc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e0_bc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1_dbc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1_dbc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1_bc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e1_bc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2_dbc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2_dbc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2_bc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e2_bc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3_dbc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3_dbc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3_bc: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, e3_bc_T: ~mrx.extraction_operators.MatrixFreeExtraction | None = None, mass_preconds: ~mrx.preconditioners.MassPreconditioners | None = None, runtime_tuning: ~mrx.operators.SequenceRuntimeTuning = <factory>, d0: ~jax.experimental.sparse.bcsr.BCSR | None = None, d0_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, d1: ~jax.experimental.sparse.bcsr.BCSR | None = None, d1_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, d2: ~jax.experimental.sparse.bcsr.BCSR | None = None, d2_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0: ~mrx.operators._MatrixFreeIncidence | None = None, g0_T: ~mrx.operators._MatrixFreeIncidence | None = None, g1: ~mrx.operators._MatrixFreeIncidence | None = None, g1_T: ~mrx.operators._MatrixFreeIncidence | None = None, g2: ~mrx.operators._MatrixFreeIncidence | None = None, g2_T: ~mrx.operators._MatrixFreeIncidence | None = None, inc_gram_inv_1: ~jax.experimental.sparse.bcsr.BCSR | None = None, inc_gram_inv_1_dbc: ~jax.experimental.sparse.bcsr.BCSR | None = None, inc_gram_inv_2: ~jax.experimental.sparse.bcsr.BCSR | None = None, inc_gram_inv_2_dbc: ~jax.experimental.sparse.bcsr.BCSR | None = None, inc_gram_inv_3: ~jax.experimental.sparse.bcsr.BCSR | None = None, inc_gram_inv_3_dbc: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_00: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_00_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_01: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_01_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_10: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_10_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_11: ~jax.experimental.sparse.bcsr.BCSR | None = None, g0_grad_11_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_00: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_00_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_01: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_01_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_10: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_10_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_11: ~jax.experimental.sparse.bcsr.BCSR | None = None, g1_curl_11_T: ~jax.experimental.sparse.bcsr.BCSR | None = None, grad_grad: ~jax.experimental.sparse.bcsr.BCSR | None = None, curl_curl: ~jax.experimental.sparse.bcsr.BCSR | None = None, div_div: ~jax.experimental.sparse.bcsr.BCSR | None = None, dd0_diaginv: object | None = None, dd1_diaginv: object | None = None, dd2_diaginv: object | None = None, dd3_diaginv: object | None = None, dd0_diaginv_dbc: object | None = None, dd1_diaginv_dbc: object | None = None, dd2_diaginv_dbc: object | None = None, dd3_diaginv_dbc: object | None = None, p21: ~jax.experimental.sparse.bcsr.BCSR | None = None, p12: ~jax.experimental.sparse.bcsr.BCSR | None = None, p03: ~jax.experimental.sparse.bcsr.BCSR | None = None, p30: ~jax.experimental.sparse.bcsr.BCSR | None = None, schur_diaginv_k1: ~jax.Array | None = None, schur_diaginv_k1_dbc: ~jax.Array | None = None, schur_diaginv_k2: ~jax.Array | None = None, schur_diaginv_k2_dbc: ~jax.Array | None = None, schur_diaginv_k3: ~jax.Array | None = None, schur_diaginv_k3_dbc: ~jax.Array | None = None, schur_diaginv_mode_k1: str | None = None, schur_diaginv_mode_k1_dbc: str | None = None, schur_diaginv_mode_k2: str | None = None, schur_diaginv_mode_k2_dbc: str | None = None, schur_diaginv_mode_k3: str | None = None, schur_diaginv_mode_k3_dbc: str | None = None, null_0: ~jax.Array | None = None, null_1: ~jax.Array | None = None, null_2: ~jax.Array | None = None, null_3: ~jax.Array | None = None, null_0_dbc: ~jax.Array | None = None, null_1_dbc: ~jax.Array | None = None, null_2_dbc: ~jax.Array | None = None, null_3_dbc: ~jax.Array | None = None, dense: ~mrx.operators.DenseSequenceOperators | None = None) → None

curl_curl: BCSR | None = None

d0: BCSR | None = None

d0_T: BCSR | None = None

d1: BCSR | None = None

d1_T: BCSR | None = None

d2: BCSR | None = None

d2_T: BCSR | None = None

dd0_diaginv: object | None = None

dd0_diaginv_dbc: object | None = None

dd1_diaginv: object | None = None

dd1_diaginv_dbc: object | None = None

dd2_diaginv: object | None = None

dd2_diaginv_dbc: object | None = None

dd3_diaginv: object | None = None

dd3_diaginv_dbc: object | None = None

dense: DenseSequenceOperators | None = None

div_div: BCSR | None = None

e0: MatrixFreeExtraction | None = None

e0_T: MatrixFreeExtraction | None = None

e0_bc: MatrixFreeExtraction | None = None

e0_bc_T: MatrixFreeExtraction | None = None

e0_dbc: MatrixFreeExtraction | None = None

e0_dbc_T: MatrixFreeExtraction | None = None

e1: MatrixFreeExtraction | None = None

e1_T: MatrixFreeExtraction | None = None

e1_bc: MatrixFreeExtraction | None = None

e1_bc_T: MatrixFreeExtraction | None = None

e1_dbc: MatrixFreeExtraction | None = None

e1_dbc_T: MatrixFreeExtraction | None = None

e2: MatrixFreeExtraction | None = None

e2_T: MatrixFreeExtraction | None = None

e2_bc: MatrixFreeExtraction | None = None

e2_bc_T: MatrixFreeExtraction | None = None

e2_dbc: MatrixFreeExtraction | None = None

e2_dbc_T: MatrixFreeExtraction | None = None

e3: MatrixFreeExtraction | None = None

e3_T: MatrixFreeExtraction | None = None

e3_bc: MatrixFreeExtraction | None = None

e3_bc_T: MatrixFreeExtraction | None = None

e3_dbc: MatrixFreeExtraction | None = None

e3_dbc_T: MatrixFreeExtraction | None = None

g0: _MatrixFreeIncidence | None = None

g0_T: _MatrixFreeIncidence | None = None

g0_grad_00: BCSR | None = None

g0_grad_00_T: BCSR | None = None

g0_grad_01: BCSR | None = None

g0_grad_01_T: BCSR | None = None

g0_grad_10: BCSR | None = None

g0_grad_10_T: BCSR | None = None

g0_grad_11: BCSR | None = None

g0_grad_11_T: BCSR | None = None

g1: _MatrixFreeIncidence | None = None

g1_T: _MatrixFreeIncidence | None = None

g1_curl_00: BCSR | None = None

g1_curl_00_T: BCSR | None = None

g1_curl_01: BCSR | None = None

g1_curl_01_T: BCSR | None = None

g1_curl_10: BCSR | None = None

g1_curl_10_T: BCSR | None = None

g1_curl_11: BCSR | None = None

g1_curl_11_T: BCSR | None = None

g2: _MatrixFreeIncidence | None = None

g2_T: _MatrixFreeIncidence | None = None

get_laplacian_diaginv(k: int, dirichlet: bool = True): Return the stored Jacobi inverse diagonal for L_k if available.

grad_grad: BCSR | None = None

inc_gram_inv_1: BCSR | None = None

inc_gram_inv_1_dbc: BCSR | None = None

inc_gram_inv_2: BCSR | None = None

inc_gram_inv_2_dbc: BCSR | None = None

inc_gram_inv_3: BCSR | None = None

inc_gram_inv_3_dbc: BCSR | None = None

k0_tensor_hodge_precond: BoundaryConditionPair | None = None

k1_tensor_stiff_model: K1TensorCurlCurlForwardModel | None = None

k1_tensor_stiff_precond: BoundaryConditionPair | None = None

k2_tensor_stiff_model: K2TensorDivDivForwardModel | None = None

k2_tensor_stiff_precond: BoundaryConditionPair | None = None

property laplacian0_diaginv

property laplacian0_diaginv_dbc

property laplacian1_diaginv

property laplacian1_diaginv_dbc

property laplacian2_diaginv

property laplacian2_diaginv_dbc

property laplacian3_diaginv

property laplacian3_diaginv_dbc

m0: BCSR | None = None

m1: BCSR | None = None

m2: BCSR | None = None

m3: BCSR | None = None

mass_preconds: MassPreconditioners | None = None

null_0: Array | None = None

null_0_dbc: Array | None = None

null_1: Array | None = None

null_1_dbc: Array | None = None

null_2: Array | None = None

null_2_dbc: Array | None = None

null_3: Array | None = None

null_3_dbc: Array | None = None

p03: BCSR | None = None

p12: BCSR | None = None

p21: BCSR | None = None

p30: BCSR | None = None

runtime_tuning: SequenceRuntimeTuning

schur_diaginv_k1: Array | None = None

schur_diaginv_k1_dbc: Array | None = None

schur_diaginv_k2: Array | None = None

schur_diaginv_k2_dbc: Array | None = None

schur_diaginv_k3: Array | None = None

schur_diaginv_k3_dbc: Array | None = None

schur_diaginv_mode_k1: str | None = None

schur_diaginv_mode_k1_dbc: str | None = None

schur_diaginv_mode_k2: str | None = None

schur_diaginv_mode_k2_dbc: str | None = None

schur_diaginv_mode_k3: str | None = None

schur_diaginv_mode_k3_dbc: str | None = None

todense(seq, operator: str, k, dirichlet: bool = True, transpose: bool = False): Return a dense matrix for one assembled operator block.

with_laplacian_diaginv(k: int, value, dirichlet: bool = True): Return a copy with updated stored Jacobi inverse diagonal for L_k.

class mrx.SequenceRuntimeTuning(mass: DegreeIterativeRuntimeTuning = <factory>, scalar_hodge: DegreeShiftedIterativeRuntimeTuning = <factory>, schur: DegreeShiftedIterativeRuntimeTuning = <factory>, diffusion: DegreeShiftedIterativeRuntimeTuning = <factory>)

Bases: Module

Dynamic runtime-tuning payload carried by SequenceOperators.

__init__(mass: ~mrx.operators.DegreeIterativeRuntimeTuning = <factory>, scalar_hodge: ~mrx.operators.DegreeShiftedIterativeRuntimeTuning = <factory>, schur: ~mrx.operators.DegreeShiftedIterativeRuntimeTuning = <factory>, diffusion: ~mrx.operators.DegreeShiftedIterativeRuntimeTuning = <factory>) → None

diffusion: DegreeShiftedIterativeRuntimeTuning

mass: DegreeIterativeRuntimeTuning

scalar_hodge: DegreeShiftedIterativeRuntimeTuning

schur: DegreeShiftedIterativeRuntimeTuning

class mrx.ShiftedIterativeRuntimeTuning(unshifted: 'IterativeRuntimeTuning' = <factory>, shifted: 'IterativeRuntimeTuning' = <factory>)

Bases: Module

__init__(unshifted: ~mrx.operators.IterativeRuntimeTuning = <factory>, shifted: ~mrx.operators.IterativeRuntimeTuning = <factory>) → None

shifted: IterativeRuntimeTuning

unshifted: IterativeRuntimeTuning

class mrx.SplineBasis(n: int, p: int, type: str, T: Array | None = None)

Bases: object

A class representing a basis of spline functions.

This class implements various types of spline bases including clamped, periodic, and constant splines of different degrees (0 to 3). The splines are evaluated using JAX for efficient computation and automatic differentiation.

n

The number of splines in the basis

Type:: int

ns

Array of spline indices

Type:: jnp.ndarray

p

The degree of the spline

Type:: int

type

The type of spline (‘clamped’, ‘periodic’, or ‘constant’)

Type:: str

T

The knot vector defining the spline basis

Type:: jnp.ndarray

T: Array

__call__(x: float, i: int) → Array: Alias for evaluate().

__getitem__(i: int) → Callable[[float], Array]: Return lambda x: self(x, i).

__init__(n: int, p: int, type: str, T: Array | None = None) → None

Initialize a spline basis.

Parameters:

n – The number of splines in the basis
p – The degree of the spline
type – The type of spline (‘clamped’, ‘periodic’, or ‘constant’)
T – Optional knot vector. If None, knots will be initialized based on type

_const_spline(x: float, t: Array) → Array

Evaluate a constant (degree 0) spline.

Parameters:

x – The point at which to evaluate
t – A vector of two elements - the start and end of the interval

Returns:

1.0 if t[0] ≤ x < t[1], 0.0 otherwise

_evaluate(x: float, i: int) → Array

Evaluate the ith spline at x using the appropriate degree-specific method.

Parameters:

x – The point at which to evaluate the spline
i – The index of the spline to evaluate

Returns:

The value of the ith spline at x

_init_knots() → Array

Initialize the knot vector based on the spline type.

Returns:: The initialized knot vector
Raises:: ValueError – If an invalid spline type is provided

_p_spline(x, t, p)

Evaluate a p-spline at point x.

Parameters:

x – The point at which to evaluate the spline
t – The knot vector
p – The degree of the spline

Returns:

The value of the p-spline at x

_safe_divide(x: Array, y: Array) → Array

Divide x by y, returning 0 wherever y is zero.

Uses a dummy denominator of 1 in the zero branch so that x / safe_y is always finite — avoiding 0 * NaN NaN-poisoning in JAX autodiff.

Parameters:

x – The numerator
y – The denominator

Returns:

x / y where y != 0, 0 elsewhere.

collocation_matrix(points: Array | None = None) → Array

Assemble the point-collocation matrix for this basis.

Parameters:: points – Evaluation points. If omitted, the Greville abscissae are used.
Returns:: Array of shape (len(points), n) where entry [k, i] is the value of the i-th basis function at points[k].

evaluate(x: float, i: int) → Array

Evaluate the ith spline at point x, handling special cases.

Parameters:

x – The point at which to evaluate the spline
i – The index of the spline to evaluate

Returns:

The value of the ith spline at x

greville_points() → Array

Return the Greville abscissae for this one-dimensional spline basis.

For degree p > 0 this uses the standard average of the interior knots of each basis function support. For degree p = 0 it falls back to support midpoints.

n: int

ns: Array

p: int

type: str

class mrx.SplineMap(coefficients: Array, extraction: Any, extraction_T: Any | None = None, basis_0: DifferentialForm = None)

Bases: Module

A logical-to-physical map represented in the scalar spline basis.

coefficients, extraction and extraction_T are dynamic pytree children, so SplineMap can be passed through jit / grad / vmap and its coefficients can be differentiated. basis_0 is a static topology object and rides along as aux data.

__init__(coefficients: Array, extraction: Any, extraction_T: Any | None = None, basis_0: DifferentialForm = None) → None

basis_0: DifferentialForm = None

coefficients: Array

extraction: Any

extraction_T: Any | None = None

with_coefficients(coefficients): Return a new spline map with updated coefficients.

class mrx.TensorBasis(bases: list[SplineBasis])

Bases: object

A class representing a tensor product of spline bases.

This class implements a multidimensional basis formed by taking tensor products of one-dimensional spline bases. It is particularly useful for constructing basis functions in higher dimensions (2D or 3D) from one-dimensional splines.

bases

List of one-dimensional spline bases

Type:: list[SplineBasis]

shape

Array containing the number of basis functions in each dimension

Type:: jnp.ndarray

n

Total number of basis functions (product of individual dimensions)

Type:: int

ns

Array of indices for all basis functions

Type:: jnp.ndarray

__call__(x: Array, i: int) → Array: Alias for evaluate().

__getitem__(i: int) → Callable[[Array], Array]: Return lambda x: self(x, i).

__init__(bases: list[SplineBasis]) → None

Initialize a tensor product basis.

The number of basis functions needs to be tracked during JAX tracing/compilation, so we store it explicitly rather than computing it from the bases.

Parameters:: bases – List of one-dimensional SplineBasis objects to form the tensor product
Raises:: ValueError – If the number of bases is not exactly 3

evaluate(x: Array, i: int) → Array

Evaluate the i-th tensor product basis function at point x.

Computes the value by taking the product of the appropriate one-dimensional basis functions in each coordinate direction.

Parameters:

x – Point at which to evaluate the basis function (array of coordinates)
i – Index of the tensor product basis function to evaluate

Returns:

Value of the i-th tensor product basis function at x

mrx.apply_derivative_matrix(seq, operators: SequenceOperators, v, k: int, dirichlet_in: bool = True, dirichlet_out: bool = True, transpose: bool = False)

Apply a weak derivative matrix from an explicit operator bundle.

D_k = M_{k+1} G_k is applied as a composition of matrix-free applies; the full D_k is never materialised.

mrx.apply_derivative_matrix_ops(seq, operators: SequenceOperators, v, k: int, dirichlet_in: bool = True, dirichlet_out: bool = True, transpose: bool = False)

Apply a weak derivative matrix from an explicit operator bundle.

D_k = M_{k+1} G_k is applied as a composition of matrix-free applies; the full D_k is never materialised.

mrx.apply_hodge_laplacian(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None)

Apply the Hodge Laplacian using explicit operator data.

This uses bundled mass, weak derivative, and stiffness operators.

mrx.apply_hodge_laplacian_approx(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

Linear approximation of the Hodge Laplacian apply.

Replaces the exact M_{k-1}^{-1} in the Schur term of L_k with one apply of the configured mass preconditioner. The result is a fully linear SPD matvec: safe to nest inside Krylov iterations and to use as a preconditioner or a diagnostic L_k-apply. It is not exactly L_k unless the metric is tensor-separable on the reference domain.

mrx.apply_hodge_laplacian_preconditioner(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, kind: str = 'auto')

Apply the Hodge-Laplacian preconditioner from an operator bundle.

kind options:

'none' — identity (no preconditioning).
'jacobi' — per-DoF diagonal of L_k; always available.
- 'tensor' — assembled surgery-plus-Schur tensor Hodge model for
  k = 0 only.
- 'auto' — picks 'tensor' when available for k = 0 and falls
  back to 'jacobi' otherwise.

mrx.apply_incidence_matrix(seq, operators: SequenceOperators, v, k: int, dirichlet_in: bool = True, dirichlet_out: bool = True, transpose: bool = False)

Apply the strong exterior-derivative G_k on extracted DoF spaces.

The raw extracted incidence is E_out^T sp E_in (sp has entries in {-1, 0, +1}). On polar sequences the extraction is non-unitary at the axis, so the raw form is NOT the topological derivative and d.d != 0. The true strong derivative is G = Gram_{k+1}^{-1} (E_out^T sp E_in) with the cached coefficient-Gram inverse (None / identity where the extraction is unitary, e.g. non-polar). The correction is a sparse matvec localised to the polar-axis DoFs, so off-axis the result is bit-identical to the raw incidence and d.d = 0 holds exactly on extracted DoFs everywhere.

mrx.apply_incidence_matrix_ops(seq, operators: SequenceOperators, v, k: int, dirichlet_in: bool = True, dirichlet_out: bool = True, transpose: bool = False)

Apply the strong exterior-derivative G_k on extracted DoF spaces.

The raw extracted incidence is E_out^T sp E_in (sp has entries in {-1, 0, +1}). On polar sequences the extraction is non-unitary at the axis, so the raw form is NOT the topological derivative and d.d != 0. The true strong derivative is G = Gram_{k+1}^{-1} (E_out^T sp E_in) with the cached coefficient-Gram inverse (None / identity where the extraction is unitary, e.g. non-polar). The correction is a sparse matvec localised to the polar-axis DoFs, so off-axis the result is bit-identical to the raw incidence and d.d = 0 holds exactly on extracted DoFs everywhere.

mrx.apply_inverse_hodge_laplacian(seq, operators: SequenceOperators, rhs, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False)

Solve with the inverse of the unshifted Hodge Laplacian L_k.

For k = 0 this uses the dedicated singular scalar-Laplacian solve directly rather than routing through the shifted eps = 0 path. For k >= 1 the saddle-point implementation remains shared with the shifted solve because the only difference is the absent mass shift.

mrx.apply_inverse_laplacian(seq, operators: SequenceOperators, rhs, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False): Alias of apply_inverse_hodge_laplacian using Laplacian naming.

mrx.apply_inverse_laplacian_ops(seq, operators: SequenceOperators, rhs, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False): Alias of apply_inverse_hodge_laplacian using Laplacian naming.

mrx.apply_inverse_mass_matrix(seq, operators: SequenceOperators, rhs, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False)

Solve with the inverse mass matrix from an explicit operator bundle.

preconditioner accepts a kind string or a MassPreconditionerSpec. When omitted, the default is tensor when assembled and Jacobi otherwise.

mrx.apply_inverse_mass_matrix_ops(seq, operators: SequenceOperators, rhs, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False)

Solve with the inverse mass matrix from an explicit operator bundle.

preconditioner accepts a kind string or a MassPreconditionerSpec. When omitted, the default is tensor when assembled and Jacobi otherwise.

mrx.apply_inverse_mass_plus_eps_laplace_matrix(seq, operators: SequenceOperators, rhs, k: int, eps: float, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False): Solve with the inverse of M_k + eps L_k using an explicit operator bundle.

mrx.apply_inverse_mass_plus_eps_laplace_matrix_ops(seq, operators: SequenceOperators, rhs, k: int, eps: float, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', return_info: bool = False): Solve with the inverse of M_k + eps L_k using an explicit operator bundle.

mrx.apply_inverse_shifted_hodge_laplacian(seq, operators: SequenceOperators, rhs, k: int, eps: float, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', use_harmonic_coarse: bool | None = None, return_info: bool = False)

Solve with the inverse of the shifted Hodge Laplacian L_k + eps M_k.

For k >= 1 the interface is preconditioner, a structured saddle-point preconditioner spec with a lower mass block, a Schur-inner mass inverse, a Schur-outer preconditioner, and an optional coupled completion. Kind strings are accepted as convenience shorthands.

mrx.apply_inverse_shifted_laplacian(seq, operators: SequenceOperators, rhs, k: int, eps: float, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', use_harmonic_coarse: bool | None = None, return_info: bool = False): Alias of apply_inverse_shifted_hodge_laplacian using Laplacian naming.

mrx.apply_inverse_shifted_laplacian_ops(seq, operators: SequenceOperators, rhs, k: int, eps: float, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None, preconditioner='auto', use_harmonic_coarse: bool | None = None, return_info: bool = False): Alias of apply_inverse_shifted_hodge_laplacian using Laplacian naming.

mrx.apply_laplacian(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None): Alias of apply_hodge_laplacian using Laplacian naming.

mrx.apply_laplacian_approx(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True): Alias of apply_hodge_laplacian_approx using Laplacian naming.

mrx.apply_laplacian_approx_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True): Alias of apply_hodge_laplacian_approx using Laplacian naming.

mrx.apply_laplacian_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, guess=None, tol: float | None = None, maxiter: int | None = None): Alias of apply_hodge_laplacian using Laplacian naming.

mrx.apply_laplacian_preconditioner(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, kind: str = 'auto'): Alias of apply_hodge_laplacian_preconditioner using Laplacian naming.

mrx.apply_laplacian_preconditioner_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, kind: str = 'auto'): Alias of apply_hodge_laplacian_preconditioner using Laplacian naming.

mrx.apply_mass_matrix(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True): Apply a mass matrix from an explicit operator bundle.

mrx.apply_mass_matrix_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True): Apply a mass matrix from an explicit operator bundle.

mrx.apply_mass_matrix_preconditioner(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, kind: str = 'auto')

Apply a mass-matrix preconditioner from an explicit operator bundle.

Parameters:: kind ({'auto', 'jacobi', 'tensor'}) – Which preconditioner to use. 'auto' picks 'tensor' when the tensor mass preconditioner is assembled and available for this k; otherwise it falls back to 'jacobi'.

mrx.apply_mass_matrix_preconditioner_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, kind: str = 'auto')

Apply a mass-matrix preconditioner from an explicit operator bundle.

Parameters:: kind ({'auto', 'jacobi', 'tensor'}) – Which preconditioner to use. 'auto' picks 'tensor' when the tensor mass preconditioner is assembled and available for this k; otherwise it falls back to 'jacobi'.

mrx.apply_mass_tensor_forward_model(seq, preconds: MassPreconditioners | None, v, k: int, dirichlet: bool = True)

mrx.apply_mass_tensor_forward_model_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

mrx.apply_mass_tensor_preconditioner(seq, preconds: MassPreconditioners | None, v, k: int, dirichlet: bool = True, *, true_block_apply_k3=None)

mrx.apply_mass_tensor_preconditioner_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

mrx.apply_projection_matrix(seq, operators: SequenceOperators, v, k_in: int, k_out: int, dirichlet_in: bool = True, dirichlet_out: bool = True): Apply a projection matrix from an explicit operator bundle.

mrx.apply_projection_matrix_ops(seq, operators: SequenceOperators, v, k_in: int, k_out: int, dirichlet_in: bool = True, dirichlet_out: bool = True): Apply a projection matrix from an explicit operator bundle.

mrx.apply_stiffness(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

Apply a stiffness matrix from an explicit operator bundle.

K_k = G_k^T M_{k+1} G_k is applied as a composition of matrix-free applies; the full K_k is never materialised.

mrx.apply_stiffness_ops(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

Apply a stiffness matrix from an explicit operator bundle.

K_k = G_k^T M_{k+1} G_k is applied as a composition of matrix-free applies; the full K_k is never materialised.

mrx.apply_stiffness_tensor_forward_model(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True, *, regular_space: bool = False)

Apply the stored tensor stiffness forward model for k = 1 or k = 2.

By default this mirrors apply_stiffness() on the extracted space. Set regular_space=True to apply the regular-space tensor model directly.

mrx.apply_stiffness_tensor_preconditioner(seq, operators: SequenceOperators, v, k: int, dirichlet: bool = True)

mrx.approx_inverse_map(y: Array, eps: float, R0: float = 1.0) → Array

Approximate inverse of toroid_map for a circular cross-section.

Parameters:

y – Cartesian coordinates (X, Y, Z).
eps – Minor radius (same as epsilon in toroid_map()).
R0 – Major radius.

mrx.assemble_all_dense_operators(seq, operators: SequenceOperators | None = None)

Materialize dense extracted operators into operators.dense.

Requires the corresponding sparse operators to already be assembled. This is intended as a courtesy/debugging path for dense inspection and direct solves, not as the default operator assembly route.

mrx.assemble_all_operators(seq, geometry, operators: SequenceOperators | None = None, include_preconditioners: bool = True)

Assemble all geometry-dependent operators.

When include_preconditioners is true, also assemble the eager preconditioner payloads that back the solver-facing convenience paths. Set it to false when only the sparse operators are needed, for example for densification or direct solves.

mrx.assemble_derivative_operators(seq, geometry, operators: SequenceOperators | None = None, ks: Sequence[int] = (0, 1, 2)): Assemble weak derivative operators for the requested form degrees.

mrx.assemble_hodge_operators(seq, geometry, operators: SequenceOperators | None = None, ks: Sequence[int] = (0, 1, 2, 3)): Assemble Hodge/stiffness operators for the requested form degrees.

mrx.assemble_incidence_operators(seq, operators: SequenceOperators | None = None, ks: Sequence[int] = (0, 1, 2))

Assemble topological incidence operators for the requested degrees.

Also caches the TRUE polar-derivative Gram^{-1} corrections for the output spaces {k+1} so apply_incidence_matrix() returns the strong derivative M^{-1} D (exact d.d = 0 on extracted DoFs) on polar sequences, while staying bit-identical to the raw incidence elsewhere.

mrx.assemble_laplacian_operators(seq, geometry, operators: SequenceOperators | None = None, ks: Sequence[int] = (0, 1, 2, 3)): Alias of assemble_hodge_operators using Laplacian naming.

mrx.assemble_leray_projection(seq): Assemble the Leray projection matrix.

mrx.assemble_mass_jacobi_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (0, 1, 2, 3))

Assemble/store Jacobi mass diagonals eagerly for requested degrees.

Reuses the existing direct diagonal extraction helper and stores only the inverse-diagonal vectors on operators.mass_preconds.jacobi (free + DBC) for reuse by Jacobi preconditioners.

mrx.assemble_mass_local(seq, k, geometry=None): Dispatch element-local mass assembly for form degree k (BCOO).

mrx.assemble_mass_operators(seq, geometry, operators: SequenceOperators | None = None, ks: Sequence[int] = (0, 1, 2, 3)): Assemble mass operators for the requested form degrees.

mrx.assemble_mass_surgery_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (0, 1, 2), precompute_coupling: bool = True)

mrx.assemble_projection_operators(seq, operators: SequenceOperators | None = None, pairs: Sequence[tuple[int, int]] = ((2, 1), (1, 2), (0, 3), (3, 0))): Assemble projection operators for the requested degree pairs.

mrx.assemble_schur_jacobi_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (1, 2, 3), dirichlet_variants: Sequence[bool] | None = None, eps: float = 0.0, schur_diag_mode: str = 'tensor_probe') → SequenceOperators

Probe and store the approximate Schur diagonal at assembly time.

For each (k, dirichlet) pair, builds the approximate Schur operator

A_k(x) = S_k x + D_{k-1} B_{k-1} D_{k-1}^T x

and probes its diagonal by O(n_k) matrix-vector products. The resulting 1/diag(A_k) is stored on the operator bundle so that the saddle-point Schur-outer Jacobi preconditioner is a cheap multiply at solve time rather than an O(n_k) probing scan.

Parameters:

seq (DeRhamSequence)
operators (SequenceOperators, optional)
ks (sequence of int) – Form degrees to assemble (must be in 1, 2, 3).
dirichlet_variants (sequence of bool, optional) – Boundary condition variants to assemble. Defaults to (True, False).
eps (float) – Shift for the stiffness term; 0 gives the unshifted Schur.
schur_diag_mode (where B_{k-1} is selected by)
'tensor_probe' (-)
'exact_probe' (-)
'diag' (-)
k-1 ('tensor_probe' requires the tensor mass preconditioner for)
first. (to be assembled)

mrx.assemble_tensor_hodge_preconditioner(seq, operators: SequenceOperators | None = None, *, rank: int | None = None, cp_maxiter: int | None = None, cp_tol: float | None = None, cp_ridge: float | None = None): Deprecated no-op; use assemble_tensor_laplacian_preconditioner().

mrx.assemble_tensor_laplacian_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (0,), rank: int | None = None, cp_kwargs: dict | None = None)

Assemble the scalar k=0 tensor Hodge-Laplacian preconditioner.

Only k=0 is supported. rank and cp_kwargs override the CP fit parameters; when None the values stored on the tensor mass preconditioner are used.

mrx.assemble_tensor_mass_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (1,), rank: int = 3, cp_kwargs: dict | None = None)

Assemble the k=0/k=1/k=2/k=3 tensor mass preconditioner on operators.

The current production tensor path implements the extracted scalar core-plus-bulk Schur model for polar k=0 and the surgery-plus-Schur model for polar k=1. For polar k=2 it uses an outer Schur split on the extracted r surgery block together with tensor-diagonal bulk block inverses. For polar k=3 it implements a direct extracted scalar tensor inverse. All use low-rank CP fits of the diagonal metric factors to build tensor block inverse applies.

mrx.assemble_tensor_stiffness_models(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (2,), rank: int = 1, cp_kwargs: dict | None = None)

Assemble the stored higher-form tensor stiffness forward models.

This stores regular-space tensor models for k = 1 curl-curl and k = 2 div-div on the operator bundle so they can be applied through a stable API rather than only via internal debug helpers.

mrx.assemble_tensor_stiffness_preconditioner(seq, operators: SequenceOperators | None = None, *, ks: Sequence[int] = (1, 2), rank: int = 1, cp_kwargs: dict | None = None)

Assemble standalone tensor stiffness preconditioners for k = 1, 2.

These are preconditioners for the semidefinite stiffness blocks curl-curl and div-div themselves. They are intentionally kept separate from the mixed saddle-point Hodge-Laplacian path.

mrx.assemble_vectorial(row_terms, col_terms, W_flat_3x3, quad_shape, comp_shapes, hw, col_comp_shapes=None)

Tensor-product assembly for vectorial DOFs with block structure.

Computes M[i,j] = Σ_{k,l} ∫ (OpΛ_i)_k · W_{kl} · (OpΛ_j)_l dx

where the operator maps each source component c to one or more output components k, each factoring as a product of 1D functions.

For mass matrices each component has a single identity term [(c, R, T, Z, +1)]. For stiffness matrices (e.g. curl-curl) each component may have multiple signed terms.

Supports rectangular matrices when col_comp_shapes is provided (e.g. derivative matrices mapping between different form degrees).

Parameters:

row_terms (list of lists) – row_terms[c] is a list of (output_idx, R, T, Z, sign) tuples.
col_terms (list of lists) – Same structure for the column operator.
W_flat_3x3 (array, shape (n_q, 3, 3)) – Weight tensor indexed by output component pair (k, l).
quad_shape (tuple (n_qt, n_qr, n_qz))
comp_shapes (list of tuples (s1, s2, s3)) – DOF grid shape per row source component.
hw (int) – Stencil half-width (polynomial degree p).
col_comp_shapes (list of tuples (s1, s2, s3), optional) – DOF grid shape per column source component. When None, defaults to comp_shapes (square matrix).

Returns:

M

Return type:

jax.experimental.sparse.BCOO

mrx.backtracking_line_search(x, direction, J_current, J_fn, *, step_init=1.0, step_min=1e-09, step_max=1000000.0, c1=0.0001, shrink=0.5, grow=2.0, max_backtracks=40, directional_derivative=None, feasible=None)

Armijo backtracking line search with an optional feasibility filter.

Finds a step s along the descent direction direction such that

J(x + s * direction) <= J(x) + c1 * s * <grad J, direction>,

and, if feasible is supplied, feasible(x + s * direction) is True. Intended to be called once per outer iteration of a Python-level descent loop; the returned step is already grown (on success) or left at the last trial value (on failure) so it can be passed back in as step_init next iteration.

Parameters:

x (array) – Current iterate.
direction (array) – Descent direction (typically -grad J).
J_current (float) – J(x).
J_fn (callable) – x_trial -> float. May return non-finite; such trials are rejected.
step_init (float) – Initial trial step, floor, and cap.
step_min (float) – Initial trial step, floor, and cap.
step_max (float) – Initial trial step, floor, and cap.
c1 (float) – Armijo sufficient-decrease constant.
shrink (float) – Multipliers applied to step on rejection / acceptance.
grow (float) – Multipliers applied to step on rejection / acceptance.
max_backtracks (int) – Maximum number of trials per call.
directional_derivative (float, optional) – <grad J, direction>. When omitted, we assume direction = -grad J and use -||direction||^2.
feasible (callable, optional) – x_trial -> bool. Trials for which this returns False are rejected without evaluating J_fn.

Returns:

result – Keys: "x" (new iterate, equals x if not accepted), "J" (J_fn at the new iterate, else J_current), "step" (next trial step to use), "accepted" (bool), "n_backtracks" (int).

Return type:

dict

mrx.bc_extraction_op(e, e_dbc, n_full: int)

Build the extraction operator for Dirichlet boundary DOFs.

Returns a MatrixFreeExtraction of shape (n_bc, n_full) that selects the DOFs present in e (unrestricted) but absent from e_dbc (DBC), i.e. the DOFs that are set to zero by the homogeneous Dirichlet BC.

Uses the identity: columns present in e but not e_dbc satisfy: (e.T @ 1 - e_dbc.T @ 1)[i] == 1

mrx.build_curl_stencil_g1(seq, xi, dirichlet_in: bool, dirichlet_out: bool)

Analytic, INVERSE-FREE polar discrete curl G_1 (V1 -> V2).

The degree-1 analog of build_grad_stencil_g0(): the true strong curl on extracted DoFs as an explicit sparse matrix from the incidence pattern and the polar coefficients xi (shape (3, 2, nt)) – coefficient differences and xi weights only, NO mass and NO matrix inverse. The closed form of Gram_2^{-1} (E_2 sp_1 E_1^T); the V2 axis-fusion inverse cancels to clean +/-1 / xi-difference stencils (verified bit-exact vs that oracle).

Full-space curl (a=s, b=chi, c=zeta -> V2 comps P,Q,R; see _apply_incidence_mf): P=-d_z b + d_t c, Q=d_z a - d_r c, R=-d_t a + d_r b. V1 input fusion is inverted by expand_v1 (the V1 analog of grad’s expand); the only fused V2 output DoFs are the comp0 surgery rows, whose stencil is the axis form of P = -d_z(chi apex) + d_t(zeta apex).

mrx.build_grad_stencil_g0(seq, xi, dirichlet_in: bool, dirichlet_out: bool)

Analytic, INVERSE-FREE polar discrete gradient G_0 (V0 -> V1).

Builds the true strong gradient on extracted DoFs as an explicit sparse matrix straight from the incidence pattern and the polar mapping coefficients xi (shape (3, 2, nt)) – coefficient differences and xi weights only, NO mass and NO matrix inverse. This is the closed form of Gram_1^{-1} (E_1 sp_0 E_0^T); the axis-fusion inverse cancels to clean +/-1 / -xi[l,1,j] stencils (verified bit-exact vs that oracle).

Layout (see extraction_operators.build_extraction k=0/k=1 branches): V0 extracted = apex (p,m) -> p*nz+m (p in 0..2) then bulk (i,j,k) -> 3 nz + ravel((i,j,k),(radial0,nt,nz)) with full radial i+2. V1 extracted = theta_surgery [0,2 nz) | zeta_surgery [2 nz, 2 nz+3 dz) | r-slice (comp0) | theta_bulk (comp1) | zeta_bulk (comp2). The full-space grad is d_r f, d_theta f (periodic), d_z f (periodic), with the near-axis full radial rows 0/1 expanded as f(0,j,k)=sum_p xi[p,0,j] apex, f(1,j,k)=sum_p xi[p,1,j] apex.

mrx.build_mass_surgery_preconditioner(seq, mass_apply, *, k: int, existing: MassSurgeryPreconditioner | None = None, dirichlet_flags: tuple[bool, ...] = (False, True), precompute_coupling: bool = True) → MassSurgeryPreconditioner

mrx.build_mass_tensor_preconditioner(seq, *, k: int, rank: int = 1, fallback_rank: int | None = None, cp_kwargs: Mapping[str, object] | None = None, existing: TensorMassPreconditioner | None = None, surgery_precond: MassSurgeryPreconditioner | None = None, dirichlet_flags: tuple[bool, ...] = (False, True), k3_true_block_apply: Mapping[bool, object] | None = None) → TensorMassPreconditioner

mrx.build_matrixfree_mass_apply(seq, k, geometry=None): Return a jitted raw-DOF-space x -> M_k x that never stores M_k.

mrx.cg(A, b, x0=None, *, tol=1e-05, atol=0.0, maxiter=None, M=None)

Use Conjugate Gradient iteration to solve Ax = b.

The numerics of JAX’s cg should exact match SciPy’s cg (up to numerical precision), but note that the interface is slightly different: you need to supply the linear operator A as a function instead of a sparse matrix or LinearOperator.

Derivatives of cg are implemented via implicit differentiation with another cg solve, rather than by differentiating through the solver. They will be accurate only if both solves converge.

Parameters:

A (ndarray, function, or matmul-compatible object) – 2D array or function that calculates the linear map (matrix-vector product) Ax when called like A(x) or A @ x. A must represent a hermitian, positive definite matrix, and must return array(s) with the same structure and shape as its argument.
b (array or tree of arrays) – Right hand side of the linear system representing a single vector. Can be stored as an array or Python container of array(s) with any shape.
x0 (array or tree of arrays) – Starting guess for the solution. Must have the same structure as b.
tol (float, optional) – Tolerances for convergence, norm(residual) <= max(tol*norm(b), atol). We do not implement SciPy’s “legacy” behavior, so JAX’s tolerance will differ from SciPy unless you explicitly pass atol to SciPy’s cg.
atol (float, optional) – Tolerances for convergence, norm(residual) <= max(tol*norm(b), atol). We do not implement SciPy’s “legacy” behavior, so JAX’s tolerance will differ from SciPy unless you explicitly pass atol to SciPy’s cg.
maxiter (integer) – Maximum number of iterations. Iteration will stop after maxiter steps even if the specified tolerance has not been achieved.
M (ndarray, function, or matmul-compatible object) – Preconditioner for A. The preconditioner should approximate the inverse of A. Effective preconditioning dramatically improves the rate of convergence, which implies that fewer iterations are needed to reach a given error tolerance.

Returns:

x (array or tree of arrays) – The converged solution. Has the same structure as b.
info (None) – Placeholder for convergence information. In the future, JAX will report the number of iterations when convergence is not achieved, like SciPy.