Skip to content

General Linear Stepper¤

exponax.stepper.generic.GeneralLinearStepper ¤

Bases: BaseStepper

Source code in exponax/stepper/generic/_linear.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
 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
class GeneralLinearStepper(BaseStepper):
    linear_coefficients: tuple[float, ...]

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        dt: float,
        *,
        linear_coefficients: tuple[float, ...] = (0.0, -0.1, 0.01),
    ):
        """
        General timestepper for a d-dimensional (`d ∈ {1, 2, 3}`) linear
        equation with an arbitrary combination of derivative terms and
        respective coefficients. To simplify the interface, only isotropicity is
        supported. For example, the advection speed is the same in all spatial
        dimensions, or diffusion is equally strong in all spatial dimensions.

        In 1d the equation is given by

        ```
            uₜ = sum_j a_j uₓˢ
        ```

        with `uₓˢ` denoting the s-th derivative of `u` with respect to `x`. The
        coefficient corresponding to this derivative is `a_j`.

        The isotropic version in higher dimensions can expressed as

        ```
            uₜ = sum_j a_j (1⋅∇ʲ)u
        ```

        with `1⋅∇ʲ` denoting the j-th repeated elementwise product of the nabla
        operator with itself in an inner product with the one vector. For
        example, `1⋅∇¹` is the collection of first derivatives, `1⋅∇²` is the
        collection of second derivatives (=Laplace operator), etc.

        The interface to this general stepper is the list of coefficients
        containing the `a_j`. Its length determines the highes occuring order of
        derivative. Note that this list starts at zero. If only one specific
        linear term is wanted, have all prior coefficients set to zero.

        The default configuration is an advection-diffusion equation with `a_0 =
        0`, `a_1 = -0.1`, and `a_2 = 0.01`.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `d`.
        - `domain_extent`: The size of the domain `L`; in higher dimensions
            the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`.
        - `num_points`: The number of points `N` used to discretize the
            domain. This **includes** the left boundary point and **excludes**
            the right boundary point. In higher dimensions; the number of points
            in each dimension is the same. Hence, the total number of degrees of
            freedom is `Nᵈ`.
        - `dt`: The timestep size `Δt` between two consecutive states.
        - `linear_coefficients` (keyword-only): The list of coefficients `a_j`
            corresponding to the derivatives. Default: `[0.0, -0.1, 0.01]`.

        **Notes:**

        - There is a repeating pattern in the effect of orders of
            derivatives:
            - Even derivatives (i.e., 0, 2, 4, 6, ...) scale the
                solution. Order 0 scales all wavenumbers/modes equally (if
                its coefficient is negative, this is also called a drag).
                Order 2 scales higher wavenumbers/modes stronger with the
                dependence on the effect on the wavenumber being
                quadratically. Order 4 also scales but stronger than order
                4. Its dependency on the wavenumber is quartically. This
                pattern continues for higher orders.
            - Odd derivatives (i.e, 1, 3, 5, 7, ...) rotate the solution in
                Fourier space. In state space, this is observed as
                advection. Order 1 rotates all wavenumbers/modes equally. In
                state space, this is observed in that the initial condition
                just moves over the domain. Order 3 rotates higher
                wavenumbers/modes stronger with the dependence on the
                wavenumber being quadratic. If certain wavenumbers are
                rotated at a different speed, there is still advection in
                state space but certain patterns in the initial condition
                are advected at different speeds. As such, it is observed
                that the shape of the initial condition dissolves. The
                effect continues for higher orders with the dependency on
                the wavenumber becoming continuously stronger.
        - Take care of the signs of coefficients. In contrast to the
            indivial linear steppers ([`exponax.stepper.Advection`][],
            [`exponax.stepper.Diffusion`][], etc.), the signs are not
            automatically taken care of to produce meaningful coefficients.
            For the general linear stepper all linear derivatives are on the
            right-hand side of the equation. This has the following effect
            based on the order of derivative (this a consequence of squaring
            the imaginary unit returning -1):
            - Zeroth-Order: A negative coeffcient is a drag and removes
                energy from the system. A positive coefficient adds energy
                to the system.
            - First-Order: A negative coefficient rotates the solution
                clockwise. A positive coefficient rotates the solution
                counter-clockwise. Hence, negative coefficients advect
                solutions to the right, positive coefficients advect
                solutions to the left.
            - Second-Order: A positive coefficient diffuses the solution
                (i.e., removes energy from the system). A negative
                coefficient adds energy to the system.
            - Third-Order: A negative coefficient rotates the solution
                counter-clockwise. A positive coefficient rotates the
                solution clockwise. Hence, negative coefficients advect
                solutions to the left, positive coefficients advect
                solutions to the right.
            - Fourth-Order: A negative coefficient diffuses the solution
                (i.e., removes energy from the system). A positive
                coefficient adds energy to the system.
            - ...
        - The stepper is unconditionally stable, no matter the choice of
            any argument because the equation is solved analytically in
            Fourier space. **However**, note if you have odd-order
            derivative terms (e.g., advection or dispersion) and your
            initial condition is **not** bandlimited (i.e., it contains
            modes beyond the Nyquist frequency of `(N//2)+1`) there is a
            chance spurious oscillations can occur.
        - Ultimately, only the factors `a_j Δt / Lʲ` affect the
            characteristic of the dynamics. See also
            [`exponax.stepper.generic.NormalizedLinearStepper`][] with
            `normalized_coefficients = [0, alpha_1, alpha_2, ...]` with
            `alpha_j = coefficients[j] * dt / domain_extent**j`. You can use
            the function
            [`exponax.stepper.generic.normalize_coefficients`][] to obtain
            the normalized coefficients.
        """
        self.linear_coefficients = linear_coefficients
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            domain_extent=domain_extent,
            num_points=num_points,
            dt=dt,
            num_channels=1,
            order=0,
        )

    def _build_linear_operator(
        self,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
    ) -> Complex[Array, "1 ... (N//2)+1"]:
        linear_operator = sum(
            jnp.sum(
                c * (derivative_operator) ** i,
                axis=0,
                keepdims=True,
            )
            for i, c in enumerate(self.linear_coefficients)
        )
        return linear_operator

    def _build_nonlinear_fun(
        self,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
    ) -> ZeroNonlinearFun:
        return ZeroNonlinearFun(
            self.num_spatial_dims,
            self.num_points,
        )
__init__ ¤
__init__(
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    linear_coefficients: tuple[float, ...] = (
        0.0,
        -0.1,
        0.01,
    )
)

General timestepper for a d-dimensional (d ∈ {1, 2, 3}) linear equation with an arbitrary combination of derivative terms and respective coefficients. To simplify the interface, only isotropicity is supported. For example, the advection speed is the same in all spatial dimensions, or diffusion is equally strong in all spatial dimensions.

In 1d the equation is given by

    uₜ = sum_j a_j uₓˢ

with uₓˢ denoting the s-th derivative of u with respect to x. The coefficient corresponding to this derivative is a_j.

The isotropic version in higher dimensions can expressed as

    uₜ = sum_j a_j (1⋅∇ʲ)u

with 1⋅∇ʲ denoting the j-th repeated elementwise product of the nabla operator with itself in an inner product with the one vector. For example, 1⋅∇¹ is the collection of first derivatives, 1⋅∇² is the collection of second derivatives (=Laplace operator), etc.

The interface to this general stepper is the list of coefficients containing the a_j. Its length determines the highes occuring order of derivative. Note that this list starts at zero. If only one specific linear term is wanted, have all prior coefficients set to zero.

The default configuration is an advection-diffusion equation with a_0 = 0, a_1 = -0.1, and a_2 = 0.01.

Arguments:

  • num_spatial_dims: The number of spatial dimensions d.
  • domain_extent: The size of the domain L; in higher dimensions the domain is assumed to be a scaled hypercube Ω = (0, L)ᵈ.
  • num_points: The number of points N used to discretize the domain. This includes the left boundary point and excludes the right boundary point. In higher dimensions; the number of points in each dimension is the same. Hence, the total number of degrees of freedom is Nᵈ.
  • dt: The timestep size Δt between two consecutive states.
  • linear_coefficients (keyword-only): The list of coefficients a_j corresponding to the derivatives. Default: [0.0, -0.1, 0.01].

Notes:

  • There is a repeating pattern in the effect of orders of derivatives:
    • Even derivatives (i.e., 0, 2, 4, 6, ...) scale the solution. Order 0 scales all wavenumbers/modes equally (if its coefficient is negative, this is also called a drag). Order 2 scales higher wavenumbers/modes stronger with the dependence on the effect on the wavenumber being quadratically. Order 4 also scales but stronger than order
      1. Its dependency on the wavenumber is quartically. This pattern continues for higher orders.
    • Odd derivatives (i.e, 1, 3, 5, 7, ...) rotate the solution in Fourier space. In state space, this is observed as advection. Order 1 rotates all wavenumbers/modes equally. In state space, this is observed in that the initial condition just moves over the domain. Order 3 rotates higher wavenumbers/modes stronger with the dependence on the wavenumber being quadratic. If certain wavenumbers are rotated at a different speed, there is still advection in state space but certain patterns in the initial condition are advected at different speeds. As such, it is observed that the shape of the initial condition dissolves. The effect continues for higher orders with the dependency on the wavenumber becoming continuously stronger.
  • Take care of the signs of coefficients. In contrast to the indivial linear steppers (exponax.stepper.Advection, exponax.stepper.Diffusion, etc.), the signs are not automatically taken care of to produce meaningful coefficients. For the general linear stepper all linear derivatives are on the right-hand side of the equation. This has the following effect based on the order of derivative (this a consequence of squaring the imaginary unit returning -1):
    • Zeroth-Order: A negative coeffcient is a drag and removes energy from the system. A positive coefficient adds energy to the system.
    • First-Order: A negative coefficient rotates the solution clockwise. A positive coefficient rotates the solution counter-clockwise. Hence, negative coefficients advect solutions to the right, positive coefficients advect solutions to the left.
    • Second-Order: A positive coefficient diffuses the solution (i.e., removes energy from the system). A negative coefficient adds energy to the system.
    • Third-Order: A negative coefficient rotates the solution counter-clockwise. A positive coefficient rotates the solution clockwise. Hence, negative coefficients advect solutions to the left, positive coefficients advect solutions to the right.
    • Fourth-Order: A negative coefficient diffuses the solution (i.e., removes energy from the system). A positive coefficient adds energy to the system.
    • ...
  • The stepper is unconditionally stable, no matter the choice of any argument because the equation is solved analytically in Fourier space. However, note if you have odd-order derivative terms (e.g., advection or dispersion) and your initial condition is not bandlimited (i.e., it contains modes beyond the Nyquist frequency of (N//2)+1) there is a chance spurious oscillations can occur.
  • Ultimately, only the factors a_j Δt / Lʲ affect the characteristic of the dynamics. See also exponax.stepper.generic.NormalizedLinearStepper with normalized_coefficients = [0, alpha_1, alpha_2, ...] with alpha_j = coefficients[j] * dt / domain_extent**j. You can use the function exponax.stepper.generic.normalize_coefficients to obtain the normalized coefficients.
Source code in exponax/stepper/generic/_linear.py 
 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def __init__(
    self,
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    linear_coefficients: tuple[float, ...] = (0.0, -0.1, 0.01),
):
    """
    General timestepper for a d-dimensional (`d ∈ {1, 2, 3}`) linear
    equation with an arbitrary combination of derivative terms and
    respective coefficients. To simplify the interface, only isotropicity is
    supported. For example, the advection speed is the same in all spatial
    dimensions, or diffusion is equally strong in all spatial dimensions.

    In 1d the equation is given by

    ```
        uₜ = sum_j a_j uₓˢ
    ```

    with `uₓˢ` denoting the s-th derivative of `u` with respect to `x`. The
    coefficient corresponding to this derivative is `a_j`.

    The isotropic version in higher dimensions can expressed as

    ```
        uₜ = sum_j a_j (1⋅∇ʲ)u
    ```

    with `1⋅∇ʲ` denoting the j-th repeated elementwise product of the nabla
    operator with itself in an inner product with the one vector. For
    example, `1⋅∇¹` is the collection of first derivatives, `1⋅∇²` is the
    collection of second derivatives (=Laplace operator), etc.

    The interface to this general stepper is the list of coefficients
    containing the `a_j`. Its length determines the highes occuring order of
    derivative. Note that this list starts at zero. If only one specific
    linear term is wanted, have all prior coefficients set to zero.

    The default configuration is an advection-diffusion equation with `a_0 =
    0`, `a_1 = -0.1`, and `a_2 = 0.01`.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `d`.
    - `domain_extent`: The size of the domain `L`; in higher dimensions
        the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`.
    - `num_points`: The number of points `N` used to discretize the
        domain. This **includes** the left boundary point and **excludes**
        the right boundary point. In higher dimensions; the number of points
        in each dimension is the same. Hence, the total number of degrees of
        freedom is `Nᵈ`.
    - `dt`: The timestep size `Δt` between two consecutive states.
    - `linear_coefficients` (keyword-only): The list of coefficients `a_j`
        corresponding to the derivatives. Default: `[0.0, -0.1, 0.01]`.

    **Notes:**

    - There is a repeating pattern in the effect of orders of
        derivatives:
        - Even derivatives (i.e., 0, 2, 4, 6, ...) scale the
            solution. Order 0 scales all wavenumbers/modes equally (if
            its coefficient is negative, this is also called a drag).
            Order 2 scales higher wavenumbers/modes stronger with the
            dependence on the effect on the wavenumber being
            quadratically. Order 4 also scales but stronger than order
            4. Its dependency on the wavenumber is quartically. This
            pattern continues for higher orders.
        - Odd derivatives (i.e, 1, 3, 5, 7, ...) rotate the solution in
            Fourier space. In state space, this is observed as
            advection. Order 1 rotates all wavenumbers/modes equally. In
            state space, this is observed in that the initial condition
            just moves over the domain. Order 3 rotates higher
            wavenumbers/modes stronger with the dependence on the
            wavenumber being quadratic. If certain wavenumbers are
            rotated at a different speed, there is still advection in
            state space but certain patterns in the initial condition
            are advected at different speeds. As such, it is observed
            that the shape of the initial condition dissolves. The
            effect continues for higher orders with the dependency on
            the wavenumber becoming continuously stronger.
    - Take care of the signs of coefficients. In contrast to the
        indivial linear steppers ([`exponax.stepper.Advection`][],
        [`exponax.stepper.Diffusion`][], etc.), the signs are not
        automatically taken care of to produce meaningful coefficients.
        For the general linear stepper all linear derivatives are on the
        right-hand side of the equation. This has the following effect
        based on the order of derivative (this a consequence of squaring
        the imaginary unit returning -1):
        - Zeroth-Order: A negative coeffcient is a drag and removes
            energy from the system. A positive coefficient adds energy
            to the system.
        - First-Order: A negative coefficient rotates the solution
            clockwise. A positive coefficient rotates the solution
            counter-clockwise. Hence, negative coefficients advect
            solutions to the right, positive coefficients advect
            solutions to the left.
        - Second-Order: A positive coefficient diffuses the solution
            (i.e., removes energy from the system). A negative
            coefficient adds energy to the system.
        - Third-Order: A negative coefficient rotates the solution
            counter-clockwise. A positive coefficient rotates the
            solution clockwise. Hence, negative coefficients advect
            solutions to the left, positive coefficients advect
            solutions to the right.
        - Fourth-Order: A negative coefficient diffuses the solution
            (i.e., removes energy from the system). A positive
            coefficient adds energy to the system.
        - ...
    - The stepper is unconditionally stable, no matter the choice of
        any argument because the equation is solved analytically in
        Fourier space. **However**, note if you have odd-order
        derivative terms (e.g., advection or dispersion) and your
        initial condition is **not** bandlimited (i.e., it contains
        modes beyond the Nyquist frequency of `(N//2)+1`) there is a
        chance spurious oscillations can occur.
    - Ultimately, only the factors `a_j Δt / Lʲ` affect the
        characteristic of the dynamics. See also
        [`exponax.stepper.generic.NormalizedLinearStepper`][] with
        `normalized_coefficients = [0, alpha_1, alpha_2, ...]` with
        `alpha_j = coefficients[j] * dt / domain_extent**j`. You can use
        the function
        [`exponax.stepper.generic.normalize_coefficients`][] to obtain
        the normalized coefficients.
    """
    self.linear_coefficients = linear_coefficients
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        domain_extent=domain_extent,
        num_points=num_points,
        dt=dt,
        num_channels=1,
        order=0,
    )
__call__ ¤
__call__(
    u: Float[Array, "C ... N"]
) -> Float[Array, "C ... N"]

Perform one step of the time integration for a single state.

Arguments:

  • u: The state vector, shape (C, ..., N,).

Returns:

  • u_next: The state vector after one step, shape (C, ..., N,).

Tip

Use this call method together with exponax.rollout to efficiently produce temporal trajectories.

Info

For batched operation, use jax.vmap on this function.

Source code in exponax/_base_stepper.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def __call__(
    self,
    u: Float[Array, "C ... N"],
) -> Float[Array, "C ... N"]:
    """
    Perform one step of the time integration for a single state.

    **Arguments:**

    - `u`: The state vector, shape `(C, ..., N,)`.

    **Returns:**

    - `u_next`: The state vector after one step, shape `(C, ..., N,)`.

    !!! tip
        Use this call method together with `exponax.rollout` to efficiently
        produce temporal trajectories.

    !!! info
        For batched operation, use `jax.vmap` on this function.
    """
    expected_shape = (self.num_channels,) + spatial_shape(
        self.num_spatial_dims, self.num_points
    )
    if u.shape != expected_shape:
        raise ValueError(
            f"Expected shape {expected_shape}, got {u.shape}. For batched operation use `jax.vmap` on this function."
        )
    return self.step(u)