astronomix.time_stepping.time_integration module

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) → 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.

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.