astronomix.data_classes.simulation_helper_data

Geometric helper data for the simulation.

Bundles the precomputed per-cell geometry (cell centers, radii, volumes, cell boundaries) used throughout the solver and the snapshot diagnostics, together with the machinery to build only the fields the active subsystems actually need (see HelperDataRequirements), optionally on the host, and to pad / shard / unpad them for the time integrator.

Module Contents

Classes

HelperData	Helper data used throughout the simulation.
HelperDataRequirements	Per-field requirements for HelperData.

Functions

get_helper_data

Generate the helper data for the simulation.

API

class astronomix.data_classes.simulation_helper_data.HelperData

Bases: typing.NamedTuple

Helper data used throughout the simulation.

geometric_centers: jax.numpy.ndarray = None

volumetric_centers: jax.numpy.ndarray = None

cell_centers_x: jax.numpy.ndarray = None

cell_centers_y: jax.numpy.ndarray = None

cell_centers_z: jax.numpy.ndarray = None

r: jax.numpy.ndarray = None

r_hat_alpha: jax.numpy.ndarray = None

cell_volumes: jax.numpy.ndarray = None

inner_cell_boundaries: jax.numpy.ndarray = None

outer_cell_boundaries: jax.numpy.ndarray = None

class astronomix.data_classes.simulation_helper_data.HelperDataRequirements

Bases: typing.NamedTuple

Per-field requirements for HelperData.

Each flag indicates whether the matching HelperData field is consumed by any active subsystem (solver, physics module, snapshot calculator). Fields whose flag is False are left unmaterialised and stay at the HelperData default of None.

needs_geometric_centers: bool = False

needs_volumetric_centers: bool = False

needs_cell_centers_x: bool = False

needs_cell_centers_y: bool = False

needs_cell_centers_z: bool = False

needs_r: bool = False

needs_r_hat_alpha: bool = False

needs_cell_volumes: bool = False

needs_inner_cell_boundaries: bool = False

needs_outer_cell_boundaries: bool = False

astronomix.data_classes.simulation_helper_data.get_helper_data(config: astronomix.option_classes.simulation_config.SimulationConfig, sharding: Union[types.NoneType, jax.NamedSharding] = None, padded: bool = False, requirements: Union[types.NoneType, astronomix.data_classes.simulation_helper_data.HelperDataRequirements] = None) → astronomix.data_classes.simulation_helper_data.HelperData

Generate the helper data for the simulation.

Only the fields requested by requirements are materialised; the rest stay at the HelperData default of None. When requirements is None (the default), all fields are built so that scripts using get_helper_data to construct initial conditions keep working. The time integrator passes a config-derived HelperDataRequirements explicitly so the in-simulation helper data is minimal.

With config.host_helper_data=True the arrays are built inside a CPU device context and only the requested fields are moved onto the target device (respecting sharding for 3D Cartesian geometric_centers). Fields needed only as intermediates (e.g. the full X, Y, Z meshgrid when only r is requested) never materialise on the accelerator.