Nonlinear¤

exponax.stepper.generic.NormalizedNonlinearStepper ¤

Bases: GeneralNonlinearStepper

Source code in exponax/stepper/generic/_nonlinear.py

class NormalizedNonlinearStepper(GeneralNonlinearStepper):
    normalized_linear_coefficients: tuple[float, ...]
    normalized_nonlinear_coefficients: tuple[float, float, float]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1),
        normalized_nonlinear_coefficients: tuple[float, float, float] = (
            0.0,
            -1.0 * 0.1,
            0.0,
        ),
        order=2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        By default Burgers.

        Timesteppr for **normalized** 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. Uses a normalized interface, i.e.,
        the domain is scaled to `Ω = (0, 1)ᵈ` and time step size is `Δt = 1.0`.

        See `exponax.stepper.generic.GeneralNonlinearStepper` for more details
        on the functional form of the PDE.

        Behaves like a single-channel Burgers equation by default under the
        following settings

        ```python

        exponax.stepper.Burgers(
            num_spatial_dims=num_spatial_dims, domain_extent=1.0,
            num_points=num_points, dt=1.0, convection_scale=1.0,
            diffusivity=0.1, single_channel=True,
        )
        ```

        Effectively, this timestepper is a combination of the
        `exponax.stepper.generic.NormalizedPolynomialStepper` (with only the
        coefficient to the quadratic polynomial being set with `b₀`), the
        `exponax.stepper.generic.NormalizedConvectionStepper` (with the
        single-channel hack activated via `single_channel=True` and the
        convection scale set with `- b₁`), and the
        `exponax.stepper.generic.NormalizedGradientNormStepper` (with the
        gradient norm scale set with `- b₂`).

        !!! warning
            In contrast to the
            `exponax.stepper.generic.NormalizedConvectionStepper` and the
            `exponax.stepper.generic.NormalizedGradientNormStepper`, 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.NormalizedPolynomialStepper` the polynomial
            nonlinearity is already on the right-hand side.)


        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `d`.
        - `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ᵈ`.
        - `normalized_linear_coefficients`: The list of coefficients `αⱼ`
            corresponding to the linear derivatives. The length of this tuple
            represents the highest occuring derivative. The default value `(0.0,
            0.0, 0.1 * 0.1)` together with the default
            `normalized_nonlinear_coefficients` corresponds to the Burgers
            equation (in single-channel mode).
        - `normalized_nonlinear_coefficients`: The list of coefficients `β₀`,
            `β₁`, and `β₂` (in this order) corresponding to the quadratic,
            (single-channel) convection, and gradient norm nonlinearity,
            respectively. The default value `(0.0, -1.0 * 0.1, 0.0)` corresponds
            to a (single-channel) convection nonlinearity scaled with `0.1`.
            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. (Also
            relevant for the gradient norm nonlinearity)
        - `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.
        """

        self.normalized_linear_coefficients = normalized_linear_coefficients
        self.normalized_nonlinear_coefficients = normalized_nonlinear_coefficients

        super().__init__(
            num_spatial_dims=num_spatial_dims,
            domain_extent=1.0,  # Derivative operator is just scaled with 2 * jnp.pi
            num_points=num_points,
            dt=1.0,
            linear_coefficients=normalized_linear_coefficients,
            nonlinear_coefficients=normalized_nonlinear_coefficients,
            order=order,
            dealiasing_fraction=dealiasing_fraction,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
        )

init ¤

__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_coefficients: tuple[float, ...] = (
        0.0,
        0.0,
        0.1 * 0.1,
    ),
    normalized_nonlinear_coefficients: tuple[
        float, float, float
    ] = (0.0, -1.0 * 0.1, 0.0),
    order=2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

By default Burgers.

Timesteppr for normalized 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. Uses a normalized interface, i.e., the domain is scaled to Ω = (0, 1)ᵈ and time step size is Δt = 1.0.

See exponax.stepper.generic.GeneralNonlinearStepper for more details on the functional form of the PDE.

Behaves like a single-channel Burgers equation by default under the following settings

exponax.stepper.Burgers(
    num_spatial_dims=num_spatial_dims, domain_extent=1.0,
    num_points=num_points, dt=1.0, convection_scale=1.0,
    diffusivity=0.1, single_channel=True,
)

Effectively, this timestepper is a combination of the exponax.stepper.generic.NormalizedPolynomialStepper (with only the coefficient to the quadratic polynomial being set with b₀), the exponax.stepper.generic.NormalizedConvectionStepper (with the single-channel hack activated via single_channel=True and the convection scale set with - b₁), and the exponax.stepper.generic.NormalizedGradientNormStepper (with the gradient norm scale set with - b₂).

Warning

In contrast to the exponax.stepper.generic.NormalizedConvectionStepper and the exponax.stepper.generic.NormalizedGradientNormStepper, 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.NormalizedPolynomialStepper the polynomial nonlinearity is already on the right-hand side.)

Arguments:

num_spatial_dims: The number of spatial dimensions d.
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ᵈ.
normalized_linear_coefficients: The list of coefficients αⱼ corresponding to the linear derivatives. The length of this tuple represents the highest occuring derivative. The default value (0.0, 0.0, 0.1 * 0.1) together with the default normalized_nonlinear_coefficients corresponds to the Burgers equation (in single-channel mode).
normalized_nonlinear_coefficients: The list of coefficients β₀, β₁, and β₂ (in this order) corresponding to the quadratic, (single-channel) convection, and gradient norm nonlinearity, respectively. The default value (0.0, -1.0 * 0.1, 0.0) corresponds to a (single-channel) convection nonlinearity scaled with 0.1. 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. (Also relevant for the gradient norm nonlinearity)
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

def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1),
    normalized_nonlinear_coefficients: tuple[float, float, float] = (
        0.0,
        -1.0 * 0.1,
        0.0,
    ),
    order=2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    By default Burgers.

    Timesteppr for **normalized** 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. Uses a normalized interface, i.e.,
    the domain is scaled to `Ω = (0, 1)ᵈ` and time step size is `Δt = 1.0`.

    See `exponax.stepper.generic.GeneralNonlinearStepper` for more details
    on the functional form of the PDE.

    Behaves like a single-channel Burgers equation by default under the
    following settings

    ```python

    exponax.stepper.Burgers(
        num_spatial_dims=num_spatial_dims, domain_extent=1.0,
        num_points=num_points, dt=1.0, convection_scale=1.0,
        diffusivity=0.1, single_channel=True,
    )
    ```

    Effectively, this timestepper is a combination of the
    `exponax.stepper.generic.NormalizedPolynomialStepper` (with only the
    coefficient to the quadratic polynomial being set with `b₀`), the
    `exponax.stepper.generic.NormalizedConvectionStepper` (with the
    single-channel hack activated via `single_channel=True` and the
    convection scale set with `- b₁`), and the
    `exponax.stepper.generic.NormalizedGradientNormStepper` (with the
    gradient norm scale set with `- b₂`).

    !!! warning
        In contrast to the
        `exponax.stepper.generic.NormalizedConvectionStepper` and the
        `exponax.stepper.generic.NormalizedGradientNormStepper`, 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.NormalizedPolynomialStepper` the polynomial
        nonlinearity is already on the right-hand side.)


    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `d`.
    - `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ᵈ`.
    - `normalized_linear_coefficients`: The list of coefficients `αⱼ`
        corresponding to the linear derivatives. The length of this tuple
        represents the highest occuring derivative. The default value `(0.0,
        0.0, 0.1 * 0.1)` together with the default
        `normalized_nonlinear_coefficients` corresponds to the Burgers
        equation (in single-channel mode).
    - `normalized_nonlinear_coefficients`: The list of coefficients `β₀`,
        `β₁`, and `β₂` (in this order) corresponding to the quadratic,
        (single-channel) convection, and gradient norm nonlinearity,
        respectively. The default value `(0.0, -1.0 * 0.1, 0.0)` corresponds
        to a (single-channel) convection nonlinearity scaled with `0.1`.
        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. (Also
        relevant for the gradient norm nonlinearity)
    - `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.
    """

    self.normalized_linear_coefficients = normalized_linear_coefficients
    self.normalized_nonlinear_coefficients = normalized_nonlinear_coefficients

    super().__init__(
        num_spatial_dims=num_spatial_dims,
        domain_extent=1.0,  # Derivative operator is just scaled with 2 * jnp.pi
        num_points=num_points,
        dt=1.0,
        linear_coefficients=normalized_linear_coefficients,
        nonlinear_coefficients=normalized_nonlinear_coefficients,
        order=order,
        dealiasing_fraction=dealiasing_fraction,
        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

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)

Nonlinear¤

exponax.stepper.generic.NormalizedNonlinearStepper ¤

__init__ ¤

__call__ ¤

init ¤

call ¤