astronomix.option_classes.simulation_config module

astronomix.option_classes.simulation_config module#

Static simulation configuration.

Defines SimulationConfig — the bundle of options that, unlike the simulation parameters, necessitate recompilation when changed — together with the integer-coded enumerations they reference (backends, solver/boundary/Riemann modes, positivity modes, …), the small geometry vector helpers, the sub-configs for gravity and positivity, and the finalize_config pass that fills in derived fields and validates the configuration.

class astronomix.option_classes.simulation_config.BoundarySettings(x: BoundarySettings1D = (0, 0), y: BoundarySettings1D = (0, 0), z: BoundarySettings1D = (0, 0))[source]#

Bases: NamedTuple

Per-axis boundary settings for the simulation.

x: BoundarySettings1D#: Alias for field number 0

y: BoundarySettings1D#: Alias for field number 1

z: BoundarySettings1D#: Alias for field number 2

class astronomix.option_classes.simulation_config.BoundarySettings1D(left_boundary: int = 0, right_boundary: int = 0)[source]#

Bases: NamedTuple

The boundary-condition type at the left and right end of a single axis.

left_boundary: int#: Alias for field number 0

right_boundary: int#: Alias for field number 1

class astronomix.option_classes.simulation_config.GravityConfig(self_gravity: bool = False, self_gravity_version: int = 2, external_potential: bool = False, poisson_manual_open_boundaries: bool = False, gravity: bool = False)[source]#

Bases: NamedTuple

Self-gravity and external-potential configuration.

external_potential: bool#: Enable an external, static gravitational potential provided via params.gravitational_potential. It is added to the self-gravity potential (if any) in _compute_total_potential.

gravity: bool#: Master gravity switch. Set automatically in finalize_config to self_gravity or external_potential; gates the gravity source-term machinery so an external potential works without self-gravity. Not set by the user directly.

poisson_manual_open_boundaries: bool#: Manual open boundary conditions in the Poisson solver.

self_gravity: bool#: Self-gravity switch (currently only for periodic / manual-open boundaries).

self_gravity_version: int#: Coupling of the self-gravity source to the hydrodynamics. One of SIMPLE_SOURCE / SECOND_ORDER_CONSERVATIVE / FOURTH_ORDER_CONSERVATIVE.

astronomix.option_classes.simulation_config.OPTIMAL_BACKEND = 2#

it is a request to pick the fastest available one at finalize_config time. It resolves to PALLAS on GPUs new enough to run the Triton kernels (compute capability >= 8.0) and falls back to NATIVE_JAX everywhere else (older GPUs, CPU, no nvidia-smi).

Type:: OPTIMAL_BACKEND is not a backend of its own

astronomix.option_classes.simulation_config.POSITIVITY_CONSERVATIVE = 3#

enforce internal-energy positivity by an antisymmetric face-flux diffusion that pulls internal energy into (near-)negative-pressure cells from their hotter neighbours (exact total-energy conservation), plus a density floor / vacuum-rest for voids and a minimal residual pressure floor as the unconditional guarantee. The smooth, conservative cousin of HARD_FLOOR: it keeps the energy-conserving self-gravity scheme stable on violent collapse without the 100%+ energy injection a bare floor causes.

Type:: CONSERVATIVE

class astronomix.option_classes.simulation_config.PositivityConfig(default_positivity_protection: bool = False, per_stage_mode: int = 0, per_step_mode: int = 0, clamp_in_estimates: bool = True, vacuum_rest: bool = False, nan_safe: bool = False, cons_coeff: float = 0.15, cons_passes: int = 16, cons_activate: float = 1.0, deepvoid_blend: bool = False, deepvoid_blend_factor: float = 8.0, preserving_flux: bool = False)[source]#

Bases: NamedTuple

Density/pressure positivity-enforcement configuration.

clamp_in_estimates: bool#: Read-only density/pressure clamp in the flux / eigenvalue / timestep estimates (NaN-safety; does NOT modify the evolved state). This is the role the old enforce_positivity bool played in those estimators. DECOUPLED from default_positivity_protection and ON by default – cheap insurance that never touches the conserved solution.

cons_activate: float#: Alias for field number 8

cons_coeff: float#: POSITIVITY_CONSERVATIVE-mode parameters (conservative internal-energy redistribution): per-axis diffusion coefficient (stability needs < 1/(2*dim)), number of Jacobi passes, and the activation margin in units of the internal-energy floor (keep ~1 – genuine near-violations only).

cons_passes: int#: Alias for field number 7

deepvoid_blend: bool#

blend the WENO interface flux toward LLF in cells near the density floor; the weight ramps from 1 at the floor to 0 at deepvoid_blend_factor * minimum_density.

Type:: Deep-void first-order flux blending (FOFC-style)

deepvoid_blend_factor: float#: Alias for field number 10

default_positivity_protection: bool#: Casual on/off switch for the per-stage / per-step STATE floors. Default False (no flooring). When True, finalize_config sets per_stage_mode and per_step_mode to HARD_FLOOR unless explicitly overridden. Does NOT affect the read-only clamp_in_estimates (always respected).

nan_safe: bool#

reset non-finite conserved entries to zero before the density/pressure floors so they become a valid floored state.

Type:: NaN/inf backstop

per_stage_mode: int#: Positivity enforcement applied inside every SSPRK/LSRK stage (on the conserved state — the CFL lever for strong shocks). One of POSITIVITY_{NONE,HARD_FLOOR,REDISTRIBUTE,CONSERVATIVE}. Default NONE; set to HARD_FLOOR by finalize when default_positivity_protection.

per_step_mode: int#: Positivity enforcement applied once per step before the evolve (on the primitive state). With turbulent forcing + vacuum_protection the conservative prot redistribution already runs once per step, so a per-step REDISTRIBUTE here is redundant and is auto-skipped.

preserving_flux: bool#

blend the WENO flux toward LLF by the largest weight keeping the LF-updated density AND pressure above their floors. Shares the unified flux-blending infrastructure with deepvoid_blend (different activation path; both may be on, the stronger blend wins). Forces the non-fused WENO+divergence path.

Type:: Positivity-preserving (Hu-Adams-Shu / Zalesak FCT) flux limiter

vacuum_rest: bool#

zero the momentum in below-floor (vacuum) cells so the recovered velocity is 0 rather than momentum/rho_floored (which spikes and drives high-Mach blow-up); lets minimum_density be lowered by orders of magnitude without instability.

Type:: Vacuum-rest velocity recovery

class astronomix.option_classes.simulation_config.SimulationConfig(backend: int = 2, pallas_block_shape: Tuple[int, int, int] = (4, 4, 8), pallas_use_triton: bool = True, pallas_interpret: bool = False, pallas_num_warps: int = 4, pallas_ct: bool = False, solver_mode: int = 1, numerical_precision: int = 0, runtime_debugging: bool = False, donate_state: bool = False, memory_analysis: bool = False, host_helper_data: bool = False, print_elapsed_time: bool = False, progress_bar: bool = False, dimensionality: int = 1, state_struct: bool = False, geometry: int = 0, random_seed: int = 42, equation_of_state: int = 0, mhd: bool = False, fv_magnetic_integrator: int = 0, positivity_config: PositivityConfig = (False, 0, 0, True, False, False, 0.15, 16, 1.0, False, 8.0, False), gravity_config: GravityConfig = (False, 2, False, False, False), diffusion: bool = False, viscosity_type: int = 1, thermal_conduction: bool = False, box_size: float | StaticFloatVector = 1.0, num_cells: int | StaticIntVector = 400, reconstruction_order: int = 1, limiter: int = 0, riemann_solver: int = 0, split: int = 0, time_integrator: int = 0, num_ghost_cells: int = 2, grid_spacing: float = 0.0025, boundary_handling: int = 0, boundary_settings: None | BoundarySettings1D | BoundarySettings = None, fixed_timestep: bool = False, exact_end_time: bool = True, source_term_aware_timestep: bool = False, num_timesteps: int = 1000, use_max_adaptive_timestep: bool = True, differentiation_mode: int = 0, num_checkpoints: int = 100, return_snapshots: bool = False, snapshot_settings: SnapshotSettings = (False, False, False, False, False, False, False, False, False, False, False, False, False, 100, 1e-10, 10000000000.0), snapshot_storage_mode: int = 0, snapshot_storage_path: str | None = None, activate_snapshot_callback: bool = False, use_specific_snapshot_timepoints: bool = False, num_snapshots: int = 10, first_order_fallback: bool = False, turbulent_forcing_config: TurbulentForcingConfig = (False, False, False), wind_config: WindConfig = (False, 10, 1, False), cosmic_ray_config: CosmicRayConfig = (False, False), cooling_config: CoolingConfig = (False, 2, (1, (None,))), frame_tracking: bool = False, neural_net_force_config: NeuralNetForceConfig = (False, None), cnn_mhd_corrector_config: CNNMHDconfig = (False, None))[source]#

Bases: NamedTuple

Configuration object for the simulation. The simulation configuration are parameters defining the simulation where changes necessitate recompilation.

activate_snapshot_callback: bool#: Call a user given function on the snapshot data, e.g. for saving or plotting. Must have signature callback(time, state, registered_variables).

backend: int#: Backend. Defaults to OPTIMAL_BACKEND, which finalize_config resolves to PALLAS on compute-capability >= 8.0 GPUs and NATIVE_JAX otherwise.

boundary_handling: int#: Explicit boundary handling mode.

boundary_settings: None | BoundarySettings1D | BoundarySettings#: Boundary settings for the simulation.

box_size: float | StaticFloatVector#: The size of the simulation box.

cnn_mhd_corrector_config: CNNMHDconfig#: Configuration of the CNN MHD corrector module.

cooling_config: CoolingConfig#: The configuration for the cooling module.

cosmic_ray_config: CosmicRayConfig#: Cosmic rays

differentiation_mode: int#: The differentiation mode one whats to use the solver in (forwards or backwards).

diffusion: bool#: Explicit diffusion term (currently only for finite difference mode)

dimensionality: int#: The number of dimensions of the simulation.

donate_state: bool#: Donate the state arrays to the time integration function to reduce memory allocations. If activated, the initial state arrays will be invalid after the simulation.

equation_of_state: int#: The equation of state for the simulation. NOTE: CURRENTLY ONLY IMPLEMENTED FOR FINITE DIFFERENCE MODE.

exact_end_time: bool#: Exactly reach the end time. In adaptive timestepping, one might otherwise overshoot.

first_order_fallback: bool#: Fallback to the first order Godunov scheme.

fixed_timestep: bool#: Enables a fixed timestep for the simulation based on the specified number of timesteps.

frame_tracking: bool#: Frame tracking in z-direction shifting the frame to follow a turbulent radiative mixing layer

fv_magnetic_integrator: int#: Integrator used for the magnetic part in the FV MHD scheme.

geometry: int#: The geometry of the simulation.

gravity_config: GravityConfig#: Self-gravity / external-potential configuration (see GravityConfig).

grid_spacing: float#: Grid spacing.

host_helper_data: bool#: Build the simulation helper data on the host (CPU) and only move the fields that are actually needed by the enabled subsystems onto the accelerator. Useful in production runs where a large meshgrid like geometric_centers is not required on device and the per-field memory footprint matters.

limiter: int#: The limiter for the reconstruction. Only affects finite volume mode.

memory_analysis: bool#: Memory analysis of the main time integration function

mhd: bool#: Magnetohydrodynamics switch.

neural_net_force_config: NeuralNetForceConfig#: Configuration of the neural network force module.

num_cells: int | StaticIntVector#: The number of cells in the simulation.

num_checkpoints: int#: The number of checkpoints used in the setup with backwards differetiability and adaptive time stepping.

num_ghost_cells: int#: The number of ghost cells.

num_snapshots: int#: The number of snapshots to return.

num_timesteps: int#: The number of timesteps for the fixed timestep mode.

numerical_precision: int#: Precision mode.

pallas_block_shape: Tuple[int, int, int]#: Alias for field number 1

pallas_ct: bool#: Toggle for the Pallas constrained-transport helpers (update_cell_center_fields, constrained_transport_rhs). Disabled by default: the staged Pallas-CT pipeline gives a clear memory win at small grids (~65% temp at N=16 on alfven_wave3D) but only marginal savings at production scale (~2% temp at N=64) while adding ~25s of one-time compile cost. Flip to True if the small-N memory profile matters; the rest of the Pallas backend stays on regardless.

pallas_interpret: bool#: Alias for field number 3

pallas_num_warps: int#: Alias for field number 4

pallas_use_triton: bool#: Alias for field number 2

positivity_config: PositivityConfig#: Density/pressure positivity-enforcement configuration (see PositivityConfig).

print_elapsed_time: bool#: Print the elapsed time of the simulation

progress_bar: bool#: Activate progress bar

random_seed: int#: The random seed for any stochastic processes in the simulation, e.g. turbulent forcing.

reconstruction_order: int#: The reconstruction order is the number of cells on each side of the cell of interest used to calculate the gradients for the reconstruction at the interfaces.

return_snapshots: bool#: Return intermediate snapshots of the time evolution instead of only the final fluid state.

riemann_solver: int#: The Riemann solver used Only for finite volume mode.

runtime_debugging: bool#: Debug runtime errors, throws exceptions on e.g. negative pressure or density. Significantly reduces performance.

snapshot_settings: SnapshotSettings#: Snapshot settings

snapshot_storage_mode: int#: Where the snapshots are stored. ON_DEVICE (default) keeps the snapshot diagnostics in preallocated device buffers and returns them at the end (the classic behaviour). TO_DISK instead streams each snapshot to disk via Orbax: the run is split into segments between the snapshot times, and the loop carry (primitive state, PRNG key, OU forcing field) plus the time is written to snapshot_storage_path after each segment. Each device writes its own shard, so this scales to multiple devices / nodes. TO_DISK is forward-mode only.

snapshot_storage_path: str | None#: Directory the Orbax checkpoints are written to / read from when snapshot_storage_mode == TO_DISK. Required in that mode.

solver_mode: int#: Basic solver mode, either finite volume or finite difference. Defaults to the finite-difference HOW-MHD scheme (Jeongbhin Seo, Dongsu Ryu, 2023), which is the recommended solver.

source_term_aware_timestep: bool#: Adds the sources with the current timestep to a hypothetical state to estimate the actual timestep. Useful for time-dependent sources, but additional computational overhead.

split: int#: Dimensional splitting / unsplit mode. Note that the UNSPLIT scheme currently interferes with energy conservation in settings with self-gravity.

state_struct: bool#: Use a struct for the state.

thermal_conduction: bool#: Explicit thermal conduction term div(kappa grad T) in the energy equation (constant conductivity params.thermal_conductivity, explicit integration). Currently only for finite difference mode.

time_integrator: int#: Time integration method.

turbulent_forcing_config: TurbulentForcingConfig#: Turbulent forcing configuration.

use_max_adaptive_timestep: bool#: Use a maximum timestep in adaptive timestep mode.

use_specific_snapshot_timepoints: bool#: Return snapshots at specific time points.

viscosity_type: int#: Viscosity type - either kinematic or dynamic viscosity.

wind_config: WindConfig#: The configuration for the stellar wind module.

class astronomix.option_classes.simulation_config.SnapshotSettings(return_states: bool = False, return_final_state: bool = False, return_total_mass: bool = False, return_total_energy: bool = False, return_internal_energy: bool = False, return_kinetic_energy: bool = False, return_gravitational_energy: bool = False, return_radial_momentum: bool = False, return_kinetic_energy_spectrum: bool = False, return_magnetic_energy_spectrum: bool = False, return_helicity_spectrum: bool = False, return_magnetic_divergence: bool = False, return_temperature_pdf: bool = False, num_temperature_bins: int = 100, temperature_pdf_min: float = 1e-10, temperature_pdf_max: float = 10000000000.0)[source]#

Bases: NamedTuple

Settings for the snapshot output of the simulation.

num_temperature_bins: int#: Alias for field number 13

return_final_state: bool#: Whether to return the final state of the simulation.

return_gravitational_energy: bool#: Whether to return gravitational energy

return_helicity_spectrum: bool#: Whether to return the helicity spectrum

return_internal_energy: bool#: Whether to return internal energy

return_kinetic_energy: bool#: Whether to return kinetic energy

return_kinetic_energy_spectrum: bool#: Whether to return the kinetic energy spectrum

return_magnetic_divergence: bool#: Whether to return the magnetic field divergence NOTE: currently only implemented for finite difference MHD

return_magnetic_energy_spectrum: bool#: Whether to return the magnetic energy spectrum

return_radial_momentum: bool#: Whether to return radial momentum

return_states: bool#: Whether to record the full primitive state at every checkpoint. This is the single biggest snapshot allocation (num_snapshots × num_vars × num_cells^d); it is opt-in. Set to True if you actually need the per-snapshot states; for the common case of only wanting a final state plus integrated diagnostics (energies, total mass, runtime, num_iterations), the default False skips the per-snapshot state allocation entirely.

return_temperature_pdf: bool#: Whether to return the temperature PDF (dV/dlogT)

return_total_energy: bool#: Whether to return the total energy at the times the snapshots were taken.

return_total_mass: bool#: Whether to return the total mass at the times the snapshots were taken.

temperature_pdf_max: float#: Alias for field number 15

temperature_pdf_min: float#: Alias for field number 14

class astronomix.option_classes.simulation_config.StaticFloatVector(x: float = -1.0, y: float = -1.0, z: float = -1.0)[source]#

Bases: NamedTuple

A static (compile-time) per-axis float triple (e.g. box size per axis).

x: float#: Alias for field number 0

y: float#: Alias for field number 1

z: float#: Alias for field number 2

class astronomix.option_classes.simulation_config.StaticIntVector(x: int = -1, y: int = -1, z: int = -1)[source]#

Bases: NamedTuple

A static (compile-time) per-axis integer triple (e.g. cells per axis).

x: int#: Alias for field number 0

y: int#: Alias for field number 1

z: int#: Alias for field number 2

astronomix.option_classes.simulation_config.config_to_string(config: SimulationConfig) → str[source]#: Return a compact one-line description of the solver configuration.

astronomix.option_classes.simulation_config.finalize_config(config: SimulationConfig, state_shape) → SimulationConfig[source]#

Fill in derived configuration fields and validate the configuration.

Resolves the values that depend on the actual state shape or on cross-field consistency: the positivity-protection defaults, the number of cells per axis, the grid spacing, the geometry- and solver-specific overrides, the master gravity switch, the boundary defaults, and the disk-snapshot requirements.

Parameters:

config – The user-supplied simulation configuration.
state_shape – The shape of the (unpadded) primitive state array, used to derive num_cells per axis.

Returns:

The finalized simulation configuration.

astronomix.option_classes.simulation_config.gpu_compute_capability_at_least_80() → bool[source]#

Return whether every visible NVIDIA GPU has compute capability >= 8.0.

Compute capability 8.0 (Ampere) is the floor for the Triton kernels the Pallas backend compiles to, so this is the predicate that decides whether OPTIMAL_BACKEND resolves to PALLAS. Any failure to query the GPUs — no nvidia-smi on the PATH (e.g. a CPU-only host) or the call erroring out — is treated as “not capable” so the safe NATIVE_JAX fallback is chosen.

Returns:: True if all visible NVIDIA GPUs report compute capability >= 8.0, False otherwise (including when no GPU could be queried).

astronomix.option_classes.simulation_config.limiter_to_string(limiter: int) → str[source]#: Return the human-readable name of a slope-limiter constant.

astronomix.option_classes.simulation_config.riemann_solver_to_string(riemann_solver: int) → str[source]#: Return the human-readable name of a Riemann-solver constant.

astronomix.option_classes.simulation_config.solver_mode_to_string(solver_mode: int) → str[source]#: Return the short label ("FV" / "FD") of a solver-mode constant.

astronomix.option_classes.simulation_config module

Contents

astronomix.option_classes.simulation_config module#