Solver

Spatial PDE solvers.

class immunowave.solver.SpatialPDETerm(vector_field)[source]

Bases: AbstractTerm

A term representing \(f(t, y(x, t), args) \mathrm{d}t\). That is to say, the term appearing on the right hand side of a PDE, in which the control is time.

vector_field should return some PyTree, with the same structure as the initial state y0, and with every leaf broadcastable to the equivalent leaf in y0.

Parameters:

vector_field (Callable[[Shaped[Array, ''], State, PyTree], State])

vector_field: Callable[[Shaped[Array, ''], State, PyTree], State]
vf(t, y, args)[source]

The vector field.

Represents a function $f(t, y(t), args)$.

Arguments:

  • t: the integration time.

  • y: the evolving state; a PyTree of structure $T$.

  • args: any static arguments as passed to [diffrax.diffeqsolve][].

Returns:

A PyTree of structure $S$.

Parameters:

  • t (Shaped[Array, ''])

  • y (State)

  • args (PyTree)

Return type:

State

static contr(t0, t1)[source]

The control.

Represents the $mathrm{d}t$ in an ODE, or the $mathrm{d}w(t)$ in an SDE, etc.

Most numerical ODE solvers work by making a step of length $Delta t = t_1 - t_0$. Likewise most numerical SDE solvers work by sampling some Brownian motion $Delta w sim mathcal{N}(0, t_1 - t_0)$.

Correspondingly a control is not defined at a point. Instead it is defined over an interval $[t_0, t_1]$.

Arguments:

  • t0: the start of the interval.

  • t1: the end of the interval.

Returns:

A PyTree of structure $U$. For a control $x$ then the result should represent $x(t_1) - x(t_0)$.

Parameters:

  • t0 (Shaped[Array, ''])

  • t1 (Shaped[Array, ''])

Return type:

Shaped[Array, ‘’]

static prod(vf, control)[source]

Determines the interaction between vector field and control.

With a solution $y$ to a differential equation with vector field $f$ and control $x$, this computes $f(t, y(t), args) Delta x(t)$ given $f(t, y(t), args)$ and $Delta x(t)$.

!!! note

This function must be bilinear.

Arguments:

  • vf: The vector field evaluation; a PyTree of structure $S$.

  • control: The control evaluated over an interval; a PyTree of structure $U$.

Returns:

The interaction between the vector field and control; a PyTree of structure $T$.

Parameters:

  • vf (State)

  • control (Shaped[Array, ''])

Return type:

PyTree

class immunowave.solver.CrankNicolson(rtol, atol)[source]

Bases: AbstractSolver

Crank-Nicolson solver for spatial PDEs, adapted from Diffrax docs.

Parameters:

  • rtol (float) – Relative tolerance.

  • atol (float) – Absolute tolerance.

rtol: float
atol: float
term_structure

alias of SpatialPDETerm

interpolation_cls

alias of LocalLinearInterpolation

order(terms)[source]

Order of the solver for solving ODEs.

init(terms, t0, t1, y0, args)[source]

Initialises any hidden state for the solver.

Arguments as [diffrax.diffeqsolve][].

Returns:

The initial solver state, which should be used the first time step is called.

step(terms, t0, t1, y0, args, solver_state, made_jump)[source]

Make a single step of the solver.

Each step is made over the specified interval $[t_0, t_1]$.

Arguments:

  • terms: The PyTree of terms representing the vector fields and controls.

  • t0: The start of the interval that the step is made over.

  • t1: The end of the interval that the step is made over.

  • y0: The current value of the solution at t0.

  • args: Any extra arguments passed to the vector field.

  • solver_state: Any evolving state for the solver itself, at t0.

  • made_jump: Whether there was a discontinuity in the vector field at t0.

    Some solvers (notably FSAL Runge–Kutta solvers) usually assume that there are no jumps and for efficiency re-use information between steps; this indicates that a jump has just occurred and this assumption is not true.

Returns:

A tuple of several objects:

  • The value of the solution at t1.

  • A local error estimate made during the step. (Used by adaptive step size

    controllers to change the step size.) May be None if no estimate was made.

  • Some dictionary of information that is passed to the solver’s interpolation

    routine to calculate dense output. (Used with SaveAt(ts=…) or SaveAt(dense=…).)

  • The value of the solver state at t1.

  • An integer (corresponding to diffrax.RESULTS) indicating whether the step

    happened successfully, or if (unusually) it failed for some reason.

func(terms, t0, y0, args)[source]

Evaluate the vector field at a point. (This is unlike [diffrax.AbstractSolver.step][], which operates over an interval.)

For most operations differential equation solvers are interval-based, so this opertion should be used sparingly. This operation is needed for things like selecting an initial step size.

Arguments: As [diffrax.diffeqsolve][]

Returns:

The evaluation of the vector field at t0, y0.