tensormesh.ode

Runge-Kutta methods

\[\begin{split}\begin{array}{c|c} \textbf c & \mathfrak A \\ \hline &\textbf b^\top \end{array} \quad = \quad \begin{array}{c|ccc} c_1 & a_{11} & \cdots & a_{1s}\\ \vdots & \vdots & \ddots & \vdots \\ c_s & a_{s1} & \cdots & a_{ss}\\\hline & b_1 & \cdots & b_s \end{array} \qquad \textbf c, \textbf b \in \mathbb R^s,\mathfrak A\in \mathbb R^{s\times s}\end{split}\]

\[\textbf k_i =\textbf f(t+c_i\tau, \textbf u +\tau \sum_{j=1}^s a_{ij}\textbf k_j)\quad \Psi^{t,t+\tau}\textbf u = \textbf u+\tau\sum_{i=1}^s b_i \textbf k_i\]

\[c_i = \sum_j a_{ij}\]

Base classes

class ExplicitRungeKutta(a, b)

Bases: object

Base class for explicit Runge-Kutta schemes.

Integrates the ODE

\[\frac{\partial u}{\partial t} = f(t, u)\]

with a user-supplied right-hand side forward(). The scheme is encoded by its Butcher tableau (a, b) and a single step() advances one time step.

Parameters:

• a (Tensor) –

  2D tensor of shape [s, s]; should be lower triangular.

  \[\begin{split}a = \begin{bmatrix} 0 & \cdots & 0 & 0 \\ a_{21} & \cdots & 0 & 0 \\ \vdots & \ddots & \vdots & \vdots \\ a_{s1} & \cdots & a_{s,s-1} & 0 \end{bmatrix}\end{split}\]

• b (Tensor) – 1D tensor of shape [s] with \(\sum_i b_i = 1\).

Examples

Solve \(\frac{\mathrm{d}u}{\mathrm{d}t} = u\) with explicit Euler:

import torch
from tensormesh.ode import ExplicitRungeKutta

class MyExplicitRungeKutta(ExplicitRungeKutta):
    def forward(self, t, u):
        return u

a = torch.zeros(1, 1)
b = torch.ones(1)
u0 = torch.rand(4)
dt = 0.1
ut = MyExplicitRungeKutta(a, b).step(0, u0, dt)

__init__(a, b)

forward(t, u)

Right-hand side of the ODE.

\[\frac{\partial u}{\partial t} = f(t, u)\]

Default returns u (i.e. \(f(t, u) = u\)); subclasses are expected to override.

Parameters:

• t (float) – Current time.

• u (Tensor) – State of shape [D] where D is the spatial dimension.

Returns:

\(f(t, u)\), same shape as u.

Return type:

Tensor

step(t0, u0, dt)

Advance one explicit Runge-Kutta step from t0 to t0 + dt.

\[\begin{split}k_i &= f\!\left(t_0 + c_i\,\tau,\ u_0 + \tau \sum_{j=1}^{s} a_{ij}\,k_j\right) \\ \Psi^{t_0,\,t_0 + \tau} u_0 &= u_0 + \tau \sum_{i=1}^{s} b_i\,k_i\end{split}\]

Parameters:

• t0 (float) – Initial time.

• u0 (Tensor) – Initial state of shape [D].

• dt (float) – Time step \(\tau\).

Returns:

State at time \(t_0 + \mathrm{d}t\), same shape as u0.

Return type:

Tensor

class ImplicitLinearRungeKutta(a, b)

Bases: object

Base class for implicit linear Runge-Kutta schemes.

Integrates the linear (in \(u\)) system

\[M(t) \frac{\partial u}{\partial t} = A(t) u + B(t)\]

where \(M(t),\,A(t) \in \mathbb{R}^{n \times n}\) and \(B(t) \in \mathbb{R}^{n}\). Each step() assembles and solves the block system

\[\begin{split}\begin{bmatrix} M_0 - A_0\tau a_{0,0} & -A_0\tau a_{0,1} & \cdots & -A_0\tau a_{0,s-1} \\ -A_1\tau a_{1,0} & M_1 - A_1\tau a_{1,1} & \cdots & -A_1\tau a_{1,s-1} \\ \vdots & \vdots & \ddots & \vdots \\ -A_{s-1}\tau a_{s-1,0} & -A_{s-1}\tau a_{s-1,1} & \cdots & M_{s-1} - A_{s-1}\tau a_{s-1,s-1} \end{bmatrix} \begin{bmatrix} k_0 \\ k_1 \\ \vdots \\ k_{s-1} \end{bmatrix} = \begin{bmatrix} B_0 + A_0 u \\ B_1 + A_1 u \\ \vdots \\ B_{s-1} + A_{s-1} u \end{bmatrix}\end{split}\]

for the stage values \(k_i\), then advances \(u_{n+1} = u_n + \tau \sum_i b_i\,k_i\).

Parameters:

• a (Tensor) – 2D tensor of shape [s, s].

• b (Tensor) – 1D tensor of shape [s] with \(\sum_i b_i = 1\).

Examples

Solve \(\frac{\mathrm{d}u}{\mathrm{d}t} = u\) with backward Euler:

import torch
from tensormesh.ode import ImplicitLinearRungeKutta

class MyImplicitLinearRungeKutta(ImplicitLinearRungeKutta):
    def forward_M(self, t):
        return torch.eye(4)

    def forward_A(self, t):
        return torch.eye(4)

    def forward_B(self, t):
        return torch.zeros(4)

a = torch.ones(1, 1)
b = torch.ones(1)
u0 = torch.rand(4).double()
dt = 0.1
ut = MyImplicitLinearRungeKutta(a, b).step(0, u0, dt)

__init__(a, b)

forward_M(t)

Mass-like operator \(M(t)\).

\[M(t) \frac{\partial u}{\partial t} = A(t) u + B(t)\]

Parameters:

t (float) – Current time.

Returns:

2D operator of shape [D, D]. If returned as int / float, \(M\) is taken to be that scalar multiple of the identity. Default returns 1.0 (i.e. \(M = I\)).

Return type:

tensormesh.sparse.SparseMatrix or torch.Tensor or float

forward_A(t)

Linear operator \(A(t)\) on the right-hand side.

\[M(t) \frac{\partial u}{\partial t} = A(t) u + B(t)\]

Parameters:

t (float) – Current time.

Returns:

2D operator of shape [D, D]. If returned as int / float, \(A\) is taken to be that scalar multiple of the identity. Default returns 1.0 (i.e. \(A = I\)).

Return type:

tensormesh.sparse.SparseMatrix or torch.Tensor or float

forward_B(t)

Source / forcing term \(B(t)\).

\[M(t) \frac{\partial u}{\partial t} = A(t) u + B(t)\]

Parameters:

t (float) – Current time.

Returns:

1D vector of shape [D]. If int / float, \(B\) is taken to be that scalar broadcast to all components. Default returns 0.0.

Return type:

Tensor or float

pre_solve_lhs(K)

Preprocess the assembled block matrix before solving.

Hook for boundary-condition condensation (or similar). Called once per [i][j] block.

Parameters:

K (torch.Tensor or tensormesh.sparse.SparseMatrix) – One block of the left-hand-side matrix.

Returns:

The (possibly condensed) block. Default returns K unchanged.

Return type:

torch.Tensor or tensormesh.sparse.SparseMatrix

pre_solve_rhs(f)

Preprocess each stage right-hand side before solving.

Hook for boundary-condition condensation (or similar). Called once per stage.

Parameters:

f (Tensor) – One stage of the right-hand-side vector.

Returns:

The (possibly condensed) vector. Default returns f unchanged.

Return type:

Tensor

recover_stage(k_i)

Lift one stage slope back from the solve space to full DOF.

Default: identity. Override when pre_solve_lhs() / pre_solve_rhs() reduce the system size — e.g. static condensation via Condenser. In that case each stage value \(k_i\) is solved in the inner-DOF subspace and must be prolonged back to the full DOF layout before being combined with \(u_0\).

For Dirichlet boundary conditions the boundary entries of \(k_i\) must be zero, regardless of the prescribed value of \(u\) there: a Dirichlet DOF is fixed, so its time-derivative is zero. Use prolong() (or an equivalent scatter-with-zero-boundary) here, not recover() (which writes the prescribed Dirichlet value into the boundary slots).

Parameters:

k_i (Tensor) – One stage slope from the solve, shape [D_solve].

Returns:

The lifted stage slope of shape [D]. Default returns k_i unchanged.

Return type:

Tensor

post_solve(u)

Postprocess the combined solution after the linear solve.

Runs after the stage slopes have already been lifted by recover_stage() and combined with \(u_0\), so u is always full-DOF here. Use this hook for things that act on the final state — clamping, normalization, projecting onto a manifold, etc. For static-condensation BCs, override recover_stage() instead.

Parameters:

u (Tensor) – Combined solution of shape [D].

Returns:

The (possibly post-processed) solution. Default returns u unchanged.

Return type:

Tensor

step(t0, u0, dt)

Advance one implicit-linear Runge-Kutta step from t0 to t0 + dt.

Builds the block system described in the class docstring, applies pre_solve_lhs() / pre_solve_rhs() to each block, solves it (via tensormesh.sparse.SparseMatrix.solve() when the operators are sparse, otherwise torch.linalg.solve()), applies post_solve(), and combines the stage values:

\[u_{n+1} = u_0 + \tau \sum_{i=1}^{s} b_i\,k_i\]

Parameters:

• t0 (float) – Initial time.

• u0 (Tensor) – Initial state of shape [D].

• dt (float) – Time step \(\tau\).

Returns:

State at time \(t_0 + \mathrm{d}t\), same shape as u0.

Return type:

Tensor

Built-in schemes

class ExplicitEuler

Bases: ExplicitRungeKutta

\[\begin{split}\begin{array}{c|c} \textbf{c} & \mathfrak{A} \\ \hline & \textbf{b}^\top \end{array} = \begin{array}{c|c} 0 & 0 \\ \hline & 1 \end{array}\end{split}\]

\[\Psi^{t,t+\tau}\textbf{u} \approx \textbf{u} + \tau \textbf{f}(t,\textbf{u})\]

Examples

\[\frac{\text{d}u}{\text{d}t} = u\]

import torch
from tensormesh.ode import ExplicitEuler

class MyExplicitEuler(ExplicitEuler):
    def forward(self, t, u):
        return u

u0 = torch.rand(4)
dt = 0.1
ut_gt = u0 + dt * u0
ut_my = MyExplicitEuler().step(0, u0, dt)
assert torch.allclose(ut_gt, ut_my)

__init__()

class ImplicitLinearEuler

Bases: ImplicitLinearRungeKutta

\[\begin{split}\begin{array}{c|c} \textbf{c} & \mathfrak{A} \\ \hline & \textbf{b}^\top \end{array} = \begin{array}{c|c} 1 & 1 \\ \hline & 1 \end{array}\end{split}\]

\[\Psi^{t,t+\tau}\textbf{u} \approx \textbf{w}\quad \textbf{w}=\textbf{u}+\tau\textbf{f}(t+\tau,\textbf{w})\]

Examples

\[\frac{\text{d}u}{\text{d}t} = u\]

import torch
from tensormesh.ode import ImplicitLinearEuler

u0 = torch.rand(4).double()
dt = 0.1
ut_gt = (1/(1-dt)) * u0
ut_my = ImplicitLinearEuler().step(0, u0, dt)
assert torch.allclose(ut_gt, ut_my), f"expected {ut_gt}, got {ut_my}"

__init__()

class MidPointLinearEuler

Bases: ImplicitLinearRungeKutta

\[\begin{split}\begin{array}{c|c} \textbf{c} & \mathfrak{A} \\ \hline & \textbf{b}^\top \end{array} = \begin{array}{c|c} \frac{1}{2} & \frac{1}{2} \\ \hline & 1 \end{array}\end{split}\]

\[\Psi^{t,t+\tau}\textbf{u} \approx \textbf{w}\quad \textbf{w} = \textbf{u} +\tau \textbf{f}\left(t+\frac{\tau}{2},\frac{\textbf{w}+\textbf{u}}{2}\right)\]

Examples

\[\frac{\text{d} u}{\text{d} t} = u\]

import torch
from tensormesh.ode import MidPointLinearEuler

u0 = torch.rand(4)
dt = 0.1
ut_gt = ((dt+2)/(2-dt)) * u0
ut_my = MidPointLinearEuler().step(0, u0, dt)
assert torch.allclose(ut_gt, ut_my), f"expected {ut_gt}, got {ut_my}"

__init__()