pyzag.nonlinear

Solving recursive nonlinear equations and computing parameter sensitivities via the adjoint method.

class pyzag.nonlinear.NonlinearFunctionOperator

Bases: ABC

Abstract operator passed to the Newton solver.

Forces and time are captured at construction. Newton only passes the state x as a BlockVector and receives back (R, J) where R is a BlockVector and J is a chunktime.BidiagonalForwardOperator.

class pyzag.nonlinear.NonlinearFunctionOperatorFactory

Bases: ABC

Factory for chunk-level NonlinearFunctionOperator objects.

The forward solve calls make_operator() once per chunk to create a stateful operator that captures the chunk’s previous-solution lookback and forces. The adjoint pass calls evaluate_raw() to recompute the residual and Jacobian; the latter as a backend-typed BlockJacobian consumed by the adjoint walk.

Subclasses must expose lookback and wrapper (typically as properties) and implement make_operator() and evaluate_raw(). The canonical ChunkOp below is a drop-in implementation that most subclasses can return from make_operator without modification.

abstract property lookback: int

Number of previous solution steps the residual depends on.

Declared abstract so a subclass that forgets it fails at instantiation rather than with a late AttributeError deep in the solver.

abstract property wrapper: Backend bridge (an ODEWrapper-like object) that wraps/unwraps raw tensors as backend BlockVector objects.

abstractmethod make_operator(prev_solution: Tensor, forces: Sequence[Tensor], inverse_operator) → NonlinearFunctionOperator: Build a NonlinearFunctionOperator for a single chunk.

abstractmethod evaluate_raw(x_full: Tensor, forces: Sequence[Tensor]) → tuple[Tensor, BlockJacobian]

Return (R, J) for adjoint reconstruction.

x_full includes the lookback steps prepended to the chunk state. R is a raw torch tensor; J is a backend-typed BlockJacobian.

class pyzag.nonlinear.ChunkOp(factory: NonlinearFunctionOperatorFactory, prev_solution: Tensor, forces: Sequence[Tensor], inverse_operator)

Bases: NonlinearFunctionOperator

Generic per-chunk operator for any NonlinearFunctionOperatorFactory.

Holds the chunk’s previous-solution lookback and forces; on each Newton call, assembles the full state, asks the factory for raw (R, J), then wraps R via the factory’s ODEWrapper and builds the chunk’s forward bidiagonal system from the BlockJacobian. Most factories can return ChunkOp(self, ...) from their make_operator without modification.

class pyzag.nonlinear.FullTrajectoryPredictor(history: Tensor)

Bases: object

Predict steps using a complete user-provided trajectory.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.ZeroPredictor

Bases: object

Predict steps just using zeros.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.PreviousStepsPredictor

Bases: object

Predict by providing the values from the previous chunk of steps.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.LastStepPredictor

Bases: object

Predict by providing the values from the previous single step.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.StepExtrapolatingPredictor

Bases: object

Predict by extrapolating using the previous chunks of steps.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.ExtrapolatingPredictor

Bases: object

Predict by extrapolating the values from the previous single steps.

predict(results: Tensor, k: int, kinc: int) → Tensor

class pyzag.nonlinear.StepGenerator(block_size: int = 1, first_block_size: int = 0)

Bases: object

Generate chunks of recursive steps to produce at once.

reverse() → StepGenerator: Reverse the iterator to yield chunks starting from the end.

class pyzag.nonlinear.InitialOffsetStepGenerator(*args, initial_steps: Sequence[int] | None = None, **kwargs)

Bases: StepGenerator

Generate chunks of recursive steps with optional user-provided initial chunks.

class pyzag.nonlinear.RecursiveNonlinearEquationSolver(func: ~pyzag.nonlinear.NonlinearFunctionOperatorFactory, step_generator: ~pyzag.nonlinear.StepGenerator = <pyzag.nonlinear.StepGenerator object>, predictor=<pyzag.nonlinear.ZeroPredictor object>, direct_solve_operator=<class 'pyzag.chunktime.BidiagonalThomasFactorization'>, nonlinear_solver: ~pyzag.chunktime.ChunkNewtonRaphson = <pyzag.chunktime.ChunkNewtonRaphson object>, callbacks: ~typing.Sequence[~typing.Callable[[~torch.Tensor], ~torch.Tensor]] | None = None, convert_nan_gradients: bool = True)

Bases: Module

Generates a time series from a recursive nonlinear equation and (optionally) uses the adjoint method to provide derivatives.

func is a NonlinearFunctionOperatorFactory. The factory knows how to create chunk-level operators (for the forward solve) and how to recompute residuals / Jacobians (for the adjoint). Built-in factories include pyzag.ode.BackwardEulerODE and pyzag.ode.ForwardEulerODE; users can also subclass NonlinearFunctionOperatorFactory directly and return ChunkOp from their make_operator.

forward(*args, **kwargs): Alias for solve.

solve(y0: Tensor, n: int, *args: Tensor, adjoint_params: Sequence[Tensor] | None = None) → Tensor: Solve the recursive equations for n steps.

block_update(prev_solution: Tensor, solution: Tensor, forces: Sequence[Tensor]) → Tensor

Solve one chunk via Newton.

Builds a chunk-level operator from the factory, wraps the initial guess as a BlockVector, runs Newton, and unwraps back to a raw tensor for the result array.

rewind(output_grad: Tensor) → tuple[tuple[Tensor, ...], Tensor]: Rewind through an adjoint pass to provide the dot product for each quantity in output_grad.

accumulate(grad_result: tuple[Tensor, ...], adjoint: BlockVector, R: Tensor, retain: bool = False) → tuple[Tensor, ...]: Accumulate the updated gradient values.

block_update_adjoint(J: BlockJacobian, grads: BlockVector, a_prev: BlockVector) → BlockVector

Do the blocked adjoint solve.

J is the chunk’s BlockJacobian in adjoint-walk order (i.e. the result of BlockJacobian.as_adjoint_walk()). The single-block inter-chunk coupling is delegated to BlockJacobian.couple_prev_chunk() so the solver never touches raw tensor indexing.

pyzag.nonlinear.solve(solver, y0, n, *forces): Solve a RecursiveNonlinearEquationSolver without the adjoint method.

pyzag.nonlinear.solve_adjoint(solver, y0, n, *forces): Apply a RecursiveNonlinearEquationSolver to solve adjoint-differentiably.