Polynomial

exponax.stepper.generic.NormalizedPolynomialStepper

Bases: GeneralPolynomialStepper

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

class NormalizedPolynomialStepper(GeneralPolynomialStepper):
    normalized_linear_coefficients: tuple[float, ...]
    normalized_polynomial_coefficients: tuple[float, ...]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_linear_coefficients: tuple[float, ...] = (
            10.0 * 0.001 / (10.0**0),
            0.0,
            1.0 * 0.001 / (10.0**2),
        ),
        normalized_polynomial_coefficients: tuple[float, ...] = (
            0.0,
            0.0,
            -10.0 * 0.001,
        ),
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Timestepper for the **normalized** d-dimensional (`d ∈ {1, 2, 3}`)
        semi-linear PDEs consisting of an arbitrary combination of polynomial
        nonlinearities and (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.GeneralPolynomialStepper` for more details
        on the functional form of the PDE.

        The default settings correspond to the Fisher-KPP equation.

        **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
            corresponds to the Fisher-KPP equation.
        - `normalized_polynomial_coefficients`: The list of scales `βₖ`
            corresponding to the polynomial contributions. The length of this
            tuple represents the highest occuring polynomial. The default value
            corresponds to the Fisher-KPP equation.
        - `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 which is sufficient if the highest occuring polynomial is
            quadratic (i.e., there are at maximum three entries in the
            `normalized_polynomial_scales` tuple).
        - `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_polynomial_coefficients = normalized_polynomial_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,
            polynomial_coefficients=normalized_polynomial_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, ...] = (
        10.0 * 0.001 / 10.0**0,
        0.0,
        1.0 * 0.001 / 10.0**2,
    ),
    normalized_polynomial_coefficients: tuple[
        float, ...
    ] = (0.0, 0.0, -10.0 * 0.001),
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

Timestepper for the normalized d-dimensional (d ∈ {1, 2, 3}) semi-linear PDEs consisting of an arbitrary combination of polynomial nonlinearities and (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.GeneralPolynomialStepper for more details on the functional form of the PDE.

The default settings correspond to the Fisher-KPP equation.

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 corresponds to the Fisher-KPP equation.
• normalized_polynomial_coefficients: The list of scales βₖ corresponding to the polynomial contributions. The length of this tuple represents the highest occuring polynomial. The default value corresponds to the Fisher-KPP equation.
• 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 which is sufficient if the highest occuring polynomial is quadratic (i.e., there are at maximum three entries in the normalized_polynomial_scales tuple).
• 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/_polynomial.py

def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_coefficients: tuple[float, ...] = (
        10.0 * 0.001 / (10.0**0),
        0.0,
        1.0 * 0.001 / (10.0**2),
    ),
    normalized_polynomial_coefficients: tuple[float, ...] = (
        0.0,
        0.0,
        -10.0 * 0.001,
    ),
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    Timestepper for the **normalized** d-dimensional (`d ∈ {1, 2, 3}`)
    semi-linear PDEs consisting of an arbitrary combination of polynomial
    nonlinearities and (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.GeneralPolynomialStepper` for more details
    on the functional form of the PDE.

    The default settings correspond to the Fisher-KPP equation.

    **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
        corresponds to the Fisher-KPP equation.
    - `normalized_polynomial_coefficients`: The list of scales `βₖ`
        corresponding to the polynomial contributions. The length of this
        tuple represents the highest occuring polynomial. The default value
        corresponds to the Fisher-KPP equation.
    - `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 which is sufficient if the highest occuring polynomial is
        quadratic (i.e., there are at maximum three entries in the
        `normalized_polynomial_scales` tuple).
    - `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_polynomial_coefficients = normalized_polynomial_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,
        polynomial_coefficients=normalized_polynomial_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)