torch_fem.ode
Runge-Kutta methods
Built-in Methods
- class ExplicitRungeKutta(a, b)[source]
Bases:
object\[\frac{\partial u}{\partial t} = f(t, u)\]- Parameters:
a (torch.Tensor) –
2D tensor of shape [s, s] shoule 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 (torch.Tensor) – 1D tensor of shape [s], \(\sum_{b_i} = 1\)
- forward(t, u)[source]
right side function \(f(t, u)\)
\[\frac{\partial u}{\partial t} = \text{forward}(t, u)\]default \(f(t, u) = u\)
- Parameters:
t (float) – time
u (torch.Tensor) – value of shape \([D]\) where D is the dimension of the problem
- Returns:
value of shape \([D]\) where D is the dimension of the problem the value of the right side function \(f(t, u)\)
- Return type:
- step(t0, u0, dt)[source]
one step of explicit Runge-Kutta method
\[\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\]- Parameters:
t0 (float) – initial time
u0 (torch.Tensor) – initial value of shape \([D]\) where D is the dimension of the problem
dt (float) – time step
- Returns:
value of shape \([D]\) where D is the dimension of the problem the value of the solution at time \(t_0 + \text{d}t\)
- Return type:
- class ImplicitLinearRungeKutta(a, b)[source]
Bases:
object\[M(t) \frac{\partial u}{\partial t} = A(t) u + B(t)\]\(M\in \mathbb R^{n\times n}\)
\(A\in \mathbb R^{n\times n}\)
\(B\in \mathbb R^{n}\)
\(u\in \mathbb R^n\)
\[\begin{split}\begin{bmatrix} M_0 - A_0\tau a_{0,0}& - A_0\tau a_{0,1}&\cdots & - A_{0}\tau a_{0,{n-1}}\\ -A_1\tau a_{1,0}& M_1-A_1\tau a_{1,1} & \cdots & - A_{1}\tau a_{1,{n-1}}\\ \vdots & \vdots &\ddots & \vdots \\ -A_{n-1}\tau a_{{n-1},0} & -A_{n-1}\tau a_{{n-1},1} & \cdots & M_{n-1} - A_{n-1}\tau a_{n-1,n-1} \end{bmatrix} \begin{bmatrix} \textbf k_0\\ \textbf k_1 \\\vdots \\\textbf k_{n-1} \end{bmatrix}= \begin{bmatrix} B_0 + A_0 u \\ B_1 + A_1 u \\ \vdots\\ B_{n-1} + A_{n-1} u \end{bmatrix}\end{split}\]- forward_M(t)[source]
left side matrix
\[M \frac{\partial u}{\partial t} = A(t)u + B(t)\]- Parameters:
t (float) – time
- Returns:
normally, 2D
torch.Tensor()ortorch_fem.sparse.SparseMatrix()of shape \([D, D]\) where \(D\) is the dimension of the problem; if returnintorfloat, the left side matrix \(M\) is assumed to be a diagonal matrix with the same value- Return type:
- forward_A(t)[source]
compute the linear mapping term \(A(t)\)
\[M \frac{\partial u}{\partial t} = A(t)u + B(t)\]- Parameters:
t (float) – time
- Returns:
2D
torch.Tensor()ortorch_fem.sparse.SparseMatrix()of shape \([D, D]\) where \(D\) is the dimension of the problem; if returnintorfloat, the linear mapping term is assumed to be a diagonal matrix with the same value- Return type:
- forward_B(t)[source]
compute the linear mapping term \(B(t)\)
\[M \frac{\partial u}{\partial t} = A(t)u + B(t)\]
- pre_solve_lhs(K)[source]
precompute something before solving the linear system, for example, do the condensation
- Parameters:
K (torch.Tensor or torch_fem.sparse.SparseMatrix) – the left side matrix
- Returns:
the left side matrix after precompute
- Return type:
- pre_solve_rhs(f)[source]
precompute something before solving the linear system, for example, do the condensation
- Parameters:
f (torch.Tensor) – the right hand side vector
- Returns:
the right hand side vector after precompute
- Return type:
- post_solve(u)[source]
postprocess after solving the linear system, for example, do the condensation recovery
- Parameters:
u (torch.Tensor) – the solution of the linear system
- Returns:
the solution after postprocess
- Return type:
- step(t0, u0, dt)[source]
- \[\]
- Parameters:
t0 (float) – initial time
u0 (torch.Tensor) – initial value of shape \([D]\) where D is the dimension of the problem
dt (float) – time step
- class ExplicitEuler[source]
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 torch_fem.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)
- class ImplicitLinearEuler[source]
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 torch_fem.ode import ImplicitLinearEuler class MyImplicitLinearEuler(ImplicitLinearEuler): pass u0 = torch.rand(4).double() dt = 0.1 ut_gt = (1/(1-dt)) * u0 ut_my = MyImplicitLinearEuler().step(0, u0, dt) assert torch.allclose(ut_gt, ut_my), f"expected {ut_gt}, got {ut_my}"
- class MidPointLinearEuler[source]
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 torch_fem.ode import MidPointLinearEuler class MyMidPointLinearEuler(MidPointLinearEuler): pass u0 = torch.rand(4) dt = 0.1 ut_gt = ((dt+2)/(2-dt)) * u0 ut_my = MyMidPointLinearEuler().step(0, u0, dt) assert torch.allclose(ut_gt, ut_my), f"expected {ut_gt}, got {ut_my}"