Matrix Ops

matrix_ops

Matrix operations and linear solvers for multiphysics modeling.

This module provides small, performance-oriented numerical utilities used by multiphysics engines:

• Construction of common 1D linear operators (e.g., upwind advection, Laplacian).
• Cached implicit solves for repeated linear systems with fixed operators.
• High-throughput aggregation utilities for large numbers of subpopulations.
• Optional Kronecker composition utilities for separable multi-axis operators.

Design notes

• Dense implicit solves use the input state's Array-API linalg.solve.
• Sparse acceleration is selected structurally from a registry containing SciPy and, when installed, CuPy adapters.
• Cache semantics: sparse factorizations are keyed by (id(left_op), id(right_op)) within each ecosystem and retain those exact operator objects to prevent stale hits after Python ID recycling. For caching to be effective, operator objects must be constructed once and reused.

Stage operator factories (IMEX/TR-BDF2 support): TR-BDF2 and similar IMEX methods can require stage-specific implicit operators that depend on: - dt (time step) - scale (method stage scalar) - t (stage time) - y (stage state) - stage (a label, e.g. "tr" or "bdf2")

This module supports dynamic base operators via a builder:
    base_builder(t, y, stage) -> Operator

Then the stage-operator factory produces (L, R) to solve:
    L @ y_next = R @ y_in

where (L, R) follow schemes like implicit Euler or trapezoidal.

DiffusionConfig(coeff, dtype=np.float64, bc='neumann') dataclass

Configuration for diffusion-like linear operators.

Attributes:

Name	Type	Description
coeff	float	Physical diffusion coefficient D (units length^2 / time).
dtype	DTypeLike	Floating dtype (e.g. np.float64).
bc	str	Boundary condition; either "neumann" or "absorbing".

__post_init__()

Validate context-free diffusion parameters.

Raises:

Type	Description
ValueError	If the coefficient or boundary condition is invalid.

Source code in src/op_engine/matrix_ops.py

def __post_init__(self) -> None:
    """Validate context-free diffusion parameters.

    Raises:
        ValueError: If the coefficient or boundary condition is invalid.
    """
    if not np.isfinite(self.coeff) or self.coeff < 0.0:
        msg = "coeff must be finite and non-negative"
        raise ValueError(msg)
    if self.bc not in {"neumann", "absorbing"}:
        raise ValueError(_UNKNOWN_BC_ERROR.format(bc=self.bc))
    np.dtype(self.dtype)

GridGeometry(n, dx) dataclass

Geometry of a 1D spatial grid.

Attributes:

Name	Type	Description
n	int	Number of grid points.
dx	float	Grid spacing.

__post_init__()

Validate context-free grid geometry.

Raises:

Type	Description
ValueError	If the grid size or spacing is invalid.

Source code in src/op_engine/matrix_ops.py

def __post_init__(self) -> None:
    """Validate context-free grid geometry.

    Raises:
        ValueError: If the grid size or spacing is invalid.
    """
    if not isinstance(self.n, Integral) or isinstance(self.n, bool) or self.n < 1:
        msg = "n must be a positive integer"
        raise ValueError(msg)
    if not np.isfinite(self.dx) or self.dx <= 0.0:
        msg = "dx must be finite and positive"
        raise ValueError(msg)

SparseAdapter(ecosystem_id, issparse, factorize, solve, cache_key) dataclass

Function bundle for one sparse-array ecosystem.

Sparse acceleration is optional: dense operators always fall back to the input array's Array-API linalg.solve implementation. Adapters are selected structurally with their own issparse predicate.

StageOperatorContext(t, y, stage=None, extra=None) dataclass

Context passed to time/state-dependent operator builders.

Attributes:

Name	Type	Description
t	float	Stage time.
y	Array	Stage state in its runtime Array-API namespace.
stage	StageName	Optional stage label (e.g. "tr", "bdf2").
extra	Any | None	Optional extra payload for future use (kept generic).

build_advection_matrix(n, dx, velocity, *, bc='absorbing', reference=None)

Build a dense first-order upwind finite-volume operator.

The returned matrix A acts on a column state as dy = A @ y. Positive velocity transports values toward increasing coordinate indices; negative velocity transports toward decreasing indices. The velocity may be an Array-API scalar, so its value stays dynamic under transformations such as JAX jit and grad.

Boundary modes are:

• absorbing: zero inflow at the upstream boundary and free outflow at the downstream boundary;
• reflecting: zero flux through the downstream boundary;
• periodic: downstream outflow wraps to the upstream boundary.

The namespace comes from reference when provided, then from velocity; plain numeric velocities use NumPy. Grid geometry and the boundary mode are static structural inputs.

Parameters:

Name	Type	Description	Default
n	int	Number of uniformly spaced finite-volume cells.	required
dx	float	Positive cell width.	required
velocity	float | Array	Signed scalar transport velocity.	required
bc	str	Boundary mode: "absorbing", "reflecting", or "periodic".	'absorbing'
reference	Array | None	Optional array whose namespace and floating dtype control the result.	None

Returns:

Type	Description
Array	Dense (n, n) operator in the selected Array-API namespace.

Raises:

Type	Description
ValueError	If the grid, velocity shape/value, or boundary mode is invalid.

Source code in src/op_engine/matrix_ops.py

def build_advection_matrix(
    n: int,
    dx: float,
    velocity: float | Array,
    *,
    bc: str = "absorbing",
    reference: Array | None = None,
) -> Array:
    """Build a dense first-order upwind finite-volume operator.

    The returned matrix ``A`` acts on a column state as ``dy = A @ y``.
    Positive velocity transports values toward increasing coordinate indices;
    negative velocity transports toward decreasing indices. The velocity may be
    an Array-API scalar, so its value stays dynamic under transformations such
    as JAX ``jit`` and ``grad``.

    Boundary modes are:

    - ``absorbing``: zero inflow at the upstream boundary and free outflow at
      the downstream boundary;
    - ``reflecting``: zero flux through the downstream boundary;
    - ``periodic``: downstream outflow wraps to the upstream boundary.

    The namespace comes from ``reference`` when provided, then from ``velocity``;
    plain numeric velocities use NumPy. Grid geometry and the boundary mode are
    static structural inputs.

    Args:
        n: Number of uniformly spaced finite-volume cells.
        dx: Positive cell width.
        velocity: Signed scalar transport velocity.
        bc: Boundary mode: ``"absorbing"``, ``"reflecting"``, or ``"periodic"``.
        reference: Optional array whose namespace and floating dtype control the
            result.

    Returns:
        Dense ``(n, n)`` operator in the selected Array-API namespace.

    Raises:
        ValueError: If the grid, velocity shape/value, or boundary mode is
            invalid.
    """
    if n < 2:
        msg = f"advection grid size must be at least 2; got {n}."
        raise ValueError(msg)
    if not np.isfinite(dx) or dx <= 0.0:
        msg = f"advection dx must be finite and positive; got {dx}."
        raise ValueError(msg)

    bc_normalized = str(bc).strip().lower()
    if bc_normalized not in {"absorbing", "reflecting", "periodic"}:
        msg = (
            f"Unknown advection bc: {bc!r}; expected absorbing, reflecting, "
            "or periodic."
        )
        raise ValueError(msg)

    namespace_source: object = reference if reference is not None else velocity
    try:
        _namespace_of(namespace_source)
    except TypeError:
        namespace_source = np.asarray(namespace_source)
    xp = _namespace_of(namespace_source)
    source_dtype = cast("Any", namespace_source).dtype
    dtype = xp.result_type(xp.asarray(0.0).dtype, source_dtype)
    velocity_array = cast("Array", xp.asarray(velocity, dtype=dtype))
    if velocity_array.shape != ():
        msg = f"advection velocity must be scalar; got shape {velocity_array.shape}."
        raise ValueError(msg)
    if isinstance(velocity_array, np.ndarray) and not np.isfinite(velocity_array).all():
        msg = "advection velocity must be finite."
        raise ValueError(msg)

    positive = -np.eye(n, dtype=np.float64)
    negative = -np.eye(n, dtype=np.float64)
    indices = np.arange(n - 1)
    positive[indices + 1, indices] = 1.0
    negative[indices, indices + 1] = 1.0

    if bc_normalized == "reflecting":
        positive[-1, -1] = 0.0
        negative[0, 0] = 0.0
    elif bc_normalized == "periodic":
        positive[0, -1] = 1.0
        negative[-1, 0] = 1.0

    positive_array = xp.asarray(positive, dtype=dtype)
    negative_array = xp.asarray(negative, dtype=dtype)
    zero = xp.asarray(0.0, dtype=dtype)
    stencil = xp.where(
        xp.greater_equal(velocity_array, zero),
        positive_array,
        negative_array,
    )
    speed = xp.divide(xp.abs(velocity_array), dx)
    return cast("Array", xp.multiply(stencil, speed))

build_crank_nicolson_operator(geom, cfg, dt)

Build Crank-Nicolson operators with dense/sparse autodispatch.

Parameters:

Name	Type	Description	Default
geom	GridGeometry	Grid geometry.	required
cfg	DiffusionConfig	Diffusion configuration.	required
dt	float	Time step.	required

Returns:

Type	Description
tuple[Operator, Operator]	Tuple of (L, R) operators for Crank-Nicolson scheme.

Source code in src/op_engine/matrix_ops.py

def build_crank_nicolson_operator(
    geom: GridGeometry,
    cfg: DiffusionConfig,
    dt: float,
) -> tuple[Operator, Operator]:
    """
    Build Crank-Nicolson operators with dense/sparse autodispatch.

    Args:
        geom: Grid geometry.
        cfg: Diffusion configuration.
        dt: Time step.

    Returns:
        Tuple of (L, R) operators for Crank-Nicolson scheme.
    """
    if geom.n < _DISPATCH_THRESHOLD:
        return _build_crank_nicolson_dense(geom, cfg, dt)
    return _build_crank_nicolson_sparse(geom, cfg, dt)

build_diffusion_matrix(n, dx, coefficient, *, grid=None, bc='neumann', reference=None)

Build a dense second-order 1D finite-volume diffusion operator.

The returned matrix A acts on a column state as dy = A @ y. Supply either a uniform cell width with dx or strictly increasing cell-center coordinates with grid. On a non-uniform grid, interior flux differences are divided by the local Voronoi-cell width. Boundary faces are inferred one half-spacing beyond the first and last centers.

The coefficient may be an Array-API scalar, so it stays dynamic under transformations such as JAX jit and grad. Static geometry is assembled with NumPy and transferred once to the namespace selected by reference or coefficient.

Boundary modes are:

• neumann or reflecting: zero flux at both boundaries;
• absorbing: zero-valued exterior ghost cells;
• periodic: the two grid ends are adjacent. A grid passed explicitly must be uniform because its coordinates do not determine wrap spacing.

Parameters:

Name	Type	Description	Default
n	int	Number of finite-volume cells.	required
dx	float | None	Positive uniform cell width, or None when grid is supplied.	required
coefficient	float | Array	Non-negative scalar diffusion coefficient.	required
grid	ArrayLike | None	Optional strictly increasing cell-center coordinates of shape (n,). Exactly one of dx and grid must be supplied.	None
bc	str	Boundary mode: neumann, reflecting, absorbing, or periodic.	'neumann'
reference	Array | None	Optional array whose namespace and floating dtype control the result.	None

Returns:

Type	Description
Array	Dense (n, n) operator in the selected Array-API namespace.

Raises:

Type	Description
ValueError	If the geometry, coefficient, or boundary mode is invalid.

Source code in src/op_engine/matrix_ops.py

def build_diffusion_matrix(  # noqa: PLR0913
    n: int,
    dx: float | None,
    coefficient: float | Array,
    *,
    grid: ArrayLike | None = None,
    bc: str = "neumann",
    reference: Array | None = None,
) -> Array:
    """Build a dense second-order 1D finite-volume diffusion operator.

    The returned matrix ``A`` acts on a column state as ``dy = A @ y``.
    Supply either a uniform cell width with ``dx`` or strictly increasing
    cell-center coordinates with ``grid``. On a non-uniform grid, interior
    flux differences are divided by the local Voronoi-cell width. Boundary
    faces are inferred one half-spacing beyond the first and last centers.

    The coefficient may be an Array-API scalar, so it stays dynamic under
    transformations such as JAX ``jit`` and ``grad``. Static geometry is
    assembled with NumPy and transferred once to the namespace selected by
    ``reference`` or ``coefficient``.

    Boundary modes are:

    - ``neumann`` or ``reflecting``: zero flux at both boundaries;
    - ``absorbing``: zero-valued exterior ghost cells;
    - ``periodic``: the two grid ends are adjacent. A grid passed explicitly
      must be uniform because its coordinates do not determine wrap spacing.

    Args:
        n: Number of finite-volume cells.
        dx: Positive uniform cell width, or ``None`` when ``grid`` is supplied.
        coefficient: Non-negative scalar diffusion coefficient.
        grid: Optional strictly increasing cell-center coordinates of shape
            ``(n,)``. Exactly one of ``dx`` and ``grid`` must be supplied.
        bc: Boundary mode: neumann, reflecting, absorbing, or periodic.
        reference: Optional array whose namespace and floating dtype control the
            result.

    Returns:
        Dense ``(n, n)`` operator in the selected Array-API namespace.

    Raises:
        ValueError: If the geometry, coefficient, or boundary mode is invalid.
    """
    lower, main, upper, periodic_corner = _diffusion_diagonals(
        n,
        dx,
        grid=grid,
        bc=bc,
    )

    namespace_source: object = reference if reference is not None else coefficient
    try:
        _namespace_of(namespace_source)
    except TypeError:
        namespace_source = np.asarray(namespace_source)
    xp = _namespace_of(namespace_source)
    source_dtype = cast("Any", namespace_source).dtype
    dtype = xp.result_type(xp.asarray(0.0).dtype, source_dtype)
    coefficient_array = cast("Array", xp.asarray(coefficient, dtype=dtype))
    if coefficient_array.shape != ():
        msg = (
            "diffusion coefficient must be scalar; "
            f"got shape {coefficient_array.shape}."
        )
        raise ValueError(msg)
    if isinstance(coefficient_array, np.ndarray):
        if not np.isfinite(coefficient_array).all():
            msg = "diffusion coefficient must be finite."
            raise ValueError(msg)
        if bool(coefficient_array < 0.0):
            msg = "diffusion coefficient must be non-negative."
            raise ValueError(msg)

    stencil = np.diag(main) + np.diag(lower, k=-1) + np.diag(upper, k=1)
    if periodic_corner:
        stencil[0, -1] += periodic_corner
        stencil[-1, 0] += periodic_corner

    stencil_array = xp.asarray(stencil, dtype=dtype)
    return cast("Array", xp.multiply(stencil_array, coefficient_array))

build_identity_operator(n, *, dtype=np.float64, prefer_sparse=None)

Build an identity operator with dense/sparse autodispatch.

Parameters:

Name	Type	Description	Default
n	int	Size of the identity operator (n x n).	required
dtype	DTypeLike	Floating dtype (e.g. np.float64).	float64
prefer_sparse	bool | None	If True, always return a sparse operator; if False, always return a dense operator; if None, autodispatch based on n.	None

Returns:

Type	Description
Operator	Identity operator of shape (n, n) as either a dense ndarray or CSR matrix.

Source code in src/op_engine/matrix_ops.py

def build_identity_operator(
    n: int,
    *,
    dtype: DTypeLike = np.float64,
    prefer_sparse: bool | None = None,
) -> Operator:
    """
    Build an identity operator with dense/sparse autodispatch.

    Args:
        n: Size of the identity operator (n x n).
        dtype: Floating dtype (e.g. np.float64).
        prefer_sparse: If True, always return a sparse operator; if False,
            always return a dense operator; if None, autodispatch based on n.

    Returns:
        Identity operator of shape (n, n) as either a dense ndarray or CSR matrix.
    """
    dtype_obj = np.dtype(dtype)

    if prefer_sparse is True:
        return identity(n, format="csr", dtype=dtype_obj)

    if prefer_sparse is False:
        return cast("DenseOperator", np.eye(n, dtype=dtype_obj))

    # Autodispatch
    if n >= _DISPATCH_THRESHOLD:
        return identity(n, format="csr", dtype=dtype_obj)

    return cast("DenseOperator", np.eye(n, dtype=dtype_obj))

build_implicit_euler_operators(base_op, dt_scale)

Build implicit Euler operators for a time-scaled linear operator.

Parameters:

Name	Type	Description	Default
base_op	Operator	Base linear operator A.	required
dt_scale	float	Time-step scaling factor (dt * scale).	required

Returns:

Type	Description
tuple[Operator, Operator]	Tuple of (L, R) operators for implicit Euler scheme.

Raises:

Type	Description
ValueError	If dt_scale is not finite.

Source code in src/op_engine/matrix_ops.py

def build_implicit_euler_operators(
    base_op: Operator,
    dt_scale: float,
) -> tuple[Operator, Operator]:
    """Build implicit Euler operators for a time-scaled linear operator.

    Args:
        base_op: Base linear operator A.
        dt_scale: Time-step scaling factor (dt * scale).

    Returns:
        Tuple of (L, R) operators for implicit Euler scheme.

    Raises:
        ValueError: If dt_scale is not finite.
    """
    if not np.isfinite(dt_scale):
        raise ValueError(_OPERATOR_SCALE_ERROR.format(scale=dt_scale))

    n = base_op.shape[0]

    if issparse(base_op):
        base_csr = base_op.tocsr()
        identity_csr = identity(n, format="csr", dtype=base_csr.dtype)
        left_csr = (identity_csr - (dt_scale * base_csr)).tocsr()
        right_csr = identity_csr.tocsr()
        return left_csr, right_csr

    base_arr = np.asarray(base_op)
    identity_arr = np.eye(n, dtype=base_arr.dtype)
    left_arr = identity_arr - (dt_scale * base_arr)
    right_arr = identity_arr
    return cast("DenseOperator", left_arr), cast("DenseOperator", right_arr)

build_laplacian_tridiag(n, dx, coeff, dtype=np.float64, bc='neumann', *, grid=None)

Build a sparse 1D finite-volume Laplacian.

The resulting operator corresponds to coeff * Δ_h. Supply either a uniform cell width with dx or strictly increasing cell-center coordinates with grid. No time-step scaling is applied here.

Parameters:

Name	Type	Description	Default
n	int	Number of finite-volume cells.	required
dx	float | None	Positive uniform cell width, or None when grid is supplied.	required
coeff	float	Physical diffusion coefficient D (units length^2 / time).	required
dtype	DTypeLike	Floating dtype (e.g. np.float64).	float64
bc	str	Boundary condition; either "neumann" or "absorbing".	'neumann'
grid	ArrayLike | None	Optional strictly increasing cell-center coordinates of shape (n,). Exactly one of dx and grid must be supplied.	None

Returns:

Type	Description
csr_matrix	Sparse CSR matrix representing the Laplacian operator.

Raises:

Type	Description
ValueError	If the geometry or boundary condition is invalid.

Source code in src/op_engine/matrix_ops.py

def build_laplacian_tridiag(  # noqa: PLR0913
    n: int,
    dx: float | None,
    coeff: float,
    dtype: DTypeLike = np.float64,
    bc: str = "neumann",
    *,
    grid: ArrayLike | None = None,
) -> csr_matrix:
    """Build a sparse 1D finite-volume Laplacian.

    The resulting operator corresponds to ``coeff * Δ_h``. Supply either a
    uniform cell width with ``dx`` or strictly increasing cell-center
    coordinates with ``grid``. No time-step scaling is applied here.

    Args:
        n: Number of finite-volume cells.
        dx: Positive uniform cell width, or ``None`` when ``grid`` is supplied.
        coeff: Physical diffusion coefficient D (units length^2 / time).
        dtype: Floating dtype (e.g. np.float64).
        bc: Boundary condition; either ``"neumann"`` or ``"absorbing"``.
        grid: Optional strictly increasing cell-center coordinates of shape
            ``(n,)``. Exactly one of ``dx`` and ``grid`` must be supplied.

    Returns:
        Sparse CSR matrix representing the Laplacian operator.

    Raises:
        ValueError: If the geometry or boundary condition is invalid.
    """
    bc_normalized = str(bc).strip().lower()
    if bc_normalized not in {"neumann", "absorbing"}:
        msg = _UNKNOWN_BC_ERROR.format(bc=bc)
        raise ValueError(msg)

    dtype_obj = np.dtype(dtype)
    lower, main, upper, _periodic_corner = _diffusion_diagonals(
        n,
        dx,
        grid=grid,
        bc=bc_normalized,
        dtype=dtype_obj,
    )
    laplacian = diags(
        [lower.tolist(), main.tolist(), upper.tolist()],
        [-1, 0, 1],
        shape=(n, n),
        dtype=dtype_obj,
    )
    return (laplacian * coeff).tocsr()

build_predictor_corrector(base_matrix)

Build predictor-corrector matrices with dense/sparse autodispatch.

Parameters:

Name	Type	Description	Default
base_matrix	DenseOperator | csr_matrix	Base linear operator A.	required

Returns:

Type	Description
tuple[Operator, Operator, Operator]	Tuple of (predictor, L, R) operators for predictor-corrector scheme.

Source code in src/op_engine/matrix_ops.py

def build_predictor_corrector(
    base_matrix: DenseOperator | csr_matrix,
) -> tuple[Operator, Operator, Operator]:
    """
    Build predictor-corrector matrices with dense/sparse autodispatch.

    Args:
        base_matrix: Base linear operator A.

    Returns:
        Tuple of (predictor, L, R) operators for predictor-corrector scheme.
    """
    n = base_matrix.shape[0]
    if issparse(base_matrix) and n >= _DISPATCH_THRESHOLD:
        return _build_predictor_corrector_sparse(base_matrix)

    if issparse(base_matrix):
        dense_base = np.asarray(base_matrix.toarray())
    else:
        dense_base = np.asarray(base_matrix)

    return _build_predictor_corrector_dense(cast("DenseOperator", dense_base))

build_trapezoidal_operators(base_op, dt_scale)

Build trapezoidal operators for a time-scaled linear operator.

Parameters:

Name	Type	Description	Default
base_op	Operator	Base linear operator A.	required
dt_scale	float	Time-step scaling factor (dt * scale).	required

Returns:

Type	Description
tuple[Operator, Operator]	Tuple of (L, R) operators for trapezoidal scheme.

Raises:

Type	Description
ValueError	If dt_scale is not finite.

Source code in src/op_engine/matrix_ops.py

def build_trapezoidal_operators(
    base_op: Operator,
    dt_scale: float,
) -> tuple[Operator, Operator]:
    """Build trapezoidal operators for a time-scaled linear operator.

    Args:
        base_op: Base linear operator A.
        dt_scale: Time-step scaling factor (dt * scale).

    Returns:
        Tuple of (L, R) operators for trapezoidal scheme.

    Raises:
        ValueError: If dt_scale is not finite.
    """
    if not np.isfinite(dt_scale):
        raise ValueError(_OPERATOR_SCALE_ERROR.format(scale=dt_scale))

    n = base_op.shape[0]
    half = 0.5 * dt_scale

    if issparse(base_op):
        base_csr = base_op.tocsr()
        identity_mat = identity(n, format="csr", dtype=base_csr.dtype)
        left_csr = (identity_mat - (half * base_csr)).tocsr()
        right_csr = (identity_mat + (half * base_csr)).tocsr()
        return left_csr, right_csr

    base_arr = np.asarray(base_op)
    identity_arr = np.eye(n, dtype=base_arr.dtype)
    left_arr = identity_arr - (half * base_arr)
    right_arr = identity_arr + (half * base_arr)
    return cast("DenseOperator", left_arr), cast("DenseOperator", right_arr)

clear_implicit_solver_cache()

Clear the internal implicit solver cache.

Source code in src/op_engine/matrix_ops.py

def clear_implicit_solver_cache() -> None:
    """Clear the internal implicit solver cache."""
    _IMPLICIT_SOLVER_CACHE.clear()

encode_groups(group_ids, n_groups, *, prefer_sparse=None, dtype=np.float64)

Encode group IDs into a one-hot group membership matrix.

Parameters:

Name	Type	Description	Default
group_ids	NDArray[integer]	1D array of integer group IDs for each item.	required
n_groups	int	Total number of groups.	required
prefer_sparse	bool | None	If True, always return a sparse matrix; if False, always return a dense array; if None, autodispatch based on n_groups.	None
dtype	DTypeLike	Data type for the output matrix.	float64

Returns:

Type	Description
csr_matrix | DenseOperator	A (n_groups, n_items) one-hot encoded group membership matrix.

Source code in src/op_engine/matrix_ops.py

def encode_groups(
    group_ids: NDArray[np.integer],
    n_groups: int,
    *,
    prefer_sparse: bool | None = None,
    dtype: DTypeLike = np.float64,
) -> csr_matrix | DenseOperator:
    """
    Encode group IDs into a one-hot group membership matrix.

    Args:
        group_ids: 1D array of integer group IDs for each item.
        n_groups: Total number of groups.
        prefer_sparse: If True, always return a sparse matrix; if False, always
            return a dense array; if None, autodispatch based on n_groups.
        dtype: Data type for the output matrix.

    Returns:
        A (n_groups, n_items) one-hot encoded group membership matrix.
    """
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)

    if prefer_sparse is True:
        return _encode_sparse_groups(group_ids_arr, n_groups, dtype=dtype)
    if prefer_sparse is False:
        return _encode_dense_groups(group_ids_arr, n_groups, dtype=dtype)

    if n_groups >= _DISPATCH_THRESHOLD:
        return _encode_sparse_groups(group_ids_arr, n_groups, dtype=dtype)
    return _encode_dense_groups(group_ids_arr, n_groups, dtype=dtype)

grouped_count_ids(group_ids, n_groups)

Perform grouped count using group IDs.

Parameters:

Name	Type	Description	Default
group_ids	NDArray[integer]	1D array of integer group IDs.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 1D array of length n_groups where each element contains the count of
NDArray[floating]	occurrences of the corresponding group ID.

Source code in src/op_engine/matrix_ops.py

def grouped_count_ids(
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped count using group IDs.

    Args:
        group_ids: 1D array of integer group IDs.
        n_groups: Total number of groups.

    Returns:
        A 1D array of length n_groups where each element contains the count of
        occurrences of the corresponding group ID.
    """
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)
    counts = np.bincount(group_ids_arr, minlength=n_groups)
    return counts.astype(float)

grouped_sum_ids(values, group_ids, n_groups)

Perform grouped sum over 1D values array using group IDs.

Parameters:

Name	Type	Description	Default
values	NDArray[floating]	1D array of values.	required
group_ids	NDArray[integer]	1D array of integer group IDs.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 1D array of length n_groups where each element contains the sum of values
NDArray[floating]	for the corresponding group ID.

Source code in src/op_engine/matrix_ops.py

def grouped_sum_ids(
    values: NDArray[np.floating],
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped sum over 1D values array using group IDs.

    Args:
        values: 1D array of values.
        group_ids: 1D array of integer group IDs.
        n_groups: Total number of groups.

    Returns:
        A 1D array of length n_groups where each element contains the sum of values
        for the corresponding group ID.
    """
    values_arr = np.asarray(values)
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)
    sums = np.bincount(group_ids_arr, weights=values_arr, minlength=n_groups)
    return sums.astype(float)

grouped_sum_ids_2d(values, group_ids, n_groups)

Perform grouped sum over 2D values array using group IDs.

Parameters:

Name	Type	Description	Default
values	NDArray[floating]	2D (N, K) array where N is num of items and K is num of features.	required
group_ids	NDArray[integer]	1D array of integer group IDs of length N.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 2D array of shape (n_groups, K) where each row contains the sum of values
NDArray[floating]	for the corresponding group ID.

Raises:

Type	Description
ValueError	If values is not 2D or if group_ids length does not match the number of items in values.

Source code in src/op_engine/matrix_ops.py

def grouped_sum_ids_2d(
    values: NDArray[np.floating],
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped sum over 2D values array using group IDs.

    Args:
        values: 2D (N, K) array where N is num of items and K is num of features.
        group_ids: 1D array of integer group IDs of length N.
        n_groups: Total number of groups.

    Returns:
        A 2D array of shape (n_groups, K) where each row contains the sum of values
        for the corresponding group ID.

    Raises:
        ValueError: If values is not 2D or if group_ids length does not match
            the number of items in values.
    """
    values_arr = np.asarray(values)
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)

    if values_arr.ndim != 2:
        raise ValueError(_VALUES_2D_ERROR)

    n_items, n_features = values_arr.shape
    if group_ids_arr.shape[0] != n_items:
        raise ValueError(_GROUP_IDS_LENGTH_ERROR)

    out = np.zeros((n_groups, n_features), dtype=values_arr.dtype)
    for feature_idx in range(n_features):
        out[:, feature_idx] = np.bincount(
            group_ids_arr,
            weights=values_arr[:, feature_idx],
            minlength=n_groups,
        )
    return out

implicit_solve(left_op, right_op, x)

Perform an implicit solve with dense fallback and cached sparse dispatch.

Parameters:

Name	Type	Description	Default
left_op	object	Left operator L in the equation L @ y = R @ x.	required
right_op	object	Right operator R in the equation L @ y = R @ x.	required
x	Array	1D or 2D array representing the input vector(s).	required

Returns:

Type	Description
Array	A 1D or 2D array containing the solution vector(s) y.

Source code in src/op_engine/matrix_ops.py

def implicit_solve(
    left_op: object,
    right_op: object,
    x: Array,
) -> Array:
    """
    Perform an implicit solve with dense fallback and cached sparse dispatch.

    Args:
        left_op: Left operator L in the equation L @ y = R @ x.
        right_op: Right operator R in the equation L @ y = R @ x.
        x: 1D or 2D array representing the input vector(s).

    Returns:
        A 1D or 2D array containing the solution vector(s) y.
    """
    _validate_solve_dimensions(left_op, right_op, x)

    adapter = _find_sparse_adapter(left_op, right_op)
    if adapter is not None:
        return _sparse_implicit_solve(adapter, left_op, right_op, x)

    xp = _namespace_of(x)
    left_dense = _dense_operator_array(left_op, xp=xp, dtype=x.dtype)
    right_dense = _dense_operator_array(right_op, xp=xp, dtype=x.dtype)
    rhs = cast("Array", xp.matmul(right_dense, x))
    return cast("Array", xp.linalg.solve(left_dense, rhs))

kron_prod(a, b)

Compute the Kronecker product of two operators.

Parameters:

Name	Type	Description	Default
a	Operator	First operator.	required
b	Operator	Second operator.	required

Returns:

Type	Description
Operator	The Kronecker product operator.

Source code in src/op_engine/matrix_ops.py

def kron_prod(a: Operator, b: Operator) -> Operator:
    """
    Compute the Kronecker product of two operators.

    Args:
        a: First operator.
        b: Second operator.

    Returns:
        The Kronecker product operator.
    """
    if issparse(a) or issparse(b):
        a_csr = a if issparse(a) else csr_matrix(np.asarray(a))
        b_csr = b if issparse(b) else csr_matrix(np.asarray(b))
        return kron(
            a_csr,
            b_csr,
            format="csr",
        )
    return cast("DenseOperator", np.kron(np.asarray(a), np.asarray(b)))

kron_sum(ops)

Compute a Kronecker sum of square operators.

Parameters:

Name	Type	Description	Default
ops	list[Operator]	List of 2D square operators.	required

Returns:

Type	Description
Operator	The Kronecker sum operator.

Raises:

Type	Description
ValueError	If ops is empty or if operators are not square or have incompatible shapes.

Source code in src/op_engine/matrix_ops.py

def kron_sum(ops: list[Operator]) -> Operator:
    """
    Compute a Kronecker sum of square operators.

    Args:
        ops: List of 2D square operators.

    Returns:
        The Kronecker sum operator.

    Raises:
        ValueError: If ops is empty or if operators are not square or
            have incompatible shapes.
    """
    if not ops:
        raise ValueError(_KRON_EMPTY_ERROR)

    shapes = [
        tuple(np.asarray(op).shape) if not issparse(op) else op.shape for op in ops
    ]

    if any(s[0] != s[1] for s in shapes):
        raise ValueError(_KRON_INCOMPATIBLE_ERROR.format(shapes=shapes))

    any_sparse = any(issparse(op) for op in ops)
    sizes = [s[0] for s in shapes]
    dtype_obj = np.result_type(*[
        (op.dtype if issparse(op) else np.asarray(op).dtype) for op in ops
    ])

    def _eye(n: int) -> Operator:
        if any_sparse:
            return identity(n, format="csr", dtype=dtype_obj)
        return cast("DenseOperator", np.eye(n, dtype=dtype_obj))

    total: Operator | None = None
    n_ops = len(ops)

    for i, op_i in enumerate(ops):
        term: Operator = op_i
        for j in range(i - 1, -1, -1):
            term = kron_prod(_eye(sizes[j]), term)
        for j in range(i + 1, n_ops):
            term = kron_prod(term, _eye(sizes[j]))
        total = term if total is None else cast("Operator", total + term)

    if total is None:
        raise ValueError(_KRON_EMPTY_ERROR)
    return total

make_constant_base_builder(operator)

Convenience: wrap a constant operator as a BaseOperatorBuilder.

Parameters:

Name	Type	Description	Default
operator	Operator	Constant operator to wrap.	required

Returns:

Type	Description
BaseOperatorBuilder	A BaseOperatorBuilder that always returns the given operator.

Source code in src/op_engine/matrix_ops.py

def make_constant_base_builder(operator: Operator) -> BaseOperatorBuilder:
    """
    Convenience: wrap a constant operator as a BaseOperatorBuilder.

    Args:
        operator: Constant operator to wrap.

    Returns:
        A BaseOperatorBuilder that always returns the given operator.
    """
    operator_0 = _ensure_operator_type(operator)

    def _builder(ctx: StageOperatorContext) -> Operator:  # noqa: ARG001
        return operator_0

    return _builder

make_stage_operator_factory(base_builder, *, scheme='implicit-euler')

Create a stage operator factory supporting time/state dependent base ops.

Parameters:

Name	Type	Description	Default
base_builder	BaseOperatorBuilder	Function that builds a base operator given stage context.	required
scheme	str	Implicit scheme; either "implicit-euler" or "trapezoidal".	'implicit-euler'

Returns:

Type	Description
StageOperatorFactory	A StageOperatorFactory that builds (L, R) operators for the given scheme.

Raises:

Type	Description
ValueError	If an unknown scheme is provided.

Source code in src/op_engine/matrix_ops.py

def make_stage_operator_factory(
    base_builder: BaseOperatorBuilder,
    *,
    scheme: str = "implicit-euler",
) -> StageOperatorFactory:
    """
    Create a stage operator factory supporting time/state dependent base ops.

    Args:
        base_builder: Function that builds a base operator given stage context.
        scheme: Implicit scheme; either "implicit-euler" or "trapezoidal".

    Returns:
        A StageOperatorFactory that builds (L, R) operators for the given scheme.

    Raises:
        ValueError: If an unknown scheme is provided.
    """
    scheme_norm = str(scheme).strip().lower()

    if scheme_norm == "implicit-euler":

        def _factory(
            dt: float, scale: float, ctx: StageOperatorContext
        ) -> tuple[Operator, Operator]:
            operator = _ensure_operator_type(base_builder(ctx))
            return build_implicit_euler_operators(operator, float(dt) * float(scale))

        return _factory

    if scheme_norm == "trapezoidal":

        def _factory(
            dt: float, scale: float, ctx: StageOperatorContext
        ) -> tuple[Operator, Operator]:
            operator = _ensure_operator_type(base_builder(ctx))
            return build_trapezoidal_operators(operator, float(dt) * float(scale))

        return _factory

    raise ValueError(_UNKNOWN_SCHEME_ERROR.format(scheme=scheme))

matrix_grouped_count(group_matrix)

Perform grouped count using a group matrix.

Parameters:

Name	Type	Description	Default
group_matrix	csr_matrix | DenseOperator	2D group matrix (csr_matrix or dense ndarray).	required

Returns:

Type	Description
NDArray[floating]	A 1D array containing the counts for each group.

Source code in src/op_engine/matrix_ops.py

def matrix_grouped_count(
    group_matrix: csr_matrix | DenseOperator,
) -> NDArray[np.floating]:
    """
    Perform grouped count using a group matrix.

    Args:
        group_matrix: 2D group matrix (csr_matrix or dense ndarray).

    Returns:
        A 1D array containing the counts for each group.
    """
    n_groups = group_matrix.shape[0]
    if issparse(group_matrix) and n_groups >= _DISPATCH_THRESHOLD:
        return _matrix_grouped_count_sparse(group_matrix)
    if issparse(group_matrix):
        dense_matrix = group_matrix.toarray()
    else:
        dense_matrix = np.asarray(group_matrix)
    return _matrix_grouped_count_dense(cast("DenseOperator", dense_matrix))

matrix_grouped_sum(group_matrix, values)

Perform grouped sum using a group matrix.

Parameters:

Name	Type	Description	Default
group_matrix	csr_matrix | DenseOperator	2D group matrix (csr_matrix or dense ndarray).	required
values	NDArray[floating]	1D array of values to be summed.	required

Returns:

Type	Description
NDArray[floating]	A 1D array containing the grouped sums.

Source code in src/op_engine/matrix_ops.py

def matrix_grouped_sum(
    group_matrix: csr_matrix | DenseOperator,
    values: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Perform grouped sum using a group matrix.

    Args:
        group_matrix: 2D group matrix (csr_matrix or dense ndarray).
        values: 1D array of values to be summed.

    Returns:
        A 1D array containing the grouped sums.
    """
    n_groups = group_matrix.shape[0]
    if issparse(group_matrix) and n_groups >= _DISPATCH_THRESHOLD:
        return _matrix_grouped_sum_sparse(group_matrix, values)
    if issparse(group_matrix):
        dense_matrix = np.asarray(group_matrix.toarray())
    else:
        dense_matrix = np.asarray(group_matrix)
    return _matrix_grouped_sum_dense(cast("DenseOperator", dense_matrix), values)

matrix_masked_sum(mask_matrix, data)

Perform masked sum using a mask matrix and data array.

Parameters:

Name	Type	Description	Default
mask_matrix	csr_matrix | DenseOperator	2D mask matrix (csr_matrix or dense ndarray).	required
data	NDArray[floating]	1D or 2D data array to be masked and summed.	required

Returns:

Type	Description
NDArray[floating]	A 1D or 2D array containing the masked sums.

Source code in src/op_engine/matrix_ops.py

def matrix_masked_sum(
    mask_matrix: csr_matrix | DenseOperator,
    data: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Perform masked sum using a mask matrix and data array.

    Args:
        mask_matrix: 2D mask matrix (csr_matrix or dense ndarray).
        data: 1D or 2D data array to be masked and summed.

    Returns:
        A 1D or 2D array containing the masked sums.
    """
    n_masks = mask_matrix.shape[0]
    if issparse(mask_matrix) and n_masks >= _DISPATCH_THRESHOLD:
        return _matrix_masked_sum_sparse(mask_matrix, data)
    if issparse(mask_matrix):
        dense_matrix = np.asarray(mask_matrix.toarray())
    else:
        dense_matrix = np.asarray(mask_matrix)
    return _matrix_masked_sum_dense(dense_matrix, data)

smooth(x, alpha=0.02, out=None)

Apply simple smoothing along the last axis.

Parameters:

Name	Type	Description	Default
x	NDArray[floating]	Input array to smooth.	required
alpha	float	Smoothing factor between 0 and 1.	0.02
out	NDArray[floating] | None	Optional output array to store the result.	None

Returns:

Type	Description
NDArray[floating]	Smoothed array with the same shape as x.

Source code in src/op_engine/matrix_ops.py

def smooth(
    x: NDArray[np.floating],
    alpha: float = 0.02,
    out: NDArray[np.floating] | None = None,
) -> NDArray[np.floating]:
    """
    Apply simple smoothing along the last axis.

    Args:
        x: Input array to smooth.
        alpha: Smoothing factor between 0 and 1.
        out: Optional output array to store the result.

    Returns:
        Smoothed array with the same shape as x.
    """
    x_arr = np.asarray(x)
    smoothed = (1.0 - alpha) * x_arr + alpha * x_arr.mean(axis=-1, keepdims=True)
    if out is not None:
        np.copyto(out, smoothed)
        return out
    return cast("NDArray[np.floating]", smoothed)