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., Laplacian, Crank-Nicolson).
• 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

• CPU-first: dense paths rely on NumPy/SciPy BLAS/LAPACK; sparse paths rely on SciPy sparse factorizations.
• Backend-friendly surface: public APIs operate on plain ndarrays or CSR matrices and avoid leaking SciPy-specific solver objects.
• Cache semantics: implicit solver caching is keyed by (id(left_op), id(right_op)). 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".

GridGeometry(n, dx) dataclass

Geometry of a 1D spatial grid.

Attributes:

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

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	NDArray[floating]	Stage state as a 1D float array (flattened along solver operator axis).
stage	StageName	Optional stage label (e.g. "tr", "bdf2").
extra	Any | None	Optional extra payload for future use (kept generic).

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_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

Raises:

Type	Description
ValueError	If dt_scale is not finite.

Returns:

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

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).

    Raises:
        ValueError: If dt_scale is not finite.

    Returns:
        Tuple of (L, R) operators for implicit Euler scheme.
    """
    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')

Build a Laplacian tridiagonal matrix for a given boundary condition.

The resulting operator corresponds to coeff * Δ_h, where Δ_h is the standard second-order central-difference Laplacian in 1D. No time-step scaling is applied here; coeff is interpreted as the physical diffusion coefficient D or a generic spatial scaling.

Parameters:

Name	Type	Description	Default
n	int	Number of grid points.	required
dx	float	Grid spacing.	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'

Raises:

Type	Description
ValueError	If an unknown boundary condition is provided.

Returns:

Type	Description
csr_matrix	Sparse CSR matrix representing the Laplacian operator.

Source code in src/op_engine/matrix_ops.py

def build_laplacian_tridiag(
    n: int,
    dx: float,
    coeff: float,
    dtype: DTypeLike = np.float64,
    bc: str = "neumann",
) -> csr_matrix:
    """Build a Laplacian tridiagonal matrix for a given boundary condition.

    The resulting operator corresponds to `coeff * Δ_h`, where `Δ_h` is the
    standard second-order central-difference Laplacian in 1D. No time-step
    scaling is applied here; `coeff` is interpreted as the physical diffusion
    coefficient `D` or a generic spatial scaling.

    Args:
        n: Number of grid points.
        dx: Grid spacing.
        coeff: Physical diffusion coefficient D (units length^2 / time).
        dtype: Floating dtype (e.g. np.float64).
        bc: Boundary condition; either "neumann" or "absorbing".

    Raises:
        ValueError: If an unknown boundary condition is provided.

    Returns:
        Sparse CSR matrix representing the Laplacian operator.
    """
    dtype_obj = np.dtype(dtype)
    factor = coeff / dx**2

    main_diag = -2.0 * np.ones(n, dtype=dtype_obj)
    off_diag = np.ones(n - 1, dtype=dtype_obj)

    if bc == "neumann":
        main_diag[0] = -1.0
        main_diag[-1] = -1.0
    elif bc == "absorbing":
        main_diag[0] = -2.0
        main_diag[-1] = -2.0
    else:
        msg = _UNKNOWN_BC_ERROR.format(bc=bc)
        raise ValueError(msg)

    laplacian = diags(
        [off_diag.tolist(), main_diag.tolist(), off_diag.tolist()],
        [-1, 0, 1],
        shape=(n, n),
        dtype=dtype_obj,
    )

    scaled = laplacian * factor
    return scaled.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

Raises:

Type	Description
ValueError	If dt_scale is not finite.

Returns:

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

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).

    Raises:
        ValueError: If dt_scale is not finite.

    Returns:
        Tuple of (L, R) operators for trapezoidal scheme.
    """
    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 cast("csr_matrix", left_csr), cast("csr_matrix", 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

Raises:

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

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.

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.

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

    Returns:
        A 2D array of shape (n_groups, K) where each row 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)

    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/sparse dispatch and caching.

Parameters:

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

Returns:

Type	Description
NDArray[floating]	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: Operator,
    right_op: Operator,
    x: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Perform an implicit solve with dense/sparse dispatch and caching.

    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.
    """
    x_arr = np.asarray(x)
    _validate_solve_dimensions(left_op, right_op, cast("NDArray[np.floating]", x_arr))

    key = (id(left_op), id(right_op))
    meta = _operator_meta(left_op, right_op)

    cached = _IMPLICIT_SOLVER_CACHE.get(key)
    if cached is not None:
        cached_meta, solver = cached
        if cached_meta != meta:
            solver = _build_implicit_solver(left_op, right_op)
            _IMPLICIT_SOLVER_CACHE[key] = (meta, solver)
        return solver(cast("NDArray[np.floating]", x_arr))

    solver = _build_implicit_solver(left_op, right_op)
    _IMPLICIT_SOLVER_CACHE[key] = (meta, solver)
    return solver(cast("NDArray[np.floating]", x_arr))

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

Raises:

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

Returns:

Type	Description
Operator	The Kronecker sum operator.

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.

    Raises:
        ValueError: If ops is empty or if operators are not square or
            have incompatible shapes.

    Returns:
        The Kronecker sum operator.
    """
    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'

Raises:

Type	Description
ValueError	If an unknown scheme is provided.

Returns:

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

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".

    Raises:
        ValueError: If an unknown scheme is provided.

    Returns:
        A StageOperatorFactory that builds (L, R) operators for the given scheme.
    """
    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)