pyzag.chunktime

Functions and objects to help with blocked/chunked time integration.

class pyzag.chunktime.ChunkNewtonRaphson(rtol: float = 1e-06, atol: float = 1e-10, miter: int = 200, throw_on_fail: bool = False, record_failed: bool = False, ignore_batches: Sequence[int] | None = None)

Bases: object

Solve a nonlinear system with Newton’s method where the residual and Jacobian are presented as chunked operators

Keyword Arguments:

  • rtol (float) – nonlinear relative tolerance

  • atol (float) – nonlinear absolute tolerance

  • miter (int) – maximum number of iterations

  • throw_on_fail (bool) – if True, throw an exception on a failed solve. If False just issue a warning.

  • record_failed (bool) – if True, store the indices of the bad batches

  • ignore_batches (list of indices) – if provided, don’t check these batches in evaluating the stopping criteria

setup(x: BlockVector) → None

Do any initialization required before solving

solve(fn: Callable[[BlockVector], tuple[BlockVector, BidiagonalForwardOperator]], x0: BlockVector) → BlockVector

Solve the nonlinear system.

Parameters:

  • fn – callable that returns (R, J) where R is a BlockVector and J is a BidiagonalForwardOperator.

  • x0 – initial guess as a BlockVector.

Returns:

solution

Return type:

BlockVector

not_converged(nR: Tensor, nR0: Tensor) → Tensor

The logical to determine if we’ve converged in a particular time/batch.

step(x: BlockVector, J: BidiagonalForwardOperator, fn: Callable[[BlockVector], tuple[BlockVector, BidiagonalForwardOperator]], R0: BlockVector, take_step: Tensor) → tuple[BlockVector, BlockVector, BidiagonalForwardOperator, Tensor]

Take a simple Newton step.

Partial step application uses the abstract BlockVector.where() primitive: the candidate x - dx is committed only for batches whose entries in final_steps are True; converged batches keep their current value.

class pyzag.chunktime.ChunkNewtonRaphsonLineSearch(*args, alpha: float = 0.5, linesearch_iter: int = 3, **kwargs)

Bases: ChunkNewtonRaphson

Newton Raphson with backtracking line search

Keyword Arguments:

  • rtol (float) – nonlinear relative tolerance

  • atol (float) – nonlinear absolute tolerance

  • miter (int) – maximum number of iterations

  • throw_on_fail (bool) – if True, throw an exception on a failed solve. If False just issue a warning.

  • record_failed (bool) – if True, store the indices of the bad batches

  • ignore_batches (list of indices) – if provided, don’t check these batches in evaluating the stopping criteria

  • alpha (float) – line search cutback

  • linesearch_iter (int) – maximum number of line search iterations

step(x: BlockVector, J: BidiagonalForwardOperator, fn: Callable[[BlockVector], tuple[BlockVector, BidiagonalForwardOperator]], R0: BlockVector, take_step: Tensor) → tuple[BlockVector, BlockVector, BidiagonalForwardOperator, Tensor]

Take a Newton step with backtracking line search.

Uses the abstract BlockVector.flat_norm(), BlockVector.scale_batches(), and BlockVector.where() primitives so no backend storage is touched directly.

class pyzag.chunktime.BidiagonalOperator(A: BlockOperator, B: BlockOperator, *args, **kwargs)

Bases: Module

Base class for block bidiagonal operators.

property dtype: dtype

The dtype of the underlying operator blocks.

property device: device

The device of the underlying operator blocks.

property batch_size: int

The batch size of the underlying operator blocks.

class pyzag.chunktime.BidiagonalInverseOperator(A: BlockOperator, B: BlockOperator, *args, **kwargs)

Bases: BidiagonalOperator

Base for bidiagonal inverse operators (Thomas / PCR / hybrid).

Applying the operator solves the bidiagonal system, i.e. it acts as the inverse. The actual factorization now lives in the SolvableBlockOperator backend (e.g. DenseBlockOperator); this class only aliases forward to matvec().

forward(v: BlockVector) → BlockVector

Apply the inverse operator (solve the system) for a vector v

pyzag.chunktime.thomas_solve(A: BlockOperator, B: BlockOperator, v: BlockVector) → BlockVector

Generic Thomas solve over block views.

All vector and operator arguments use the abstract block interfaces.

class pyzag.chunktime.BidiagonalThomasFactorization(A: BlockOperator, B: BlockOperator, *args, **kwargs)

Bases: BidiagonalInverseOperator

Manages the data needed to solve our bidiagonal system via Thomas factorization.

matvec(v: BlockVector) → BlockVector

Apply the Thomas factorization.

class pyzag.chunktime.BidiagonalPCRFactorization(A: BlockOperator, B: BlockOperator, *args, **kwargs)

Bases: BidiagonalInverseOperator

PCR factorization — algorithm lives here, backend provides pcr_init/reduce_level/finalize.

matvec(v: BlockVector) → BlockVector

Apply the PCR factorization.

class pyzag.chunktime.BidiagonalHybridFactorizationImpl(*args, min_size: int = 0, **kwargs)

Bases: BidiagonalPCRFactorization

Hybrid PCR/Thomas factorization.

matvec(v: BlockVector) → BlockVector

Apply the hybrid PCR/Thomas factorization.

pyzag.chunktime.BidiagonalHybridFactorization(min_size: int = 1)

Factory wrapper for the hybrid factorization with a given min_size.

class pyzag.chunktime.BidiagonalForwardOperator(*args, inverse_operator=<class 'pyzag.chunktime.BidiagonalThomasFactorization'>, **kwargs)

Bases: BidiagonalOperator

Forward bidiagonal operator that wraps an inverse-operator factory.

forward(v: BlockVector) → BlockVector

Apply the forward bidiagonal operator to a vector v.

matvec(v: BlockVector) → BlockVector

Return the matrix-vector product of the bidiagonal operator with v.

vecmat(v: BlockVector) → BlockVector

Return the transpose matrix-vector product of the operator with v.

inverse() → BidiagonalInverseOperator

Return the inverse operator built via the configured factory.

class pyzag.chunktime.SquareBatchedBlockDiagonalMatrix(data, diags)

Bases: object

Utility for converting block-diagonal data into dense / sparse representations.

property dtype: dtype

The dtype of the block-diagonal data.

property device: device

The device of the block-diagonal data.

property n: int

Size of the unbatched square matrix.

property shape: tuple[int, int, int]

Logical shape of the dense array.

property nnz: int

Number of logical non-zeros (not counting the batch dimension).

to_dense() → Tensor

Convert the representation to a dense tensor.

to_batched_coo() → Tensor

Convert to a torch sparse batched COO tensor.

to_unrolled_csr() → list[Tensor]

Return a list of CSR tensors with length equal to the batch size.