astronomix.initial_condition_generation.construct_primitive_state module#
Assemble the primitive state array from individual primitive fields.
Given the per-variable fields (density, velocities, magnetic field components,
pressures, …) this stacks them into the single state array used throughout
the solver, placing each field at the index dictated by registered_variables
for the active configuration (dimensionality, MHD, solver mode, equation of
state, cosmic rays).
- astronomix.initial_condition_generation.construct_primitive_state.construct_primitive_state(config: SimulationConfig, registered_variables: RegisteredVariables, density: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'], velocity_x: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, velocity_y: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, velocity_z: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, magnetic_field_x: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, magnetic_field_y: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, magnetic_field_z: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, interface_magnetic_field_x: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, interface_magnetic_field_y: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, interface_magnetic_field_z: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, gas_pressure: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, cosmic_ray_pressure: Float[Array, 'num_cells_x'] | Float[Array, 'num_cells_x num_cells_y'] | Float[Array, 'num_cells_x num_cells_y num_cells_z'] | None = None, sharding=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'][source]#
Stack the primitive variables into the state array, checking for NaNs.
Thin wrapper around the jitted
_assemble_primitive_state()that adds a one-time sanity check: a NaN in the freshly built primitive state almost always means a bad initial condition (a divide-by-zero, an out-of-range log, a mismatched field shape, …) that would otherwise only surface much later as an opaque solver blow-up. Catching it here points the user straight at their setup. The check reads back a single scalar, so it is cheap; it is skipped when the state is a tracer so differentiating or jitting through the setup still works.See
_assemble_primitive_state()for the argument semantics.- Returns:
The state array.
- Raises:
ValueError – If the assembled primitive state contains NaNs.