electrostatics.solver — the linear solve

tempura/electrostatics/solver.py is the sparse linear-algebra core. It is deliberately independent of geometry: given the finalized system matrix and the region indices, it builds the right-hand sides, applies the Dirichlet constraints, and factors the system once with MUMPS for reuse across many solves.

The discrete system

pescado discretizes the Poisson problem into a sparse linear system

\[ A\,x = b, \]

where \(A\) is the finalized capacitance_matrix and \(x\) is the potential at the mesh sites. For a fixed device \(A\) is built once; different gate voltages change only the right-hand side \(b\). That single observation is what makes the gate-basis solve cheap (see _pescado_solve).

Symmetry-preserving Dirichlet elimination

Gate and boundary nodes are Dirichlet-constrained: their potentials are prescribed, not solved for. Naively striking those rows would break the symmetry of \(A\) and prevent a symmetric factorization. Instead Tempura rewrites the system into an equivalent symmetric form. For the Dirichlet index set \(D\) with prescribed values \(x_D\):

\[ \tilde A_{ij} = \begin{cases} \delta_{ij}, & i \in D \ \text{or}\ j \in D,\\[2pt] A_{ij}, & \text{otherwise,} \end{cases} \qquad \tilde b_i = \begin{cases} x_i, & i \in D,\\[2pt] b_i - \displaystyle\sum_{j \in D} A_{ij}\,x_j, & i \notin D. \end{cases} \]

In words: zero the Dirichlet rows and columns, put \(1\) on their diagonal, and move the removed coupling \(\sum_{j \in D} A_{ij} x_j\) into the right-hand side. The solution of \(\tilde A\,x = \tilde b\) matches the constrained problem while keeping \(\tilde A\) symmetric.

Factor once, solve many

setup_solver(...) factors \(\tilde A\) with MUMPS (optionally with block low-rank compression), and solve_linear_system(...) reuses that factorization for each right-hand side. build_rhs_vectors(...) assembles the multi-column RHS for a set of gates: without a distributed charge term each column is a sparse indicator on one gate's nodes; with a charge term the interior rows carry the charge and the result is dense. This is the machinery the gate-basis solve drives in blocks.

Optional dependency

The solve path requires python-mumps. The helpers raise a clear error if it is unavailable.

API

solver

Linear-system helpers for Tempura electrostatics.

This module is intentionally independent from Tempura's geometry assembly. It contains the sparse constrained-system helpers used for the Poisson basis solves. Self-consistent Thomas-Fermi workflows should use Pescado directly with the finalized problem returned by Tempura's public electrostatics workflow.

build_rhs_vectors(region_inds, gate_names, *, size=None, charge=None, dtype=DTYPE)

Build the multi-RHS matrix for the ordered gate basis solves.

Parameters:

Name	Type	Description	Default
region_inds	IndexMap	Mapping of region name to solver indices.	required
gate_names	list[str]	Gate names in the solve order to encode into the RHS.	required
size	int | None	Optional total row count. When omitted, it is inferred from the largest region index.	None
charge	ndarray | None	Optional distributed charge term added to the interior rows.	None
dtype	DTypeLike	Numeric dtype used for the returned matrix.	DTYPE

Returns:

Type	Description
csc_matrix | ndarray	Sparse CSC gate-basis RHS when charge is omitted, otherwise a dense
csc_matrix | ndarray	(n_rows, n_gates) array containing both the gate basis and charge
csc_matrix | ndarray	terms.

Source code in tempura/electrostatics/solver.py

def build_rhs_vectors(
    region_inds: IndexMap,
    gate_names: list[str],
    *,
    size: int | None = None,
    charge: np.ndarray | None = None,
    dtype: DTypeLike = DTYPE,
) -> sp.csc_matrix | np.ndarray:
    """Build the multi-RHS matrix for the ordered gate basis solves.

    Args:
        region_inds: Mapping of region name to solver indices.
        gate_names: Gate names in the solve order to encode into the RHS.
        size: Optional total row count. When omitted, it is inferred from the
            largest region index.
        charge: Optional distributed charge term added to the interior rows.
        dtype: Numeric dtype used for the returned matrix.

    Returns:
        Sparse CSC gate-basis RHS when ``charge`` is omitted, otherwise a dense
        ``(n_rows, n_gates)`` array containing both the gate basis and charge
        terms.
    """
    gate_inds = _ordered_gate_inds(region_inds, gate_names)
    n_gates = len(gate_inds)
    n_rows = _region_system_size(region_inds, gate_inds) if size is None else size
    if n_gates == 0:
        return sp.csc_matrix((n_rows, 0), dtype=np.dtype(dtype))

    gate_sizes = np.fromiter((inds.size for inds in gate_inds), dtype=int, count=n_gates)
    gate_row_inds = (
        np.concatenate(gate_inds) if gate_sizes.sum() else np.empty(0, dtype=int)
    )
    gate_col_inds = np.repeat(np.arange(n_gates), gate_sizes)

    # Without distributed charge, each RHS column is just a sparse gate basis
    # vector, so keep the representation sparse as long as possible.
    if charge is None:
        return sp.csc_matrix(
            (
                np.ones(gate_row_inds.size, dtype=np.dtype(dtype)),
                (gate_row_inds, gate_col_inds),
            ),
            shape=(n_rows, n_gates),
        )

    charge_matrix = np.asarray(charge, dtype=np.dtype(dtype))
    if charge_matrix.ndim == 1:
        charge_matrix = charge_matrix[:, None]
    if charge_matrix.shape[0] != n_rows:
        raise ValueError(f"charge must have {n_rows} rows; got {charge_matrix.shape[0]}")
    if charge_matrix.shape[1] == 1 and n_gates != 1:
        charge_matrix = np.repeat(charge_matrix, n_gates, axis=1)
    elif charge_matrix.shape[1] != n_gates:
        raise ValueError(
            f"charge must have {n_gates} columns; got {charge_matrix.shape[1]}"
        )

    # Charge terms populate interior rows for every solve, so the dense path is
    # simpler than building a mixed sparse/dense block structure.
    rhs_matrix = np.zeros((n_rows, n_gates), dtype=np.dtype(dtype), order="F")
    interior_inds = np.asarray(region_inds["Interior"], dtype=int)
    if interior_inds.size:
        rhs_matrix[interior_inds, :] = charge_matrix[interior_inds, :]
    if gate_row_inds.size:
        rhs_matrix[gate_row_inds, gate_col_inds] = 1.0
    return rhs_matrix

setup_solver(matrix, *, dtype=DTYPE, blr=BLR, eps_blr=EPS_BLR, symmetric=False)

Create and factor a MUMPS solver for a sparse linear system.

Parameters:

Name	Type	Description	Default
matrix	Any	Real-valued system matrix to factor.	required
dtype	DTypeLike	Real numeric dtype used to cast matrix before factorization. Tempura's electrostatics solve path does not preserve imaginary components.	DTYPE
blr	bool	Whether to enable MUMPS block low-rank factorization.	BLR
eps_blr	float	BLR compression threshold used when blr is enabled.	EPS_BLR
symmetric	bool	Whether matrix should be factorized as symmetric.	False

Returns:

Type	Description
Any	Factored MUMPS context that can be reused across solves.

Raises:

Type	Description
RuntimeError	If python-mumps is not available.

Source code in tempura/electrostatics/solver.py

def setup_solver(
    matrix: Any,
    *,
    dtype: DTypeLike = DTYPE,
    blr: bool = BLR,
    eps_blr: float = EPS_BLR,
    symmetric: bool = False,
) -> Any:
    """Create and factor a MUMPS solver for a sparse linear system.

    Args:
        matrix: Real-valued system matrix to factor.
        dtype: Real numeric dtype used to cast ``matrix`` before factorization.
            Tempura's electrostatics solve path does not preserve imaginary
            components.
        blr: Whether to enable MUMPS block low-rank factorization.
        eps_blr: BLR compression threshold used when ``blr`` is enabled.
        symmetric: Whether ``matrix`` should be factorized as symmetric.

    Returns:
        Factored MUMPS context that can be reused across solves.

    Raises:
        RuntimeError: If ``python-mumps`` is not available.
    """
    solver = _require_mumps().Context()
    _factorize_solver(
        solver,
        matrix,
        dtype=dtype,
        blr=blr,
        eps_blr=eps_blr,
        symmetric=symmetric,
    )
    return solver

solve_linear_system(matrix, rhs, *, solver=None, dtype=DTYPE, blr=BLR, eps_blr=EPS_BLR, symmetric=False)

Solve matrix x = rhs with a factored MUMPS solver.

Parameters:

Name	Type	Description	Default
matrix	Any	Real-valued system matrix to solve.	required
rhs	Any	Dense or sparse right-hand side array.	required
solver	Any | None	Optional existing factored MUMPS context to reuse.	None
dtype	DTypeLike	Real numeric dtype used to cast the matrix and RHS. Tempura's electrostatics solve path does not preserve imaginary components.	DTYPE
blr	bool	Whether to enable MUMPS block low-rank factorization.	BLR
eps_blr	float	BLR compression threshold used when blr is enabled.	EPS_BLR
symmetric	bool	Whether matrix should be solved as symmetric.	False

Returns:

Type	Description
tuple[ndarray, Any]	Tuple of (solution, solver).

Raises:

Type	Description
RuntimeError	If python-mumps is not available and solver is not provided.

Source code in tempura/electrostatics/solver.py

def solve_linear_system(
    matrix: Any,
    rhs: Any,
    *,
    solver: Any | None = None,
    dtype: DTypeLike = DTYPE,
    blr: bool = BLR,
    eps_blr: float = EPS_BLR,
    symmetric: bool = False,
) -> tuple[np.ndarray, Any]:
    """Solve ``matrix x = rhs`` with a factored MUMPS solver.

    Args:
        matrix: Real-valued system matrix to solve.
        rhs: Dense or sparse right-hand side array.
        solver: Optional existing factored MUMPS context to reuse.
        dtype: Real numeric dtype used to cast the matrix and RHS. Tempura's
            electrostatics solve path does not preserve imaginary components.
        blr: Whether to enable MUMPS block low-rank factorization.
        eps_blr: BLR compression threshold used when ``blr`` is enabled.
        symmetric: Whether ``matrix`` should be solved as symmetric.

    Returns:
        Tuple of ``(solution, solver)``.

    Raises:
        RuntimeError: If ``python-mumps`` is not available and ``solver`` is
            not provided.
    """
    if solver is None:
        solver = _require_mumps().Context()

    _factorize_solver(
        solver,
        matrix,
        dtype=dtype,
        blr=blr,
        eps_blr=eps_blr,
        symmetric=symmetric,
    )

    if sp.issparse(rhs):
        rhs_fortran: Any = sp.csc_matrix(rhs, dtype=np.dtype(dtype))
    else:
        rhs_fortran = np.asfortranarray(np.asarray(rhs, dtype=np.dtype(dtype)))
    solution = np.asarray(solver.solve(rhs_fortran))
    if np.issubdtype(np.dtype(dtype), np.complexfloating):
        solution = solution.real
    return solution, solver