Exponax

https://fkoehler.site/exponax
Admin
Exponax is a suite for building efficient and differentiable Fourier spectral ETDRK time-steppers...
Tokens:479,214
Snippets:2,234
Trust Score:5.7
Update:3 months ago
Show doc for...
Tokens:
Raw
### 1D Kuramoto-Sivashinsky Equation Quickstart

Source: https://fkoehler.site/exponax

This example demonstrates how to set up and run a simulation for the 1D Kuramoto-Sivashinsky equation using Exponax. It initializes a random state and plots the trajectory.

```python
import jax
import exponax as ex
import matplotlib.pyplot as plt

ks_stepper = ex.stepper.KuramotoSivashinskyConservative(
    num_spatial_dims=1, domain_extent=100.0,
    num_points=200, dt=0.1,
)

u_0 = ex.ic.RandomTruncatedFourierSeries(
    num_spatial_dims=1, cutoff=5
)(num_points=200, key=jax.random.PRNGKey(0))

trajectory = ex.rollout(ks_stepper, 500, include_init=True)(u_0)

plt.imshow(trajectory[:, 0, :].T, aspect='auto', cmap='RdBu', vmin=-2, vmax=2, origin="lower")
plt.xlabel("Time"); plt.ylabel("Space"); plt.show()

```

--------------------------------

### Visualize 3D Fourier Coefficients

Source: https://fkoehler.site/exponax/api/utilities/spectral

Example of visualizing 3D Fourier coefficients using `exponax.viz.volume_render_state_3d` and logarithmic scaling. Requires `vape4d` to be installed.

```python
state_3d = ...  # shape (1, N, N, N)

coef_3d = exponax.spectral.get_fourier_coefficients(
    state_3d, round=None,
)

images = ex.viz.volume_render_state_3d(
    jnp.log10(jnp.abs(coef_3d)), vlim=(-8, 2),
)

plt.imshow(images[0])
```

--------------------------------

### Install Exponax

Source: https://fkoehler.site/exponax

Install the Exponax package using pip. Requires Python 3.10+ and JAX 0.4.13+.

```bash
pip install exponax

```

--------------------------------

### Burgers Equation Stepper Setup

Source: https://fkoehler.site/exponax/examples/understanding_general_and_normalized_stepper

Sets up a stepper for the 1D Burgers equation with specified parameters for domain extent, number of points, time step, diffusivity, and convection scale.

```python
burgers_stepper = ex.stepper.Burgers(
    1,  # number of dimensions (set to 1)
    domain_extent=5.0,
    num_points=100,
    dt=0.7,
    diffusivity=0.05,
    convection_scale=0.25,  # Let's be brave and set it different from its default 1.0
)

```

--------------------------------

### Build Initial Condition Set

Source: https://fkoehler.site/exponax/api/utilities/initial_conditions/helper

Generates a set of initial conditions by sampling from a given distribution. Use this to create multiple starting states for simulations.

```python
def build_ic_set(
    ic_generator: Callable[[int, PRNGKeyArray], Float[Array, "C ... N"]],
    *,
    num_points: int,
    num_samples: int,
    key: PRNGKeyArray,
) -> Float[Array, "S C ... N"]:
    """
    Generate a set of initial conditions by sampling from a given initial
    condition distribution and evaluating the function on the given grid.

    **Arguments**:

    - `ic_generator`: A function that takes a number of points and a PRNGKey
        and returns an array representing the discrete state of an initial
        condition. The shape of the returned array is `(C, ..., N)`.
    - `num_samples`: The number of initial conditions to sample.
    - `key`: The PRNGKey to use for sampling.

    **Returns**:

    - `ic_set`: The set of initial conditions. Shape: `(S, C, ..., N)`.
        `S = num_samples`.
    """

    def scan_fn(k, _):
        k, sub_k = jr.split(k)
        ic = ic_generator(num_points, key=sub_k)
        return k, ic

    _, ic_set = jax.lax.scan(scan_fn, key, None, length=num_samples)

    return ic_set

```

--------------------------------

### Rollout with Auxiliary Input

Source: https://fkoehler.site/exponax/api/utilities/rollout_and_repeat

This example demonstrates using `exponax.repeat` with a stepper function that accepts an auxiliary input. The auxiliary input can be constant or vary over the trajectory.

```python
from jax import numpy as jnp

def stepper_fn_aux(u, aux):
    return u + aux['force']

# Constant auxiliary input
rollout_stepper_fn_const_aux = exponax.repeat(
    stepper_fn_aux, n=5, takes_aux=True, constant_aux=True
)
u_0 = jnp.array([0.0])
aux_const = {'force': 1.0}
trajectory_const = rollout_stepper_fn_const_aux(u_0, aux_const)

print(trajectory_const.shape) # (5, 1)
print(trajectory_const) # [[1.]
                      #  [2.]
                      #  [3.]
                      #  [4.]
                      #  [5.]]

# Variable auxiliary input
rollout_stepper_fn_var_aux = exponax.repeat(
    stepper_fn_aux, n=5, takes_aux=True, constant_aux=False
)
aux_var = {'force': jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])}
trajectory_var = rollout_stepper_fn_var_aux(u_0, aux_var)

print(trajectory_var.shape) # (5, 1)
print(trajectory_var) # [[1.]
                    #  [3.]
                    #  [6.]
                    #  [10.]
                    #  [15.]]
```

--------------------------------

### Initialize SimplePeriodicConvNet

Source: https://fkoehler.site/exponax/examples/learning_burgers_autoregressive_neural_operator

Provides a sample initialization for the `SimplePeriodicConvNet` model. This is useful for sanity checks and to obtain an initial state of the network before training.

```python
neural_stepper = SimplePeriodicConvNet(1, 1, 32, 5, key=jr.PRNGKey(INIT_SEED))
neural_stepper
```

--------------------------------

### Example 1D Fourier Coefficients

Source: https://fkoehler.site/exponax/api/utilities/spectral

An example of the output format for 1D Fourier coefficients, representing a signal with a constant offset and sine/cosine modes.

```python
array([[3.0 + 0.0j, 0.0 - 1.5j, 0.3 + 0.8j, 0.0 + 0.0j,]])
```

--------------------------------

### Initialize Exponax Stepper and Initial Conditions

Source: https://fkoehler.site/exponax/examples/performance_hints

Sets up the Burgers stepper and generates initial conditions for simulation. Includes both a single initial condition and a set of multiple initial conditions.

```python
NUM_SPATIAL_DIMS = 1
DOMAIN_EXTENT = 3.0
NUM_POINTS = 100
DT = 0.1

burgers_stepper = ex.stepper.Burgers(NUM_SPATIAL_DIMS, DOMAIN_EXTENT, NUM_POINTS, DT)

ic_gen = ex.ic.RandomTruncatedFourierSeries(
    NUM_SPATIAL_DIMS,
    cutoff=5,
    max_one=True,
)

one_u_0 = ic_gen(NUM_POINTS, key=jax.random.PRNGKey(0))

multiple_u_0 = ex.build_ic_set(
    ic_gen, num_points=NUM_POINTS, num_samples=10, key=jax.random.PRNGKey(0)
)

```

--------------------------------

### Setup Reference and Predicted States

Source: https://fkoehler.site/exponax/examples/on_metrics_simple

Creates two 1D fields with shared large-scale structure but different smoothing levels using DiffusedNoise. This setup is for comparing states using metrics.

```python
NUM_POINTS = 100

u_ref = ex.ic.DiffusedNoise(1, intensity=0.001)(NUM_POINTS, key=jax.random.PRNGKey(0))
u_pred = ex.ic.DiffusedNoise(1, intensity=0.0005)(NUM_POINTS, key=jax.random.PRNGKey(0))

print("u_ref shape:", u_ref.shape)
print("u_pred shape:", u_pred.shape)
```

--------------------------------

### KdV Equation Stepper Initialization

Source: https://fkoehler.site/exponax/api/stepper/nonlinear/kdv

Initializes the KdV stepper for d-dimensional simulations with periodic boundary conditions. Configure parameters like domain extent, grid points, time step, and various physical coefficients for convection, diffusion, and dispersion.

```python
def __init__(
    self,
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    convection_scale: float = -6.0,
    diffusivity: float = 0.0,
    dispersivity: float = 1.0,
    hyper_diffusivity: float = 0.01,
    advect_over_diffuse: bool = False,
    diffuse_over_diffuse: bool = False,
    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,
):
    """
    Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) (hyper-viscous)
    Korteweg-de Vries equation on periodic boundary conditions.

    In 1d, the Korteweg-de Vries equation is given by

    ```
        uₜ + b₁ 1/2 (u²)ₓ + a₃ uₓₓₓ = ν uₓₓ - ζ uₓₓₓₓ
    ```

    with `b₁` the convection coefficient, `a₃` the dispersion coefficient
    and `ν` the diffusivity. Oftentimes `b₁ = -6` and `ν = 0`. The
    nonlinearity is similar to the Burgers equation and the number of
    channels grow with the number of spatial dimensions. In higher
    dimensions, the equation reads (using vector format for the channels)

    ```
        uₜ + b₁ 1/2 ∇ ⋅ (u ⊗ u) + a₃ 1 ⋅ (∇⊙∇⊙(∇u)) = ν Δu - ζ ((∇⊙∇) ⋅ (∇⊙∇)) u
    ```

    or

    ```
        uₜ + b₁ 1/2 ∇ ⋅ (u ⊗ u) + a₃ ∇ ⋅ ∇(Δu) = ν Δu - μ Δ(Δu)
    ```

    if `advect_over_diffuse` is `True` and `diffuse_on_diffuse` is `True`,

    In 1d, the expected temporal behavior is that the initial condition
    breaks into soliton waves that propagate at a speed depending on their
    height. They interact with other soliton waves by being spatially
    displaced but having an unchanged shape and propagation speed. If either
    the `diffusivity` or the `hyper_diffusivity` is non-zero (by default, 
    the latter is active), the solution will decay over time with the
    produced small-scale features decaying the fastest.

    **Arguments:**

    
```

--------------------------------

### Interpolate State on Unstructured Coordinates

Source: https://fkoehler.site/exponax/examples/additional_features

Applies the interpolator to a set of unstructured coordinates to get the state values at those points.

```python
values = jax.vmap(interpolator)(unstructured_coordinates)

```

--------------------------------

### Rollout for General Convection Stepper

Source: https://fkoehler.site/exponax/examples/understanding_general_and_normalized_stepper

Performs a rollout simulation using the GeneralConvectionStepper with the same initial condition and number of steps as previous examples.

```python
general_trj = ex.rollout(burgers_general, 4, include_init=True)(u_0)
```

--------------------------------

### Visualize 2D Fourier Coefficients

Source: https://fkoehler.site/exponax/api/utilities/spectral

Example of visualizing 2D Fourier coefficients using `plt.imshow` and logarithmic scaling of absolute values.

```python
state_2d = ...  # shape (1, N, N)

coef_2d = exponax.spectral.get_fourier_coefficients(state_2d)

# shape (1, N, (N//2)+1)

plt.imshow(
    jnp.log10(jnp.abs(coef_2d[0])),
)
```

--------------------------------

### Get Fourier Coefficients

Source: https://fkoehler.site/exponax/api/utilities/spectral

Extracts Fourier coefficients from a state. Allows specifying scaling compensation mode, rounding precision, and indexing.

```python
get_fourier_coefficients(
    state: Float[Array, "C ... N"],
    *,
    scaling_compensation_mode: (
        Literal[
            "norm_compensation",
            "reconstruction",
            "coef_extraction",
        ]
        | None
    ) = "coef_extraction",
    round: int | None = 5,
    indexing: str = "ij"
) -> Complex[Array, "C ... (N//2)+1"]

```

--------------------------------

### MultiChannelIC Evaluation

Source: https://fkoehler.site/exponax/api/utilities/initial_conditions/helper

Evaluates the combined multi-channel initial condition across grid points. This is called internally to get the initial state for the simulation.

```python
def __call__(self, x: Float[Array, "D ... N"]) -> Float[Array, "C ... N"]:
    """
    Evaluate the initial condition.

    **Arguments**:

    - `x`: The grid points.

    **Returns**:

    - `u`: The initial condition evaluated at the grid points.
    """
    return jnp.concatenate([ic(x) for ic in self.initial_conditions], axis=0)

```

--------------------------------

### FisherKPP Initialization

Source: https://fkoehler.site/exponax/api/stepper/reaction/fisher_kpp

Initializes the FisherKPP stepper for simulating the d-dimensional Fisher-KPP equation with periodic boundary conditions.

```APIDOC
## exponax.stepper.reaction.FisherKPP

### Description
Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) Fisher-KPP equation on periodic boundary conditions. This reaction-diffusion equation is related to logistic growth and describes the spread of a population.

In 1d, the Fisher-KPP equation is given by:

```
uₜ = ν uₓₓ + r u (1 - u)
```

with `ν` the diffusivity and `r` the reactivity. In 1d, the state `u` has only one channel. As such the discretized state is represented by a tensor of shape `(1, num_points)`. For higher dimensions, the number of channels will be constant 1, no matter the dimension. The higher-dimensional equation reads:

```
uₜ = ν Δu + r u (1 - u)
```

with `Δ` the Laplacian. The dynamics require initial conditions in the range `[0, 1]`. Then, the expected temporal behavior is a collective spread and growth. The limit of the solution is the constant state `1`.

### Method Signature
```python
def __init__(
    self,
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *, 
    diffusivity: float = 0.01,
    reactivity=1.0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    ...
```

### Arguments
- `num_spatial_dims` (int): The number of spatial dimensions `d`.
- `domain_extent` (float): The size of the domain `L`; in higher dimensions the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`.
- `num_points` (int): 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ᵈ`.
- `dt` (float): The timestep size `Δt` between two consecutive states.
- `diffusivity` (float, optional): The diffusivity `ν`. Defaults to `0.01`.
- `reactivity` (float, optional): The reactivity `r`. Defaults to `1.0`.
- `order` (int, optional): 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. Defaults to `2`.
- `dealiasing_fraction` (float, optional): 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. Defaults to `2/3`.
- `num_circle_points` (int, optional): How many points to use in the complex contour integral method to compute the coefficients of the exponential time differencing Runge Kutta method. Defaults to `16`.
- `circle_radius` (float, optional): The radius of the contour used to compute the coefficients of the exponential time differencing Runge Kutta method. Defaults to `1.0`.

### Notes
- The dynamics require initial conditions in the range `[0, 1]`. This can be achieved by combining any of the available IC generators with the [`exponax.ic.ClampingICGenerator`]. Alternatively, a good choice is also the [`exponax.ic.GaussianBlobs`]

### Good Values
- Use the `ClampingICGenerator` on `RandomTruncatedFourierSeries` with limits `[0, 1]` to generate initial conditions. Set `domain_extent = 1.0`, `num_points = 100`, `dt = 0.001`, and produce a trajectory of 500 steps. The final state of almost constant `1` will be reached after 200-400 steps.
```

--------------------------------

### Exponax Poisson Solver Initialization

Source: https://fkoehler.site/exponax/api/stepper/additional/poisson

Initializes the Exponax Poisson solver. This constructor sets up the necessary parameters to solve the Poisson equation, including spatial dimensions, domain extent, number of points, and the order of the equation.

```APIDOC
## __init__

### Description
Exactly solve the Poisson equation with periodic boundary conditions.

This "stepper" is different from all other steppers in this package in that it does not solve a time-dependent PDE. Instead, it solves the Poisson equation

$$ u_{xx} = - f $$

for a given right hand side $f$.

It is included for completion.

### Parameters

#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters:
- **num_spatial_dims** (int) - Required - The number of spatial dimensions.
- **domain_extent** (float) - Required - The extent of the domain.
- **num_points** (int) - Required - The number of points in each spatial dimension.
- **order** (int) - Optional - The order of the Poisson equation. Defaults to 2. You can also set `order=4` for the biharmonic equation.

### Request Example
None

### Response
None

### Response Example
None
```

--------------------------------

### BaseRandomICGenerator.__call__ Method

Source: https://fkoehler.site/exponax/api/utilities/initial_conditions/base_initial_condition

Generates a random initial condition on a grid. It uses gen_ic_fun to get the initial condition function and then evaluates it on a generated grid.

```python
def __call__(
    self,
    num_points: int,
    *, 
    key: PRNGKeyArray,
) -> Float[Array, "1 ... N"]:
    """
    Generate a random initial condition on a grid with `num_points` points.

    **Arguments**:

    - `num_points`: The number of grid points in each dimension.
    - `key`: A jax random key.

    **Returns**:

    - `u`: The initial condition evaluated at the grid points.
    """
    ic_fun = self.gen_ic_fun(key=key)
    grid = make_grid(
        self.num_spatial_dims,
        self.domain_extent,
        num_points,
        indexing=self.indexing,
    )
    return ic_fun(grid)
```

--------------------------------

### 2D Diffusion Solver

Source: https://fkoehler.site/exponax/examples/solver_showcase_2d

Simulates the diffusion equation ∂u/∂t = νΔu. Set up the domain, grid, diffusivity, and time step. This example uses isotropic diffusion.

```python
DOMAIN_EXTENT = 1.0
NUM_POINTS = 100
DT = 0.01
# See next section for anisotropic diffusion
NU = 0.01

anisotropic_diffusion_stepper = ex.stepper.Diffusion(
    2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU
)

grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)
u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(
    2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT
)

diffusion_trj = ex.rollout(anisotropic_diffusion_stepper, 40, include_init=True)(u_0)

diffusion_ani = ex.viz.animate_state_2d(diffusion_trj)

HTML(diffusion_ani.to_html5_video())
```

--------------------------------

### Initialize NavierStokesVelocity Stepper

Source: https://fkoehler.site/exponax/api/stepper/nonlinear/navier_stokes

Instantiate the `NavierStokesVelocity` stepper for 3D simulations. Requires spatial dimensions, domain extent, number of points, and timestep. Optional parameters control diffusivity, drag, and numerical method order.

```python
class NavierStokesVelocity(BaseStepper):
    diffusivity: float
    drag: float
    dealiasing_fraction: float

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        dt: float,
        *,
        diffusivity: float = 0.01,
        drag: float = 0.0,
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Timestepper for the 3d incompressible Navier-Stokes equations on
        periodic boundary conditions in velocity formulation. The equation
        reads

        ```
            uₜ = ν Δu + λu + 𝒫(u × ω)
        ```

        with `u` the three-channel velocity field and `ω = ∇ × u` the
        vorticity. The term `𝒫` denotes the Leray projection which enforces
        the incompressibility constraint `∇ ⋅ u = 0` by removing gradient
        components (pressure and kinetic energy gradient). The first term on
        the right-hand side is a diffusion with coefficient `ν` and the second
        term is an optional drag with coefficient `λ`.

        The nonlinear term uses the rotational form which exploits the vector
        identity `(u ⋅ ∇)u = ∇(|u|²/2) + ω × u` and the fact that the Leray
        projection annihilates all gradient fields.

        The Reynolds number of the problem is `Re = U L / ν` with `U` a
        characteristic velocity scale and `L` the `domain_extent`.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `d`. Must be
            `3`.
        - `domain_extent`: The size of the domain `L`; the domain is assumed
            to be a scaled hypercube `Ω = (0, L)³`.
        - `num_points`: The number of points `N` used to discretize the
            domain. This **includes** the left boundary point and **excludes**
            the right boundary point. The number of points in each dimension
            is the same. Hence, the total number of degrees of freedom is
            `N³`.
        - `dt`: The timestep size `Δt` between two consecutive states.
        - `diffusivity`: The diffusivity (viscosity) coefficient `ν`. This
            affects the Reynolds number. The lower the diffusivity, the "more
            turbulent". Default is `0.01`.
        - `drag`: The drag coefficient `λ`. Default is `0.0`.
        - `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.

        **Notes:**


```

--------------------------------

### Create full grid and wrap boundary conditions

Source: https://fkoehler.site/exponax/examples/simple_advection_example_2d

Generates a grid that includes boundary points and applies boundary conditions to the initial condition for visualization.

```python
full_grid = ex.make_grid(
    2,
    DOMAIN_EXTENT,
    NUM_POINTS,
    full=True,
)
full_ic = ex.wrap_bc(ic)
plt.contourf(
    full_grid[0],
    full_grid[1],
    full_ic[0],
    cmap="RdBu_r",
    levels=jnp.linspace(-1, 1, 21),
)
plt.xlim(0.0, DOMAIN_EXTENT)
plt.ylim(0.0, DOMAIN_EXTENT)
plt.colorbar()
plt.grid()
plt.title("Initial condition on full grid");
```

--------------------------------

### Simulate Single and Batched Initial Conditions

Source: https://fkoehler.site/exponax/examples/performance_hints

Performs a single simulation step for one initial condition and a batched step for multiple initial conditions using `jax.vmap`.

```python
one_u_1 = burgers_stepper(one_u_0)
# This will fail because the vanilla timestepper is single-batch only
# burgers_stepper(mutliple_u_0)
multiple_u_1 = jax.vmap(burgers_stepper)(multiple_u_0)

```

--------------------------------

### Initialize Burgers Stepper and Random Initial Condition

Source: https://fkoehler.site/exponax/examples/performance_hints

Sets up the Burgers equation stepper and generates a random initial condition in Fourier space. Ensure JAX is imported and configured.

```python
NUM_SPATIAL_DIMS = 1
DOMAIN_EXTENT = 3.0
NUM_POINTS = 100
DT = 0.1

burgers_stepper = ex.stepper.Burgers(NUM_SPATIAL_DIMS, DOMAIN_EXTENT, NUM_POINTS, DT)

u_0 = ex.ic.RandomTruncatedFourierSeries(
    NUM_SPATIAL_DIMS,
    cutoff=5,
    max_one=True,
)(NUM_POINTS, key=jax.random.PRNGKey(0))

ex.viz.plot_state_1d(u_0)
```

--------------------------------

### Burgers Equation Equivalence

Source: https://fkoehler.site/exponax/api/stepper/generic/normalized/normalized_convection

Under default settings, the NormalizedConvectionStepper behaves like the Burgers equation with specified parameters. This example shows how to instantiate a Burgers stepper for comparison.

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

--------------------------------

### 3D Burgers (Conservative) Simulation

Source: https://fkoehler.site/exponax/examples/solver_showcase_3d

Simulates the conservative Burgers equation in 3D. Similar setup to the non-conservative version but with `conservative=True`. Uses the same initial condition and visualization method.

```python
DOMAIN_EXTENT = 1.0
NUM_POINTS = 100
DT = 0.01
NU = 0.01

burgers_stepper_conservative = ex.stepper.Burgers(
    3, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU, conservative=True
)

grid = ex.make_grid(3, DOMAIN_EXTENT, NUM_POINTS)

# Burgers has two channels!
u_0 = jnp.concatenate(
    [
        jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)
        * jnp.cos(2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT)
        * jnp.sin(3 * 2 * jnp.pi * grid[2:3] / DOMAIN_EXTENT),
        jnp.cos(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)
        * jnp.sin(2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT)
        * jnp.cos(3 * 2 * jnp.pi * grid[2:3] / DOMAIN_EXTENT),
        jnp.sin(3 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)
        * jnp.cos(3 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT)
        * jnp.sin(4 * 2 * jnp.pi * grid[2:3] / DOMAIN_EXTENT),
    ]
)

burgers_trj_conservative = ex.rollout(
    burgers_stepper_conservative, 40, include_init=True
)(u_0)

burgers_ani_conservative = ex.viz.animate_state_3d_facet(
    burgers_trj_conservative,
    grid=(1, 3),
    figsize=(8, 3),
)

HTML(burgers_ani_conservative.to_html5_video())

```

--------------------------------

### Importing necessary libraries for Exponax

Source: https://fkoehler.site/exponax/examples/creating_your_own_solvers_1d

Import the core Exponax library and the BaseStepper class for creating custom solvers. Also import BaseNonlinearFun and ConvectionNonlinearFun for defining nonlinear terms.

```python
import exponax as ex
from exponax import BaseStepper
from exponax.nonlin_fun import BaseNonlinearFun, ConvectionNonlinearFun
```

--------------------------------

### Build Derivative Operator for Heat Equation

Source: https://fkoehler.site/exponax/api/etdrk_backbone

Example demonstrating how to construct a linear operator for the heat equation using spectral methods. This operator is used for initializing ETDRK methods.

```python
import jax.numpy as jnp
import exponax as ex

# Define the linear operator
N = 256
L = 5.0  # The domain size
D = 1  # Being in 1D

derivative_operator = 1j * ex.spectral.build_derivative_operator(
    D,
    L,
    N,
)

print(derivative_operator.shape)  # (1, (N//2)+1)

nu = 0.01 # The diffusion coefficient

linear_operator = nu * derivative_operator**2

```

--------------------------------

### Simulate Multiple Initial Conditions with Ensemble Steppers

Source: https://fkoehler.site/exponax/examples/performance_hints

Combines `jax.vmap` for multiple initial conditions and `eqx.filter_vmap` for an ensemble of steppers.

```python
ensembled_multiple_u_1 = eqx.filter_vmap(
    lambda stepper: jax.vmap(stepper)(multiple_u_0)
)(burgers_stepper_ensemble)

```

--------------------------------

### Wave Stepper Initialization

Source: https://fkoehler.site/exponax/api/stepper/linear/wave

Initializes the Wave stepper for the d-dimensional wave equation with periodic boundary conditions. It requires the number of spatial dimensions, domain extent, number of points, and timestep size. The speed of sound can also be specified.

```python
class Wave(BaseStepper):
    speed_of_sound: float
    wavenumber_norm: Float[Array, " 1 ... (N//2)+1"]

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        dt: float,
        *,
        speed_of_sound: float = 1.0,
    ):
        """
        Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) wave equation on
        periodic boundary conditions.

        In 1d, the wave equation is given by

        ```
            uₜₜ = c² uₓₓ
        ```

        with `c ∈ ℝ` being the speed of sound (or wave speed).

        In higher dimensions, the wave equation is written using the Laplacian

        ```
            uₜₜ = c² Δu
        ```

        Internally, the second-order equation is rewritten as a first-order
        system of two coupled fields — height `h` and velocity `v = hₜ`:

        ```
            hₜ = v
            vₜ = c² Δh
        ```

        As a result, the state has **two channels**: `u[0]` is the height field
        `h` and `u[1]` is the velocity field `v`.

        **Diagonalization:**

        The general solution of the wave equation is a superposition of
        right-traveling and left-traveling waves (d'Alembert's decomposition).
        This stepper exploits that structure: rather than time-stepping the
        coupled `(h, v)` system directly, it transforms into independent
        traveling-wave modes that each evolve as a simple phase rotation.

        In Fourier space, each wavenumber `k` gives a 2×2 ODE for `(ĥ, v̂)`
        that oscillates at frequency `ω = c|k|` — analogous to a harmonic
        oscillator trading potential and kinetic energy. Three steps
        diagonalize it:

        1. **Rescale** — `h` and `v` live on different scales (displacement vs.
           rate). Defining `w = iωĥ` puts them on equal footing. The coupled
           system becomes symmetric: `wₜ = iω v̂`, `v̂ₜ = iω w`.

        2. **Rotate** — Taking the sum and difference
           `pos = (w + v̂)/√2`, `neg = (w − v̂)/√2` decouples the system
           into two independent modes: `posₜ = +iω · pos` and
           `negₜ = −iω · neg`. Physically, `pos` is the right-traveling
           wave and `neg` the left-traveling wave.

        3. **Exponentiate** — Each decoupled mode evolves as a pure phase
           rotation: `pos(t+Δt) = exp(+iωΔt) · pos(t)`. This is what the
           ETDRK0 integrator computes exactly.

        After the exponential step, the inverse rotation and unscaling recover
        the updated `(h, v)`.

        At `k = 0` (the spatial mean), the frequency is zero and the two modes
        collapse — the system is no longer diagonalizable. There, the exact
        update is simply `h_mean += Δt · v_mean`, which is applied as a
        separate correction.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `d`.
        - `domain_extent`: The size of the domain `L`; in higher dimensions
            the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`.
        - `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ᵈ`.
        - `dt`: The timestep size `Δt` between two consecutive states.
        - `speed_of_sound` (keyword-only): The wave speed `c`. Default: `1.0`.

        **Notes:**

        
```

--------------------------------

### Kuramoto-Sivashinsky Equation Default Behavior

Source: https://fkoehler.site/exponax/api/stepper/generic/normalized/normalized_gradient_norm

Under default settings, the NormalizedGradientNormStepper behaves like the Kuramoto-Sivashinsky equation on a domain of size L=60.0. This example shows how to instantiate the KuramotoSivashinsky stepper for comparison.

```python
exponax.stepper.KuramotoSivashinsky(
    num_spatial_dims=D, domain_extent=60.0, num_points=N, dt=0.1,
    gradient_norm_scale=1.0, second_order_diffusivity=1.0,
    fourth_order_diffusivity=1.0,
)
```

--------------------------------

### Initialize Gray-Scott Stepper

Source: https://fkoehler.site/exponax/api/stepper/reaction/gray_scott

Instantiate the GrayScott stepper for simulating the reaction-diffusion system. Configure parameters like spatial dimensions, domain extent, discretization, timestep, diffusivities, feed rate, kill rate, and numerical integration order.

```python
class GrayScott(BaseStepper):
    diffusivity_1: float
    diffusivity_2: float
    feed_rate: float
    kill_rate: float
    dealiasing_fraction: float

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        dt: float,
        *,
        diffusivity_1: float = 2e-5,
        diffusivity_2: float = 1e-5,
        feed_rate: float = 0.04,
        kill_rate: float = 0.06,
        order: int = 2,
        # Needs lower value due to cubic nonlinearity
        dealiasing_fraction: float = 1 / 2,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) Gray-Scott reaction
        diffusion equation on periodic boundary conditions. This
        reaction-diffusion equation models the interaction of two chemical
        species u & v.

        In 1d, the Gray-Scott equation is given by

        ```
            uₜ = ν₁ uₓₓ + f(1 - u) - u v²

            vₜ = ν₂ vₓₓ - (f + k) v + u v²
        ```

        with `ν₁` and `ν₂` the diffusivities, `f` the feed rate, and `k` the
        kill rate. No matter the spatial dimension, this dynamics always has two
        channels, referring to the two chemical species. In higher dimensions,
        the equations read

        ```
            uₜ = ν₁ Δu + f(1 - u) - u v²

            vₜ = ν₂ Δv - (f + k) v + u v²
        ```

        with `Δ` the Laplacian.

        The Gray-Scott equation is known to produce a variety of patterns, such
        as spots, stripes, and spirals. The expected temporal behavior is highly
        dependent on the values of the feed and kill rates, see also this paper:
        https://www.ljll.fr/hecht/ftp/ff++/2015-cimpa-IIT/edp-tuto/Pearson.pdf

        IMPORTANT: Both channels are expected to have values in the range `[0,
        1]`.

        **Arguments**:

        - `num_spatial_dims`: The number of spatial dimensions `d`.
        - `domain_extent`: The size of the domain `L`; in higher dimensions
            the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`.
        - `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ᵈ`.
        - `dt`: The timestep size `Δt` between two consecutive states.
        - `diffusivity_1`: The diffusivity `ν₁` of the first species.
            Default is `2e-5`.
        - `diffusivity_2`: The diffusivity `ν₂` of the second species.
            Default is `1e-5`.
        - `feed_rate`: The feed rate `f`. Default is `0.04`.
        - `kill_rate`: The kill rate `k`. Default is `0.06`.
        - `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. Default: 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. 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.
        
```

--------------------------------

### Korteweg-de Vries Stepper Constructor

Source: https://fkoehler.site/exponax/api/stepper/nonlinear/kdv

Initializes the Korteweg-de Vries (KdV) equation stepper for d-dimensional problems with periodic boundary conditions. Allows configuration of convection, diffusion, dispersion, and hyper-diffusion parameters.

```python
__init__(
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    dt: float,
    *,
    convection_scale: float = -6.0,
    diffusivity: float = 0.0,
    dispersivity: float = 1.0,
    hyper_diffusivity: float = 0.01,
    advect_over_diffuse: bool = False,
    diffuse_over_diffuse: bool = False,
    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
)

```

--------------------------------

### BaseStepper Class

Source: https://fkoehler.site/exponax/api/stepper/base_stepper

The BaseStepper class is an abstract base class for time integrators. It handles the setup of linear and nonlinear operators and provides methods for stepping in time, both in real and Fourier space.

```APIDOC
## Class exponax.BaseStepper

This class is an abstract base class for time integrators.

### Methods

#### `_build_linear_operator(derivative_operator)`

Assemble the L operator in Fourier space.

- **Arguments**:
  - `derivative_operator` (Complex[Array, "D ... (N//2)+1"]): The derivative operator, shape `( D, ..., N//2+1 )`.

- **Returns**:
  - `L` (Complex[Array, "C ... (N//2)+1"]): The linear operator, shape `( C, ..., N//2+1 )`.

#### `_build_nonlinear_fun(derivative_operator)`

Build the function that evaluates nonlinearity in physical space, transforms to Fourier space, and evaluates derivatives there.

- **Arguments**:
  - `derivative_operator` (Complex[Array, "D ... (N//2)+1"]): The derivative operator, shape `( D, ..., N//2+1 )`.

- **Returns**:
  - `nonlinear_fun` (BaseNonlinearFun): A function that evaluates the nonlinearities in time space, transforms to Fourier space, and evaluates the derivatives there. Should be a subclass of `BaseNonlinearFun`.

#### `step(u)`

Perform one step of the time integration.

- **Arguments**:
  - `u` (Float[Array, "C ... N"]): The state vector, shape `(C, ..., N,)`.

- **Returns**:
  - `u_next` (Float[Array, "C ... N"]): The state vector after one step, shape `(C, ..., N,)`.

#### `step_fourier(u_hat)`

Perform one step of the time integration in Fourier space. Oftentimes, this is more efficient than `step` since it avoids back and forth transforms.

- **Arguments**:
  - `u_hat` (Complex[Array, "C ... (N//2)+1"]): The (real) Fourier transform of the state vector.

- **Returns**:
  - `u_next_hat` (Complex[Array, "C ... (N//2)+1"]): The (real) Fourier transform of the state vector after one step.

#### `__call__(u)`

Perform one step of the time integration for a single state.

- **Arguments**:
  - `u` (Float[Array, "C ... N"]): The state vector, shape `(C, ..., N,)`.

- **Returns**:
  - (Float[Array, "C ... N"]): The state vector after one step, shape `(C, ..., N,)`.
```

--------------------------------

### Initialize Multiple Diffusion Steppers with Varied Parameters

Source: https://fkoehler.site/exponax/examples/understanding_general_and_normalized_stepper

Initializes three additional 1D Diffusion steppers with different configurations of domain extent, time step, and diffusivity to observe their effects.

```python
diffusion_stepper_2 = ex.stepper.Diffusion(
    1,  # number of dimensions (set to 1)
    domain_extent=5.0,
    num_points=100,
    dt=0.1,
    diffusivity=1.0,
)
diffusion_stepper_3 = ex.stepper.Diffusion(
    1,  # number of dimensions (set to 1)
    domain_extent=2.5,
    num_points=100,
    dt=0.5,
    diffusivity=0.05,
)
diffusion_stepper_4 = ex.stepper.Diffusion(
    1,  # number of dimensions (set to 1)
    domain_extent=2.5,
    num_points=100,
    dt=0.125,
    diffusivity=0.2,
)
```
Tokens:
Raw