astronomix.time_stepping.time_integration

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.

Module Contents

Classes

LoopState

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

Functions

time_integration

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

API

class astronomix.time_stepping.time_integration.LoopState

Bases: typing.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.

Attributes:

primitive_state: The (padded) primitive fluid state being evolved. key: The PRNG key advanced by stochastic per-step modules (forcing, …). forcing: The persistent OU forcing field f (shape (3, nx, ny, nz)),

or None when OU forcing is inactive.

primitive_state: Any = None

key: Any = None

forcing: Any = None

astronomix.time_stepping.time_integration.time_integration(primitive_state: astronomix.option_classes.simulation_config.STATE_TYPE, config: astronomix.option_classes.simulation_config.SimulationConfig, params: astronomix.option_classes.simulation_params.SimulationParams, registered_variables: astronomix.variable_registry.registered_variables.RegisteredVariables, snapshot_callable=None, sharding: Union[types.NoneType, jax.NamedSharding] = None, restart_state: Union[types.NoneType, astronomix.time_stepping.time_integration.LoopState] = None) → Union[astronomix.option_classes.simulation_config.STATE_TYPE, astronomix.data_classes.simulation_snapshot_data.SnapshotData]

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

Args:

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.