Linear

exponax.stepper.generic.DifficultyLinearStepper

Bases: NormalizedLinearStepper

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

class DifficultyLinearStepper(NormalizedLinearStepper):
    linear_difficulties: tuple[float, ...]

    def __init__(
        self,
        num_spatial_dims: int = 1,
        num_points: int = 48,
        *,
        linear_difficulties: tuple[float, ...] = (0.0, -2.0),
    ):
        """
        Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
        boundary conditions with normalized dynamics in a difficulty-based
        interface.

        Different to `NormalizedLinearStepper`, 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.

        This interface is more natural 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.

        In the default configuration of this timestepper, the PDE is an
        advection equation with CFL number 2 solved in 1d with 48 resolution
        points to discretize the domain.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `d`. Default is
            1.
        - `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ᵈ`. Default is 48.
        - `difficulties`: The difficulties of the normalized dynamics. This must
            be a tuple of floats. The length of the tuple defines the highest
            occuring linear derivative in the PDE. Default is `(0.0, -2.0)`.
        """
        self.linear_difficulties = linear_difficulties
        normalized_coefficients = extract_normalized_coefficients_from_difficulty(
            linear_difficulties,
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
        )

        super().__init__(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            normalized_linear_coefficients=normalized_coefficients,
        )

__init__

__init__(
    num_spatial_dims: int = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (0.0, -2.0)
)

Timestepper for d-dimensional (d ∈ {1, 2, 3}) linear PDEs on periodic boundary conditions with normalized dynamics in a difficulty-based interface.

Different to NormalizedLinearStepper, 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.

This interface is more natural 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.

In the default configuration of this timestepper, the PDE is an advection equation with CFL number 2 solved in 1d with 48 resolution points to discretize the domain.

Arguments:

• num_spatial_dims: The number of spatial dimensions d. Default is 1.
• 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ᵈ. Default is 48.
• difficulties: The difficulties of the normalized dynamics. This must be a tuple of floats. The length of the tuple defines the highest occuring linear derivative in the PDE. Default is (0.0, -2.0).

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

def __init__(
    self,
    num_spatial_dims: int = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (0.0, -2.0),
):
    """
    Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
    boundary conditions with normalized dynamics in a difficulty-based
    interface.

    Different to `NormalizedLinearStepper`, 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.

    This interface is more natural 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.

    In the default configuration of this timestepper, the PDE is an
    advection equation with CFL number 2 solved in 1d with 48 resolution
    points to discretize the domain.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `d`. Default is
        1.
    - `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ᵈ`. Default is 48.
    - `difficulties`: The difficulties of the normalized dynamics. This must
        be a tuple of floats. The length of the tuple defines the highest
        occuring linear derivative in the PDE. Default is `(0.0, -2.0)`.
    """
    self.linear_difficulties = linear_difficulties
    normalized_coefficients = extract_normalized_coefficients_from_difficulty(
        linear_difficulties,
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
    )

    super().__init__(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
        normalized_linear_coefficients=normalized_coefficients,
    )

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