Skip to content

General Nonlinear Stepper¤

exponax.stepper.generic.GeneralNonlinearStepper ¤

Bases: BaseStepper

Source code in exponax/stepper/generic/_nonlinear.py 
 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
class GeneralNonlinearStepper(BaseStepper):
    linear_coefficients: tuple[float, ...]
    nonlinear_coefficients: tuple[float, float, float]
    dealiasing_fraction: float

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        dt: float,
        *,
        linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
        nonlinear_coefficients: tuple[float, float, float] = (0.0, -1.0, 0.0),
        order=2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) semi-linear PDEs
        consisting of a quadratic, a single-channel convection, and a gradient
        norm nonlinearity together with an arbitrary combination of (isotropic)
        linear derivatives.

        In 1d, the PDE is of the form

        ```
            uₜ = b₀ u² + b₁ 1/2 (u²)ₓ + b₂ 1/2 (uₓ)² + sum_j a_j uₓʲ
        ```

        where `b₀`, `b₁`, `b₂` are the coefficients of the quadratic,
        convection, and gradient norm nonlinearity, respectively, and `a_j` are
        the coefficients of the linear derivatives. Effectively, this
        timestepper is a combination of the
        `exponax.stepper.generic.GeneralPolynomialStepper` (with only the
        coefficient to the quadratic polynomial being set with `b₀`), the
        `exponax.stepper.generic.GeneralConvectionStepper` (with the
        single-channel hack activated via `single_channel=True` and the
        convection scale set with `- b₁`), and the
        `exponax.stepper.generic.GeneralGradientNormStepper` (with the gradient
        norm scale set with `- b₂`).

        !!! warning
            In contrast to the
            `exponax.stepper.generic.GeneralConvectionStepper` and the
            `exponax.stepper.generic.GeneralGradientNormStepper`, the nonlinear
            terms are no considered to be on right-hand side of the PDE. Hence,
            in order to get the same dynamics as in the other steppers, the
            coefficients must be negated. (This is not relevant for the
            coefficient of the quadratic polynomial because in the
            `exponax.stepper.generic.GeneralPolynomialStepper` the polynomial
            nonlinearity is already on the right-hand side.)

        The higher-dimensional generalization is

        ```
            uₜ = b₀ u² + b₁ 1/2 (1⃗ ⋅ ∇)(u²) + b₂ 1/2 ‖ ∇u ‖₂² + sum_j a_j uₓˢ
        ```

        Under the default configuration, this correspons to a Burgers equation
        in single-channel mode.

        **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`: The list of coefficients `a_j` corresponding to
            the derivatives. The length of this tuple represents the highest
            occuring derivative. The default value `(0.0, 0.0, 0.01)` together
            with the default `nonlinear_coefficients` corresponds to the Burgers
            equation.
        - `nonlinear_coefficients`: The list of coefficients `b₀`, `b₁`, `b₂`
            (in this order). The default value `(0.0, -1.0, 0.0)` corresponds to
            a (single-channel) convection nonlinearity scaled with `1.0`. Note
            that all nonlinear contributions are considered to be on the
            right-hand side of the PDE. Hence, in order to get the "correct
            convection" dynamics, the coefficients must be negated.
        - `order`: The order of the ETDRK method to use. Must be one of {0, 1, 2,
            3, 4}. The option `0` only solves the linear part of the equation.
            Use higher values for higher accuracy and stability. The default
            choice of `2` is a good compromise for single precision floats.
        - `dealiasing_fraction`: The fraction of the wavenumbers to keep before
            evaluating the nonlinearity. The default value `2/3` corresponds to
            Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2.
        - `num_circle_points`: How many points to use in the complex contour
            integral method to compute the coefficients of the exponential time
            differencing Runge Kutta method.
        - `circle_radius`: The radius of the contour used to compute the
            coefficients of the exponential time differencing Runge Kutta method.
        """
        if len(nonlinear_coefficients) != 3:
            raise ValueError(
                "The nonlinear coefficients list must have exactly 3 elements"
            )
        self.linear_coefficients = linear_coefficients
        self.nonlinear_coefficients = nonlinear_coefficients
        self.dealiasing_fraction = dealiasing_fraction

        super().__init__(
            num_spatial_dims=num_spatial_dims,
            domain_extent=domain_extent,
            num_points=num_points,
            dt=dt,
            num_channels=1,
            order=order,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
        )

    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"],
    ) -> GeneralNonlinearFun:
        return GeneralNonlinearFun(
            self.num_spatial_dims,
            self.num_points,
            derivative_operator=derivative_operator,
            dealiasing_fraction=self.dealiasing_fraction,
            scale_list=self.nonlinear_coefficients,
            zero_mode_fix=True,  # ToDo: check this
        )
__init__ ¤
__init__(
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    linear_coefficients: tuple[float, ...] = (
        0.0,
        0.0,
        0.01,
    ),
    nonlinear_coefficients: tuple[float, float, float] = (
        0.0,
        -1.0,
        0.0,
    ),
    order=2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Timestepper for d-dimensional (d ∈ {1, 2, 3}) semi-linear PDEs consisting of a quadratic, a single-channel convection, and a gradient norm nonlinearity together with an arbitrary combination of (isotropic) linear derivatives.

In 1d, the PDE is of the form

    uₜ = b₀ u² + b₁ 1/2 (u²)ₓ + b₂ 1/2 (uₓ)² + sum_j a_j uₓʲ

where b₀, b₁, b₂ are the coefficients of the quadratic, convection, and gradient norm nonlinearity, respectively, and a_j are the coefficients of the linear derivatives. Effectively, this timestepper is a combination of the exponax.stepper.generic.GeneralPolynomialStepper (with only the coefficient to the quadratic polynomial being set with b₀), the exponax.stepper.generic.GeneralConvectionStepper (with the single-channel hack activated via single_channel=True and the convection scale set with - b₁), and the exponax.stepper.generic.GeneralGradientNormStepper (with the gradient norm scale set with - b₂).

Warning

In contrast to the exponax.stepper.generic.GeneralConvectionStepper and the exponax.stepper.generic.GeneralGradientNormStepper, the nonlinear terms are no considered to be on right-hand side of the PDE. Hence, in order to get the same dynamics as in the other steppers, the coefficients must be negated. (This is not relevant for the coefficient of the quadratic polynomial because in the exponax.stepper.generic.GeneralPolynomialStepper the polynomial nonlinearity is already on the right-hand side.)

The higher-dimensional generalization is

    uₜ = b₀ u² + b₁ 1/2 (1⃗ ⋅ ∇)(u²) + b₂ 1/2 ‖ ∇u ‖₂² + sum_j a_j uₓˢ

Under the default configuration, this correspons to a Burgers equation in single-channel mode.

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: The list of coefficients a_j corresponding to the derivatives. The length of this tuple represents the highest occuring derivative. The default value (0.0, 0.0, 0.01) together with the default nonlinear_coefficients corresponds to the Burgers equation.
  • nonlinear_coefficients: The list of coefficients b₀, b₁, b₂ (in this order). The default value (0.0, -1.0, 0.0) corresponds to a (single-channel) convection nonlinearity scaled with 1.0. Note that all nonlinear contributions are considered to be on the right-hand side of the PDE. Hence, in order to get the "correct convection" dynamics, the coefficients must be negated.
  • order: The order of the ETDRK method to use. Must be one of {0, 1, 2, 3, 4}. The option 0 only solves the linear part of the equation. Use higher values for higher accuracy and stability. The default choice of 2 is a good compromise for single precision floats.
  • dealiasing_fraction: The fraction of the wavenumbers to keep before evaluating the nonlinearity. The default value 2/3 corresponds to Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2.
  • num_circle_points: How many points to use in the complex contour integral method to compute the coefficients of the exponential time differencing Runge Kutta method.
  • circle_radius: The radius of the contour used to compute the coefficients of the exponential time differencing Runge Kutta method.
Source code in exponax/stepper/generic/_nonlinear.py 
 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
def __init__(
    self,
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
    nonlinear_coefficients: tuple[float, float, float] = (0.0, -1.0, 0.0),
    order=2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) semi-linear PDEs
    consisting of a quadratic, a single-channel convection, and a gradient
    norm nonlinearity together with an arbitrary combination of (isotropic)
    linear derivatives.

    In 1d, the PDE is of the form

    ```
        uₜ = b₀ u² + b₁ 1/2 (u²)ₓ + b₂ 1/2 (uₓ)² + sum_j a_j uₓʲ
    ```

    where `b₀`, `b₁`, `b₂` are the coefficients of the quadratic,
    convection, and gradient norm nonlinearity, respectively, and `a_j` are
    the coefficients of the linear derivatives. Effectively, this
    timestepper is a combination of the
    `exponax.stepper.generic.GeneralPolynomialStepper` (with only the
    coefficient to the quadratic polynomial being set with `b₀`), the
    `exponax.stepper.generic.GeneralConvectionStepper` (with the
    single-channel hack activated via `single_channel=True` and the
    convection scale set with `- b₁`), and the
    `exponax.stepper.generic.GeneralGradientNormStepper` (with the gradient
    norm scale set with `- b₂`).

    !!! warning
        In contrast to the
        `exponax.stepper.generic.GeneralConvectionStepper` and the
        `exponax.stepper.generic.GeneralGradientNormStepper`, the nonlinear
        terms are no considered to be on right-hand side of the PDE. Hence,
        in order to get the same dynamics as in the other steppers, the
        coefficients must be negated. (This is not relevant for the
        coefficient of the quadratic polynomial because in the
        `exponax.stepper.generic.GeneralPolynomialStepper` the polynomial
        nonlinearity is already on the right-hand side.)

    The higher-dimensional generalization is

    ```
        uₜ = b₀ u² + b₁ 1/2 (1⃗ ⋅ ∇)(u²) + b₂ 1/2 ‖ ∇u ‖₂² + sum_j a_j uₓˢ
    ```

    Under the default configuration, this correspons to a Burgers equation
    in single-channel mode.

    **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`: The list of coefficients `a_j` corresponding to
        the derivatives. The length of this tuple represents the highest
        occuring derivative. The default value `(0.0, 0.0, 0.01)` together
        with the default `nonlinear_coefficients` corresponds to the Burgers
        equation.
    - `nonlinear_coefficients`: The list of coefficients `b₀`, `b₁`, `b₂`
        (in this order). The default value `(0.0, -1.0, 0.0)` corresponds to
        a (single-channel) convection nonlinearity scaled with `1.0`. Note
        that all nonlinear contributions are considered to be on the
        right-hand side of the PDE. Hence, in order to get the "correct
        convection" dynamics, the coefficients must be negated.
    - `order`: The order of the ETDRK method to use. Must be one of {0, 1, 2,
        3, 4}. The option `0` only solves the linear part of the equation.
        Use higher values for higher accuracy and stability. The default
        choice of `2` is a good compromise for single precision floats.
    - `dealiasing_fraction`: The fraction of the wavenumbers to keep before
        evaluating the nonlinearity. The default value `2/3` corresponds to
        Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2.
    - `num_circle_points`: How many points to use in the complex contour
        integral method to compute the coefficients of the exponential time
        differencing Runge Kutta method.
    - `circle_radius`: The radius of the contour used to compute the
        coefficients of the exponential time differencing Runge Kutta method.
    """
    if len(nonlinear_coefficients) != 3:
        raise ValueError(
            "The nonlinear coefficients list must have exactly 3 elements"
        )
    self.linear_coefficients = linear_coefficients
    self.nonlinear_coefficients = nonlinear_coefficients
    self.dealiasing_fraction = dealiasing_fraction

    super().__init__(
        num_spatial_dims=num_spatial_dims,
        domain_extent=domain_extent,
        num_points=num_points,
        dt=dt,
        num_channels=1,
        order=order,
        num_circle_points=num_circle_points,
        circle_radius=circle_radius,
    )
__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)