jaxfluids.time_integration package

Submodules

jaxfluids.time_integration.RK2 module

class jaxfluids.time_integration.RK2.RungeKutta2(nh: int, inactive_axis: List)[source]

Bases: TimeIntegrator

2nd-order TVD RK2 scheme

integrate(cons: Array, rhs: Array, timestep: float, stage: int) → Array[source]

Wrapper function around integrate_conservatives. Adjusts the timestep according to current RK stage and calls integrate_conservatives. Implementation in child class.

Parameters:

  • cons (jnp.ndarray) – conservative variables buffer before integration

  • rhs (jnp.ndarray) – right-hand side buffer

  • timestep (float) – timestep to be integrated

Returns:

conservative variables buffer after integration

Return type:

DeviceArray

prepare_buffer_for_integration(cons: Array, init: Array, stage: int) → Array[source]

u_cons = 0.5 u^n + 0.5 u^*

jaxfluids.time_integration.RK3 module

class jaxfluids.time_integration.RK3.RungeKutta3(nh: int, inactive_axis: List)[source]

Bases: TimeIntegrator

3rd-order TVD RK3 scheme

integrate(cons: Array, rhs: Array, timestep: float, stage: int) → Array[source]

Wrapper function around integrate_conservatives. Adjusts the timestep according to current RK stage and calls integrate_conservatives. Implementation in child class.

Parameters:

  • cons (jnp.ndarray) – conservative variables buffer before integration

  • rhs (jnp.ndarray) – right-hand side buffer

  • timestep (float) – timestep to be integrated

Returns:

conservative variables buffer after integration

Return type:

DeviceArray

prepare_buffer_for_integration(cons: Array, init: Array, stage: int) → Array[source]

stage 1: u_cons = 3/4 u^n + 1/4 u^* stage 2: u_cons = 1/3 u^n + 2/3 u^**

jaxfluids.time_integration.euler module

class jaxfluids.time_integration.euler.Euler(nh: int, inactive_axis: List)[source]

Bases: TimeIntegrator

First-order explicit Euler time integration scheme

integrate(cons: Array, rhs: Array, timestep: float, stage: int) → Array[source]

Wrapper function around integrate_conservatives. Adjusts the timestep according to current RK stage and calls integrate_conservatives. Implementation in child class.

Parameters:

  • cons (jnp.ndarray) – conservative variables buffer before integration

  • rhs (jnp.ndarray) – right-hand side buffer

  • timestep (float) – timestep to be integrated

Returns:

conservative variables buffer after integration

Return type:

DeviceArray

jaxfluids.time_integration.time_integrator module

class jaxfluids.time_integration.time_integrator.TimeIntegrator(nh: int, inactive_axis: List)[source]

Bases: ABC

Abstract base class for explicit time integration schemes. All time intergration schemes are derived from TimeIntegrator.

abstract integrate(cons: Array, rhs: Array, timestep: float, stage: int) → Array[source]

Wrapper function around integrate_conservatives. Adjusts the timestep according to current RK stage and calls integrate_conservatives. Implementation in child class.

Parameters:

  • cons (jnp.ndarray) – conservative variables buffer before integration

  • rhs (jnp.ndarray) – right-hand side buffer

  • timestep (float) – timestep to be integrated

Returns:

conservative variables buffer after integration

Return type:

DeviceArray

integrate_conservatives(cons: Array, rhs: Array, timestep: float) → Array[source]

Integrates the conservative variables.

Parameters:

  • cons (jnp.ndarray) – conservative variables buffer before integration

  • rhs (jnp.ndarray) – right-hand side buffer

  • timestep (float) – timestep adjusted according to sub-stage in Runge-Kutta

Returns:

conservative variables buffer after integration

Return type:

DeviceArray

prepare_buffer_for_integration(cons: Array, init: Array, stage: int) → Array[source]

In multi-stage Runge-Kutta methods, prepares the buffer for integration. Implementation in child class.

Parameters:

  • cons (jnp.ndarray) – Buffer of conservative variables.

  • init (jnp.ndarray) – Initial conservative buffer.

  • stage (int) – Current stage of the RK time integrator.

Returns:

Sum of initial buffer and current buffer.

Return type:

jnp.ndarray

Module contents