astronomix.initial_condition_generation.construct_primitive_state

astronomix.initial_condition_generation.construct_primitive_state#

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

Module Contents#

Functions#

construct_primitive_state

Stack the primitive variables into the state array, checking for NaNs.

API#

astronomix.initial_condition_generation.construct_primitive_state.construct_primitive_state(config: astronomix.option_classes.simulation_config.SimulationConfig, registered_variables: astronomix.variable_registry.registered_variables.RegisteredVariables, density: astronomix.option_classes.simulation_config.FIELD_TYPE, velocity_x: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, velocity_y: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, velocity_z: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, magnetic_field_x: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, magnetic_field_y: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, magnetic_field_z: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, interface_magnetic_field_x: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, interface_magnetic_field_y: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, interface_magnetic_field_z: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, gas_pressure: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, cosmic_ray_pressure: Union[astronomix.option_classes.simulation_config.FIELD_TYPE, types.NoneType] = None, sharding=None) → astronomix.option_classes.simulation_config.STATE_TYPE[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.

Contents