astronomix.time_stepping.time_integration module

Time integration of the fluid equations.

This module wires together everything needed to advance a primitive state in time: it prepares the helper data and sharding, optionally compiles for memory analysis or runtime debugging, and then drives the per-step update through the generic loop driver (fixed-step / adaptive while / checkpointed). The snapshot machinery that records diagnostics along the way also lives here. For the available options see the simulation configuration and the simulation parameters.

class astronomix.time_stepping.time_integration.LoopState(primitive_state: Any, key: Any, forcing: Any = None)

Bases: NamedTuple

The physics evolving-state threaded through the generic time-loop driver.

The driver (astronomix.time_stepping._time_loop.integrate) treats this as an opaque pytree; only the closures here unpack it. Holding the PRNG key and the persistent Ornstein-Uhlenbeck forcing field here (rather than overloading a bare tuple slot) keeps the carry explicit and extensible.

primitive_state

The (padded) primitive fluid state being evolved.

Type:

Any

key

The PRNG key advanced by stochastic per-step modules (forcing, …).

Type:

Any

forcing

The persistent OU forcing field f (shape (3, nx, ny, nz)), or None when OU forcing is inactive.

Type:

Any

forcing: Any

Alias for field number 2

key: Any

Alias for field number 1

primitive_state: Any

Alias for field number 0

astronomix.time_stepping.time_integration.time_integration(primitive_state: Float[Array, 'num_vars num_cells_x'] | Float[Array, 'num_vars num_cells_x num_cells_y'] | Float[Array, 'num_vars num_cells_x num_cells_y num_cells_z'], config: SimulationConfig, params: SimulationParams, registered_variables: RegisteredVariables, snapshot_callable=None, sharding: None | NamedSharding = None, restart_state: None | LoopState = None) → Float[Array, 'num_vars num_cells_x'] | Float[Array, 'num_vars num_cells_x num_cells_y'] | Float[Array, 'num_vars num_cells_x num_cells_y num_cells_z'] | SnapshotData

Integrate the fluid equations in time. For the options of the time integration see the simulation configuration and the simulation parameters.

Parameters:

• primitive_state – The primitive state array.

• config – The simulation configuration.

• params – The simulation parameters.

• registered_variables – The registered variables.

• snapshot_callable –

  A callable which is called at certain time points if config.activate_snapshot_callback is True. The callable must have the signature

  callable(time: float, state: STATE_TYPE, registered_variables: RegisteredVariables) -> None

  and can be used to e.g. output the current state to disk or directly produce intermediate plots. Note that inside the callable, to pass data to memory, one must use

  jax.debug.callback(

  function, args…

  )

  To avoid moving large amounts of data to the host, only pass the necessary data to the function in the jax.debug.callback call, e.g. only the slice or summary statistics you need.

• sharding – The sharding to use for the padded helper data. If None, no sharding is applied.

• restart_state – An optional LoopState providing the PRNG key and persistent OU forcing field to resume from (typically obtained from astronomix.setup_helpers.restart_from_latest_checkpoint). Its primitive_state slot is ignored; pass the restored state as the primitive_state argument. Mainly used together with snapshot_storage_mode == TO_DISK and params.t_start.

Returns:

Depending on the configuration (return_snapshots, num_snapshots) either the final state of the fluid after the time integration of snapshots of the time evolution.