Structured direct solver API

Block-tridiagonal

Structured direct solvers: block-tridiagonal Schur elimination (block Thomas).

For a system with blocks L_k x_{k-1} + D_k x_k + U_k x_{k+1} = b_k (k = 0..N-1), a Schur-complement sweep from the last block down to the first,

Delta_{N-1} = D_{N-1} Delta_k = D_k - U_k Delta_{k+1}^{-1} L_{k+1} sigma_k = b_k - U_k Delta_{k+1}^{-1} sigma_{k+1}

followed by substitution upward from block 0:

x_0 = Delta_0^{-1} sigma_0 x_k = Delta_k^{-1} (sigma_k - L_k x_{k-1})

Each step costs one dense LU solve plus one matrix product — never an explicit inverse. Cost is O(N m^3) time; the factor/solve split lets several right-hand sides share one elimination. block_thomas_truncated additionally exploits a common kinetic-equation structure: when the right-hand side vanishes for k >= K and only the lowest K blocks of the solution are needed (e.g. velocity moments touching only the first few spectral modes), the upward substitution can stop at block K and the downward sweep needs no storage above it, so peak memory is O(K m^2), independent of N.

Stability note: block LU without pivoting is guaranteed stable only for block-diagonally-dominant systems (Demmel, Higham & Schreiber, Numer. Linear Algebra Appl. 2, 173 (1995)); each block here is factored with partial pivoting, and callers should monitor conditioning in weakly-dominant regimes (see solvax.refine for iterative-refinement fallbacks).

References

      1. Golub & C. F. Van Loan, Matrix Computations, 4th ed., section 4.5.

  • F. J. Escoto, PhD thesis (2025), https://arxiv.org/abs/2510.27513 — block-tridiagonal elimination over Legendre modes for kinetic equations, including the truncated-storage observation.

class solvax.direct.BlockTridiagFactors(delta_lu, delta_piv, lower, upper)

Reusable elimination state from block_thomas_factor().

Variables:

  • delta_lu (jax.Array) – LU factors of every Schur complement Delta_k, shape (n_blocks, m, m).

  • delta_piv (jax.Array) – matching pivot indices, shape (n_blocks, m).

  • lower (jax.Array) – sub-diagonal blocks L_k (lower[0] unused).

  • upper (jax.Array) – super-diagonal blocks U_k (upper[-1] unused).

Parameters:
delta_lu: Array

Alias for field number 0

delta_piv: Array

Alias for field number 1

lower: Array

Alias for field number 2

upper: Array

Alias for field number 3

solvax.direct.block_thomas_factor(lower, diag, upper, factor_dtype=None)

Run the downward Schur sweep once, for reuse across right-hand sides.

Parameters:

  • lower (Array) – sub-diagonal blocks L_k, shape (n_blocks, m, m); lower[0] is ignored.

  • diag (Array) – diagonal blocks D_k, shape (n_blocks, m, m).

  • upper (Array) – super-diagonal blocks U_k, shape (n_blocks, m, m); upper[-1] is ignored.

  • factor_dtype – if given, the Schur-complement LU factorizations and their triangular solves run in this lower precision, while the block products U_k Delta^{-1} L_k and the stored off-diagonal bands stay in the working precision of diag. The returned delta_lu is then low precision; pair with block_thomas_solve() under iterative_refinement() (see mixed_precision_block_thomas()) to recover working-precision accuracy on fast low-precision hardware. Practically this is jnp.float32: lu_factor dispatches to LAPACK/cuSOLVER getrf, which has float32 and float64 kernels but no half precision, so bfloat16/float16 raise NotImplementedError.

Returns:

Factors for block_thomas_solve().

Return type:

BlockTridiagFactors

solvax.direct.block_thomas_factor_fn(block_fn, n_blocks, factor_dtype=None)

Factor generated block rows once for reusable primal/transpose solves.

Unlike block_thomas_factor(), this entry point never materializes the diagonal band. block_fn is evaluated exactly once per block index; the returned state stores Schur LU factors and the two off-diagonal bands needed by block_thomas_solve().

Parameters:

  • block_fn (Callable[[Array], tuple[Array, Array, Array]]) – maps a traced int32 index to (lower, diagonal, upper) blocks of identical square shape.

  • n_blocks (int) – static positive number of block rows.

  • factor_dtype – optional lower precision for Schur LU factorizations, with the same contract as block_thomas_factor().

Returns:

Reusable factors accepted by block_thomas_solve(), including its exact transpose=True path.

Return type:

BlockTridiagFactors

solvax.direct.block_thomas_solve(factors, rhs, transpose=False)

Solve using precomputed factors.

With transpose=True this solves the transposed system A^T x = rhs reusing the same factors: for the same elimination order the Schur complements of A^T are exactly Delta_k^T (inductively, Delta'_{N-1} = D_{N-1}^T and Delta'_k = D_k^T - L_{k+1}^T Delta_{k+1}^{-T} U_k^T = Delta_k^T), so the stored LU factors serve both directions via trans=1 triangular solves. The off-diagonal blocks swap roles and transpose: the downward sweep uses L_{k+1}^T where the forward solve used U_k, and the upward substitution uses U_{k-1}^T where it used L_k. One elimination thus covers the forward and the adjoint solve — exactly what implicit differentiation needs.

Parameters:
Returns:

Solution with the same shape as rhs.

Return type:

Array

solvax.direct.block_thomas(lower, diag, upper, rhs)

Solve a block-tridiagonal system by Schur-complement elimination.

Convenience wrapper: block_thomas_factor() then block_thomas_solve(). For repeated solves with the same matrix, call the two stages directly and reuse the factors.

Parameters:

  • lower (Array) – L_k blocks, (n_blocks, m, m); lower[0] ignored.

  • diag (Array) – D_k blocks, (n_blocks, m, m).

  • upper (Array) – U_k blocks, (n_blocks, m, m); upper[-1] ignored.

  • rhs (Array) – (n_blocks, m) or (n_blocks, m, n_rhs).

Returns:

x with the same shape as rhs.

Return type:

Array

solvax.direct.block_tridiag_matvec(lower, diag, upper, x)

Apply a block-tridiagonal operator without forming a dense matrix.

(A x)_k = L_k x_{k-1} + D_k x_k + U_k x_{k+1}, evaluated for every block at once. x and the result share the layout of the right-hand side, (n_blocks, m) or (n_blocks, m, n_rhs). This independent operator action is also used by residual diagnostics.

Parameters:
Return type:

Array

solvax.direct.block_tridiag_relative_residual(lower, diag, upper, solution, rhs)

Return ||b - A x||_2 / max(||b||_2, tiny) per right-hand side.

A vector right-hand side returns a scalar. Multiple right-hand sides return one value per final column. Every block row is included, so this diagnostic cannot silently omit a truncated high-mode tail.

Parameters:
Return type:

Array

solvax.direct.mixed_precision_block_thomas(lower, diag, upper, rhs, *, factor_dtype=<class 'jax.numpy.float32'>, refine_steps=2, implicit_adjoint=False)

Block-tridiagonal solve with a low-precision factorization + refinement.

Factors once with block_thomas_factor() in factor_dtype — the dense Schur-complement LU factorizations, the dominant cost of the sweep, then run at (e.g.) float32 throughput, up to 32x that of float64 on consumer GPUs — and recovers working-precision accuracy with refine_steps sweeps of solvax.refine.iterative_refinement(). Each sweep forms the residual with the working-precision operator and corrects it with one low-precision block_thomas_solve(), so the result matches the full-precision solve to working-precision accuracy whenever kappa(A) * u_low < 1 (Carson & Higham 2018; see solvax.refine).

This composes the existing factor/solve with iterative refinement — no parallel scan — and stays jit/vmap/grad-transparent like the rest of the module.

Parameters:

  • lower (Array) – L_k blocks, (n_blocks, m, m); lower[0] ignored.

  • diag (Array) – D_k blocks, (n_blocks, m, m).

  • upper (Array) – U_k blocks, (n_blocks, m, m); upper[-1] ignored.

  • rhs (Array) – (n_blocks, m) or (n_blocks, m, n_rhs).

  • factor_dtype – precision for the LU factorizations (default float32, the fast low precision supported by the LAPACK/cuSOLVER getrf backend; half precision is not — see block_thomas_factor()).

  • refine_steps (int) – number of refinement sweeps (static int); 0 returns the bare low-precision solve.

  • implicit_adjoint (bool) – if True, reverse mode uses a custom VJP that solves the adjoint system A^T lam = cotangent by the same refinement reusing the transposed low-precision factors — zero additional factorizations, no differentiation through the factorization, and the gradient inherits the refined working-precision forward error rather than the factorization precision. If False (default), reverse mode differentiates through the refinement loop directly.

Returns:

x with the same shape and precision as rhs.

Return type:

Array

solvax.direct.block_thomas_truncated(lower, diag, upper, rhs_low, keep_lowest, *, adjoint_window=None)

Block-tridiagonal solve returning only the lowest keep_lowest blocks.

Requires the right-hand side to vanish for k >= keep_lowest (rhs_low holds the nonzero head). The downward Schur sweep runs over all blocks but stores nothing above keep_lowest; the upward substitution stops there. Peak memory O(keep_lowest * m^2), independent of n_blocks.

Parameters:

  • lower (Array) – as in block_thomas().

  • diag (Array) – as in block_thomas().

  • upper (Array) – as in block_thomas().

  • rhs_low (Array) – nonzero head of the right-hand side, shape (keep_lowest, m) or (keep_lowest, m, n_rhs).

  • keep_lowest (int) – static number of solution blocks to compute (1 <= keep_lowest <= n_blocks; equality recovers the full solve).

  • adjoint_window (int | None) – if None (default), reverse mode differentiates the elimination directly, taping the sweep at O(n_blocks * m^2). If an integer w, a structure-preserving custom VJP is used: the right-hand-side gradient is the exact transposed truncated solve and the band gradients come from a leading (keep_lowest + w)-block re-solve, so the differentiated solve also runs at O((keep_lowest + w) m^2) memory, independent of n_blocks. Band-gradient error decays geometrically in w for block diagonally dominant systems; w >= n_blocks reproduces the exact gradient.

Returns:

The lowest keep_lowest solution blocks, same layout as rhs_low.

Return type:

Array

solvax.direct.block_thomas_truncated_fn(block_fn, n_blocks, rhs_low, keep_lowest, *, params=None, adjoint_window=None)

Truncated block-tridiagonal solve with on-the-fly block assembly.

The full band arrays are never materialized, and each block index is assembled once. See block_thomas_truncated_fn_with_residual() when an algebraic residual of the retained Schur system is also required.

By default block_fn maps an index to blocks, k -> (L_k, D_k, U_k), and reverse mode differentiates through the generated sweeps (taping them at O(n_blocks m^2)). Passing params (any pytree) switches block_fn to the explicit form (params, k) -> (L_k, D_k, U_k) and selects a structure-preserving custom VJP governed by adjoint_window: the right-hand-side gradient is an exactly generated truncated solve of A^T (three block assemblies per index), and the params gradient pulls the windowed band cotangents back through block_fn’s own VJP at each of the leading keep_lowest + adjoint_window indices. Forward and reverse then both run at memory independent of n_blocks — the band arrays are never materialized in either direction.

Parameters:

  • block_fn (Callable[[...], tuple[Array, Array, Array]]) – k -> (L_k, D_k, U_k), or (params, k) -> ... when params is given. L_0 and U_{n_blocks-1} are ignored.

  • n_blocks (int) – static total number of blocks (>= 1).

  • rhs_low (Array) – nonzero head of the right-hand side, shape (keep_lowest, m) or (keep_lowest, m, n_rhs).

  • keep_lowest (int) – static number of solution blocks to compute.

  • params – optional differentiable parameters consumed by block_fn.

  • adjoint_window (int | None) – window w for the params gradient (required with params); band-gradient error decays as O(rho^{2w}) for block diagonally dominant systems, exactly as for block_thomas_truncated().

Returns:

The lowest keep_lowest solution blocks, same layout as rhs_low.

Return type:

Array

solvax.direct.block_thomas_truncated_fn_with_residual(block_fn, n_blocks, rhs_low, keep_lowest, *, residual_rhs_index=None)

Return the generated truncated solution and its Schur-system RMS residual.

The residual is evaluated from the retained pivoted LU factors as L @ (U @ x) - P @ rhs. It therefore includes the eliminated high-mode tail without reconstructing another solution block or materializing the original diagonal band.

Parameters:
Return type:

tuple[Array, Array]

Banded and periodic-banded

Banded LU solvers: non-pivoted factorization and periodic (circulant-banded) systems.

Storage convention (identical to scipy.linalg.solve_banded()): a matrix A with lower_bw sub-diagonals and upper_bw super-diagonals is held in a bands array of shape (n_diags, n) with n_diags = lower_bw + upper_bw + 1, where row r holds the diagonal with offset upper_bw - r:

bands[upper_bw + i - j, j] = A[i, j] for max(0, j - upper_bw) <= i <= min(n - 1, j + lower_bw)

Entries of bands outside that range are ignored. The LU factorization is Doolittle elimination without pivoting, carried out column by column in banded storage with a jax.lax.scan (static shapes, so everything is jit/vmap/grad-transparent — XLA handles row pivoting poorly, and avoiding it is the point of this module). Two safeguards substitute for pivoting:

  • row equilibration: each row is pre-scaled by 1 / max|row|;

  • static pivoting: any pivot with |pivot| < floor is clamped to sign(pivot) * floor, and the number of clamps is recorded in the factors so callers can detect near-singularity and fall back to iterative refinement (see solvax.refine).

LU without pivoting is guaranteed backward-stable only for diagonally dominant (or block-dominant) systems (Demmel, Higham & Schreiber, Numer. Linear Algebra Appl. 2, 173 (1995)); with equilibration and static pivoting it is a robust practical choice for the advection-dominated periodic 1-D operators these routines target, but callers should monitor the clamp counter in weakly-dominant regimes.

Periodic (circulant-banded) systems A = B + U V^T — a banded core B plus wrap-around corner blocks expressed as a low-rank update — are solved with the Sherman-Morrison-Woodbury capacitance-matrix identity

(B + U V^T)^{-1} = B^{-1} - B^{-1} U (I + V^T B^{-1} U)^{-1} V^T B^{-1},

where the small dense capacitance matrix I + V^T B^{-1} U is LU-factored once (with partial pivoting — it is tiny) and B^{-1} U is precomputed with the banded factorization, so each periodic solve costs one banded solve plus O(bw) dense work.

References

  • G. H. Golub & C. F. Van Loan, Matrix Computations, 4th ed., section 4.3 (banded Gaussian elimination) and section 2.1.4 (Sherman-Morrison-Woodbury).

  • J. W. Demmel, N. J. Higham & R. S. Schreiber, “Stability of block LU factorization”, Numer. Linear Algebra Appl. 2, 173 (1995) — stability caveats for elimination without pivoting.

solvax.banded.banded_matvec(bands, lower_bw, upper_bw, x)

Matrix-vector product A @ x with A in banded storage.

Vectorized over the matrix dimension (the only Python loop is over the n_diags diagonals, which is static).

Parameters:

  • bands (Array) – banded storage of A, shape (lower_bw + upper_bw + 1, n) (see module docstring); out-of-range entries are ignored.

  • lower_bw (int) – number of sub-diagonals.

  • upper_bw (int) – number of super-diagonals.

  • x (Array) – vector (n,) or block of vectors (n, k).

Returns:

A @ x with the same shape as x.

Return type:

Array

class solvax.banded.BandedLUFactors(l_bands, u_bands, row_scale, n_clamped)

Packed non-pivoted banded LU from lu_factor_banded().

The bandwidths are recovered from the array shapes, so the factors stay jit/vmap-transparent (no static metadata crosses transform boundaries).

Variables:

  • l_bands (jax.Array) – unit-lower-triangular multipliers, shape (lower_bw, n); l_bands[t - 1, k] holds L[k + t, k].

  • u_bands (jax.Array) – upper factor in banded storage, shape (upper_bw + 1, n); u_bands[upper_bw + i - j, j] holds U[i, j].

  • row_scale (jax.Array) – equilibration scales applied to the rows of A (all ones when equilibrate=False), shape (n,).

  • n_clamped (jax.Array) – int32 count of pivots clamped by static pivoting; a nonzero value signals near-singularity of the (equilibrated) core.

Parameters:
l_bands: Array

Alias for field number 0

u_bands: Array

Alias for field number 1

row_scale: Array

Alias for field number 2

n_clamped: Array

Alias for field number 3

solvax.banded.lu_factor_banded(bands, lower_bw, upper_bw, *, equilibrate=True, static_pivot_floor=None)

Non-pivoted (Doolittle) LU of a banded matrix, in banded storage.

Columns are eliminated left to right with a jax.lax.scan; the carry holds the multipliers of the previous upper_bw steps, so shapes are static and the factorization is jit/vmap/grad-transparent. Row equilibration and static pivoting (see module docstring) substitute for the row pivoting that XLA handles poorly.

Parameters:

  • bands (Array) – banded storage of A, shape (lower_bw + upper_bw + 1, n); out-of-range entries are ignored.

  • lower_bw (int) – number of sub-diagonals.

  • upper_bw (int) – number of super-diagonals.

  • equilibrate (bool) – scale each row by 1 / max|row| before factoring (the scales are stored and applied to the right-hand side by lu_solve_banded()).

  • static_pivot_floor (float | None) – clamp threshold for pivots; None (default) uses sqrt(machine eps) * max|bands| of the (equilibrated) matrix.

Returns:

Factors for lu_solve_banded(), including the clamp counter.

Return type:

BandedLUFactors

solvax.banded.lu_solve_banded(factors, b)

Solve A x = b using precomputed banded LU factors.

Forward substitution with the unit-lower multipliers, then backward substitution with the banded upper factor, each as a jax.lax.scan carrying a bandwidth-sized window of the solution.

Parameters:
Returns:

Solution with the same shape as b.

Return type:

Array

class solvax.banded.PeriodicBandedLUFactors(core, z_ul, z_lr, cap_lu, cap_piv)

Woodbury factors for a periodic banded system, from lu_factor_banded_periodic().

Variables:

  • core (solvax.banded.BandedLUFactors) – banded LU of the (non-periodic) banded core B.

  • z_ul (jax.Array) – B^{-1} U columns generated by the top-right corner block, shape (n, bw_ul).

  • z_lr (jax.Array) – B^{-1} U columns generated by the bottom-left corner block, shape (n, bw_lr).

  • cap_lu (jax.Array) – dense LU of the capacitance matrix I + V^T B^{-1} U, shape (bw_ul + bw_lr, bw_ul + bw_lr).

  • cap_piv (jax.Array) – matching pivot indices.

Parameters:
core: BandedLUFactors

Alias for field number 0

z_ul: Array

Alias for field number 1

z_lr: Array

Alias for field number 2

cap_lu: Array

Alias for field number 3

cap_piv: Array

Alias for field number 4

solvax.banded.lu_factor_banded_periodic(bands, lower_bw, upper_bw, corner_ul, corner_lr, *, equilibrate=True, static_pivot_floor=None)

Factor a periodic (circulant-banded) matrix via the capacitance method.

The matrix is A = B + U V^T: a banded core B (given in bands) plus the periodic wrap-around corners as a rank-(bw_ul + bw_lr) update, where U carries the corner blocks and V selects the coupled columns. B is factored with the non-pivoted banded LU, Z = B^{-1} U is precomputed, and the small dense capacitance matrix I + V^T Z is LU-factored so lu_solve_banded_periodic() can apply the Woodbury identity (module docstring) at the cost of one banded solve.

Parameters:

  • bands (Array) – banded storage of the core B, shape (lower_bw + upper_bw + 1, n).

  • lower_bw (int) – number of sub-diagonals of B.

  • upper_bw (int) – number of super-diagonals of B.

  • corner_ul (Array) – top-right corner block A[:bw, n - bw:] coupling the first rows to the last columns, shape (bw, bw).

  • corner_lr (Array) – bottom-left corner block A[n - bw:, :bw] coupling the last rows to the first columns, shape (bw, bw).

  • equilibrate (bool) – passed to lu_factor_banded() for the core.

  • static_pivot_floor (float | None) – passed to lu_factor_banded() for the core.

Returns:

Factors for lu_solve_banded_periodic().

Return type:

PeriodicBandedLUFactors

solvax.banded.lu_solve_banded_periodic(factors, b)

Solve a periodic banded system using precomputed Woodbury factors.

Applies x = B^{-1} b - Z (I + V^T Z)^{-1} V^T B^{-1} b with Z = B^{-1} U from lu_factor_banded_periodic().

Parameters:
Returns:

Solution with the same shape as b.

Return type:

Array

Tridiagonal

Backend-aware batched tridiagonal solve (Thomas / cuSPARSE fast paths).

A scalar tridiagonal system lower[j] x[j-1] + diag[j] x[j] + upper[j] x[j+1] = rhs[j] is the banded special case lower_bw = upper_bw = 1 of solvax.banded, but it is common and structured enough to deserve a dedicated, hardware-aware fast path — this module specializes it. The system lives on the leading axis; every trailing axis of rhs (extra columns, stacked fields, batch dimensions) is solved simultaneously in one call, which is exactly the layout that maps a stack of independent tridiagonal systems onto the vendor batched kernels without an outer jax.vmap().

Two backends, selected per lowering platform at trace time:

  • Thomas (method="thomas"): the classic two-sweep elimination

    c’[0] = upper[0] / diag[0], d’[0] = rhs[0] / diag[0], c’[j] = upper[j] / (diag[j] - lower[j] c’[j-1]), d’[j] = (rhs[j] - lower[j] d’[j-1]) / (diag[j] - lower[j] c’[j-1]), x[n-1] = d’[n-1], x[j] = d’[j] - c’[j] x[j+1],

    implemented as two jax.lax.scan sweeps over the radial axis with a eps = 1e-12 guard on vanishing pivots. Fully jit/vmap/grad-transparent and, because the arithmetic is fixed, bitwise reproducible — the CPU path.

  • Fused (method="lax"): XLA’s batched jax.lax.linalg.tridiagonal_solve() (cuSPARSE gtsv2 on CUDA, LAPACK gtsv on CPU). On a GPU the n sequential scan steps of Thomas serialize into n kernel launches (latency-bound, independent of how many columns ride along), so the single fused kernel is dramatically faster there. Numerically equivalent to Thomas (same solution to roundoff) but not bit-identical.

method="auto" (default) uses jax.lax.platform_dependent() to pick Thomas when the code lowers for CPU (bit parity, honouring a JAX_PLATFORMS=cpu / jax.default_device() pin even on an accelerator host) and the fused kernel everywhere else. Systems with fewer than three rows always use Thomas — the cuSPARSE kernel requires n >= 3.

This solver is a drop-in preconditioner core: hand lambda r: tridiagonal_solve(lower, diag, upper, r) to solvax.precond. coarse_operator() or use it as a line solve inside solvax.precond.line_smoother().

References

  • L. H. Thomas, Elliptic Problems in Linear Difference Equations over a Network, Watson Sci. Comput. Lab. Report, Columbia University (1949) — the tridiagonal elimination algorithm.

  • G. H. Golub & C. F. Van Loan, Matrix Computations, 4th ed., section 4.3 (tridiagonal / banded Gaussian elimination).

  • NVIDIA cuSPARSE gtsv2 batched tridiagonal solver, exposed in JAX as jax.lax.linalg.tridiagonal_solve().

  • W. H. Press et al., Numerical Recipes, 3rd ed., section 2.7.3 — cyclic tridiagonal systems as a rank-one correction.

solvax.tridiagonal.tridiagonal_solve(lower, diag, upper, rhs, *, method='auto')

Solve a batched scalar tridiagonal system along the leading axis.

Solves lower[j] x[j-1] + diag[j] x[j] + upper[j] x[j+1] = rhs[j] for j = 0 .. n-1 (with lower[0] and upper[-1] ignored), for every trailing column of rhs at once. The band arrays lower, diag, upper share the system shape (leading axis n plus any batch columns); rhs may carry extra trailing axes beyond that shape — stacked right-hand sides / fields — which are all solved in a single pass. Real bands with a complex rhs use independent real and imaginary solves, preserving real storage and fused accelerator kernels. Genuinely complex bands use the portable Thomas implementation.

Parameters:

  • lower (Array) – sub-diagonal, couples row j to j-1; lower[0] ignored. Shape broadcastable to diag.

  • diag (Array) – main diagonal; its shape (n, *columns) defines the system extent.

  • upper (Array) – super-diagonal, couples row j to j+1; upper[-1] ignored. Shape broadcastable to diag.

  • rhs (Array) – right-hand side, shape diag.shape or diag.shape + fields.

  • method (str) –

    backend selection.

    Complex coefficient matrices use Thomas for every method because the fused JAX primitive is real-only on supported JAX releases.

Returns:

The solution with the same shape as rhs.

Raises:

ValueError – if method is not one of "auto", "thomas", "lax".

Return type:

Array

solvax.tridiagonal.cyclic_tridiagonal_solve(lower, diag, upper, rhs, *, method='auto')

Solve periodic tridiagonal systems along the leading axis.

lower[0] couples row zero to the last unknown and upper[-1] couples the last row to unknown zero. A Sherman–Morrison correction turns the periodic system into one ordinary tridiagonal solve with two stacked right-hand sides, retaining the backend selected by tridiagonal_solve().

The system arrays have shape (n, *columns). rhs may append extra trailing field axes, which are solved simultaneously.

Parameters:
Return type:

Array