linalg

triwarp.linalg

Sparse linear-algebra helpers shared by the parametrization and smoothing solvers.

Two layers:

• min_quad_with_fixed and its two steps (free_partition, assemble_interior_system) port igl::min_quad_with_fixed: minimize a quadratic form with some degrees of freedom pinned, by eliminating the fixed rows/columns into a reduced symmetric system.
• solve_spd_columns and spd_column_solver solve one symmetric positive-definite operator against several right-hand-side columns in a single conjugate-gradient call.

Why batch the columns. Solving a k-column system as a Python loop of k independent cg calls costs the sum of the per-column iteration counts, and every CG iteration then launches k separate sets of reduction and AXPY kernels. replicated_operator instead presents the same operator as k independent subproblems over one flat k * n vector: the solver then advances all columns together and stops on the worst-case residual, so the cost becomes max of the per-column iteration counts with one set of vector kernels per iteration. The sparse matrix is never replicated in memory — the matvec issues k bsr_mv calls against the single operator.

Whose conjugate gradient. One column goes to warp.optim.linear.cg; more than one goes to this module's own _BatchedCg, which runs the same iteration and the same stopping rule. The split is not a preference — it is the one input for which Warp's reduction degrades: batching is expressed to Warp as batch_offsets, and that is exactly what makes its dot-product reduction take a per-column path whose cost grows with the vector length rather than staying flat. Reducing per column with a real two-stage tree keeps the batching and drops that cost; see _BatchedCg for the mechanism.

Determinism. Build each operator natively at its final dtype in a single warp.sparse.bsr_from_triplets and never recast or rebuild it. A rebuild re-sorts and duplicate-accumulates an order the CSR already has, which is wasted work; and if the rebuild's triplet buffers are ever sized off BsrMatrix.nnz (a stale capacity, not the true entry count — see nnz_sync()) the buffers' tail reaches bsr_from_triplets uninitialized. Nothing here recasts an operator, and Q_uu is assembled as a CSR directly, without any triplet build, so it carries an exact nnz.

Why Jacobi. Every solve here preconditions with warp.optim.linear.preconditioner(A, "diag"). On a cotangent Laplacian a preconditioner costing k mat-vecs per iteration cuts the iteration count by only about sqrt(k), so total work scales as k / sqrt(k) = sqrt(k) — single-level preconditioning loses on this operator class, and only a multilevel method escapes it. IC(0) would help if it were available, but Warp has no sparse triangular solve — warp.sparse exposes only bsr_from_triplets and warp.optim.linear only the Krylov methods and a diagonal preconditioner — and no parallel substitute for a triangular solve keeps its advantage: a fully parallel Jacobi-sweep approximation, an exact apply parallelized by graph coloring (whose uncoalesced access and extra launches eat most of its own iteration win), and natural-ordering level sets (whose level count swings wildly with the input's vertex numbering) are all worse than Jacobi in practice. Chebyshev as a single-level preconditioner is a further alternative and is also worse; Chebyshev as the multigrid V-cycle's smoother is a separate question — see _MULTIGRID_SWEEPS — and neither decline covers the other.

Two further obstacles are specific to this repository. Obtuse triangles give negative cotangent weights (triwarp/kernels/laplacian.py), so -L is often not the M-matrix that IC(0) existence requires; and both heat_geodesic and heat_signed_distance solve a -L with a genuine constant null space, where IC(0) hits a zero pivot on the last row of every connected component.

The multilevel option is multigrid_preconditioner: smoothed aggregation, the one scheme that breaks the conjugate-gradient iteration count's growth with problem size rather than paying it down by a constant factor. It is not the default, and the reason is the setup rather than the cycle: building the hierarchy (one aggregation, one power iteration, a bsr_transposed and three bsr_mm per level) has a real fixed cost, so the V-cycle wins exactly where the solve it replaces is long enough to amortize that setup, and loses on a well-conditioned or already-fast-converging system. See CG_MULTIGRID_DOMINANCE for how the gate decides which systems clear that bar.

A GPU sparse direct solver (cuDSS through nvmath-python and CuPy) was also considered and declined for the same reason the multilevel preconditioner is gated: its cost is a host-side symbolic factorization plan that is flat regardless of how well-conditioned the system is, so it wins exactly on the systems the multigrid gate already routes to a hierarchy and loses everywhere CG converges quickly. It would also add an optional CUDA-only dependency, and a direct solver is singular on the empty rows CG tolerates (unreferenced free vertices, which a caller would have to pin and restore). One idea from that investigation is worth keeping in mind for a future factor- once, solve-many caller: reusing a factorization plan across solves of one sparsity pattern is comparatively cheap, so a caller that resolves the same operator repeatedly (as arap already does with its own preconditioner) is the shape that would benefit.

Nothing cheap predicts in advance which side of the multigrid-vs-Jacobi line a system falls on, which is why preconditioner="auto" uses a capped Jacobi probe rather than a heuristic predictor of the iteration count; see CG_PROBE_ITERATIONS.

The aggregation keeps a strength-of-connection threshold (_MULTIGRID_THETA) rather than every off-diagonal, because on a graded (anisotropic) patch keeping every off-diagonal aggregates across the weak direction and the hierarchy converges far more slowly.

Routing through "auto" decides from the operator, not the caller, and the axis that matters is conditioning: only a system whose off-diagonal dominance clears CG_MULTIGRID_DOMINANCE benefits from a hierarchy. A plain Laplacian (tutte, min_quad_with_fixed on a raw cotangent matrix) and lscm's coupled u/v system sit below that bar and a forced hierarchy regresses them, where a squared operator (harmonic at k >= 2) sits comfortably above it. So harmonic passes "auto" at k >= 2 and "diag" below, and min_quad_with_fixed keeps "diag" as its default rather than becoming a second "auto" caller.

harmonic at k=2 on a strongly graded patch may not converge under Jacobi

A k=2 biharmonic operator on a strongly graded patch can be outside what Jacobi-preconditioned conjugate gradient reaches in float64 at all: the iteration can hit its cap and return a residual several orders of magnitude above the requested tolerance, with nothing but a UserWarning to say so, producing a visibly wrong UV map. A caller who needs that combination should pass a stronger preconditioner and check the warning.

That heat system M - tL heat_geodesic solves needs no multilevel help of its own: it converges in a small, size-independent number of iterations, because t = h**2 makes it a small perturbation of the mass matrix.

CG_TOLERANCE module-attribute

CG_TOLERANCE = 1e-10

CG_MAXITER_FACTOR module-attribute

CG_MAXITER_FACTOR = 10

CG_CHECK_EVERY module-attribute

CG_CHECK_EVERY = 0

CG_CHECK_EVERY_FALLBACK module-attribute

CG_CHECK_EVERY_FALLBACK = 10

CG_PROBE_ITERATIONS module-attribute

CG_PROBE_ITERATIONS = 2000

CG_MULTIGRID_MIN_UNKNOWNS module-attribute

CG_MULTIGRID_MIN_UNKNOWNS = 1500

CG_MULTIGRID_LARGE_UNKNOWNS module-attribute

CG_MULTIGRID_LARGE_UNKNOWNS = 7000

CG_MULTIGRID_SIZE_FLOOR module-attribute

CG_MULTIGRID_SIZE_FLOOR = 1.9

CG_MULTIGRID_DOMINANCE module-attribute

CG_MULTIGRID_DOMINANCE = 2.15

min_quad_with_fixed

min_quad_with_fixed(
    q: BsrMatrix[float64],
    fixed_mask: array[bool],
    fixed_values: Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> tuple[twt.Array2dFloat, wp.array[wp.int32], int]

Minimize a quadratic form with pinned degrees of freedom.

Eliminates the fixed degrees of freedom from 0.5 x' Q x and solves the reduced symmetric system Q_uu x_u = -Q_ub bc for every right-hand-side column at once. Handles both shapes the parametrization solvers need: several independent columns over n_vertices unknowns (harmonic / tutte, one column per UV coordinate) and a single coupled column over 2 * n_vertices unknowns (lscm).

Parameters:

Name	Type	Description	Default
q	BsrMatrix[float64]	(n_dofs, n_dofs) symmetric positive-semi-definite operator, float64.	required
fixed_mask	array[bool]	Length-n_dofs mask: True marks a pinned degree of freedom.	required
fixed_values	Array2dFloat	(n_rhs, n_dofs) prescribed values; only the entries where fixed_mask is True are read.	required
tol	float	Relative residual tolerance for the conjugate-gradient solve.	CG_TOLERANCE
check_every	int	Iterations between residual tests, forwarded to solve_spd_columns. The default tests on device every iteration; see there for the tradeoff. This entry point's own return type does not depend on it — the solver's (iterations, residual, atol) triple is not surfaced here.	CG_CHECK_EVERY
preconditioner	str	Forwarded to solve_spd_columns. "diag" (the default) because the eliminated operator is usually a plain Laplacian, whose rows nearly sum to zero and which a V-cycle does not help. The exception is a caller that squares the operator, and harmonic passes "auto" at k >= 2 for exactly that reason.	'diag'

Returns:

Name	Type	Description
solution	Array2dFloat	(n_rhs, n_free) solved values for the unpinned degrees of freedom.
free_map	array[int32]	Length-n_dofs compact remap of the unpinned degrees of freedom.
n_free	int	Number of unpinned degrees of freedom.

Raises:

Type	Description
RuntimeError	If fixed_mask and fixed_values are not all on one device.

Notes

The elimination is igl::min_quad_with_fixed's, without its Aeq linear-equality block -- no caller here needs one.

When every degree of freedom is pinned (n_free == 0) the prescribed values are the whole answer: an empty solution is returned without a solve, so this case also works on the CPU. Reconstructing the full field from solution / free_map is left to the caller's scatter.

See Also

free_partition assemble_interior_system solve_spd_columns

Source code in triwarp/linalg.py

def min_quad_with_fixed(
    q: wps.BsrMatrix[wp.float64],
    fixed_mask: wp.array[wp.bool],
    fixed_values: twt.Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> tuple[twt.Array2dFloat, wp.array[wp.int32], int]:
    """
    Minimize a quadratic form with pinned degrees of freedom.

    Eliminates the fixed degrees of freedom from ``0.5 x' Q x`` and solves the reduced symmetric
    system ``Q_uu x_u = -Q_ub bc`` for every right-hand-side column at once. Handles both shapes the
    parametrization solvers need: several *independent* columns over ``n_vertices`` unknowns
    (``harmonic`` / ``tutte``, one column per UV coordinate) and a single *coupled* column over
    ``2 * n_vertices`` unknowns (``lscm``).

    Parameters
    ----------
    q
        ``(n_dofs, n_dofs)`` symmetric positive-semi-definite operator, ``float64``.
    fixed_mask
        Length-``n_dofs`` mask: ``True`` marks a pinned degree of freedom.
    fixed_values
        ``(n_rhs, n_dofs)`` prescribed values; only the entries where ``fixed_mask`` is ``True`` are
        read.
    tol
        Relative residual tolerance for the conjugate-gradient solve.
    check_every
        Iterations between residual tests, forwarded to
        [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]. The default tests on device every
        iteration; see there for the tradeoff. This entry point's own return type does not depend
        on it — the solver's ``(iterations, residual, atol)`` triple is not surfaced here.
    preconditioner
        Forwarded to [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]. ``"diag"`` (the
        default) because the eliminated operator is usually a plain Laplacian, whose rows nearly sum
        to zero and which a V-cycle does not help. The exception is a caller that squares the
        operator, and [`harmonic`][triwarp.parametrization.harmonic] passes ``"auto"`` at ``k >= 2``
        for exactly that reason.

    Returns
    -------
    solution : twt.Array2dFloat
        ``(n_rhs, n_free)`` solved values for the unpinned degrees of freedom.
    free_map : wp.array[wp.int32]
        Length-``n_dofs`` compact remap of the unpinned degrees of freedom.
    n_free : int
        Number of unpinned degrees of freedom.

    Raises
    ------
    RuntimeError
        If ``fixed_mask`` and ``fixed_values`` are not all on one device.

    Notes
    -----
    The elimination is ``igl::min_quad_with_fixed``'s, without its ``Aeq`` linear-equality block --
    no caller here needs one.

    When every degree of freedom is pinned (``n_free == 0``) the prescribed values are the whole
    answer: an empty solution is returned without a solve, so this case also works on the CPU.
    Reconstructing the full field from ``solution`` / ``free_map`` is left to the caller's scatter.

    See Also
    --------
    [`free_partition`][triwarp.linalg.free_partition]
    [`assemble_interior_system`][triwarp.linalg.assemble_interior_system]
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]
    """
    require_same_device(fixed_mask=fixed_mask, fixed_values=fixed_values)
    device = fixed_mask.device
    n_rhs = int(fixed_values.shape[0])
    free_map, n_free = free_partition(fixed_mask)

    solution = wp.zeros((n_rhs, n_free), dtype=wp.float64, device=device)
    if n_free == 0:
        return twt.as_array2d(solution, wp.float64), free_map, n_free

    q_uu, rhs = assemble_interior_system(q, fixed_mask, free_map, fixed_values, n_free)
    solve_spd_columns(
        q_uu,
        rhs,
        twt.as_array2d(solution, wp.float64),
        tol=tol,
        check_every=check_every,
        preconditioner=preconditioner,
    )
    return twt.as_array2d(solution, wp.float64), free_map, n_free

free_partition

free_partition(
    fixed_mask: array[bool],
) -> tuple[wp.array[wp.int32], int]

Compact remap of the unpinned degrees of freedom, plus their count.

Thin inversion of mask_to_compact_ranks: fixed_mask marks the constrained degrees of freedom, and the returned map indexes the reduced system built over the complement.

Parameters:

Name	Type	Description	Default
fixed_mask	array[bool]	Length-n_dofs mask: True marks a pinned degree of freedom.	required

Returns:

Name	Type	Description
free_map	array[int32]	Length-n_dofs; free_map[i] is the reduced-system index of degree of freedom i, meaningful only where fixed_mask[i] is False.
n_free	int	Number of unpinned degrees of freedom.

See Also

mask_to_compact_ranks

Source code in triwarp/linalg.py

def free_partition(fixed_mask: wp.array[wp.bool]) -> tuple[wp.array[wp.int32], int]:
    """
    Compact remap of the *unpinned* degrees of freedom, plus their count.

    Thin inversion of [`mask_to_compact_ranks`][triwarp.array.mask_to_compact_ranks]:
    ``fixed_mask`` marks the constrained degrees of freedom, and the returned map indexes the
    reduced system built over the complement.

    Parameters
    ----------
    fixed_mask
        Length-``n_dofs`` mask: ``True`` marks a pinned degree of freedom.

    Returns
    -------
    free_map : wp.array[wp.int32]
        Length-``n_dofs``; ``free_map[i]`` is the reduced-system index of degree of freedom ``i``,
        meaningful only where ``fixed_mask[i]`` is ``False``.
    n_free : int
        Number of unpinned degrees of freedom.

    See Also
    --------
    [`mask_to_compact_ranks`][triwarp.array.mask_to_compact_ranks]
    """
    return tw.array.mask_to_compact_ranks(fixed_mask, invert=True)

assemble_interior_system

assemble_interior_system(
    q: BsrMatrix[float64],
    fixed_mask: array[bool],
    free_map: array[int32],
    fixed_values: Array2dFloat,
    n_free: int,
) -> tuple[wps.BsrMatrix[wp.float64], twt.Array2dFloat]

Extract the free-free block Q_uu and the constant right-hand side -Q_ub bc.

A CSR-to-CSR extraction in two launches over the n_dofs rows of q: the first counts each free row's surviving entries (and accumulates the pinned-column contributions into the n_rhs right-hand-side rows, which needs the same row scan), a prefix sum turns those counts into row offsets, and the second fills columns / values in place, remapped through free_map. The result is handed to a warp.sparse.BsrMatrix directly.

Why not bsr_from_triplets

Because the input is already a CSR. Emitting q.nnz triplets and letting bsr_from_triplets sort and duplicate-accumulate them re-derives an order the input already had, at real cost for a large operator. Row order comes from the launch index and column order from q's own rows through the monotone free_map, so the extracted matrix is sorted by construction.

It also avoids a second, quieter cost: bsr_from_triplets leaves nnz at the triplet count, which is q.nnz — an upper bound on the true entry count of a lightly-pinned Q_uu — so every downstream bsr_mv would be dimensioned for the unreduced matrix. The count here is exact.

Parameters:

Name	Type	Description	Default
q	BsrMatrix[float64]	(n_dofs, n_dofs) symmetric positive-semi-definite operator, float64.	required
fixed_mask	array[bool]	Length-n_dofs mask: True marks a pinned degree of freedom.	required
free_map	array[int32]	Compact remap from free_partition.	required
fixed_values	Array2dFloat	(n_rhs, n_dofs) prescribed values.	required
n_free	int	Number of unpinned degrees of freedom.	required

Returns:

Name	Type	Description
q_uu	``warp.sparse.BsrMatrix``	(n_free, n_free) reduced operator.
rhs	Array2dFloat	(n_rhs, n_free) constant right-hand side. arap adds its per-iteration rotation term to this buffer rather than reassembling it.

Raises:

Type	Description
RuntimeError	If fixed_mask, free_map and fixed_values are not all on one device.

See Also

min_quad_with_fixed

Source code in triwarp/linalg.py

def assemble_interior_system(
    q: wps.BsrMatrix[wp.float64],
    fixed_mask: wp.array[wp.bool],
    free_map: wp.array[wp.int32],
    fixed_values: twt.Array2dFloat,
    n_free: int,
) -> tuple[wps.BsrMatrix[wp.float64], twt.Array2dFloat]:
    """
    Extract the free-free block ``Q_uu`` and the constant right-hand side ``-Q_ub bc``.

    A **CSR-to-CSR** extraction in two launches over the ``n_dofs`` rows of ``q``: the first counts
    each free row's surviving entries (and accumulates the pinned-column contributions into the
    ``n_rhs`` right-hand-side rows, which needs the same row scan), a prefix sum turns those counts
    into row offsets, and the second fills ``columns`` / ``values`` in place, remapped through
    ``free_map``. The result is handed to a ``warp.sparse.BsrMatrix`` directly.

    !!! note "Why not ``bsr_from_triplets``"
        Because the input is *already* a CSR. Emitting ``q.nnz`` triplets and letting
        ``bsr_from_triplets`` sort and duplicate-accumulate them re-derives an order the input
        already had, at real cost for a large operator. Row order comes from the launch index and
        column order from ``q``'s own rows through the monotone ``free_map``, so the extracted
        matrix is sorted by construction.

        It also avoids a second, quieter cost: ``bsr_from_triplets`` leaves ``nnz`` at the
        *triplet* count, which is ``q.nnz`` — an upper bound on the true entry count of a
        lightly-pinned ``Q_uu`` — so every downstream ``bsr_mv`` would be dimensioned for the
        unreduced matrix. The count here is exact.

    Parameters
    ----------
    q
        ``(n_dofs, n_dofs)`` symmetric positive-semi-definite operator, ``float64``.
    fixed_mask
        Length-``n_dofs`` mask: ``True`` marks a pinned degree of freedom.
    free_map
        Compact remap from [`free_partition`][triwarp.linalg.free_partition].
    fixed_values
        ``(n_rhs, n_dofs)`` prescribed values.
    n_free
        Number of unpinned degrees of freedom.

    Returns
    -------
    q_uu : ``warp.sparse.BsrMatrix``
        ``(n_free, n_free)`` reduced operator.
    rhs : twt.Array2dFloat
        ``(n_rhs, n_free)`` constant right-hand side. ``arap`` adds its per-iteration rotation term
        to this buffer rather than reassembling it.

    Raises
    ------
    RuntimeError
        If ``fixed_mask``, ``free_map`` and ``fixed_values`` are not all on one device.

    See Also
    --------
    [`min_quad_with_fixed`][triwarp.linalg.min_quad_with_fixed]
    """
    require_same_device(fixed_mask=fixed_mask, free_map=free_map, fixed_values=fixed_values)
    device = fixed_mask.device
    n_dofs = int(fixed_mask.shape[0])
    n_rhs = int(fixed_values.shape[0])
    # Every free row writes its own count exactly once (``free_map`` is a bijection onto
    # ``[0, n_free)``), so there is nothing to pre-zero.
    counts = wp.empty(n_free, dtype=wp.int32, device=device)
    rhs = wp.zeros((n_rhs, n_free), dtype=wp.float64, device=device)
    wp.launch(
        kernel_linalg.interior_row_counts,
        dim=n_dofs,
        inputs=[q.offsets, q.columns, q.values, fixed_mask, free_map, fixed_values, counts, rhs],
        device=device,
    )
    # The total-terminated form *is* the CSR offsets array, and the one host read it costs is what
    # sizes ``columns`` / ``values`` for their final use at allocation time.
    row_offsets, nnz_uu = tw.array.counts_to_offsets(counts, include_total=True)
    columns = wp.empty(nnz_uu, dtype=wp.int32, device=device)
    values = wp.empty(nnz_uu, dtype=wp.float64, device=device)
    wp.launch(
        kernel_linalg.interior_system_csr,
        dim=n_dofs,
        inputs=[q.offsets, q.columns, q.values, fixed_mask, free_map, row_offsets, columns, values],
        device=device,
    )
    # Hand the finished CSR to a compact ``BsrMatrix`` directly. ``notify_nnz_changed`` is Warp's
    # documented entry point for exactly this -- storage metadata assigned from outside
    # ``warp.sparse`` -- and ``bsr_zeros`` leaves ``row_counts`` at ``None``, which is the compact
    # topology these arrays describe.
    q_uu = wps.bsr_zeros(n_free, n_free, wp.float64, device=device)
    q_uu.offsets = row_offsets
    q_uu.columns = columns
    q_uu.values = values
    q_uu.notify_nnz_changed(nnz=nnz_uu)
    return q_uu, twt.as_array2d(rhs, wp.float64)

solve_spd

solve_spd(
    matrix: BsrMatrix,
    rhs: array,
    solution: array,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY_FALLBACK,
    preconditioner: LinearOperator | None = None,
    name: str = "solve_spd",
) -> tuple[int, float, float]

Solve one symmetric positive-definite system by preconditioned conjugate gradient.

The single-right-hand-side form of solve_spd_columns, and deliberately dtype-agnostic: it carries scalar float64 systems as well as the wp.mat22d-block operators the tangent-field solvers build, whose right-hand side is an array of wp.vec2d. solution is the initial guess and is overwritten in place, so warm starting is free.

This exists so the Jacobi preconditioner and the iteration cap are decided in one place rather than re-derived at every call site.

check_every defaults to the host-side cadence here

Unlike solve_spd_columns, which defaults to CG_CHECK_EVERY (the device-side test), this defaults to CG_CHECK_EVERY_FALLBACK — Warp's own default — so the return values stay host scalars. The device-side check trades a host readback for a conditional-graph loop, which only pays off when the solve runs long enough to amortize it; the short solves that call this (tens of iterations) run slower with it. Pass check_every=0 where a solve is known to be long.

Parameters:

Name	Type	Description	Default
matrix	BsrMatrix	Symmetric positive-(semi-)definite operator. Scalar or block dtype.	required
rhs	array	Right-hand side, with as many rows as matrix.	required
solution	array	Initial guess, overwritten with the result. Same shape and dtype as rhs.	required
tol	float	Relative residual tolerance.	CG_TOLERANCE
maxiter	int | None	Iteration cap. When None, uses CG_MAXITER_FACTOR times the number of rows.	None
check_every	int	Residual-test cadence; 0 tests on device every iteration and returns device arrays.	CG_CHECK_EVERY_FALLBACK
preconditioner	LinearOperator | None	Optional Jacobi preconditioner for matrix, built here when None. Pass one to hoist its construction out of a loop that solves against the same operator repeatedly.	None
name	str	Caller name, used in the non-convergence warning.	'solve_spd'

Returns:

Type	Description
tuple[int, float, float]	Whatever warp.optim.linear.cg returns: iteration count, residual and tolerance. Device 1-element arrays instead of host scalars when check_every=0.

Raises:

Type	Description
RuntimeError	If rhs and solution are not all on one device.

Warns:

Type	Description
UserWarning	When the solve exhausts maxiter without reaching tol. The returned solution is then whatever the last iterate happened to be, not an answer -- silence here is how a diverging solve reaches a caller looking like a converged one. Only detectable when check_every > 0; under check_every=0 the counts stay on device and testing them would reintroduce the readback that setting exists to avoid.

See Also

solve_spd_columns spd_column_solver

Source code in triwarp/linalg.py

def solve_spd(
    matrix: wps.BsrMatrix,
    rhs: wp.array,
    solution: wp.array,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY_FALLBACK,
    preconditioner: wpl.LinearOperator | None = None,
    name: str = "solve_spd",
) -> tuple[int, float, float]:
    """
    Solve one symmetric positive-definite system by preconditioned conjugate gradient.

    The single-right-hand-side form of
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns], and deliberately dtype-agnostic: it
    carries scalar ``float64`` systems as well as the ``wp.mat22d``-block operators the
    tangent-field solvers build, whose right-hand side is an array of ``wp.vec2d``. ``solution`` is
    the initial guess and is overwritten in place, so warm starting is free.

    This exists so the Jacobi preconditioner and the iteration cap are decided in one place rather
    than re-derived at every call site.

    !!! note "``check_every`` defaults to the host-side cadence here"
        Unlike [`solve_spd_columns`][triwarp.linalg.solve_spd_columns], which defaults to
        [`CG_CHECK_EVERY`][triwarp.linalg.CG_CHECK_EVERY] (the device-side test), this defaults to
        [`CG_CHECK_EVERY_FALLBACK`][triwarp.linalg.CG_CHECK_EVERY_FALLBACK] — Warp's own default —
        so the return values stay host scalars. The device-side check trades a host readback for a
        conditional-graph loop, which only pays off when the solve runs long enough to amortize it;
        the short solves that call this (tens of iterations) run slower with it. Pass
        ``check_every=0`` where a solve is known to be long.

    Parameters
    ----------
    matrix
        Symmetric positive-(semi-)definite operator. Scalar or block dtype.
    rhs
        Right-hand side, with as many rows as ``matrix``.
    solution
        Initial guess, overwritten with the result. Same shape and dtype as ``rhs``.
    tol
        Relative residual tolerance.
    maxiter
        Iteration cap. When ``None``, uses
        [`CG_MAXITER_FACTOR`][triwarp.linalg.CG_MAXITER_FACTOR] times the number of rows.
    check_every
        Residual-test cadence; ``0`` tests on device every iteration and returns device arrays.
    preconditioner
        Optional Jacobi preconditioner for ``matrix``, built here when ``None``. Pass one to hoist
        its construction out of a loop that solves against the same operator repeatedly.
    name
        Caller name, used in the non-convergence warning.

    Returns
    -------
    tuple[int, float, float]
        Whatever ``warp.optim.linear.cg`` returns: iteration count, residual and tolerance. Device
        1-element arrays instead of host scalars when ``check_every=0``.

    Raises
    ------
    RuntimeError
        If ``rhs`` and ``solution`` are not all on one device.

    Warns
    -----
    UserWarning
        When the solve exhausts ``maxiter`` without reaching ``tol``. The returned ``solution`` is
        then whatever the last iterate happened to be, not an answer -- silence here is how a
        diverging solve reaches a caller looking like a converged one. Only detectable when
        ``check_every > 0``; under ``check_every=0`` the counts stay on device and testing them
        would reintroduce the readback that setting exists to avoid.

    See Also
    --------
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]
    [`spd_column_solver`][triwarp.linalg.spd_column_solver]
    """
    require_same_device(rhs=rhs, solution=solution)
    n_rows = int(rhs.shape[0])
    iteration_cap = CG_MAXITER_FACTOR * n_rows if maxiter is None else maxiter
    result = wpl.cg(
        matrix,
        rhs,
        solution,
        tol=tol,
        # ``atol=0.0`` rather than omitted: ``warp.optim.linear``'s own tolerance resolution sets
        # ``atol := tol`` whenever ``atol`` is left ``None``, silently turning this "relative
        # residual tolerance" into an *absolute* floor of the same numeric value -- a right-hand
        # side whose own norm falls below ``tol`` then "converges" at the untouched initial guess
        # in zero iterations, which looks like a correct answer rather than a failure (CLAUDE.md
        # section 12.7). ``_BatchedCg`` / ``_BlockCg2`` below already pass an explicit zero for the
        # same reason.
        atol=0.0,
        maxiter=iteration_cap,
        M=wpl.preconditioner(matrix, "diag") if preconditioner is None else preconditioner,
        check_every=_supported_check_every(check_every),
    )
    _warn_if_not_converged(result, iteration_cap, name)
    return result

solve_spd_columns

solve_spd_columns(
    matrix: BsrMatrix[float64],
    rhs: Array2dFloat,
    solution: Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> tuple[int, float, float]

Solve one symmetric positive-definite operator against several right-hand-side columns.

Single batched conjugate-gradient call: all n_rhs columns advance together and convergence uses the worst-case residual, so the iteration count is the maximum over columns rather than their sum. solution is used as the initial guess (warm starting is therefore free) and is overwritten in place.

check_every=0 returns device arrays, not host scalars

Under the default check_every=0 nothing is ever read back, so the three returned values are 1-element device arrays — and the second and third are the squared residual norm and squared absolute tolerance, matching warp.optim.linear.cg. Code that formats or compares them must call .numpy() first, which reintroduces the host sync the setting exists to avoid. Pass a positive check_every to get host scalars back. solution is unaffected either way.

On a device without conditional CUDA graphs there is nowhere to put the device-side test, so check_every=0 is replaced by CG_CHECK_EVERY_FALLBACK and host scalars come back instead. Warp's own behaviour in that case is to run every solve to maxiter — CG_MAXITER_FACTOR * n iterations of guaranteed waste — so substituting the cadence is the only usable reading of the request.

Parameters:

Name	Type	Description	Default
matrix	BsrMatrix[float64]	(n, n) symmetric positive-definite operator, float64.	required
rhs	Array2dFloat	(n_rhs, n) right-hand sides, one per row. Must be contiguous.	required
solution	Array2dFloat	(n_rhs, n) initial guess, overwritten with the result. Must be contiguous.	required
tol	float	Relative residual tolerance, as a ratio of the right-hand-side norm.	CG_TOLERANCE
maxiter	int | None	Iteration cap. Defaults to CG_MAXITER_FACTOR * n.	None
check_every	int	How many iterations run between residual tests. 0 (the default) tests every iteration on device; see Notes for the tradeoff and the warning above for what it does to the return type.	CG_CHECK_EVERY
preconditioner	str	"diag" (the default) for the Jacobi preconditioner, "multigrid" for the smoothed-aggregation V-cycle multigrid_preconditioner builds, or "auto" to let the operator decide. "auto" first tests the assembled system against CG_MULTIGRID_DOMINANCE and its size floor: a system that clears the gate builds the V-cycle outright, and one that does not runs Jacobi under CG_PROBE_ITERATIONS and escalates only if that has not converged. The V-cycle costs a setup pass and pays for itself only where the solve dominates the call, so "auto" is the setting for a caller whose systems vary -- it cannot regress a solve that was already short, though it can forgo a small win on a small ill-conditioned system that the size floor declines. See that function's Notes, and both constants.	'diag'

Returns:

Type	Description
tuple[int, float, float]	(iterations, residual_norm, absolute_tolerance) on warp.optim.linear.cg's terms, with the residual taken over the worst column. Device arrays rather than host scalars under check_every=0; see the warning above. A one-column solve returns Warp's own values and a batched one returns the same three from this module's solver.

Raises:

Type	Description
ValueError	If preconditioner is not one of the three names above. It is rejected rather than treated as "diag", so a misspelling cannot turn into a silently slower solve.
RuntimeError	If rhs and solution are not all on one device.

Warns:

Type	Description
UserWarning	When the solve exhausts maxiter without reaching tol, on the same terms as solve_spd. A batched solve converges on its worst column, so hitting the cap here means at least one column is unsolved. Only detectable when check_every > 0.

Notes

check_every is a pure performance knob on every path but one — it cannot change the converged answer, only how far past the tolerance the solver may overshoot before it notices. The exception is the exactly-two-column, "diag"-preconditioned path (this module's internal block conjugate gradient): it shares one Krylov subspace across both columns, so a positive check_every can let a batch of iterations apply a real, coupled update to a column that already crossed its own tolerance before the next readback catches up. Its device-side default (0) scales with how much work a single cg call does, and the regimes disagree, so read the one that matches the caller:

• Cold single solves — one cg call from a zero initial guess, the shape harmonic / tutte / smooth_region take. This is the regime the default is set for, and it is a clear win there.
• Warm-started solves inside an iteration loop — arap, whose right-hand side changes every iteration so each solve still runs tens of CG iterations from the previous answer. The default is neutral to mildly positive here.
• Repeated near-converged solves over one [spd_column_solver] [triwarp.linalg.spd_column_solver] state — the same right-hand side re-solved back to back, so every call after the first converges in one or two iterations. Here the device-side check's fixed per-call cost dominates a solve that short, so it is a real loss; pass a positive check_every to a state driven that way.
• Raising check_every well above the default is a loss in every regime: the readback it saves is cheap, while the extra iterations it can cause on top of the tolerance are real work — the smaller the solve, the worse the trade.

The preconditioner is not a knob worth turning beyond "diag" / "multigrid" / "auto" above: IC(0) and Chebyshev both come out a wash or a loss against the Jacobi preconditioner used here. See "Why Jacobi" in the triwarp.linalg module documentation.

The per-iteration cost of batching was a separate lever and has been taken: with more than one column this runs triwarp's own conjugate gradient rather than warp.optim.linear's, because Warp's reduction degrades on precisely the batched input the worst-case stopping rule needs. See "Whose conjugate gradient" in the triwarp.linalg module documentation.

See Also

spd_column_solver replicated_operator

Source code in triwarp/linalg.py

def solve_spd_columns(
    matrix: wps.BsrMatrix[wp.float64],
    rhs: twt.Array2dFloat,
    solution: twt.Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> tuple[int, float, float]:
    """
    Solve one symmetric positive-definite operator against several right-hand-side columns.

    Single batched conjugate-gradient call: all ``n_rhs`` columns advance together and convergence
    uses the worst-case residual, so the iteration count is the *maximum* over columns rather than
    their sum. ``solution`` is used as the initial guess (warm starting is therefore free) and is
    overwritten in place.

    !!! warning "``check_every=0`` returns device arrays, not host scalars"
        Under the default ``check_every=0`` nothing is ever read back, so the three returned values
        are 1-element **device arrays** — and the second and third are the *squared* residual norm
        and *squared* absolute tolerance, matching ``warp.optim.linear.cg``. Code that formats or
        compares them must call ``.numpy()`` first, which reintroduces the host sync the setting
        exists to avoid. Pass a positive ``check_every`` to get host scalars back. ``solution`` is
        unaffected either way.

        On a device without conditional CUDA graphs there is nowhere to put the device-side test,
        so ``check_every=0`` is replaced by
        [`CG_CHECK_EVERY_FALLBACK`][triwarp.linalg.CG_CHECK_EVERY_FALLBACK] and host scalars come
        back instead. Warp's own behaviour in that case is to run every solve to ``maxiter`` —
        ``CG_MAXITER_FACTOR * n`` iterations of guaranteed waste — so substituting the cadence is
        the only usable reading of the request.

    Parameters
    ----------
    matrix
        ``(n, n)`` symmetric positive-definite operator, ``float64``.
    rhs
        ``(n_rhs, n)`` right-hand sides, one per row. Must be contiguous.
    solution
        ``(n_rhs, n)`` initial guess, overwritten with the result. Must be contiguous.
    tol
        Relative residual tolerance, as a ratio of the right-hand-side norm.
    maxiter
        Iteration cap. Defaults to ``CG_MAXITER_FACTOR * n``.
    check_every
        How many iterations run between residual tests. ``0`` (the default) tests every iteration
        on device; see Notes for the tradeoff and the warning above for what it does to the return
        type.
    preconditioner
        ``"diag"`` (the default) for the Jacobi preconditioner, ``"multigrid"`` for the
        smoothed-aggregation V-cycle [`multigrid_preconditioner`]
        [triwarp.linalg.multigrid_preconditioner] builds, or ``"auto"`` to let the operator decide.
        ``"auto"`` first tests the assembled system against
        [`CG_MULTIGRID_DOMINANCE`][triwarp.linalg.CG_MULTIGRID_DOMINANCE] and its size floor: a
        system that clears the gate builds the V-cycle outright, and one that does not runs Jacobi
        under [`CG_PROBE_ITERATIONS`][triwarp.linalg.CG_PROBE_ITERATIONS] and escalates only if
        that has not converged. The V-cycle costs a setup pass and pays for itself only where the
        solve dominates the call, so ``"auto"`` is the setting for a caller whose systems vary --
        it cannot regress a solve that was already short, though it can forgo a small win on a
        small ill-conditioned system that the size floor declines. See that function's Notes, and
        both constants.

    Returns
    -------
    tuple[int, float, float]
        ``(iterations, residual_norm, absolute_tolerance)`` on ``warp.optim.linear.cg``'s terms,
        with the residual taken over the worst column. Device arrays rather than host scalars under
        ``check_every=0``; see the warning above. A one-column solve returns Warp's own values and
        a batched one returns the same three from this module's solver.

    Raises
    ------
    ValueError
        If ``preconditioner`` is not one of the three names above. It is rejected rather than
        treated as ``"diag"``, so a misspelling cannot turn into a silently slower solve.
    RuntimeError
        If ``rhs`` and ``solution`` are not all on one device.

    Warns
    -----
    UserWarning
        When the solve exhausts ``maxiter`` without reaching ``tol``, on the same terms as
        [`solve_spd`][triwarp.linalg.solve_spd]. A batched solve converges on its *worst* column,
        so hitting the cap here means at least one column is unsolved. Only detectable when
        ``check_every > 0``.

    Notes
    -----
    ``check_every`` is a pure performance knob on every path but one — it cannot change the
    converged answer, only how far past the tolerance the solver may overshoot before it notices.
    The exception is the exactly-two-column, ``"diag"``-preconditioned path (this module's internal
    block conjugate gradient): it shares one Krylov subspace across both columns, so a positive
    ``check_every`` can let a batch of iterations apply a real, coupled update to a column that
    already crossed its own tolerance before the next readback catches up. Its device-side default
    (``0``) scales with how much work a single ``cg`` call does, and the regimes disagree, so read
    the one that matches the caller:

    - **Cold single solves** — one ``cg`` call from a zero initial guess, the shape
      ``harmonic`` / ``tutte`` / ``smooth_region`` take. This is the regime the default is set for,
      and it is a clear win there.
    - **Warm-started solves inside an iteration loop** — ``arap``, whose right-hand side changes
      every iteration so each solve still runs tens of CG iterations from the previous answer. The
      default is neutral to mildly positive here.
    - **Repeated near-converged solves over one [`spd_column_solver`]
      [triwarp.linalg.spd_column_solver] state** — the same right-hand side re-solved back to back,
      so every call after the first converges in one or two iterations. Here the device-side check's
      fixed per-call cost dominates a solve that short, so it is a real loss; pass a positive
      ``check_every`` to a state driven that way.
    - **Raising ``check_every`` well above the default is a loss in every regime**: the readback it
      saves is cheap, while the extra iterations it can cause on top of the tolerance are real work
      — the smaller the solve, the worse the trade.

    The *preconditioner* is not a knob worth turning beyond ``"diag"`` / ``"multigrid"`` /
    ``"auto"`` above: IC(0) and Chebyshev both come out a wash or a loss against the Jacobi
    preconditioner used here. See "Why Jacobi" in the [`triwarp.linalg`][triwarp.linalg] module
    documentation.

    The *per-iteration* cost of batching was a separate lever and has been taken: with more than one
    column this runs triwarp's own conjugate gradient rather than ``warp.optim.linear``'s, because
    Warp's reduction degrades on precisely the batched input the worst-case stopping rule needs. See
    "Whose conjugate gradient" in the [`triwarp.linalg`][triwarp.linalg] module documentation.

    See Also
    --------
    [`spd_column_solver`][triwarp.linalg.spd_column_solver]
    [`replicated_operator`][triwarp.linalg.replicated_operator]
    """
    require_same_device(rhs=rhs, solution=solution)
    result = _cg_columns(
        matrix,
        rhs,
        solution,
        tol=tol,
        maxiter=maxiter,
        check_every=check_every,
        preconditioner=preconditioner,
        run=True,
        caller="solve_spd_columns",
    )
    _warn_if_not_converged(
        result,
        maxiter if maxiter is not None else CG_MAXITER_FACTOR * int(rhs.shape[1]),
        "solve_spd_columns",
    )
    return result

spd_column_solver

spd_column_solver(
    matrix: BsrMatrix[float64],
    rhs: Array2dFloat,
    solution: Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> wpl.LinearSolverState | _BatchedCg | _BlockCg2

Pre-allocated batched conjugate-gradient state, for repeated solves of one operator.

Same solve as solve_spd_columns, but the temporary buffers and the batch layout are allocated once and the returned state is callable. Build it outside a per-iteration loop and call it inside: solve_spd_columns would otherwise reallocate its temporaries and rebuild the operator wrappers on every iteration.

rhs and solution are captured at construction, so a loop that overwrites those same buffers in place can call the state with no arguments.

Parameters:

Name	Type	Description	Default
matrix	BsrMatrix[float64]	(n, n) symmetric positive-definite operator, float64. Held by the returned state.	required
rhs	Array2dFloat	(n_rhs, n) right-hand sides; re-read from this buffer on every call.	required
solution	Array2dFloat	(n_rhs, n) initial guess and result; warm-started from its current contents each call.	required
tol	float	Relative residual tolerance.	CG_TOLERANCE
maxiter	int | None	Iteration cap. Defaults to CG_MAXITER_FACTOR * n.	None
check_every	int	Iterations between residual tests; see solve_spd_columns for the tradeoff and for what the default 0 does to the values each call returns.	CG_CHECK_EVERY
preconditioner	str	"diag" or "multigrid", as in solve_spd_columns. Built once here and reused by every call against this state, which is the shape the V-cycle's setup cost wants. "auto" is not accepted: it decides on the first solve -- from that system's operator, and from the probe when the operator does not settle it -- and a hoisted state exists to be driven many times.	'diag'

Returns:

Type

Description

``warp.optim.linear.LinearSolverState`` | internal batched state

Callable solver state: a single-column solve returns Warp's own LinearSolverState, and a multi-column one returns this module's own _BatchedCg or (at exactly two columns under "diag") _BlockCg2 -- a drop-in with the identical call-and-return contract, kept internal because a caller never has to name the difference (see solve_spd_columns's "Whose conjugate gradient"). Substituted operands must match the construction-time shape, dtype, device and batch layout. Each call returns what solve_spd_columns returns, including its check_every=0 device arrays.

Raises:

Type	Description
ValueError	If preconditioner is "auto", for the reason above, or is not one of the two names this accepts.
RuntimeError	If rhs and solution are not all on one device.

Examples:

solver = spd_column_solver(q_uu, rhs, solution, tol=1e-10)
for _ in range(max_iterations):
    ...  # rewrite rhs in place
    solver()

See Also

solve_spd_columns

Source code in triwarp/linalg.py

def spd_column_solver(
    matrix: wps.BsrMatrix[wp.float64],
    rhs: twt.Array2dFloat,
    solution: twt.Array2dFloat,
    *,
    tol: float = CG_TOLERANCE,
    maxiter: int | None = None,
    check_every: int = CG_CHECK_EVERY,
    preconditioner: str = "diag",
) -> wpl.LinearSolverState | _BatchedCg | _BlockCg2:
    """
    Pre-allocated batched conjugate-gradient state, for repeated solves of one operator.

    Same solve as [`solve_spd_columns`][triwarp.linalg.solve_spd_columns], but the temporary buffers
    and the batch layout are allocated once and the returned state is callable. Build it *outside* a
    per-iteration loop and call it inside: ``solve_spd_columns`` would otherwise reallocate its
    temporaries and rebuild the operator wrappers on every iteration.

    ``rhs`` and ``solution`` are captured at construction, so a loop that overwrites those same
    buffers in place can call the state with no arguments.

    Parameters
    ----------
    matrix
        ``(n, n)`` symmetric positive-definite operator, ``float64``. Held by the returned state.
    rhs
        ``(n_rhs, n)`` right-hand sides; re-read from this buffer on every call.
    solution
        ``(n_rhs, n)`` initial guess and result; warm-started from its current contents each call.
    tol
        Relative residual tolerance.
    maxiter
        Iteration cap. Defaults to ``CG_MAXITER_FACTOR * n``.
    check_every
        Iterations between residual tests; see
        [`solve_spd_columns`][triwarp.linalg.solve_spd_columns] for the tradeoff and for what the
        default ``0`` does to the values each call returns.
    preconditioner
        ``"diag"`` or ``"multigrid"``, as in
        [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]. Built once here and reused by every
        call against this state, which is the shape the V-cycle's setup cost wants. ``"auto"`` is
        **not** accepted: it decides on the first solve -- from that system's operator, and from
        the probe when the operator does not settle it -- and a hoisted state exists to be driven
        many times.

    Returns
    -------
    ``warp.optim.linear.LinearSolverState`` | internal batched state
        Callable solver state: a single-column solve returns Warp's own
        ``LinearSolverState``, and a multi-column one returns this module's own ``_BatchedCg`` or
        (at exactly two columns under ``"diag"``) ``_BlockCg2`` -- a drop-in with the identical
        call-and-return contract, kept internal because a caller never has to name the difference
        (see [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]'s "Whose conjugate gradient").
        Substituted operands must match the construction-time shape, dtype, device and batch
        layout. Each call returns what
        [`solve_spd_columns`][triwarp.linalg.solve_spd_columns] returns, including its
        ``check_every=0`` device arrays.

    Raises
    ------
    ValueError
        If ``preconditioner`` is ``"auto"``, for the reason above, or is not one of the two names
        this accepts.
    RuntimeError
        If ``rhs`` and ``solution`` are not all on one device.

    Examples
    --------
    ```python
    solver = spd_column_solver(q_uu, rhs, solution, tol=1e-10)
    for _ in range(max_iterations):
        ...  # rewrite rhs in place
        solver()
    ```

    See Also
    --------
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]
    """
    require_same_device(rhs=rhs, solution=solution)
    return _cg_columns(
        matrix,
        rhs,
        solution,
        tol=tol,
        maxiter=maxiter,
        check_every=check_every,
        preconditioner=preconditioner,
        run=False,
        caller="spd_column_solver",
    )

replicated_operator

replicated_operator(
    matrix: BsrMatrix[float64] | LinearOperator,
    n_columns: int,
) -> wpl.LinearOperator

Present one (n, n) operator as n_columns independent subproblems over a flat vector.

The returned operator acts on length-n_columns * n vectors laid out as n_columns contiguous blocks — exactly the memory of a contiguous (n_columns, n) array — and carries the batch_offsets that make warp.optim.linear's solvers iterate all blocks together and converge on the worst-case residual. matrix is not replicated in memory: the matvec applies the single operator to each block in turn.

Parameters:

Name	Type	Description	Default
matrix	BsrMatrix[float64] | LinearOperator	(n, n) operator: a warp.sparse.BsrMatrix, a 1-D array (read as a diagonal), or an existing warp.optim.linear.LinearOperator such as a preconditioner.	required
n_columns	int	Number of independent right-hand-side columns.	required

Returns:

Type	Description
``warp.optim.linear.LinearOperator``	Batched operator of shape (n_columns * n, n_columns * n).

Notes

A block-diagonal matrix would give the same batching at n_columns times the operator memory, and 2x2 blocks holding a scalar multiple of the identity would cost four times the storage and flops per scalar entry. Replicating only the matvec avoids both.

batch_offsets is attached and max_batch_length deliberately is not, although this function knows the value exactly (every subproblem is n scalars): passing it to skip redundant reduction levels for a known maximum subproblem size made no measurable difference here, so it would be one more argument to explain for nothing. warp.optim.linear's restart option was also considered and is worse: resetting the search direction throws away the Krylov space CG has built, costing more iterations rather than buying better-conditioned ones.

See Also

solve_spd_columns

Source code in triwarp/linalg.py

def replicated_operator(
    matrix: wps.BsrMatrix[wp.float64] | wpl.LinearOperator, n_columns: int
) -> wpl.LinearOperator:
    """
    Present one ``(n, n)`` operator as ``n_columns`` independent subproblems over a flat vector.

    The returned operator acts on length-``n_columns * n`` vectors laid out as ``n_columns``
    contiguous blocks — exactly the memory of a contiguous ``(n_columns, n)`` array — and carries
    the ``batch_offsets`` that make ``warp.optim.linear``'s solvers iterate all blocks together and
    converge on the worst-case residual. ``matrix`` is **not** replicated in memory: the ``matvec``
    applies the single operator to each block in turn.

    Parameters
    ----------
    matrix
        ``(n, n)`` operator: a ``warp.sparse.BsrMatrix``, a 1-D array (read as a diagonal), or an
        existing ``warp.optim.linear.LinearOperator`` such as a preconditioner.
    n_columns
        Number of independent right-hand-side columns.

    Returns
    -------
    ``warp.optim.linear.LinearOperator``
        Batched operator of shape ``(n_columns * n, n_columns * n)``.

    Notes
    -----
    A block-diagonal *matrix* would give the same batching at ``n_columns`` times the operator
    memory, and 2x2 blocks holding a scalar multiple of the identity would cost four times the
    storage and flops per scalar entry. Replicating only the ``matvec`` avoids both.

    ``batch_offsets`` is attached and ``max_batch_length`` deliberately is not, although this
    function knows the value exactly (every subproblem is ``n`` scalars): passing it to skip
    redundant reduction levels for a known maximum subproblem size made no measurable difference
    here, so it would be one more argument to explain for nothing. ``warp.optim.linear``'s
    ``restart`` option was also considered and is worse: resetting the search direction throws away
    the Krylov space CG has built, costing more iterations rather than buying better-conditioned
    ones.

    See Also
    --------
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]
    """
    base = wpl.aslinearoperator(matrix)
    if n_columns == 1:
        # Nothing to batch: hand back the operator itself rather than wrapping every matvec in a
        # one-iteration Python loop. ``lscm`` solves a single coupled column and would otherwise pay
        # the wrapper cost for no benefit.
        return base
    n = int(base.shape[0])

    # Per-block views, cached by buffer address. A solver calls ``matvec`` once per iteration with
    # the *same* handful of temporaries, so slicing them fresh every time would add thousands of
    # Python-level array constructions to a solve — enough to erase the batching win outright, and
    # enough to keep the iteration from being captured as a CUDA graph. The cache is keyed on the
    # full memory layout, so a different buffer (or a re-slice of one) never aliases a stale view.
    block_views: dict[tuple[int, tuple[int, ...], tuple[int, ...]], list[wp.array]] = {}

    def blocks(array: wp.array) -> list[wp.array]:
        key = (array.ptr, tuple(array.shape), tuple(array.strides))
        views = block_views.get(key)
        if views is None:
            views = [array[column * n : (column + 1) * n] for column in range(n_columns)]
            block_views[key] = views
        return views

    def matvec(x: wp.array, y: wp.array, z: wp.array, alpha: float, beta: float) -> None:
        x_blocks, y_blocks, z_blocks = blocks(x), blocks(y), blocks(z)
        for column in range(n_columns):
            base.matvec(x_blocks[column], y_blocks[column], z_blocks[column], alpha, beta)

    offsets = wp.array(
        [column * n for column in range(n_columns + 1)], dtype=wp.int32, device=base.device
    )
    total = n_columns * n
    return wpl.LinearOperator(
        (total, total), base.dtype, base.device, matvec, batch_offsets=offsets
    )

multigrid_preconditioner

multigrid_preconditioner(
    matrix: BsrMatrix[float64],
    n_columns: int = 1,
    *,
    seed: int = 0,
) -> wpl.LinearOperator

Smoothed-aggregation multigrid preconditioner for a symmetric positive-semi-definite operator.

A single-level preconditioner cannot break conjugate gradient's growth in the mesh size -- see "Why Jacobi" in the triwarp.linalg module documentation for the alternatives that were tried and rejected -- because cutting the iteration count by sqrt(k) at k mat-vecs per apply leaves the total work growing. A multigrid V-cycle attacks the low-frequency error the smoother cannot see, so the iteration count grows much more slowly with the problem size than it does under Jacobi.

The returned operator acts on length-n_columns * n vectors laid out as n_columns contiguous blocks, exactly as replicated_operator does, so it drops into M= beside a Jacobi preconditioner without any other change.

Parameters:

Name	Type	Description	Default
matrix	BsrMatrix[float64]	(n, n) symmetric positive-(semi-)definite operator, float64. Zero diagonal entries are allowed: such a row is identically zero for a semi-definite operator, and the whole chain -- smoother, coarse solve and all -- leaves those unknowns at zero.	required
n_columns	int	Number of independent right-hand-side columns the operator will be applied to.	1
seed	int	Seeds the aggregation's randomized priorities and the power iteration's start vector, so a hierarchy is a deterministic function of the operator and this number.	0

Returns:

Type	Description
``warp.optim.linear.LinearOperator``	The V-cycle, of shape (n_columns * n, n_columns * n). A Jacobi preconditioner instead when the operator does not coarsen -- see Notes.

Raises:

Type	Description
ValueError	If the returned operator's matvec is called with alpha != 1 or beta != 0. A preconditioner apply is always z = M x, and the general form would cost a pass that no caller needs.

Notes

Setup is not free and it is not amortized over anything: building the hierarchy costs one aggregation, one power iteration, two bsr_mm and one bsr_transposed per level. Ask for this where the solve dominates the call, and pass the same operator's preconditioner into a loop rather than rebuilding it -- solve_spd's preconditioner parameter exists for exactly that.

The setup is host cost and it is per level, not per unknown, so it cannot be cut by making the operator smaller. A level's cost is dominated by host-side overhead in warp.sparse.bsr_mm, which makes device-to-host readbacks to size its output and so costs roughly the same regardless of the level's size. That is Warp's, not this package's, so "make the hierarchy cheap enough to run unconditionally" is not a lever available here; the reachable version is to build fewer levels, which is what this module's _MULTIGRID_MAX_COARSE is set for.

Coarsening stops at 128 rows, or earlier if a level fails to shrink; the coarsest operator is then inverted densely on the host, which is exact and is a single launch inside the cycle where an iterative coarse solve would be a data-dependent loop. When coarsening stalls while the level is still too large to factor, there is no usable hierarchy and this hands back warp.optim.linear.preconditioner(matrix, "diag") rather than a cycle whose coarse solve is a guess -- so a caller never has to branch on the operator's shape.

That fallback is silent, and the strength threshold can trigger it

A stalled hierarchy is indistinguishable from a weak one at the call site: the solve simply runs at its Jacobi iteration count. _MULTIGRID_THETA is what decides how easily it happens -- raising it makes more off-diagonals weak, and past some threshold an operator can stop coarsening entirely and come back at exactly the Jacobi count. So an aggregation change that "did nothing" should be checked against the level count before it is read as a change that did not help.

See Also

solve_spd solve_spd_columns replicated_operator

Source code in triwarp/linalg.py

def multigrid_preconditioner(
    matrix: wps.BsrMatrix[wp.float64], n_columns: int = 1, *, seed: int = 0
) -> wpl.LinearOperator:
    """
    Smoothed-aggregation multigrid preconditioner for a symmetric positive-semi-definite operator.

    A single-level preconditioner cannot break conjugate gradient's growth in the mesh size -- see
    "Why Jacobi" in the [`triwarp.linalg`][triwarp.linalg] module documentation for the alternatives
    that were tried and rejected -- because cutting the iteration count by ``sqrt(k)`` at ``k``
    mat-vecs per apply leaves the total work growing. A multigrid V-cycle attacks the low-frequency
    error the smoother cannot see, so the iteration count grows much more slowly with the problem
    size than it does under Jacobi.

    The returned operator acts on length-``n_columns * n`` vectors laid out as ``n_columns``
    contiguous blocks, exactly as [`replicated_operator`][triwarp.linalg.replicated_operator] does,
    so it drops into ``M=`` beside a Jacobi preconditioner without any other change.

    Parameters
    ----------
    matrix
        ``(n, n)`` symmetric positive-(semi-)definite operator, ``float64``. Zero diagonal entries
        are allowed: such a row is identically zero for a semi-definite operator, and the whole
        chain -- smoother, coarse solve and all -- leaves those unknowns at zero.
    n_columns
        Number of independent right-hand-side columns the operator will be applied to.
    seed
        Seeds the aggregation's randomized priorities and the power iteration's start vector, so a
        hierarchy is a deterministic function of the operator and this number.

    Returns
    -------
    ``warp.optim.linear.LinearOperator``
        The V-cycle, of shape ``(n_columns * n, n_columns * n)``. A **Jacobi** preconditioner
        instead when the operator does not coarsen -- see Notes.

    Raises
    ------
    ValueError
        If the returned operator's ``matvec`` is called with ``alpha != 1`` or ``beta != 0``. A
        preconditioner apply is always ``z = M x``, and the general form would cost a pass that no
        caller needs.

    Notes
    -----
    Setup is not free and it is not amortized over anything: building the hierarchy costs one
    aggregation, one power iteration, two ``bsr_mm`` and one ``bsr_transposed`` per level. **Ask for
    this where the solve dominates the call**, and pass the same operator's preconditioner into a
    loop rather than rebuilding it -- [`solve_spd`][triwarp.linalg.solve_spd]'s ``preconditioner``
    parameter exists for exactly that.

    **The setup is host cost and it is per *level*, not per unknown, so it cannot be cut by making
    the operator smaller.** A level's cost is dominated by host-side overhead in
    ``warp.sparse.bsr_mm``, which makes device-to-host readbacks to size its output and so costs
    roughly the same regardless of the level's size. That is Warp's, not this package's, so "make
    the hierarchy cheap enough to run unconditionally" is not a lever available here; the reachable
    version is to build **fewer levels**, which is what this module's ``_MULTIGRID_MAX_COARSE`` is
    set for.

    Coarsening stops at 128 rows, or earlier if a level fails to shrink; the coarsest operator is
    then inverted densely on the host, which is exact and is a single launch inside the cycle where
    an iterative coarse solve would be a data-dependent loop. When coarsening stalls while the level
    is still too large to factor, there is no usable hierarchy and this hands back
    ``warp.optim.linear.preconditioner(matrix, "diag")`` rather than a cycle whose coarse solve is a
    guess -- so a caller never has to branch on the operator's shape.

    !!! warning "That fallback is silent, and the strength threshold can trigger it"
        A stalled hierarchy is indistinguishable from a weak one at the call site: the solve simply
        runs at its Jacobi iteration count. ``_MULTIGRID_THETA`` is what decides how easily it
        happens -- raising it makes more off-diagonals weak, and past some threshold an operator can
        stop coarsening entirely and come back at *exactly* the Jacobi count. So an aggregation
        change that "did nothing" should be checked against the level count before it is read as a
        change that did not help.

    See Also
    --------
    [`solve_spd`][triwarp.linalg.solve_spd]
    [`solve_spd_columns`][triwarp.linalg.solve_spd_columns]
    [`replicated_operator`][triwarp.linalg.replicated_operator]
    """
    base = wpl.aslinearoperator(matrix)
    n = int(matrix.nrow)
    hierarchy = _multigrid_hierarchy(matrix, seed)
    if hierarchy is None:
        return replicated_operator(wpl.preconditioner(matrix, "diag"), n_columns)
    levels, coarse_inverse = hierarchy
    cycle = _MultigridCycle(levels, coarse_inverse, n_columns=n_columns, stride=n)

    def matvec(x: wp.array, y: wp.array, z: wp.array, alpha: float, beta: float) -> None:
        if alpha != 1.0 or beta != 0.0:
            raise ValueError(
                "multigrid_preconditioner's operator only implements z = M x "
                f"(got alpha={alpha}, beta={beta})"
            )
        cycle.apply(x, z)

    total = n_columns * n
    return wpl.LinearOperator((total, total), base.dtype, base.device, matvec)