Linear

exponax.normalized.NormalizedLinearStepper

Bases: BaseStepper

Source code in exponax/normalized/_linear.py

class NormalizedLinearStepper(BaseStepper):
    normalized_coefficients: tuple[float, ...]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01),
    ):
        """
        Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
        boundary conditions with normalized dynamics.

        If the PDE in physical description is

            uₜ = ∑ᵢ aᵢ (∂ₓ)ⁱ u

        with `aᵢ` the coefficients on `domain_extent` `L` with time step `dt`,
        the normalized coefficients are

            αᵢ = (aᵢ Δt)/(Lⁱ)

        Important: note that the `domain_extent` is raised to the order of
        linear derivative `i`.

        One can also think of normalized dynamics, as a PDE on `domain_extent`
        `1.0` with time step `dt=1.0`.

        Take care of the signs!

        In the defaulf configuration of this timestepper, the PDE is an
        advection-diffusion equation with normalized advection of 0.5 and
        normalized diffusion of 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_coefficients`: The coefficients of the normalized
            dynamics. This must a tuple of floats. The length of the tuple
            defines the highest occuring linear derivative in the PDE.
        """
        self.normalized_coefficients = normalized_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,
            num_channels=1,
            order=0,
        )

    def _build_linear_operator(self, derivative_operator: Array) -> Array:
        linear_operator = sum(
            jnp.sum(
                c * (derivative_operator) ** i,
                axis=0,
                keepdims=True,
            )
            for i, c in enumerate(self.normalized_coefficients)
        )
        return linear_operator

    def _build_nonlinear_fun(self, derivative_operator: Array):
        return ZeroNonlinearFun(
            num_spatial_dims=self.num_spatial_dims,
            num_points=self.num_points,
        )

__init__

__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_coefficients: tuple[float, ...] = (
        0.0,
        -0.5,
        0.01,
    )
)

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

If the PDE in physical description is

uₜ = ∑ᵢ aᵢ (∂ₓ)ⁱ u

with aᵢ the coefficients on domain_extent L with time step dt, the normalized coefficients are

αᵢ = (aᵢ Δt)/(Lⁱ)

Important: note that the domain_extent is raised to the order of linear derivative i.

One can also think of normalized dynamics, as a PDE on domain_extent 1.0 with time step dt=1.0.

Take care of the signs!

In the defaulf configuration of this timestepper, the PDE is an advection-diffusion equation with normalized advection of 0.5 and normalized diffusion of 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_coefficients: The coefficients of the normalized dynamics. This must a tuple of floats. The length of the tuple defines the highest occuring linear derivative in the PDE.

Source code in exponax/normalized/_linear.py

def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01),
):
    """
    Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
    boundary conditions with normalized dynamics.

    If the PDE in physical description is

        uₜ = ∑ᵢ aᵢ (∂ₓ)ⁱ u

    with `aᵢ` the coefficients on `domain_extent` `L` with time step `dt`,
    the normalized coefficients are

        αᵢ = (aᵢ Δt)/(Lⁱ)

    Important: note that the `domain_extent` is raised to the order of
    linear derivative `i`.

    One can also think of normalized dynamics, as a PDE on `domain_extent`
    `1.0` with time step `dt=1.0`.

    Take care of the signs!

    In the defaulf configuration of this timestepper, the PDE is an
    advection-diffusion equation with normalized advection of 0.5 and
    normalized diffusion of 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_coefficients`: The coefficients of the normalized
        dynamics. This must a tuple of floats. The length of the tuple
        defines the highest occuring linear derivative in the PDE.
    """
    self.normalized_coefficients = normalized_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,
        num_channels=1,
        order=0,
    )

__call__

__call__(
    u: Float[Array, "C ... N"]
) -> Float[Array, "C ... N"]

Performs a check

Source code in exponax/_base_stepper.py

def __call__(
    self,
    u: Float[Array, "C ... N"],
) -> Float[Array, "C ... N"]:
    """
    Performs a check
    """
    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)