Convection

exponax.stepper.generic.NormalizedConvectionStepper

Bases: GeneralConvectionStepper

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

class NormalizedConvectionStepper(GeneralConvectionStepper):
    normalized_linear_coefficients: tuple[float, ...]
    normalized_convection_scale: float

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1),
        normalized_convection_scale: float = 1.0 * 0.1,
        single_channel: bool = False,
        conservative: bool = False,
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Time stepper for the **normalized** d-dimensional (`d ∈ {1, 2, 3}`)
        semi-linear PDEs consisting of a convection nonlinearity and 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.GeneralConvectionStepper` for more details
        on the functional form of the PDE.

        In the default configuration, the number of channel grows with the
        number of spatial dimensions. Setting the flag `single_channel=True`
        activates a single-channel hack.

        Under the default settings, it behaves like the Burgers equation under
        the following settings

        ```python

        exponax.stepper.Burgers(
            D=D, L=1, N=N, dt=0.1, diffusivity=0.01,
        )
        ```

        **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
            `α_j` corresponding to the derivatives. The length of this tuple
            represents the highest occuring derivative. The default value `(0.0,
            0.0, 0.01)` corresponds to the Burgers equation (because of the
            diffusion contribution). Note that these coefficients are normalized
            on the unit domain and unit time step size.
        - `normalized_convection_scale`: The scale `β` of the convection term.
            Default is `1.0`.
        - `single_channel`: Whether to use the single channel mode in higher
            dimensions. In this case the the convection is `β (∇ ⋅ 1)(u²)`. In
            this case, the state always has a single channel, no matter the
            spatial dimension. Default: False.
        - `conservative`: Whether to use the conservative form of the convection
            term. Default: False.
        - `order`: The order of the Exponential Time Differencing Runge
            Kutta method. 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 2/3 corresponds to
            Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2. Default:
            2/3.
        - `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. Default: 16.
        - `circle_radius`: The radius of the contour used to compute the
            coefficients of the exponential time differencing Runge Kutta
            method. Default: 1.0.
        """
        self.normalized_linear_coefficients = normalized_linear_coefficients
        self.normalized_convection_scale = normalized_convection_scale
        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,
            convection_scale=normalized_convection_scale,
            order=order,
            dealiasing_fraction=dealiasing_fraction,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
            single_channel=single_channel,
            conservative=conservative,
        )

__init__

__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_coefficients: tuple[float, ...] = (
        0.0,
        0.0,
        0.01 * 0.1,
    ),
    normalized_convection_scale: float = 1.0 * 0.1,
    single_channel: bool = False,
    conservative: bool = False,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Time stepper for the normalized d-dimensional (d ∈ {1, 2, 3}) semi-linear PDEs consisting of a convection nonlinearity and 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.GeneralConvectionStepper for more details on the functional form of the PDE.

In the default configuration, the number of channel grows with the number of spatial dimensions. Setting the flag single_channel=True activates a single-channel hack.

Under the default settings, it behaves like the Burgers equation under the following settings

exponax.stepper.Burgers(
    D=D, L=1, N=N, dt=0.1, diffusivity=0.01,
)

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 α_j corresponding to the derivatives. The length of this tuple represents the highest occuring derivative. The default value (0.0, 0.0, 0.01) corresponds to the Burgers equation (because of the diffusion contribution). Note that these coefficients are normalized on the unit domain and unit time step size.
• normalized_convection_scale: The scale β of the convection term. Default is 1.0.
• single_channel: Whether to use the single channel mode in higher dimensions. In this case the the convection is β (∇ ⋅ 1)(u²). In this case, the state always has a single channel, no matter the spatial dimension. Default: False.
• conservative: Whether to use the conservative form of the convection term. Default: False.
• order: The order of the Exponential Time Differencing Runge Kutta method. 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 2/3 corresponds to Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2. Default: 2/3.
• 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. Default: 16.
• circle_radius: The radius of the contour used to compute the coefficients of the exponential time differencing Runge Kutta method. Default: 1.0.

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

def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1),
    normalized_convection_scale: float = 1.0 * 0.1,
    single_channel: bool = False,
    conservative: bool = False,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    Time stepper for the **normalized** d-dimensional (`d ∈ {1, 2, 3}`)
    semi-linear PDEs consisting of a convection nonlinearity and 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.GeneralConvectionStepper` for more details
    on the functional form of the PDE.

    In the default configuration, the number of channel grows with the
    number of spatial dimensions. Setting the flag `single_channel=True`
    activates a single-channel hack.

    Under the default settings, it behaves like the Burgers equation under
    the following settings

    ```python

    exponax.stepper.Burgers(
        D=D, L=1, N=N, dt=0.1, diffusivity=0.01,
    )
    ```

    **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
        `α_j` corresponding to the derivatives. The length of this tuple
        represents the highest occuring derivative. The default value `(0.0,
        0.0, 0.01)` corresponds to the Burgers equation (because of the
        diffusion contribution). Note that these coefficients are normalized
        on the unit domain and unit time step size.
    - `normalized_convection_scale`: The scale `β` of the convection term.
        Default is `1.0`.
    - `single_channel`: Whether to use the single channel mode in higher
        dimensions. In this case the the convection is `β (∇ ⋅ 1)(u²)`. In
        this case, the state always has a single channel, no matter the
        spatial dimension. Default: False.
    - `conservative`: Whether to use the conservative form of the convection
        term. Default: False.
    - `order`: The order of the Exponential Time Differencing Runge
        Kutta method. 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 2/3 corresponds to
        Orszag's 2/3 rule. To fully eliminate aliasing, use 1/2. Default:
        2/3.
    - `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. Default: 16.
    - `circle_radius`: The radius of the contour used to compute the
        coefficients of the exponential time differencing Runge Kutta
        method. Default: 1.0.
    """
    self.normalized_linear_coefficients = normalized_linear_coefficients
    self.normalized_convection_scale = normalized_convection_scale
    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,
        convection_scale=normalized_convection_scale,
        order=order,
        dealiasing_fraction=dealiasing_fraction,
        num_circle_points=num_circle_points,
        circle_radius=circle_radius,
        single_channel=single_channel,
        conservative=conservative,
    )

__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)