astronomix.time_stepping.time_integration#
Module Contents#
Functions#
Integrate the fluid equations in time. For the options of the time integration see the simulation configuration and the simulation parameters. |
API#
- 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) Union[astronomix.option_classes.simulation_config.STATE_TYPE, astronomix.data_classes.simulation_snapshot_data.SnapshotData][source]#
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.
- 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.