Skip to content

ETDRK Backbone¤

Core clases that implement the Exponential Time Differencing Runge-Kutta (ETDRK) method for solving semi-linear PDEs in form of timesteppers. Require supplying the time step size \(\Delta t\), the linear operator in Fourier space \(\hat{\mathcal{L}}_h\), and the non-linear operator in Fourier space \(\hat{\mathcal{N}}_h\).

exponax.etdrk.ETDRK0 ¤

Bases: BaseETDRK

Source code in exponax/etdrk/_etdrk_0.py 
 6
 7
 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
class ETDRK0(BaseETDRK):
    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
    ):
        r"""
        Exactly solve a linear PDE in Fourier space.

        $$
            \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot
            \hat{u}_h^{[t]}
        $$

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.
        """
        super().__init__(dt, linear_operator)

    def step_fourier(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        return self._exp_term * u_hat
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
)

Exactly solve a linear PDE in Fourier space.

\[ \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]} \]

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.
Source code in exponax/etdrk/_etdrk_0.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
):
    r"""
    Exactly solve a linear PDE in Fourier space.

    $$
        \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot
        \hat{u}_h^{[t]}
    $$

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.
    """
    super().__init__(dt, linear_operator)
step_fourier ¤
step_fourier(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]
Source code in exponax/etdrk/_etdrk_0.py 
30
31
32
33
34
def step_fourier(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    return self._exp_term * u_hat

exponax.etdrk.ETDRK1 ¤

Bases: BaseETDRK

Source code in exponax/etdrk/_etdrk_1.py 
 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
class ETDRK1(BaseETDRK):
    _nonlinear_fun: BaseNonlinearFun
    _coef_1: Complex[Array, "E ... (N//2)+1"]

    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
        nonlinear_fun: BaseNonlinearFun,
        *,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        r"""
        Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
        with a **first order approximation**.

        Adapted from Eq. (4) of [Cox and Matthews
        (2002)](https://doi.org/10.1006/jcph.2002.6995):

        $$
            \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot
            \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) -
            1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]})
        $$

        where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
        the nonlinear differential operator.

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.
        - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
            nonlinear differential operator.
        - `num_circle_points`: The number of points on the unit circle used to
            approximate the numerically challenging coefficients.
        - `circle_radius`: The radius of the circle used to approximate the
            numerically challenging coefficients.

        !!! warning
            The nonlinear function must take care of proper dealiasing.

        !!! note
            The numerically stable evaluation of the coefficients follows
            [Kassam and Trefethen
            (2005)](https://doi.org/10.1137/S1064827502410633).
        """
        super().__init__(dt, linear_operator)
        self._nonlinear_fun = nonlinear_fun

        LR = (
            circle_radius * roots_of_unity(num_circle_points)
            + linear_operator[..., jnp.newaxis] * dt
        )

        self._coef_1 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real

    def step_fourier(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        return self._exp_term * u_hat + self._coef_1 * self._nonlinear_fun(u_hat)
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta with a first order approximation.

Adapted from Eq. (4) of Cox and Matthews (2002):

\[ \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \]

where \(\hat{\mathcal{N}}_h\) is the Fourier pseudo-spectral treatment of the nonlinear differential operator.

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.
  • nonlinear_fun: The Fourier pseudo-spectral treatment of the nonlinear differential operator.
  • num_circle_points: The number of points on the unit circle used to approximate the numerically challenging coefficients.
  • circle_radius: The radius of the circle used to approximate the numerically challenging coefficients.

Warning

The nonlinear function must take care of proper dealiasing.

Note

The numerically stable evaluation of the coefficients follows Kassam and Trefethen (2005).

Source code in exponax/etdrk/_etdrk_1.py 
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
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    r"""
    Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
    with a **first order approximation**.

    Adapted from Eq. (4) of [Cox and Matthews
    (2002)](https://doi.org/10.1006/jcph.2002.6995):

    $$
        \hat{u}_h^{[t+1]} = \exp(\hat{\mathcal{L}}_h \Delta t) \odot
        \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) -
        1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]})
    $$

    where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
    the nonlinear differential operator.

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.
    - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
        nonlinear differential operator.
    - `num_circle_points`: The number of points on the unit circle used to
        approximate the numerically challenging coefficients.
    - `circle_radius`: The radius of the circle used to approximate the
        numerically challenging coefficients.

    !!! warning
        The nonlinear function must take care of proper dealiasing.

    !!! note
        The numerically stable evaluation of the coefficients follows
        [Kassam and Trefethen
        (2005)](https://doi.org/10.1137/S1064827502410633).
    """
    super().__init__(dt, linear_operator)
    self._nonlinear_fun = nonlinear_fun

    LR = (
        circle_radius * roots_of_unity(num_circle_points)
        + linear_operator[..., jnp.newaxis] * dt
    )

    self._coef_1 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real
step_fourier ¤
step_fourier(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]
Source code in exponax/etdrk/_etdrk_1.py 
70
71
72
73
74
def step_fourier(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    return self._exp_term * u_hat + self._coef_1 * self._nonlinear_fun(u_hat)

exponax.etdrk.ETDRK2 ¤

Bases: BaseETDRK

Source code in exponax/etdrk/_etdrk_2.py 
 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
87
88
89
90
91
92
class ETDRK2(BaseETDRK):
    _nonlinear_fun: BaseNonlinearFun
    _coef_1: Complex[Array, "E ... (N//2)+1"]
    _coef_2: Complex[Array, "E ... (N//2)+1"]

    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
        nonlinear_fun: BaseNonlinearFun,
        *,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        r"""
        Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
        with a **second order approximation**.

        Adopted from Eq. (22) of [Cox and Matthews
        (2002)](https://doi.org/10.1006/jcph.2002.6995):

        $$
            \begin{aligned}
                \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot
                \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) -
                1}{\hat{\mathcal{L}}_h} \odot
                \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ \hat{u}_h^{[t+1]} &=
                \hat{u}_h^* + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1 -
                \hat{\mathcal{L}}_h \Delta t}{\hat{\mathcal{L}}_h^2 \Delta t}
                \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) -
                \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right)
            \end{aligned}
        $$

        where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
        the nonlinear differential operator.

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.
        - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
            nonlinear differential operator. ! The operator must take care of
            proper dealiasing.
        - `num_circle_points`: The number of points on the unit circle used to
            approximate the numerically challenging coefficients.
        - `circle_radius`: The radius of the circle used to approximate the
            numerically challenging coefficients.

        !!! warning
            The nonlinear function must take care of proper dealiasing.

        !!! note
            The numerically stable evaluation of the coefficients follows
            [Kassam and Trefethen
            (2005)](https://doi.org/10.1137/S1064827502410633).
        """
        super().__init__(dt, linear_operator)
        self._nonlinear_fun = nonlinear_fun

        LR = (
            circle_radius * roots_of_unity(num_circle_points)
            + linear_operator[..., jnp.newaxis] * dt
        )

        self._coef_1 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real

        self._coef_2 = dt * jnp.mean((jnp.exp(LR) - 1 - LR) / LR**2, axis=-1).real

    def step_fourier(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        u_nonlin_hat = self._nonlinear_fun(u_hat)
        u_stage_1_hat = self._exp_term * u_hat + self._coef_1 * u_nonlin_hat

        u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
        u_next_hat = u_stage_1_hat + self._coef_2 * (
            u_stage_1_nonlin_hat - u_nonlin_hat
        )
        return u_next_hat
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta with a second order approximation.

Adopted from Eq. (22) of Cox and Matthews (2002):

\[ \begin{aligned} \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ \hat{u}_h^{[t+1]} &= \hat{u}_h^* + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1 - \hat{\mathcal{L}}_h \Delta t}{\hat{\mathcal{L}}_h^2 \Delta t} \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right) \end{aligned} \]

where \(\hat{\mathcal{N}}_h\) is the Fourier pseudo-spectral treatment of the nonlinear differential operator.

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.
  • nonlinear_fun: The Fourier pseudo-spectral treatment of the nonlinear differential operator. ! The operator must take care of proper dealiasing.
  • num_circle_points: The number of points on the unit circle used to approximate the numerically challenging coefficients.
  • circle_radius: The radius of the circle used to approximate the numerically challenging coefficients.

Warning

The nonlinear function must take care of proper dealiasing.

Note

The numerically stable evaluation of the coefficients follows Kassam and Trefethen (2005).

Source code in exponax/etdrk/_etdrk_2.py 
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
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    r"""
    Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
    with a **second order approximation**.

    Adopted from Eq. (22) of [Cox and Matthews
    (2002)](https://doi.org/10.1006/jcph.2002.6995):

    $$
        \begin{aligned}
            \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot
            \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) -
            1}{\hat{\mathcal{L}}_h} \odot
            \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ \hat{u}_h^{[t+1]} &=
            \hat{u}_h^* + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1 -
            \hat{\mathcal{L}}_h \Delta t}{\hat{\mathcal{L}}_h^2 \Delta t}
            \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) -
            \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right)
        \end{aligned}
    $$

    where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
    the nonlinear differential operator.

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.
    - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
        nonlinear differential operator. ! The operator must take care of
        proper dealiasing.
    - `num_circle_points`: The number of points on the unit circle used to
        approximate the numerically challenging coefficients.
    - `circle_radius`: The radius of the circle used to approximate the
        numerically challenging coefficients.

    !!! warning
        The nonlinear function must take care of proper dealiasing.

    !!! note
        The numerically stable evaluation of the coefficients follows
        [Kassam and Trefethen
        (2005)](https://doi.org/10.1137/S1064827502410633).
    """
    super().__init__(dt, linear_operator)
    self._nonlinear_fun = nonlinear_fun

    LR = (
        circle_radius * roots_of_unity(num_circle_points)
        + linear_operator[..., jnp.newaxis] * dt
    )

    self._coef_1 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real

    self._coef_2 = dt * jnp.mean((jnp.exp(LR) - 1 - LR) / LR**2, axis=-1).real
step_fourier ¤
step_fourier(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]
Source code in exponax/etdrk/_etdrk_2.py 
81
82
83
84
85
86
87
88
89
90
91
92
def step_fourier(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    u_nonlin_hat = self._nonlinear_fun(u_hat)
    u_stage_1_hat = self._exp_term * u_hat + self._coef_1 * u_nonlin_hat

    u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
    u_next_hat = u_stage_1_hat + self._coef_2 * (
        u_stage_1_nonlin_hat - u_nonlin_hat
    )
    return u_next_hat

exponax.etdrk.ETDRK3 ¤

Bases: BaseETDRK

Source code in exponax/etdrk/_etdrk_3.py 
  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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class ETDRK3(BaseETDRK):
    _nonlinear_fun: BaseNonlinearFun
    _half_exp_term: Complex[Array, "E ... (N//2)+1"]
    _coef_1: Complex[Array, "E ... (N//2)+1"]
    _coef_2: Complex[Array, "E ... (N//2)+1"]
    _coef_3: Complex[Array, "E ... (N//2)+1"]
    _coef_4: Complex[Array, "E ... (N//2)+1"]
    _coef_5: Complex[Array, "E ... (N//2)+1"]

    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
        nonlinear_fun: BaseNonlinearFun,
        *,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        r"""
        Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
        with a **third order approximation**.

        Adapted from Eq. (23-25) of [Cox and Matthews
        (2002)](https://doi.org/10.1006/jcph.2002.6995):

        $$
            \begin{aligned}
                \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
                \\
                \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^*) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right).
                \\
                \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]}
                \\
                &+ \frac{-4 - \exp(\hat{\mathcal{L}}_h \Delta t) + \exp(\hat{\mathcal{L}}_h \Delta) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
                \\
                &+ 4 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*)
                \\
                &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{**})
            \end{aligned}
        $$

        where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
        the nonlinear differential operator.

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.
        - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
            nonlinear differential operator. ! The operator must take care of
            proper dealiasing.
        - `num_circle_points`: The number of points on the unit circle used to
            approximate the numerically challenging coefficients.
        - `circle_radius`: The radius of the circle used to approximate the
            numerically challenging coefficients.

        !!! warning
            The nonlinear function must take care of proper dealiasing.

        !!! note
            The numerically stable evaluation of the coefficients follows
            [Kassam and Trefethen
            (2005)](https://doi.org/10.1137/S1064827502410633).
        """
        super().__init__(dt, linear_operator)
        self._nonlinear_fun = nonlinear_fun
        self._half_exp_term = jnp.exp(0.5 * dt * linear_operator)

        LR = (
            circle_radius * roots_of_unity(num_circle_points)
            + linear_operator[..., jnp.newaxis] * dt
        )

        self._coef_1 = dt * jnp.mean((jnp.exp(LR / 2) - 1) / LR, axis=-1).real

        self._coef_2 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real

        self._coef_3 = (
            dt
            * jnp.mean(
                (-4 - LR + jnp.exp(LR) * (4 - 3 * LR + LR**2)) / (LR**3), axis=-1
            ).real
        )

        self._coef_4 = (
            dt
            * jnp.mean(
                (4.0 * (2.0 + LR + jnp.exp(LR) * (-2 + LR))) / (LR**3), axis=-1
            ).real
        )

        self._coef_5 = (
            dt
            * jnp.mean(
                (-4 - 3 * LR - LR**2 + jnp.exp(LR) * (4 - LR)) / (LR**3), axis=-1
            ).real
        )

    def step_fourier(
        self,
        u_hat: Complex[Array, "E ... (N//2)+1"],
    ) -> Complex[Array, "E ... (N//2)+1"]:
        u_nonlin_hat = self._nonlinear_fun(u_hat)
        u_stage_1_hat = self._half_exp_term * u_hat + self._coef_1 * u_nonlin_hat

        u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
        u_stage_2_hat = self._exp_term * u_hat + self._coef_2 * (
            2 * u_stage_1_nonlin_hat - u_nonlin_hat
        )

        u_stage_2_nonlin_hat = self._nonlinear_fun(u_stage_2_hat)

        u_next_hat = (
            self._exp_term * u_hat
            + self._coef_3 * u_nonlin_hat
            + self._coef_4 * u_stage_1_nonlin_hat
            + self._coef_5 * u_stage_2_nonlin_hat
        )

        return u_next_hat
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta with a third order approximation.

Adapted from Eq. (23-25) of Cox and Matthews (2002):

\[ \begin{aligned} \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^*) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right). \\ \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]} \\ &+ \frac{-4 - \exp(\hat{\mathcal{L}}_h \Delta t) + \exp(\hat{\mathcal{L}}_h \Delta) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ &+ 4 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*) \\ &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{**}) \end{aligned} \]

where \(\hat{\mathcal{N}}_h\) is the Fourier pseudo-spectral treatment of the nonlinear differential operator.

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.
  • nonlinear_fun: The Fourier pseudo-spectral treatment of the nonlinear differential operator. ! The operator must take care of proper dealiasing.
  • num_circle_points: The number of points on the unit circle used to approximate the numerically challenging coefficients.
  • circle_radius: The radius of the circle used to approximate the numerically challenging coefficients.

Warning

The nonlinear function must take care of proper dealiasing.

Note

The numerically stable evaluation of the coefficients follows Kassam and Trefethen (2005).

Source code in exponax/etdrk/_etdrk_3.py 
 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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    r"""
    Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
    with a **third order approximation**.

    Adapted from Eq. (23-25) of [Cox and Matthews
    (2002)](https://doi.org/10.1006/jcph.2002.6995):

    $$
        \begin{aligned}
            \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
            \\
            \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^*) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right).
            \\
            \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]}
            \\
            &+ \frac{-4 - \exp(\hat{\mathcal{L}}_h \Delta t) + \exp(\hat{\mathcal{L}}_h \Delta) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
            \\
            &+ 4 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*)
            \\
            &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{**})
        \end{aligned}
    $$

    where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
    the nonlinear differential operator.

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.
    - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
        nonlinear differential operator. ! The operator must take care of
        proper dealiasing.
    - `num_circle_points`: The number of points on the unit circle used to
        approximate the numerically challenging coefficients.
    - `circle_radius`: The radius of the circle used to approximate the
        numerically challenging coefficients.

    !!! warning
        The nonlinear function must take care of proper dealiasing.

    !!! note
        The numerically stable evaluation of the coefficients follows
        [Kassam and Trefethen
        (2005)](https://doi.org/10.1137/S1064827502410633).
    """
    super().__init__(dt, linear_operator)
    self._nonlinear_fun = nonlinear_fun
    self._half_exp_term = jnp.exp(0.5 * dt * linear_operator)

    LR = (
        circle_radius * roots_of_unity(num_circle_points)
        + linear_operator[..., jnp.newaxis] * dt
    )

    self._coef_1 = dt * jnp.mean((jnp.exp(LR / 2) - 1) / LR, axis=-1).real

    self._coef_2 = dt * jnp.mean((jnp.exp(LR) - 1) / LR, axis=-1).real

    self._coef_3 = (
        dt
        * jnp.mean(
            (-4 - LR + jnp.exp(LR) * (4 - 3 * LR + LR**2)) / (LR**3), axis=-1
        ).real
    )

    self._coef_4 = (
        dt
        * jnp.mean(
            (4.0 * (2.0 + LR + jnp.exp(LR) * (-2 + LR))) / (LR**3), axis=-1
        ).real
    )

    self._coef_5 = (
        dt
        * jnp.mean(
            (-4 - 3 * LR - LR**2 + jnp.exp(LR) * (4 - LR)) / (LR**3), axis=-1
        ).real
    )
step_fourier ¤
step_fourier(
    u_hat: Complex[Array, "E ... (N//2)+1"]
) -> Complex[Array, "E ... (N//2)+1"]
Source code in exponax/etdrk/_etdrk_3.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def step_fourier(
    self,
    u_hat: Complex[Array, "E ... (N//2)+1"],
) -> Complex[Array, "E ... (N//2)+1"]:
    u_nonlin_hat = self._nonlinear_fun(u_hat)
    u_stage_1_hat = self._half_exp_term * u_hat + self._coef_1 * u_nonlin_hat

    u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
    u_stage_2_hat = self._exp_term * u_hat + self._coef_2 * (
        2 * u_stage_1_nonlin_hat - u_nonlin_hat
    )

    u_stage_2_nonlin_hat = self._nonlinear_fun(u_stage_2_hat)

    u_next_hat = (
        self._exp_term * u_hat
        + self._coef_3 * u_nonlin_hat
        + self._coef_4 * u_stage_1_nonlin_hat
        + self._coef_5 * u_stage_2_nonlin_hat
    )

    return u_next_hat

exponax.etdrk.ETDRK4 ¤

Bases: BaseETDRK

Source code in exponax/etdrk/_etdrk_4.py 
  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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class ETDRK4(BaseETDRK):
    _nonlinear_fun: BaseNonlinearFun
    _half_exp_term: Complex[Array, "E ... (N//2)+1"]
    _coef_1: Complex[Array, "E ... (N//2)+1"]
    _coef_2: Complex[Array, "E ... (N//2)+1"]
    _coef_3: Complex[Array, "E ... (N//2)+1"]
    _coef_4: Complex[Array, "E ... (N//2)+1"]
    _coef_5: Complex[Array, "E ... (N//2)+1"]
    _coef_6: Complex[Array, "E ... (N//2)+1"]

    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
        nonlinear_fun: BaseNonlinearFun,
        *,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        r"""
        Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
        with a **fourth order approximation**.

        Adapted from Eq. (26-29) of [Cox and Matthews
        (2002)](https://doi.org/10.1006/jcph.2002.6995):

        $$
            \begin{aligned}
                \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
                \\
                \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t / 2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*).
                \\
                \hat{u}_h^{***} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{*} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^{**}) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right).
                \\
                \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]}
                \\
                &+ \frac{-4 - \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]})
                \\
                &+ 2 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) + \hat{\mathcal{N}}_h(\hat{u}_h^{**}) \right)
                \\
                &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{***})
            \end{aligned}
        $$

        where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
        the nonlinear differential operator.

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.
        - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
            nonlinear differential operator. ! The operator must take care of
            proper dealiasing.
        - `num_circle_points`: The number of points on the unit circle used to
            approximate the numerically challenging coefficients.
        - `circle_radius`: The radius of the circle used to approximate the
            numerically challenging coefficients.

        !!! warning
            The nonlinear function must take care of proper dealiasing.

        !!! note
            The numerically stable evaluation of the coefficients follows
            [Kassam and Trefethen
            (2005)](https://doi.org/10.1137/S1064827502410633).
        """
        super().__init__(dt, linear_operator)
        self._nonlinear_fun = nonlinear_fun
        self._half_exp_term = jnp.exp(0.5 * dt * linear_operator)

        LR = (
            circle_radius * roots_of_unity(num_circle_points)
            + linear_operator[..., jnp.newaxis] * dt
        )

        self._coef_1 = dt * jnp.mean((jnp.exp(LR / 2) - 1) / LR, axis=-1).real

        self._coef_2 = self._coef_1
        self._coef_3 = self._coef_1

        self._coef_4 = (
            dt
            * jnp.mean(
                (-4 - LR + jnp.exp(LR) * (4 - 3 * LR + LR**2)) / (LR**3), axis=-1
            ).real
        )

        self._coef_5 = (
            dt * jnp.mean((2 + LR + jnp.exp(LR) * (-2 + LR)) / (LR**3), axis=-1).real
        )

        self._coef_6 = (
            dt
            * jnp.mean(
                (-4 - 3 * LR - LR**2 + jnp.exp(LR) * (4 - LR)) / (LR**3), axis=-1
            ).real
        )

    def step_fourier(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        u_nonlin_hat = self._nonlinear_fun(u_hat)
        u_stage_1_hat = self._half_exp_term * u_hat + self._coef_1 * u_nonlin_hat

        u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
        u_stage_2_hat = (
            self._half_exp_term * u_hat + self._coef_2 * u_stage_1_nonlin_hat
        )

        u_stage_2_nonlin_hat = self._nonlinear_fun(u_stage_2_hat)
        u_stage_3_hat = self._half_exp_term * u_stage_1_hat + self._coef_3 * (
            2 * u_stage_2_nonlin_hat - u_nonlin_hat
        )

        u_stage_3_nonlin_hat = self._nonlinear_fun(u_stage_3_hat)

        u_next_hat = (
            self._exp_term * u_hat
            + self._coef_4 * u_nonlin_hat
            + self._coef_5 * 2 * (u_stage_1_nonlin_hat + u_stage_2_nonlin_hat)
            + self._coef_6 * u_stage_3_nonlin_hat
        )

        return u_next_hat
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta with a fourth order approximation.

Adapted from Eq. (26-29) of Cox and Matthews (2002):

\[ \begin{aligned} \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}). \\ \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t / 2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*). \\ \hat{u}_h^{***} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{*} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^{**}) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right). \\ \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]} \\ &+ \frac{-4 - \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \\ &+ 2 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) + \hat{\mathcal{N}}_h(\hat{u}_h^{**}) \right) \\ &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{***}) \end{aligned} \]

where \(\hat{\mathcal{N}}_h\) is the Fourier pseudo-spectral treatment of the nonlinear differential operator.

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.
  • nonlinear_fun: The Fourier pseudo-spectral treatment of the nonlinear differential operator. ! The operator must take care of proper dealiasing.
  • num_circle_points: The number of points on the unit circle used to approximate the numerically challenging coefficients.
  • circle_radius: The radius of the circle used to approximate the numerically challenging coefficients.

Warning

The nonlinear function must take care of proper dealiasing.

Note

The numerically stable evaluation of the coefficients follows Kassam and Trefethen (2005).

Source code in exponax/etdrk/_etdrk_4.py 
 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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
    nonlinear_fun: BaseNonlinearFun,
    *,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    r"""
    Solve a semi-linear PDE using Exponential Time Differencing Runge-Kutta
    with a **fourth order approximation**.

    Adapted from Eq. (26-29) of [Cox and Matthews
    (2002)](https://doi.org/10.1006/jcph.2002.6995):

    $$
        \begin{aligned}
            \hat{u}_h^* &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}).
            \\
            \hat{u}_h^{**} &= \exp(\hat{\mathcal{L}}_h \Delta t / 2) \odot \hat{u}_h^{[t]} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t / 2) - 1}{\hat{\mathcal{L}}_h} \odot \hat{\mathcal{N}}_h(\hat{u}_h^*).
            \\
            \hat{u}_h^{***} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{*} + \frac{\exp(\hat{\mathcal{L}}_h \Delta t/2) - 1}{\hat{\mathcal{L}}_h} \odot \left( 2 \hat{\mathcal{N}}_h(\hat{u}_h^{**}) - \hat{\mathcal{N}}_h(\hat{u}_h^{[t]}) \right).
            \\
            \hat{u}_h^{[t+1]} &= \exp(\hat{\mathcal{L}}_h \Delta t) \odot \hat{u}_h^{[t]}
            \\
            &+ \frac{-4 - \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - 3 \hat{\mathcal{L}}_h \Delta t + \left(\hat{\mathcal{L}}_h \Delta t\right)^2 \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{[t]})
            \\
            &+ 2 \frac{2 + \hat{\mathcal{L}}_h \Delta t + \exp(\hat{\mathcal{L}}_h \Delta t) \left( -2 + \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \left( \hat{\mathcal{N}}_h(\hat{u}_h^*) + \hat{\mathcal{N}}_h(\hat{u}_h^{**}) \right)
            \\
            &+ \frac{-4 - 3 \hat{\mathcal{L}}_h \Delta t - \left( \hat{\mathcal{L}}_h \Delta t \right)^2 + \exp(\hat{\mathcal{L}}_h \Delta t) \left( 4 - \hat{\mathcal{L}}_h \Delta t \right)}{\hat{\mathcal{L}}_h^3 (\Delta t)^2} \odot \hat{\mathcal{N}}_h(\hat{u}_h^{***})
        \end{aligned}
    $$

    where $\hat{\mathcal{N}}_h$ is the Fourier pseudo-spectral treatment of
    the nonlinear differential operator.

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.
    - `nonlinear_fun`: The Fourier pseudo-spectral treatment of the
        nonlinear differential operator. ! The operator must take care of
        proper dealiasing.
    - `num_circle_points`: The number of points on the unit circle used to
        approximate the numerically challenging coefficients.
    - `circle_radius`: The radius of the circle used to approximate the
        numerically challenging coefficients.

    !!! warning
        The nonlinear function must take care of proper dealiasing.

    !!! note
        The numerically stable evaluation of the coefficients follows
        [Kassam and Trefethen
        (2005)](https://doi.org/10.1137/S1064827502410633).
    """
    super().__init__(dt, linear_operator)
    self._nonlinear_fun = nonlinear_fun
    self._half_exp_term = jnp.exp(0.5 * dt * linear_operator)

    LR = (
        circle_radius * roots_of_unity(num_circle_points)
        + linear_operator[..., jnp.newaxis] * dt
    )

    self._coef_1 = dt * jnp.mean((jnp.exp(LR / 2) - 1) / LR, axis=-1).real

    self._coef_2 = self._coef_1
    self._coef_3 = self._coef_1

    self._coef_4 = (
        dt
        * jnp.mean(
            (-4 - LR + jnp.exp(LR) * (4 - 3 * LR + LR**2)) / (LR**3), axis=-1
        ).real
    )

    self._coef_5 = (
        dt * jnp.mean((2 + LR + jnp.exp(LR) * (-2 + LR)) / (LR**3), axis=-1).real
    )

    self._coef_6 = (
        dt
        * jnp.mean(
            (-4 - 3 * LR - LR**2 + jnp.exp(LR) * (4 - LR)) / (LR**3), axis=-1
        ).real
    )
step_fourier ¤
step_fourier(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]
Source code in exponax/etdrk/_etdrk_4.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def step_fourier(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    u_nonlin_hat = self._nonlinear_fun(u_hat)
    u_stage_1_hat = self._half_exp_term * u_hat + self._coef_1 * u_nonlin_hat

    u_stage_1_nonlin_hat = self._nonlinear_fun(u_stage_1_hat)
    u_stage_2_hat = (
        self._half_exp_term * u_hat + self._coef_2 * u_stage_1_nonlin_hat
    )

    u_stage_2_nonlin_hat = self._nonlinear_fun(u_stage_2_hat)
    u_stage_3_hat = self._half_exp_term * u_stage_1_hat + self._coef_3 * (
        2 * u_stage_2_nonlin_hat - u_nonlin_hat
    )

    u_stage_3_nonlin_hat = self._nonlinear_fun(u_stage_3_hat)

    u_next_hat = (
        self._exp_term * u_hat
        + self._coef_4 * u_nonlin_hat
        + self._coef_5 * 2 * (u_stage_1_nonlin_hat + u_stage_2_nonlin_hat)
        + self._coef_6 * u_stage_3_nonlin_hat
    )

    return u_next_hat

exponax.etdrk.BaseETDRK ¤

Bases: Module, ABC

Source code in exponax/etdrk/_base_etdrk.py 
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
class BaseETDRK(eqx.Module, ABC):
    dt: float
    _exp_term: Complex[Array, "E ... (N//2)+1"]

    def __init__(
        self,
        dt: float,
        linear_operator: Complex[Array, "E ... (N//2)+1"],
    ):
        """
        Base class for exponential time differencing Runge-Kutta methods.

        **Arguments:**

        - `dt`: The time step size.
        - `linear_operator`: The linear operator of the PDE. Must have a leading
            channel axis, followed by one, two or three spatial axes whereas the
            last axis must be of size `(N//2)+1` where `N` is the number of
            dimensions in the former spatial axes.

        !!! Example
            Below is an example how to get the linear operator for
            the heat equation.

            ```python
            import jax.numpy as jnp
            import exponax as ex

            # Define the linear operator
            N = 256
            L = 5.0  # The domain size
            D = 1  # Being in 1D

            derivative_operator = 1j * ex.spectral.build_derivative_operator(
                D,
                L,
                N,
            )

            print(derivative_operator.shape)  # (1, (N//2)+1)

            nu = 0.01 # The diffusion coefficient

            linear_operator = nu * derivative_operator**2
            ```
        """
        self.dt = dt
        self._exp_term = jnp.exp(self.dt * linear_operator)

    @abstractmethod
    def step_fourier(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        """
        Advance the state in Fourier space.

        **Arguments:**

        - `u_hat`: The previous state in Fourier space.

        **Returns:**

        - The next state in Fourier space, i.e., `self.dt` time units later.
        """
        pass
__init__ ¤
__init__(
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
)

Base class for exponential time differencing Runge-Kutta methods.

Arguments:

  • dt: The time step size.
  • linear_operator: The linear operator of the PDE. Must have a leading channel axis, followed by one, two or three spatial axes whereas the last axis must be of size (N//2)+1 where N is the number of dimensions in the former spatial axes.

Example

Below is an example how to get the linear operator for the heat equation.

import jax.numpy as jnp
import exponax as ex

# Define the linear operator
N = 256
L = 5.0  # The domain size
D = 1  # Being in 1D

derivative_operator = 1j * ex.spectral.build_derivative_operator(
    D,
    L,
    N,
)

print(derivative_operator.shape)  # (1, (N//2)+1)

nu = 0.01 # The diffusion coefficient

linear_operator = nu * derivative_operator**2
Source code in exponax/etdrk/_base_etdrk.py 
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
def __init__(
    self,
    dt: float,
    linear_operator: Complex[Array, "E ... (N//2)+1"],
):
    """
    Base class for exponential time differencing Runge-Kutta methods.

    **Arguments:**

    - `dt`: The time step size.
    - `linear_operator`: The linear operator of the PDE. Must have a leading
        channel axis, followed by one, two or three spatial axes whereas the
        last axis must be of size `(N//2)+1` where `N` is the number of
        dimensions in the former spatial axes.

    !!! Example
        Below is an example how to get the linear operator for
        the heat equation.

        ```python
        import jax.numpy as jnp
        import exponax as ex

        # Define the linear operator
        N = 256
        L = 5.0  # The domain size
        D = 1  # Being in 1D

        derivative_operator = 1j * ex.spectral.build_derivative_operator(
            D,
            L,
            N,
        )

        print(derivative_operator.shape)  # (1, (N//2)+1)

        nu = 0.01 # The diffusion coefficient

        linear_operator = nu * derivative_operator**2
        ```
    """
    self.dt = dt
    self._exp_term = jnp.exp(self.dt * linear_operator)
step_fourier abstractmethod ¤
step_fourier(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]

Advance the state in Fourier space.

Arguments:

  • u_hat: The previous state in Fourier space.

Returns:

  • The next state in Fourier space, i.e., self.dt time units later.
Source code in exponax/etdrk/_base_etdrk.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@abstractmethod
def step_fourier(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    """
    Advance the state in Fourier space.

    **Arguments:**

    - `u_hat`: The previous state in Fourier space.

    **Returns:**

    - The next state in Fourier space, i.e., `self.dt` time units later.
    """
    pass

exponax.etdrk.roots_of_unity ¤

roots_of_unity(M: int) -> Complex[Array, M]

Return (complex-valued) array with M roots of unity. Useful to perform contour integrals in the complex plane.

Arguments:

  • M: The number of roots of unity.

Returns:

  • roots: The M roots of unity in an array of shape (M,).
Source code in exponax/etdrk/_utils.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def roots_of_unity(M: int) -> Complex[Array, "M"]:
    """
    Return (complex-valued) array with M roots of unity. Useful to perform
    contour integrals in the complex plane.

    **Arguments:**

    - `M`: The number of roots of unity.

    **Returns:**

    - `roots`: The M roots of unity in an array of shape `(M,)`.
    """
    # return jnp.exp(1j * jnp.pi * (jnp.arange(1, M+1) - 0.5) / M)
    return jnp.exp(2j * jnp.pi * (jnp.arange(1, M + 1) - 0.5) / M)