Nonlinear

exponax.stepper.generic.DifficultyNonlinearStepper

Bases: NormalizedNonlinearStepper

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

class DifficultyNonlinearStepper(NormalizedNonlinearStepper):
    linear_difficulties: tuple[float, ...]
    nonlinear_difficulties: tuple[float, float, float]

    def __init__(
        self,
        num_spatial_dims: int = 1,
        num_points: int = 48,
        *,
        linear_difficulties: tuple[float, ...] = (
            0.0,
            0.0,
            0.1 * 0.1 / 1.0 * 48**2 * 2,
        ),
        nonlinear_difficulties: tuple[float, float, float] = (
            0.0,
            -1.0 * 0.1 / 1.0 * 48,
            0.0,
        ),
        maximum_absolute: float = 1.0,
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Timestepper for **difficulty-based** 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 difficulty-based interface
        where the "intensity" of the dynamics reduces with increasing
        resolution. This is intended such that emulator learning problems on two
        resolutions are comparibly difficult.

        Different to `exponax.stepper.generic.NormalizedNonlinearStepper`, the
        dynamics are defined by difficulties. The difficulties are a different
        combination of normalized dynamics, `num_spatial_dims`, and
        `num_points`.

            γᵢ = αᵢ Nⁱ 2ⁱ⁻¹ d

        with `d` the number of spatial dimensions, `N` the number of points, and
        `αᵢ` the normalized coefficient.

        The difficulties of the nonlinear terms are

            δ₀ = β₀

            δ₁ = β₁ * M * N * D

            δ₂ = β₂ * M * N² * D

        with `βᵢ` the normalized coefficient and `M` the maximum absolute value
        of the input state (typically `1.0` if one uses the `exponax.ic` random
        generators with the `max_one=True` argument).

        This interface is more natural than the normalized interface because the
        difficulties for all orders (given by `i`) are around 1.0. Additionally,
        they relate to stability condition of explicit Finite Difference schemes
        for the particular equations. For example, for advection (`i=1`), the
        absolute of the difficulty is the Courant-Friedrichs-Lewy (CFL) number.

        Under the default settings, this timestep corresponds to a Burgers
        equation in single-channel mode.

        **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ᵈ`.
        - `linear_difficulties`: The list of difficulties `γᵢ` 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
            / 1.0 * 48**2 * 2)` together with the default `nonlinear_difficulties`
            corresponds to the Burgers equation.
        - `nonlinear_difficulties`: The list of difficulties `δ₀`, `δ₁`, 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 / 1.0 * 48, 0.0)` corresponds to a
            (single-channel) convection nonlinearity. Note that all nonlinear
            contributions are considered to be on the right-hand side of the PDE.
        - `maximum_absolute`: The maximum absolute value of the input state. This
            is used to scale the nonlinear difficulties.
        - `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.linear_difficulties = linear_difficulties
        self.nonlinear_difficulties = nonlinear_difficulties

        normalized_coefficients_linear = (
            extract_normalized_coefficients_from_difficulty(
                linear_difficulties,
                num_spatial_dims=num_spatial_dims,
                num_points=num_points,
            )
        )
        normalized_coefficients_nonlinear = (
            extract_normalized_nonlinear_scales_from_difficulty(
                nonlinear_difficulties,
                num_spatial_dims=num_spatial_dims,
                num_points=num_points,
                maximum_absolute=maximum_absolute,
            )
        )

        super().__init__(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            normalized_linear_coefficients=normalized_coefficients_linear,
            normalized_nonlinear_coefficients=normalized_coefficients_nonlinear,
            order=order,
            dealiasing_fraction=dealiasing_fraction,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
        )

__init__

__init__(
    num_spatial_dims: int = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (
        0.0,
        0.0,
        0.1 * 0.1 / 1.0 * 48**2 * 2,
    ),
    nonlinear_difficulties: tuple[float, float, float] = (
        0.0,
        -1.0 * 0.1 / 1.0 * 48,
        0.0,
    ),
    maximum_absolute: float = 1.0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Timestepper for difficulty-based 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 difficulty-based interface where the "intensity" of the dynamics reduces with increasing resolution. This is intended such that emulator learning problems on two resolutions are comparibly difficult.

Different to exponax.stepper.generic.NormalizedNonlinearStepper, the dynamics are defined by difficulties. The difficulties are a different combination of normalized dynamics, num_spatial_dims, and num_points.

γᵢ = αᵢ Nⁱ 2ⁱ⁻¹ d

with d the number of spatial dimensions, N the number of points, and αᵢ the normalized coefficient.

The difficulties of the nonlinear terms are

δ₀ = β₀

δ₁ = β₁ * M * N * D

δ₂ = β₂ * M * N² * D

with βᵢ the normalized coefficient and M the maximum absolute value of the input state (typically 1.0 if one uses the exponax.ic random generators with the max_one=True argument).

This interface is more natural than the normalized interface because the difficulties for all orders (given by i) are around 1.0. Additionally, they relate to stability condition of explicit Finite Difference schemes for the particular equations. For example, for advection (i=1), the absolute of the difficulty is the Courant-Friedrichs-Lewy (CFL) number.

Under the default settings, this timestep corresponds to a Burgers equation in single-channel mode.

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ᵈ.
• linear_difficulties: The list of difficulties γᵢ 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 / 1.0 * 48**2 * 2) together with the default nonlinear_difficulties corresponds to the Burgers equation.
• nonlinear_difficulties: The list of difficulties δ₀, δ₁, 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 / 1.0 * 48, 0.0) corresponds to a (single-channel) convection nonlinearity. Note that all nonlinear contributions are considered to be on the right-hand side of the PDE.
• maximum_absolute: The maximum absolute value of the input state. This is used to scale the nonlinear difficulties.
• 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 = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (
        0.0,
        0.0,
        0.1 * 0.1 / 1.0 * 48**2 * 2,
    ),
    nonlinear_difficulties: tuple[float, float, float] = (
        0.0,
        -1.0 * 0.1 / 1.0 * 48,
        0.0,
    ),
    maximum_absolute: float = 1.0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    Timestepper for **difficulty-based** 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 difficulty-based interface
    where the "intensity" of the dynamics reduces with increasing
    resolution. This is intended such that emulator learning problems on two
    resolutions are comparibly difficult.

    Different to `exponax.stepper.generic.NormalizedNonlinearStepper`, the
    dynamics are defined by difficulties. The difficulties are a different
    combination of normalized dynamics, `num_spatial_dims`, and
    `num_points`.

        γᵢ = αᵢ Nⁱ 2ⁱ⁻¹ d

    with `d` the number of spatial dimensions, `N` the number of points, and
    `αᵢ` the normalized coefficient.

    The difficulties of the nonlinear terms are

        δ₀ = β₀

        δ₁ = β₁ * M * N * D

        δ₂ = β₂ * M * N² * D

    with `βᵢ` the normalized coefficient and `M` the maximum absolute value
    of the input state (typically `1.0` if one uses the `exponax.ic` random
    generators with the `max_one=True` argument).

    This interface is more natural than the normalized interface because the
    difficulties for all orders (given by `i`) are around 1.0. Additionally,
    they relate to stability condition of explicit Finite Difference schemes
    for the particular equations. For example, for advection (`i=1`), the
    absolute of the difficulty is the Courant-Friedrichs-Lewy (CFL) number.

    Under the default settings, this timestep corresponds to a Burgers
    equation in single-channel mode.

    **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ᵈ`.
    - `linear_difficulties`: The list of difficulties `γᵢ` 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
        / 1.0 * 48**2 * 2)` together with the default `nonlinear_difficulties`
        corresponds to the Burgers equation.
    - `nonlinear_difficulties`: The list of difficulties `δ₀`, `δ₁`, 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 / 1.0 * 48, 0.0)` corresponds to a
        (single-channel) convection nonlinearity. Note that all nonlinear
        contributions are considered to be on the right-hand side of the PDE.
    - `maximum_absolute`: The maximum absolute value of the input state. This
        is used to scale the nonlinear difficulties.
    - `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.linear_difficulties = linear_difficulties
    self.nonlinear_difficulties = nonlinear_difficulties

    normalized_coefficients_linear = (
        extract_normalized_coefficients_from_difficulty(
            linear_difficulties,
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
        )
    )
    normalized_coefficients_nonlinear = (
        extract_normalized_nonlinear_scales_from_difficulty(
            nonlinear_difficulties,
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            maximum_absolute=maximum_absolute,
        )
    )

    super().__init__(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
        normalized_linear_coefficients=normalized_coefficients_linear,
        normalized_nonlinear_coefficients=normalized_coefficients_nonlinear,
        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)