Skip to content

torchebm.integrators.rk4

Classical 4th-order Runge-Kutta integrator (fixed step).

RK4Integrator

Bases: BaseRungeKuttaIntegrator

Classical 4th-order Runge-Kutta integrator.

A 4-stage, 4th-order explicit method. Fixed-step only (no embedded error pair).

The update rule computes four stages and combines them:

\[ k_1 = f(x_n,\, t_n) \]
\[ k_2 = f\!\bigl(x_n + \tfrac{h}{2}\,k_1,\; t_n + \tfrac{h}{2}\bigr) \]
\[ k_3 = f\!\bigl(x_n + \tfrac{h}{2}\,k_2,\; t_n + \tfrac{h}{2}\bigr) \]
\[ k_4 = f(x_n + h\,k_3,\; t_n + h) \]
\[ x_{n+1} = x_n + \tfrac{h}{6}\bigl(k_1 + 2k_2 + 2k_3 + k_4\bigr) \]

The Butcher tableau:

\[ \begin{array}{c|cccc} 0 \\ \tfrac{1}{2} & \tfrac{1}{2} \\ \tfrac{1}{2} & 0 & \tfrac{1}{2} \\ 1 & 0 & 0 & 1 \\ \hline & \tfrac{1}{6} & \tfrac{1}{3} & \tfrac{1}{3} & \tfrac{1}{6} \end{array} \]
Reference

Kutta, W. (1901). Beitrag zur naherungsweisen Integration totaler Differentialgleichungen. Z. Math. Phys., 46, 435--453.

Parameters:

Name Type Description Default
device Optional[device]

Device for computations.

 None
dtype Optional[dtype]

Data type for computations.

 None
Example 
1
2
3
4
5
6
7
8
9
from torchebm.integrators import RK4Integrator
import torch

integrator = RK4Integrator()
state = {"x": torch.randn(100, 2)}
drift = lambda x, t: -x
result = integrator.integrate(
    state, step_size=0.01, n_steps=100, drift=drift,
)
Source code in torchebm/integrators/rk4.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class RK4Integrator(BaseRungeKuttaIntegrator):
    r"""Classical 4th-order Runge-Kutta integrator.

    A 4-stage, 4th-order explicit method. Fixed-step only (no embedded
    error pair).

    The update rule computes four stages and combines them:

    \[
    k_1 = f(x_n,\, t_n)
    \]

    \[
    k_2 = f\!\bigl(x_n + \tfrac{h}{2}\,k_1,\; t_n + \tfrac{h}{2}\bigr)
    \]

    \[
    k_3 = f\!\bigl(x_n + \tfrac{h}{2}\,k_2,\; t_n + \tfrac{h}{2}\bigr)
    \]

    \[
    k_4 = f(x_n + h\,k_3,\; t_n + h)
    \]

    \[
    x_{n+1} = x_n + \tfrac{h}{6}\bigl(k_1 + 2k_2 + 2k_3 + k_4\bigr)
    \]

    The Butcher tableau:

    \[
    \begin{array}{c|cccc}
    0 \\
    \tfrac{1}{2} & \tfrac{1}{2} \\
    \tfrac{1}{2} & 0 & \tfrac{1}{2} \\
    1 & 0 & 0 & 1 \\
    \hline
    & \tfrac{1}{6} & \tfrac{1}{3} & \tfrac{1}{3} & \tfrac{1}{6}
    \end{array}
    \]

    Reference:
        Kutta, W. (1901). Beitrag zur naherungsweisen Integration
        totaler Differentialgleichungen. Z. Math. Phys., 46, 435--453.

    Args:
        device: Device for computations.
        dtype: Data type for computations.

    Example:
        ```python
        from torchebm.integrators import RK4Integrator
        import torch

        integrator = RK4Integrator()
        state = {"x": torch.randn(100, 2)}
        drift = lambda x, t: -x
        result = integrator.integrate(
            state, step_size=0.01, n_steps=100, drift=drift,
        )
        ```
    """

    @property
    def tableau_a(self) -> Tuple[Tuple[float, ...], ...]:
        return (
            (),
            (1 / 2,),
            (0.0, 1 / 2),
            (0.0, 0.0, 1.0),
        )

    @property
    def tableau_b(self) -> Tuple[float, ...]:
        return (1 / 6, 1 / 3, 1 / 3, 1 / 6)

    @property
    def tableau_c(self) -> Tuple[float, ...]:
        return (0.0, 1 / 2, 1 / 2, 1.0)