Skip to content

Linear

exponax.stepper.generic.NormalizedLinearStepper ¤

Bases: GeneralLinearStepper

Source code in exponax/stepper/generic/_linear.py
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class NormalizedLinearStepper(GeneralLinearStepper):
    normalized_linear_coefficients: tuple[float, ...]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_linear_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_linear_coefficients = normalized_linear_coefficients
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            domain_extent=1.0,
            num_points=num_points,
            dt=1.0,
            linear_coefficients=normalized_linear_coefficients,
        )
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_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/stepper/generic/_linear.py
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_linear_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_linear_coefficients = normalized_linear_coefficients
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        domain_extent=1.0,
        num_points=num_points,
        dt=1.0,
        linear_coefficients=normalized_linear_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
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
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)