pyzag.operators

The pyzag.operators subpackage defines the abstract block-operator interfaces the solver runs against, plus the concrete dense backend. New in pyzag 2.0.

Abstract interfaces

Abstract block operator and block vector interfaces.

class pyzag.operators.base.BlockVector

Bases: ABC

Abstract interface for a logical packed block vector.

Logical conventions

A block vector consists of nblk logical blocks. Solver code treats block vectors as block-major: axis 0 indexes logical blocks.

For the current bidiagonal solvers, the logical shape is

(nblk, batch_size, block_size)

where:

axis 0 is the logical block / time index
axis 1 is the batch index
axis 2 is the local vector entries for one block

Backend freedom

A backend may store data in any representation it wants (dense, structured, sparse, factored, etc.) as long as:

nblk, batch_size, block_size describe the logical shape

all methods preserve the same logical block ordering

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk: int: Number of logical blocks in this vector.

abstract property batch_size: int: Logical batch size.

abstract property block_size: int: Logical size of one block (last axis).

abstractmethod clone() → BlockVector: Return a safe copy of the vector.

abstractmethod norm(dim: int = -1) → Tensor: Compute the norm along dim. Returns a raw tensor (used for scalar convergence checks; the result has a different shape than a block vector and should not be wrapped).

abstractmethod flat_norm() → Tensor: Cross-block flattened L2 norm per batch element. Returns a raw tensor of shape (batch_size,). Counterpart to per-block norm(); used for convergence metrics that need a single scalar per batch (e.g. line search).

abstractmethod where(mask: Tensor, other: BlockVector) → BlockVector

Batch-axis conditional combination, torch.where convention: returns self where mask is True, else other.

mask is a raw tensor of shape (batch_size,); backends broadcast it over the block and state axes. Used by the Newton step for partial updates.

abstractmethod scale_batches(factor: Tensor) → BlockVector: Per-batch scalar broadcast: multiply each batch element by the corresponding entry of factor (raw tensor of shape (batch_size,), broadcast over block and state axes). Used by line-search backtracking.

abstractmethod flip(dim: int) → BlockVector: Return a new vector with axis dim reversed.

neg() → BlockVector: Return the negated vector (alias for -self).

abstractmethod classmethod cat(vectors: Sequence[BlockVector], dim: int = 0) → BlockVector: Concatenate a sequence of compatible block vectors along dim.

abstractmethod classmethod zeros_like(other: BlockVector) → BlockVector: Construct a zero-filled block vector with the same shape as other.

class pyzag.operators.base.BlockOperator

Bases: ABC

Abstract interface for a logical packed block operator.

Logical conventions

The operator consists of nblk logical blocks. The solver treats block vectors passed to this operator as block-major.

Backend freedom

A backend may store its blocks in any representation it wants as long as:

nblk reports the correct logical number of blocks
batch_size describes the logical block action
all methods preserve the same logical block ordering

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk: int: Number of logical blocks in this operator.

abstract property batch_size: int: Logical batch size expected in block-major vector inputs.

abstractmethod matvec(x: BlockVector) → BlockVector

Apply the operator to a block vector x.

Required:: x.nblk == self.nblk

abstractmethod t_matvec(x: BlockVector) → BlockVector

Apply the transpose of the operator to a block vector x.

Required:: x.nblk == self.nblk

abstractmethod clone() → BlockOperator: Return a safe copy of the operator.

abstractmethod pad_front(n: int = 1) → BlockOperator: Return an operator with n leading dummy logical blocks (creates new data, not a view).

abstractmethod trim_front(n: int = 1) → BlockOperator: Return an operator with the first n logical blocks removed.

class pyzag.operators.base.PCRState

Bases: ABC

Opaque PCR working state managed by the backend across levels.

Created by SolvableBlockOperator.pcr_init(), updated by SolvableBlockOperator.pcr_reduce_level(), and consumed by SolvableBlockOperator.pcr_finalize(). Callers in chunktime.py treat this as an opaque handle and never inspect its contents.

class pyzag.operators.base.SolvableBlockOperator

Bases: BlockOperator

Block operator supporting direct block solves and PCR-based reduction.

abstractmethod solve(rhs: BlockVector) → BlockVector

Solve the block system A x = rhs.

Required:: rhs.nblk == self.nblk

abstractmethod pcr_init(B: BlockOperator, v: BlockVector) → PCRState

Initialise the backend-native PCR working state for a power-of-two window.

self is the diagonal operator (A), B is the subdiagonal (already padded with one dummy leading block so B.nblk == self.nblk), and v is the RHS slice for this window.

The backend allocates internal working tensors (e.g. adds the extra leading dimension used by the Dense cyclic-shift trick) and returns an opaque PCRState.

abstractmethod pcr_reduce_level(state: PCRState, level: int) → PCRState

Apply one PCR level to the working state.

Updates the RHS vector and subdiagonal via A-inverse products, then applies the backend-native cyclic interleaving (as_strided for Dense). Returns the updated PCRState ready for the next level or for pcr_finalize().

level (0-based) lets the backend compute the correct stride pattern without the caller knowing about the internal working shape.

abstractmethod pcr_finalize(state: PCRState) → tuple[BlockOperator, BlockVector]

Extract the reduced (B_red, v_red) from the final PCR state.

Returns a pair with nblk == window_size - 1, suitable for writing back into the full B and v_work arrays in pyzag.chunktime.BidiagonalPCRFactorization.

class pyzag.operators.base.BlockJacobian

Bases: ABC

Abstract per-chunk Jacobian for a recursive nonlinear system.

Logical model

For lookback = 1 (the only case the solver currently supports), the chunk Jacobian represents the linearization of R[k] = f(x[k-1], x[k]) over k = 1..nblk_steps. The two structural pieces are:

the diagonal: dR[k]/dx[k] for k = 1..nblk_steps

the subdiagonal: dR[k]/dx[k-1] for k = 1..nblk_steps

The boundary subdiagonal blocks couple a chunk to its neighbours. In forward time order, sub[0] (k = 1, dR[1]/dx[0]) couples the chunk’s first residual to the lookback / previous chunk’s last state; the remaining subdiagonal blocks are internal to the chunk’s bidiagonal system. The adjoint pass walks time in reverse (see as_adjoint_walk()) and couples to the previously-processed (adjoint-order) chunk through index 0 of the walk-order subdiagonal – which is the original last forward block, not sub[0].

Time order is forward (low index = early time). The adjoint walk methods internalize all reversal – callers use couple_prev_chunk() and must NOT reach into storage to flip, reorder, or pick a boundary block themselves.

Backend freedom

Backends may store the diagonal/subdiagonal in any layout (dense per block, structured/arrowhead, sparse, factored). The contract is purely behavioral.

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk_steps: int: Number of residual rows in the chunk.

abstract property batch_size: int: Logical batch size.

abstract property block_size: int: Per-step state size (the user-facing n).

abstractmethod forward_system(inverse_operator) → BidiagonalForwardOperator

Build the chunk’s forward bidiagonal system, ready for Newton.

Returns a pyzag.chunktime.BidiagonalForwardOperator whose diagonal A has nblk == nblk_steps and whose subdiagonal B has nblk == nblk_steps - 1.

Parameters:: inverse_operator – factory used to build the inverse (e.g. pyzag.chunktime.BidiagonalThomasFactorization).

abstractmethod adjoint_system(inverse_operator)

Build the chunk’s adjoint bidiagonal solve operator, in adjoint-walk order.

Transposes are baked in. The first-row / first+last-col slicing (J[1, 1:].T and J[0, 1:-1].T after flip(1)) is also baked in. Unlike forward_system(), which returns a BidiagonalForwardOperator (so Newton can call .inverse() between iterations), this returns the inverse / solve operator directly: applying .matvec(rhs) on the returned object yields the adjoint solution. This matches how block_update_adjoint consumes the result with a single linear solve per chunk.

This method should be called on a BlockJacobian returned by as_adjoint_walk().

abstractmethod solve_terminal_adjoint(g_terminal: Tensor) → BlockVector

Compute -A_terminal^{-T} @ g_terminal for the very last forward-time step of the trajectory.

Returns a single-block BlockVector (nblk == 1 for lookback == 1; for higher lookback this would be a lookback-block vector). The returned vector must not alias g_terminal.

abstractmethod couple_prev_chunk(a_first: BlockVector) → BlockVector

Compute the inter-chunk adjoint coupling B_boundary^T @ a_first.

a_first is a single-block BlockVector holding the previous chunk’s adjoint tail (in adjoint-walk order). Returns a single-block BlockVector to be subtracted into the current chunk’s RHS first row (nblk == 1 for lookback == 1; lookback-block for higher lookback).

abstractmethod as_adjoint_walk() → BlockJacobian: Return a BlockJacobian whose forward time order is the reverse of this one. Backends are free to implement lazily (e.g. via a flag) so that no copy happens unless storage requires it. Replaces direct flip calls in solver code.

Dense backend

The dense backend is the default implementation and reproduces the pyzag 1.x behavior. Block vectors are stored as (nblk, batch, state) tensors and block operators as (nblk, sbat, sblk, sblk) tensors.

Packed block operators and vectors with dense tensor storage.

pyzag.operators.dense.batch_lu_solve(lu: Tensor, pivots: Tensor, rhs: Tensor) → Tensor: Batched version of torch.linalg.lu_solve that accepts separate LU and pivot tensors.

class pyzag.operators.dense.DenseBlockVector(data: Tensor)

Bases: BlockVector

Dense tensor-backed packed block vector.

Parameters:: data (torch.Tensor) – shape (nblk, sbat, sblk)

property device: device: Execution device for torch-backed implementations.

property dtype: dtype: Data type for torch-backed implementations.

property nblk: int: Number of logical blocks in this vector.

property batch_size: int: Logical batch size.

property block_size: int: Logical size of one block (last axis).

clone() → DenseBlockVector: Return a safe copy of the vector.

norm(dim: int = -1) → Tensor: Compute the norm along dim. Returns a raw tensor (used for scalar convergence checks; the result has a different shape than a block vector and should not be wrapped).

flat_norm() → Tensor: Cross-block flattened L2 norm per batch element. Returns a raw tensor of shape (batch_size,). Counterpart to per-block norm(); used for convergence metrics that need a single scalar per batch (e.g. line search).

where(mask: Tensor, other: BlockVector) → DenseBlockVector

Batch-axis conditional combination, torch.where convention: returns self where mask is True, else other.

mask is a raw tensor of shape (batch_size,); backends broadcast it over the block and state axes. Used by the Newton step for partial updates.

scale_batches(factor: Tensor) → DenseBlockVector: Per-batch scalar broadcast: multiply each batch element by the corresponding entry of factor (raw tensor of shape (batch_size,), broadcast over block and state axes). Used by line-search backtracking.

flip(dim: int) → DenseBlockVector: Return a new vector with axis dim reversed.

classmethod cat(vectors: Sequence[BlockVector], dim: int = 0) → DenseBlockVector: Concatenate a sequence of compatible block vectors along dim.

classmethod zeros(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockVector: Return a zero-filled DenseBlockVector of the given shape.

classmethod zeros_like(other: BlockVector) → DenseBlockVector: Construct a zero-filled block vector with the same shape as other.

classmethod empty(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockVector: Return an uninitialized DenseBlockVector of the given shape.

class pyzag.operators.dense.DensePCRState(lu: Tensor, pivots: Tensor, b: Tensor, v: Tensor)

Bases: PCRState

Dense backend PCR working state.

Holds the four working tensors used by the Dense cyclic-shift PCR kernel. Each tensor has an extra leading dimension prepended by pcr_init() so that _dense_pcr_cyclic_shift() can double it at each level.

lu

shape (1, nblk, sbat, sblk, sblk) initially.

Type:: torch.Tensor

pivots

shape (1, nblk, sbat, sblk) initially.

Type:: torch.Tensor

b

shape (1, nblk, sbat, sblk, sblk) initially.

Type:: torch.Tensor

v

shape (1, nblk, sbat, sblk, 1) initially.

Type:: torch.Tensor

class pyzag.operators.dense.DenseBlockOperator(data: Tensor, lu: Tensor | None = None, pivots: Tensor | None = None)

Bases: SolvableBlockOperator

Dense tensor-backed packed block operator.

Implements SolvableBlockOperator (LU-based solve and PCR primitives pcr_init / pcr_reduce_level / pcr_finalize).

The cached LU factorization (self.lu / self.pivots) is materialized lazily on the first solve or pcr_init call. Use factored() to construct with eager factorization for cases where you know LU will be needed (e.g., the diagonal of a Newton system). Slicing, cloning, and in-place assignment preserve the cached LU when present, so chained Thomas-style A[i:i+1].solve(...) calls don’t re-factor every block.

Parameters:

data (torch.Tensor) – shape (nblk, sbat, sblk, sblk).
lu (torch.Tensor, optional) – pre-computed LU factor with the same shape as data.
pivots (torch.Tensor, optional) – pre-computed pivots, shape (nblk, sbat, sblk). Must be paired with lu.

classmethod factored(data: Tensor) → DenseBlockOperator: Construct with eager LU factorization. Use when the caller knows solve or pcr_init will be invoked.

property device: device: Execution device for torch-backed implementations.

property dtype: dtype: Data type for torch-backed implementations.

property nblk: int: Number of logical blocks in this operator.

property batch_size: int: Logical batch size expected in block-major vector inputs.

matvec(x: BlockVector) → DenseBlockVector

Apply the operator to a block vector x.

Required:: x.nblk == self.nblk

t_matvec(x: BlockVector) → DenseBlockVector

Apply the transpose of the operator to a block vector x.

Required:: x.nblk == self.nblk

solve(rhs: BlockVector) → DenseBlockVector

Solve the block system A x = rhs.

Required:: rhs.nblk == self.nblk

clone() → DenseBlockOperator: Return a safe copy of the operator.

pad_front(n: int = 1) → DenseBlockOperator: Return an operator with n leading dummy logical blocks (creates new data, not a view).

trim_front(n: int = 1) → DenseBlockOperator: Return an operator with the first n logical blocks removed.

pcr_init(B: BlockOperator, v: BlockVector) → DensePCRState

Initialise Dense PCR working state for a power-of-two window.

The state owns its b/v working buffers (cloned here), so the in-place reductions in pcr_reduce_level() never mutate the caller’s operator/vector – honouring the base-class contract that the backend allocates internal working tensors.

pcr_reduce_level(state: PCRState, level: int) → DensePCRState: Apply one Dense PCR level: update v and B, then cyclic-shift all four tensors.

pcr_finalize(state: PCRState) → tuple[DenseBlockOperator, DenseBlockVector]: Extract (B_red, v_red) with nblk = window_size - 1 from the final PCR state.

classmethod identity(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockOperator: Return an identity DenseBlockOperator of the given shape.

classmethod from_diagonal(data: Tensor) → DenseBlockOperator: Return a DenseBlockOperator built from per-block diagonal data.

class pyzag.operators.dense.DenseBlockJacobian(diag: Tensor, sub: Tensor, _reversed: bool = False)

Bases: BlockJacobian

Dense tensor-backed per-chunk Jacobian.

Storage is canonical (forward-time order). The as_adjoint_walk() method returns a sibling instance that shares the same underlying tensors and only flips a private flag; the four adjoint methods read the data with appropriate slicing/indexing under that flag, so a physical flip never materializes.

Parameters:

diag (torch.Tensor) – per-step diagonal blocks dR[k]/dx[k], shape (nblk_steps, batch, n, n).
sub (torch.Tensor) – per-step subdiagonal blocks dR[k]/dx[k-1], shape (nblk_steps, batch, n, n). The first block sub[0] is the inter-chunk boundary coupling to the lookback / previous chunk’s last state.

classmethod from_stacked(J: Tensor) → DenseBlockJacobian: Construct from a stacked tensor of shape (2, nblk_steps, batch, n, n) where J[1] is the diagonal and J[0] is the subdiagonal (with J[0, 0] being the boundary block).

property device: device: Execution device for torch-backed implementations.

property dtype: dtype: Data type for torch-backed implementations.

property nblk_steps: int: Number of residual rows in the chunk.

property batch_size: int: Logical batch size.

property block_size: int: Per-step state size (the user-facing n).

forward_system(inverse_operator)

Build the chunk’s forward bidiagonal system, ready for Newton.

Returns a pyzag.chunktime.BidiagonalForwardOperator whose diagonal A has nblk == nblk_steps and whose subdiagonal B has nblk == nblk_steps - 1.

Parameters:: inverse_operator – factory used to build the inverse (e.g. pyzag.chunktime.BidiagonalThomasFactorization).

adjoint_system(inverse_operator)

Build the chunk’s adjoint bidiagonal solve operator, in adjoint-walk order.

Transposes are baked in. The first-row / first+last-col slicing (J[1, 1:].T and J[0, 1:-1].T after flip(1)) is also baked in. Unlike forward_system(), which returns a BidiagonalForwardOperator (so Newton can call .inverse() between iterations), this returns the inverse / solve operator directly: applying .matvec(rhs) on the returned object yields the adjoint solution. This matches how block_update_adjoint consumes the result with a single linear solve per chunk.

This method should be called on a BlockJacobian returned by as_adjoint_walk().

solve_terminal_adjoint(g_terminal: Tensor) → DenseBlockVector

Compute -A_terminal^{-T} @ g_terminal for the very last forward-time step of the trajectory.

Returns a single-block BlockVector (nblk == 1 for lookback == 1; for higher lookback this would be a lookback-block vector). The returned vector must not alias g_terminal.

couple_prev_chunk(a_first: BlockVector) → DenseBlockVector

Compute the inter-chunk adjoint coupling B_boundary^T @ a_first.

a_first is a single-block BlockVector holding the previous chunk’s adjoint tail (in adjoint-walk order). Returns a single-block BlockVector to be subtracted into the current chunk’s RHS first row (nblk == 1 for lookback == 1; lookback-block for higher lookback).

as_adjoint_walk() → DenseBlockJacobian: Return a BlockJacobian whose forward time order is the reverse of this one. Backends are free to implement lazily (e.g. via a flag) so that no copy happens unless storage requires it. Replaces direct flip calls in solver code.