From a0800d94667817064fdd2de1400a4532e946446e Mon Sep 17 00:00:00 2001 From: Lorenzo Branca Date: Tue, 14 Jul 2026 14:59:55 +0200 Subject: [PATCH 1/3] WIP: full carbox + KROME-derived chemistry/thermochemistry coupling Complete working implementation (species advection + per-cell carbox reaction + KROME-grade thermochemistry). GPL-encumbered (KROME-derived cooling + carbox dependency), so NOT suitable for MIT astronomix main -- preserved here; the MIT-clean mechanism is on feature/chemistry-species-hook. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01UhwF9zJfmnzuLeABHu75QW --- astronomix/_modules/_chemistry/__init__.py | 8 + astronomix/_modules/_chemistry/_chemistry.py | 418 +++++++++ .../_modules/_chemistry/_thermochemistry.py | 868 ++++++++++++++++++ .../_modules/_chemistry/chemistry_options.py | 157 ++++ .../_modules/_iteration_level_updates.py | 12 + .../option_classes/simulation_config.py | 4 + .../option_classes/simulation_params.py | 4 + astronomix/setup_helpers/__init__.py | 2 + astronomix/setup_helpers/chemistry_setup.py | 211 +++++ .../variable_registry/registered_variables.py | 28 + 10 files changed, 1712 insertions(+) create mode 100644 astronomix/_modules/_chemistry/__init__.py create mode 100644 astronomix/_modules/_chemistry/_chemistry.py create mode 100644 astronomix/_modules/_chemistry/_thermochemistry.py create mode 100644 astronomix/_modules/_chemistry/chemistry_options.py create mode 100644 astronomix/setup_helpers/chemistry_setup.py diff --git a/astronomix/_modules/_chemistry/__init__.py b/astronomix/_modules/_chemistry/__init__.py new file mode 100644 index 00000000..b6c517be --- /dev/null +++ b/astronomix/_modules/_chemistry/__init__.py @@ -0,0 +1,8 @@ +""" +Astrochemical reaction-network coupling. + +Advects chemical species as scalar fields on the fluid grid and, once per hydro +step, reacts them per cell using a carbox reaction network integrated with a +stiff Diffrax solver. See ``_chemistry.py`` for the driver and +``chemistry_options.py`` for the configuration / parameter containers. +""" diff --git a/astronomix/_modules/_chemistry/_chemistry.py b/astronomix/_modules/_chemistry/_chemistry.py new file mode 100644 index 00000000..44ac5c0f --- /dev/null +++ b/astronomix/_modules/_chemistry/_chemistry.py @@ -0,0 +1,418 @@ +""" +Per-cell astrochemistry as an operator-split source term. + +Once per hydro step this module reacts the chemical species carried in the state +array. The species advect with the flow as passive volumetric densities (handled +by the finite-volume solver); here we freeze the fluid and, in every cell, +integrate a carbox reaction network over the hydro time step with a stiff Diffrax +solver. Concretely the driver: + + 1. reads the species block and the local temperature from the primitive state, + 2. crosses the unit boundary from astronomix code units into the CGS units the + network expects (number densities in cm^-3, temperature in Kelvin, time in + seconds), + 3. integrates the network per cell (vmapped over the grid), and + 4. writes the updated species back into the state array. + +The carbox network is reassembled from the static structure held in the +configuration and the dynamic array leaves held in the parameters, mirroring how +the neural-network cooling curve is embedded (see ``_modules/_cooling``). + +When ``chemistry_config.thermochemistry`` is enabled the temperature is evolved +*together* with the abundances (an extra ODE variable per cell), and the evolved +temperature is converted back into the pressure field — so chemistry and hydro +exchange energy. When it is disabled the species react at a fixed temperature and +the energy field is left untouched. The finite-difference source-term path is a +deliberate follow-up. + +WARNING: The network integration uses Diffrax, which is an optional dependency. +It is imported lazily inside the driver so that ``import astronomix`` does not +require Diffrax (or carbox) when the chemistry module is inactive. +""" + +# general +from functools import partial + +# jax +import jax +import jax.numpy as jnp + +# astronomix constants +from astronomix._modules._chemistry.chemistry_options import KVAERNO5 +from astronomix.option_classes.simulation_config import STATE_TYPE + +# astronomix containers +from astronomix._modules._chemistry.chemistry_options import ChemistryConfig +from astronomix.option_classes.simulation_params import SimulationParams +from astronomix.variable_registry.registered_variables import RegisteredVariables + +# astronomix functions +from astronomix._modules._chemistry._thermochemistry import temperature_derivative +from astronomix._modules._cooling._cooling import ( + get_pressure_from_temperature, + get_temperature_from_pressure, +) + + +def _advance_single_cell( + cell_abundances, + cell_temperature_kelvin, + time_step_seconds, + reaction_network, + rate_modifier_a, + rate_modifier_b, + cosmic_ray_rate, + fuv_field, + visual_extinction, + absolute_tolerance, + relative_tolerance, + solver, + max_steps, + thermochemistry, + adiabatic_index, + dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + atomic_oxygen_index, + ionized_hydrogen_index, + helium_index, + ionized_carbon_index, + co_cooling, + carbon_monoxide_index, + co_cooling_table, + co_cooling_bounds, +): + """Advance one cell's chemistry over ``time_step_seconds`` with a stiff solve. + + This is the 0-D box integration carbox performs, reduced to a single short + step so it can serve as an operator-split source term (no snapshot grid, just + integrate from zero to the hydro time step and read the end state). It is + written for a single cell and vmapped over the grid by the driver. + + With ``thermochemistry`` enabled the integrated state is the abundances with + the temperature (Kelvin) appended as a final element, so the temperature + evolves under net heating/cooling while the reaction rates see the changing + temperature. Otherwise only the abundances are integrated, at fixed + temperature. + + Args: + cell_abundances: Absolute species number densities [cm^-3] for this cell. + cell_temperature_kelvin: Gas temperature in Kelvin for this cell. + time_step_seconds: Chemistry sub-step length in seconds (the hydro step). + reaction_network: The reassembled carbox ``JNetwork`` right-hand side. + rate_modifier_a: Per-reaction multiplicative rate modifier. + rate_modifier_b: Per-reaction additive rate modifier. + cosmic_ray_rate: Cosmic-ray ionization rate [s^-1]. + fuv_field: FUV radiation field strength (Draine units). + visual_extinction: Visual extinction Av [mag]. + absolute_tolerance: Absolute tolerance for the stiff solver. + relative_tolerance: Relative tolerance for the stiff solver. + solver: The Diffrax solver instance to integrate with. + max_steps: Maximum internal Diffrax steps. + thermochemistry: When True, evolve the temperature alongside the species. + adiabatic_index: The adiabatic index gamma (for dT/dt). + dust_to_gas_ratio: The dust-to-gas mass ratio (for photoelectric heating). + hydrogen_molecule_formation_rate_coefficient: H + H -> H2 grain formation + rate coefficient [cm^3 s^-1] (for H2 formation heating). + hydrogen_index: Index of atomic H in the abundances (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + electron_index: Index of the electron (or -1). + atomic_oxygen_index: Index of atomic O (or -1). + ionized_hydrogen_index: Index of H+ (or -1). + helium_index: Index of He (or -1). + ionized_carbon_index: Index of C+ (or -1). + co_cooling: Whether to add tabulated CO rotational cooling. + carbon_monoxide_index: Index of CO (or -1). + co_cooling_table: The 3D CO cooling table. + co_cooling_bounds: The CO table grid limits. + + Returns: + The end-of-step cell state: the updated abundances [cm^-3] alone, or the + abundances with the updated temperature [K] appended when + ``thermochemistry`` is enabled. + """ + + # Diffrax is an optional dependency; keep the import local so the package + # imports without it when chemistry is inactive. + import diffrax as dx + + # The chemistry right-hand side; the physical parameters shared across cells + # are closed over rather than threaded through Diffrax ``args``. + def chemistry_derivative(time, abundances, temperature_kelvin): + return reaction_network( + time, + abundances, + temperature_kelvin, + cosmic_ray_rate, + fuv_field, + visual_extinction, + rate_modifier_a, + rate_modifier_b, + ) + + if thermochemistry: + # State = [abundances, temperature]; the temperature evolves under net + # heating/cooling and feeds back into the reaction rates. + def vector_field(time, state, args): + abundances = state[:-1] + temperature_kelvin = state[-1] + abundance_derivative = chemistry_derivative( + time, + abundances, + temperature_kelvin, + ) + temperature_rate = temperature_derivative( + abundances, + temperature_kelvin, + adiabatic_index, + cosmic_ray_rate, + fuv_field, + dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + atomic_oxygen_index, + ionized_hydrogen_index, + helium_index, + ionized_carbon_index, + co_cooling, + carbon_monoxide_index, + co_cooling_table, + co_cooling_bounds, + ) + return jnp.concatenate( + [abundance_derivative, jnp.reshape(temperature_rate, (1,))] + ) + + initial_state = jnp.concatenate( + [cell_abundances, jnp.reshape(cell_temperature_kelvin, (1,))] + ) + else: + # Temperature is a fixed parameter of the reaction rates for this step. + def vector_field(time, abundances, args): + return chemistry_derivative(time, abundances, cell_temperature_kelvin) + + initial_state = cell_abundances + + ode_term = dx.ODETerm(vector_field) + + solution = dx.diffeqsolve( + ode_term, + solver, + t0=0.0, + t1=time_step_seconds, + dt0=jnp.minimum(1e-6, time_step_seconds * 1e-3), + y0=initial_state, + args=None, + stepsize_controller=dx.PIDController( + rtol=relative_tolerance, + atol=absolute_tolerance, + ), + saveat=dx.SaveAt(t1=True), + max_steps=max_steps, + ) + + # Only the end state matters for an operator-split sub-step. + return solution.ys[-1] + + +@partial(jax.jit, static_argnames=("chemistry_config", "registered_variables")) +def update_chemistry( + primitive_state: STATE_TYPE, + registered_variables: RegisteredVariables, + chemistry_config: ChemistryConfig, + simulation_params: SimulationParams, + time_step: float, +) -> STATE_TYPE: + """React the chemical species of the primitive state for one time step. + + Reads the species block and the local temperature, integrates the reaction + network in every cell over ``time_step`` and writes the updated species + densities back. When ``chemistry_config.thermochemistry`` is enabled the + temperature is evolved alongside the species and the resulting pressure is + written back into the energy field; otherwise the energy field is untouched. + + Args: + primitive_state: The primitive state array. + registered_variables: The registered variables (locate the species block). + chemistry_config: The static chemistry configuration (network structure, + species count, solver, thermochemistry flag, species-role indices). + simulation_params: The simulation parameters (network leaves, unit + factors, physical parameters, tolerances, gamma). + time_step: The hydro time step, in code units. + + Returns: + The primitive state with the species block advanced by one reaction step + (and, with thermochemistry, the pressure updated by heating/cooling). + """ + + chemistry_params = simulation_params.chemistry_params + + # ------------------------------------------------------------- + # =============== ↓ Pick the reaction network and solver ↓ ==== + # ------------------------------------------------------------- + + # The network rides in the (dynamic) parameters rather than the static + # configuration, because a carbox ``JNetwork`` holds a Python list of + # reactions and so is not hashable. + reaction_network = chemistry_params.network + + if chemistry_config.solver == KVAERNO5: + import diffrax as dx + + solver = dx.Kvaerno5() + else: + raise ValueError( + f"Unsupported chemistry solver tag: {chemistry_config.solver}" + ) + + # ------------------------------------------------------------- + # =============== ↑ Pick the reaction network and solver ↑ ==== + # ------------------------------------------------------------- + + # ------------------------------------------------------------- + # =============== ↓ Cross into CGS units ↓ ==================== + # ------------------------------------------------------------- + + density = primitive_state[registered_variables.density_index] + pressure = primitive_state[registered_variables.pressure_index] + + # The pressure-to-temperature relation returns the rescaled code temperature; + # scale it into Kelvin for the rate expressions. + code_temperature = get_temperature_from_pressure( + density, + pressure, + chemistry_params.hydrogen_mass_fraction, + chemistry_params.metal_mass_fraction, + ) + temperature_kelvin = code_temperature * chemistry_params.temperature_unit_kelvin + + time_step_seconds = time_step * chemistry_params.time_unit_seconds + + # The species block holds code-unit number densities; the network works in + # absolute number densities [cm^-3]. + species_start = registered_variables.chemistry_species_index + number_of_species = registered_variables.num_chemical_species + species_code_units = primitive_state[ + species_start : species_start + number_of_species + ] + species_number_densities = ( + species_code_units * chemistry_params.number_density_unit_cgs + ) + + # ------------------------------------------------------------- + # =============== ↑ Cross into CGS units ↑ ==================== + # ------------------------------------------------------------- + + # ------------------------------------------------------------- + # =============== ↓ React every cell (vmapped) ↓ ============== + # ------------------------------------------------------------- + + # Flatten the spatial grid so each cell is one row: the species array goes + # from ``(num_species, *grid)`` to ``(num_cells, num_species)`` and the + # temperature from ``(*grid,)`` to ``(num_cells,)``. + grid_shape = temperature_kelvin.shape + number_of_cells = temperature_kelvin.size + + species_per_cell = species_number_densities.reshape( + number_of_species, number_of_cells + ).T + temperature_per_cell = temperature_kelvin.reshape(number_of_cells) + + # Bind everything that is shared across cells; vmap only over the per-cell + # abundances and temperature. + react_one_cell = partial( + _advance_single_cell, + time_step_seconds=time_step_seconds, + reaction_network=reaction_network, + rate_modifier_a=chemistry_params.rate_modifier_a, + rate_modifier_b=chemistry_params.rate_modifier_b, + cosmic_ray_rate=chemistry_params.cosmic_ray_rate, + fuv_field=chemistry_params.fuv_field, + visual_extinction=chemistry_params.visual_extinction, + absolute_tolerance=chemistry_params.atol, + relative_tolerance=chemistry_params.rtol, + solver=solver, + max_steps=chemistry_config.max_steps, + thermochemistry=chemistry_config.thermochemistry, + adiabatic_index=simulation_params.gamma, + dust_to_gas_ratio=chemistry_params.dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient=( + chemistry_params.hydrogen_molecule_formation_rate_coefficient + ), + hydrogen_index=chemistry_config.hydrogen_index, + molecular_hydrogen_index=chemistry_config.molecular_hydrogen_index, + electron_index=chemistry_config.electron_index, + atomic_oxygen_index=chemistry_config.atomic_oxygen_index, + ionized_hydrogen_index=chemistry_config.ionized_hydrogen_index, + helium_index=chemistry_config.helium_index, + ionized_carbon_index=chemistry_config.ionized_carbon_index, + co_cooling=chemistry_config.co_cooling, + carbon_monoxide_index=chemistry_config.carbon_monoxide_index, + co_cooling_table=chemistry_params.co_cooling_table, + co_cooling_bounds=chemistry_params.co_cooling_bounds, + ) + reacted_per_cell = jax.vmap(react_one_cell)( + species_per_cell, + temperature_per_cell, + ) + + # ------------------------------------------------------------- + # =============== ↑ React every cell (vmapped) ↑ ============== + # ------------------------------------------------------------- + + # ------------------------------------------------------------- + # =============== ↓ Cross back and write the state ↓ ========== + # ------------------------------------------------------------- + + # The first ``number_of_species`` columns are the reacted abundances; with + # thermochemistry a final column carries the evolved temperature (Kelvin). + reacted_species_per_cell = reacted_per_cell[:, :number_of_species] + + # Undo the flatten and the CGS scaling, then write the species block back. + reacted_species_number_densities = reacted_species_per_cell.T.reshape( + number_of_species, *grid_shape + ) + reacted_species_code_units = ( + reacted_species_number_densities / chemistry_params.number_density_unit_cgs + ) + primitive_state = primitive_state.at[ + species_start : species_start + number_of_species + ].set(reacted_species_code_units) + + # With thermochemistry, convert the evolved temperature back into a pressure + # and write the energy field, mirroring the cooling module's round-trip. + if chemistry_config.thermochemistry: + reacted_temperature_kelvin = reacted_per_cell[:, number_of_species].reshape( + grid_shape + ) + + # Never let heating/cooling push the temperature below the floor. + reacted_temperature_kelvin = jnp.maximum( + reacted_temperature_kelvin, + chemistry_params.floor_temperature, + ) + + # Kelvin -> rescaled code temperature (inverse of the read above), then + # code temperature -> pressure via the shared equation of state. + reacted_code_temperature = ( + reacted_temperature_kelvin / chemistry_params.temperature_unit_kelvin + ) + new_pressure = get_pressure_from_temperature( + density, + reacted_code_temperature, + chemistry_params.hydrogen_mass_fraction, + chemistry_params.metal_mass_fraction, + ) + primitive_state = primitive_state.at[ + registered_variables.pressure_index + ].set(new_pressure) + + # ------------------------------------------------------------- + # =============== ↑ Cross back and write the state ↑ ========== + # ------------------------------------------------------------- + + return primitive_state diff --git a/astronomix/_modules/_chemistry/_thermochemistry.py b/astronomix/_modules/_chemistry/_thermochemistry.py new file mode 100644 index 00000000..789bbd71 --- /dev/null +++ b/astronomix/_modules/_chemistry/_thermochemistry.py @@ -0,0 +1,868 @@ +""" +Heating and cooling for the chemistry module (thermochemistry). + +Provides the net gas heating minus cooling rate and the corresponding +temperature derivative used to evolve the gas temperature alongside the chemical +abundances. Everything works in CGS with the temperature in Kelvin, matching the +units the per-cell reaction network is integrated in. Species are located by +index rather than a hard-coded layout: any referenced species that is absent +from the network (index ``-1``) drops out of its term. + +Provenance of the physics: + + * heating (cosmic-ray, grain photoelectric, H2-formation) is ported from the + carbox ``latent_tgas`` example; + * H2 line cooling is the **Glover & Abel (2008)** multi-collider form (H, H+, + H2, He, e), ported verbatim from KROME's ``krome_cooling.f90::cooling_H2`` + (a substantial upgrade over the single-collider example fit); + * fine-structure metal cooling — [C II] 158 um and [O I] 63 um — are the + dominant coolants of the cold neutral gas that KROME includes but the carbox + example omits. They are implemented here as standard two-level atoms. The + line data (Einstein A, level spacing) are exact; the collisional + de-excitation rate coefficients are representative literature fits + (Draine 2011; Wolfire et al. 2003; Barinovs et al. 2005), not KROME's + tabulated collision data — good enough to capture the thermal balance, to be + refined against KROME later. +""" + +# jax +import jax.numpy as jnp + +# Boltzmann constant in CGS (erg / K). +BOLTZMANN_CONSTANT_CGS = 1.38064852e-16 + +# H2 binding energy released per molecule formed (4.48 eV) in erg. +ELECTRON_VOLT_ERG = 1.602176634e-12 +HYDROGEN_MOLECULE_BINDING_ERG = 4.48 * ELECTRON_VOLT_ERG + + +def _species_density(abundances, species_index): + """Return the number density of one species, or zero if it is absent. + + Args: + abundances: The per-cell absolute abundances [cm^-3]. + species_index: The static index of the species, or -1 when the network + does not contain it. + + Returns: + The species number density, or ``0.0`` for a missing species so its + contribution vanishes. + """ + # ``species_index`` is a static (configuration) integer, so this branch is + # resolved at trace time and adds no runtime cost. + if species_index < 0: + return 0.0 + return abundances[species_index] + + +# ============================================================= +# ======================= ↓ Heating ↓ ========================= +# ============================================================= + + +def _grain_photoelectric_heating( + abundances, + temperature_kelvin, + fuv_field, + dust_to_gas_ratio, + hydrogen_index, + electron_index, +): + """Net grain photoelectric heating minus grain-recombination cooling. + + Ported from the carbox example ``chem_heating.py:get_photoelectric_heating`` + (Bakes & Tielens / Wolfire form). The heating scales with the total number + density and the FUV field; the recombination-cooling correction depends on + the grain charging parameter ``psi``. + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + fuv_field: The FUV radiation field (Draine units). + dust_to_gas_ratio: The dust-to-gas mass ratio. + hydrogen_index: Index of atomic H (or -1). + electron_index: Index of the electron (or -1). + + Returns: + The net photoelectric heating rate. + """ + total_number_density = jnp.sum(abundances) + electron_density = _species_density(abundances, electron_index) + hydrogen_density = _species_density(abundances, hydrogen_index) + + # Charging parameter of the grains; the small floor avoids a divide-by-zero + # where the electron abundance underflows. + charging_parameter = ( + fuv_field * jnp.sqrt(temperature_kelvin) / (electron_density + 1e-20) + ) + + recombination_exponent = 0.735 * temperature_kelvin ** (-0.068) + grain_recombination_cooling = ( + 4.65e-30 + * temperature_kelvin**0.94 + * charging_parameter**recombination_exponent + * electron_density + * hydrogen_density + ) + + heating_efficiency = 4.9e-2 / ( + 1.0 + 4e-3 * charging_parameter**0.73 + ) + 3.7e-2 * (temperature_kelvin * 1e-4) ** 0.7 / ( + 1.0 + 2e-4 * charging_parameter + ) + + return ( + 1.3e-24 * heating_efficiency * fuv_field * total_number_density + - grain_recombination_cooling + ) * dust_to_gas_ratio + + +def _hydrogen_molecule_formation_heating( + abundances, + temperature_kelvin, + formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, +): + """H2 grain-formation heating with KROME's critical-density partition. + + Each H2 formed on a grain releases 4.48 eV; the fraction that thermalises + (rather than escaping as internal excitation / radiation) is set by the + critical-density factor ``h2heatfac`` ported verbatim from KROME's + ``heatingChem`` (``krome_heating.f90``). The formation rate is tied to the + network's H + H -> H2 reaction, ``formation_rate_coefficient * n_H^2``. + + NB the carbox example instead multiplied the *FUV photodissociation* rate by + n(H2), an unshielded photo term mislabelled as formation heating; a 0-D + benchmark against KROME showed it overestimated the heating by ~5 dex. This + replaces it. (True FUV photodissociation heating, with self-shielding, + belongs in a separate photo module.) + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + formation_rate_coefficient: The H + H -> H2 grain formation rate + coefficient [cm^3 s^-1]. + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + + Returns: + The H2 formation heating rate [erg cm^-3 s^-1]. + """ + if hydrogen_index < 0: + return 0.0 + + hydrogen_density = abundances[hydrogen_index] + molecular_hydrogen_density = _species_density( + abundances, molecular_hydrogen_index + ) + # Total hydrogen nuclei density (KROME's ``get_Hnuclei``, dominant terms). + hydrogen_nuclei_density = hydrogen_density + 2.0 * molecular_hydrogen_density + + # Critical-density partition (KROME heatingChem): the fraction of the 4.48 eV + # that thermalises rises toward unity above the critical density ``ncr``. + critical_numerator = 1.0e6 * temperature_kelvin ** (-0.5) + critical_denominator_hydrogen = 1.6 * jnp.exp(-((4.0e2 / temperature_kelvin) ** 2)) + critical_denominator_molecular = 1.4 * jnp.exp( + -1.2e4 / (temperature_kelvin + 1.2e3) + ) + hydrogen_fraction = hydrogen_density / hydrogen_nuclei_density + molecular_fraction = molecular_hydrogen_density / hydrogen_nuclei_density + critical_density = critical_numerator / ( + critical_denominator_hydrogen * hydrogen_fraction + + critical_denominator_molecular * molecular_fraction + ) + thermalised_fraction = 1.0 / (1.0 + critical_density / hydrogen_nuclei_density) + + formation_rate = formation_rate_coefficient * hydrogen_density * hydrogen_density + return ( + HYDROGEN_MOLECULE_BINDING_ERG * thermalised_fraction * formation_rate + ) + + +def heating_rate( + abundances, + temperature_kelvin, + cosmic_ray_rate, + fuv_field, + dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + electron_index, +): + """Total gas heating rate [erg cm^-3 s^-1]. + + Sums cosmic-ray heating, grain photoelectric heating and H2 grain-formation + heating (the last with KROME's critical-density partition). + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + cosmic_ray_rate: The cosmic-ray ionization rate [s^-1]. + fuv_field: The FUV radiation field (Draine units). + dust_to_gas_ratio: The dust-to-gas mass ratio. + hydrogen_molecule_formation_rate_coefficient: The H + H -> H2 grain + formation rate coefficient [cm^3 s^-1]. + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + electron_index: Index of the electron (or -1). + + Returns: + The total heating rate. + """ + hydrogen_density = _species_density(abundances, hydrogen_index) + molecular_hydrogen_density = _species_density( + abundances, molecular_hydrogen_index + ) + + cosmic_ray_heating = cosmic_ray_rate * ( + 5.5e-12 * hydrogen_density + 2.5e-11 * molecular_hydrogen_density + ) + + photoelectric_heating = _grain_photoelectric_heating( + abundances, + temperature_kelvin, + fuv_field, + dust_to_gas_ratio, + hydrogen_index, + electron_index, + ) + + h2_formation_heating = _hydrogen_molecule_formation_heating( + abundances, + temperature_kelvin, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + ) + + return cosmic_ray_heating + photoelectric_heating + h2_formation_heating + + +# ============================================================= +# ======================= ↓ Cooling ↓ ========================= +# ============================================================= + + +def _sigmoid(argument, midpoint, steepness): + """KROME's ``sigmoid`` helper: ``10 / (10 + exp(-s (x - x0)))``.""" + return 10.0 / (10.0 + jnp.exp(-steepness * (argument - midpoint))) + + +def _cooling_window(log_temperature, log_temperature_low, log_temperature_high): + """KROME's ``wCool`` smoothing window in [0, 1]. + + Ported from ``krome_cooling.f90::wCool``. Blends the piecewise H2 cooling + fits smoothly to zero outside their fitted temperature range. + + Args: + log_temperature: ``log10(T)``. + log_temperature_low: Lower edge of the window in ``log10(T)``. + log_temperature_high: Upper edge of the window in ``log10(T)``. + + Returns: + The window value in [0, 1]. + """ + scaled = (log_temperature - log_temperature_low) / ( + log_temperature_high - log_temperature_low + ) + window = 10.0 ** ( + 200.0 * (_sigmoid(scaled, -0.2, 50.0) * _sigmoid(-scaled, -1.2, 50.0) - 1.0) + ) + return jnp.where(window < 1e-199, 0.0, window) + + +def _ten_to_the(exponent): + """``10 ** exponent`` with the exponent capped to avoid overflow. + + The Glover & Abel fits are only valid in their own temperature window; the + branches evaluated outside it (masked by ``jnp.where``) can otherwise produce + a non-finite value that would poison gradients. Capping keeps every branch + finite; the in-window values (log-rates around -16 to -25) are untouched. + """ + return 10.0 ** jnp.minimum(exponent, 30.0) + + +def _glover_abel_h2_cooling( + abundances, + temperature_kelvin, + hydrogen_index, + molecular_hydrogen_index, + ionized_hydrogen_index, + helium_index, + electron_index, +): + """Glover & Abel (2008) multi-collider H2 line cooling. + + Ported from ``krome_cooling.f90::cooling_H2``. The low-density limit is the + sum of per-collider excitation rates (H, H+, H2, e, He), each a piecewise + polynomial in ``log10(T/1000)``; the high-density (LTE) limit follows + Hollenbach/Glover. The effective rate is the harmonic combination + ``n(H2) / (1/HDL + 1/LDL)``. + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + ionized_hydrogen_index: Index of H+ (or -1). + helium_index: Index of He (or -1). + electron_index: Index of the electron (or -1). + + Returns: + The H2 cooling rate, or ``0.0`` if the network has no H2. + """ + if molecular_hydrogen_index < 0: + return 0.0 + + molecular_hydrogen_density = abundances[molecular_hydrogen_index] + hydrogen_density = _species_density(abundances, hydrogen_index) + ionized_hydrogen_density = _species_density(abundances, ionized_hydrogen_index) + helium_density = _species_density(abundances, helium_index) + electron_density = _species_density(abundances, electron_index) + + temperature = temperature_kelvin + scaled_temperature = temperature * 1e-3 + log_t3 = jnp.log10(scaled_temperature) + log_temperature = jnp.log10(temperature) + + def polynomial(coefficients): + return sum( + coefficient * log_t3**power + for power, coefficient in enumerate(coefficients) + ) + + window_1_4 = _cooling_window(log_temperature, 1.0, 4.0) + window_2_4 = _cooling_window(log_temperature, 2.0, 4.0) + + # --- H2-H (Glover & Abel), piecewise in T --- + h2_h = hydrogen_density * jnp.where( + temperature <= 1e2, + _ten_to_the( + polynomial([-16.818342, 37.383713, 58.145166, 48.656103, 20.159831, 3.8479610]) + ), + jnp.where( + temperature <= 1e3, + _ten_to_the( + polynomial([-24.311209, 3.5692468, -11.332860, -27.850082, -21.328264, -4.2519023]) + ), + jnp.where( + temperature <= 6e3, + _ten_to_the( + polynomial([-24.311209, 4.6450521, -3.7209846, 5.9369081, -5.5108049, 1.5538288]) + ), + 1.862314467912518e-22 + * _cooling_window(log_temperature, 1.0, jnp.log10(6e3)), + ), + ), + ) + + # --- H2-H+ --- + h2_ion = ionized_hydrogen_density * jnp.where( + (temperature > 1e1) & (temperature <= 1e4), + _ten_to_the( + polynomial([-22.089523, 1.5714711, 0.015391166, -0.23619985, -0.51002221, 0.32168730]) + ), + 1.182509139382060e-21 * window_1_4, + ) + + # --- H2-H2 --- + h2_h2 = molecular_hydrogen_density * window_2_4 * _ten_to_the( + polynomial([-23.962112, 2.09433740, -0.77151436, 0.43693353, -0.14913216, -0.033638326]) + ) + + # --- H2-e --- + electron_rate = jnp.where( + temperature <= 5e2, + _ten_to_the( + polynomial( + [-21.928796, 16.815730, 96.743155, 343.19180, 734.71651, 983.67576, 802.01247, 364.14446, 70.609154] + ) + ), + _ten_to_the( + polynomial( + [-22.921189, 1.6802758, 0.93310622, 4.0406627, -4.7274036, -8.8077017, 8.9167183, 6.4380698, -6.3701156] + ) + ), + ) + h2_electron = electron_density * electron_rate * window_2_4 + + # --- H2-He --- + h2_helium = helium_density * jnp.where( + (temperature > 1e1) & (temperature <= 1e4), + _ten_to_the( + polynomial([-23.689237, 2.1892372, -0.81520438, 0.29036281, -0.16596184, 0.19191375]) + ), + 1.002560385050777e-22 * window_1_4, + ) + + low_density_limit = h2_h + h2_ion + h2_h2 + h2_electron + h2_helium + + # High-density (LTE) limit: Hollenbach below 2000 K, Glover fit to 1e4 K, a + # smooth cut-off above. + high_density_rotational = ( + (9.5e-22 * scaled_temperature**3.76) + / (1.0 + 0.12 * scaled_temperature**2.1) + * jnp.exp(-((0.13 / scaled_temperature) ** 3)) + + 3.0e-24 * jnp.exp(-0.51 / scaled_temperature) + ) + high_density_vibrational = 6.7e-19 * jnp.exp(-5.86 / scaled_temperature) + 1.6e-18 * jnp.exp( + -11.7 / scaled_temperature + ) + high_density_low_temperature = high_density_rotational + high_density_vibrational + high_density_mid_temperature = _ten_to_the( + polynomial( + [-20.584225, 5.0194035, -1.5738805, -4.7155769, 2.4714161, 5.4710750, -3.9467356, -2.2148338, 1.8161874] + ) + ) + cutoff = 1.0 / (1.0 + jnp.exp(jnp.minimum((temperature - 3e4) * 2e-4, 3e2))) + high_density_limit = jnp.where( + temperature < 2e3, + high_density_low_temperature, + jnp.where( + temperature <= 1e4, + high_density_mid_temperature, + 5.531333679406485e-19 * cutoff, + ), + ) + + # Harmonic combination; guard the degenerate limits. + return jnp.where( + low_density_limit <= 0.0, + 0.0, + molecular_hydrogen_density + / (1.0 / (high_density_limit + 1e-100) + 1.0 / (low_density_limit + 1e-100)), + ) + + +def _two_level_line_cooling( + species_density, + temperature_kelvin, + level_spacing_kelvin, + einstein_coefficient, + upper_over_lower_degeneracy, + collisional_deexcitation_rate, +): + """Cooling from a single two-level (fine-structure) line. + + The upper-level fraction follows from statistical balance between collisional + excitation/de-excitation and spontaneous radiative decay (optically thin): + ``f_upper = C_lu / (C_lu + C_ul + A)`` with + ``C_lu = C_ul (g_u/g_l) exp(-dE/kT)``. The emitted power per ion is + ``A · dE · f_upper``. + + Args: + species_density: Number density of the emitting species [cm^-3]. + temperature_kelvin: The gas temperature [K]. + level_spacing_kelvin: Upper-lower level spacing ``dE/k`` [K]. + einstein_coefficient: Spontaneous decay rate ``A_ul`` [s^-1]. + upper_over_lower_degeneracy: ``g_upper / g_lower``. + collisional_deexcitation_rate: ``C_ul = sum_c n_c gamma_c(T)`` [s^-1]. + + Returns: + The line cooling rate [erg cm^-3 s^-1]. + """ + boltzmann_factor = jnp.exp(-level_spacing_kelvin / temperature_kelvin) + collisional_excitation = ( + collisional_deexcitation_rate * upper_over_lower_degeneracy * boltzmann_factor + ) + upper_fraction = collisional_excitation / ( + collisional_excitation + collisional_deexcitation_rate + einstein_coefficient + ) + line_energy_erg = level_spacing_kelvin * BOLTZMANN_CONSTANT_CGS + return species_density * einstein_coefficient * line_energy_erg * upper_fraction + + +def _ionized_carbon_cooling( + abundances, + temperature_kelvin, + ionized_carbon_index, + electron_index, + hydrogen_index, + molecular_hydrogen_index, +): + """[C II] 158 um fine-structure cooling (2P3/2 - 2P1/2). + + The dominant coolant of the cold/warm neutral medium. Line data are exact; + collision rate coefficients (e, H, H2) are representative literature fits. + + Returns: + The [C II] cooling rate, or ``0.0`` if the network has no C+. + """ + if ionized_carbon_index < 0: + return 0.0 + + carbon_ion_density = abundances[ionized_carbon_index] + electron_density = _species_density(abundances, electron_index) + hydrogen_density = _species_density(abundances, hydrogen_index) + molecular_hydrogen_density = _species_density(abundances, molecular_hydrogen_index) + scaled_temperature = temperature_kelvin / 100.0 + + # Collisional de-excitation rates [cm^3 s^-1]: electrons via the collision + # strength (Omega ~ 2, g_upper = 4); H and H2 via power-law fits. + electron_rate = 4.3e-6 / jnp.sqrt(temperature_kelvin) + hydrogen_rate = 8.0e-10 * scaled_temperature**0.07 + molecular_hydrogen_rate = 3.8e-10 * scaled_temperature**0.14 + collisional_deexcitation_rate = ( + electron_density * electron_rate + + hydrogen_density * hydrogen_rate + + molecular_hydrogen_density * molecular_hydrogen_rate + ) + + return _two_level_line_cooling( + carbon_ion_density, + temperature_kelvin, + level_spacing_kelvin=91.21, + einstein_coefficient=2.29e-6, + upper_over_lower_degeneracy=2.0, + collisional_deexcitation_rate=collisional_deexcitation_rate, + ) + + +def _atomic_oxygen_cooling( + abundances, + temperature_kelvin, + atomic_oxygen_index, + hydrogen_index, + molecular_hydrogen_index, +): + """[O I] 63 um fine-structure cooling (3P1 - 3P2). + + A major coolant of dense neutral gas. Line data exact; H and H2 collisional + de-excitation rates are representative literature fits (electron collisions + with neutral O are weak and neglected). + + Returns: + The [O I] cooling rate, or ``0.0`` if the network has no atomic O. + """ + if atomic_oxygen_index < 0: + return 0.0 + + oxygen_density = abundances[atomic_oxygen_index] + hydrogen_density = _species_density(abundances, hydrogen_index) + molecular_hydrogen_density = _species_density(abundances, molecular_hydrogen_index) + scaled_temperature = temperature_kelvin / 100.0 + + hydrogen_rate = 9.2e-11 * scaled_temperature**0.67 + molecular_hydrogen_rate = 3.0e-11 * scaled_temperature**0.1 + collisional_deexcitation_rate = ( + hydrogen_density * hydrogen_rate + + molecular_hydrogen_density * molecular_hydrogen_rate + ) + + return _two_level_line_cooling( + oxygen_density, + temperature_kelvin, + level_spacing_kelvin=227.7, + einstein_coefficient=8.91e-5, + upper_over_lower_degeneracy=0.6, + collisional_deexcitation_rate=collisional_deexcitation_rate, + ) + + +# CO Jeans column-density coefficient (KROME num2col default): +# N_CO = CO_COLUMN_COEFFICIENT * (n_CO * 1e-3) ** (2/3) [cm^-2]. +CO_COLUMN_COEFFICIENT = 1.87e21 + + +def _trilinear_table_lookup(table, bounds, coordinate_1, coordinate_2, coordinate_3): + """Trilinearly interpolate a value on a uniform 3D log grid. + + The grid is uniform in each axis between the bounds; coordinates outside the + upper edge are pulled just inside it (as KROME does), and a flag marks + coordinates below the lower edge so the caller can zero the contribution. + + Args: + table: Grid values, shape ``(n1, n2, n3)``. + bounds: ``[x1_min, x1_max, x2_min, x2_max, x3_min, x3_max]``. + coordinate_1: Query position along axis 1. + coordinate_2: Query position along axis 2. + coordinate_3: Query position along axis 3. + + Returns: + A tuple ``(interpolated_value, below_grid)`` where ``below_grid`` is True + when any coordinate fell below the grid's lower edge. + """ + grid_size_1, grid_size_2, grid_size_3 = table.shape + x1_min, x1_max, x2_min, x2_max, x3_min, x3_max = ( + bounds[0], bounds[1], bounds[2], bounds[3], bounds[4], bounds[5] + ) + edge_tolerance = 1e-5 + + coordinate_1 = jnp.where( + coordinate_1 >= x1_max, x1_max * (1.0 - edge_tolerance), coordinate_1 + ) + coordinate_2 = jnp.where( + coordinate_2 >= x2_max, x2_max * (1.0 - edge_tolerance), coordinate_2 + ) + coordinate_3 = jnp.where( + coordinate_3 >= x3_max, x3_max * (1.0 - edge_tolerance), coordinate_3 + ) + below_grid = ( + (coordinate_1 < x1_min) | (coordinate_2 < x2_min) | (coordinate_3 < x3_min) + ) + + def axis_position(coordinate, minimum, maximum, grid_size): + fractional = (coordinate - minimum) / (maximum - minimum) * (grid_size - 1) + lower = jnp.clip(jnp.floor(fractional).astype(jnp.int32), 0, grid_size - 2) + weight = fractional - lower + return lower, weight + + lower_1, weight_1 = axis_position(coordinate_1, x1_min, x1_max, grid_size_1) + lower_2, weight_2 = axis_position(coordinate_2, x2_min, x2_max, grid_size_2) + lower_3, weight_3 = axis_position(coordinate_3, x3_min, x3_max, grid_size_3) + + def corner(offset_1, offset_2, offset_3): + return table[lower_1 + offset_1, lower_2 + offset_2, lower_3 + offset_3] + + # Interpolate along axis 1, then 2, then 3. + face_low_3 = ( + (1 - weight_2) + * ((1 - weight_1) * corner(0, 0, 0) + weight_1 * corner(1, 0, 0)) + + weight_2 + * ((1 - weight_1) * corner(0, 1, 0) + weight_1 * corner(1, 1, 0)) + ) + face_high_3 = ( + (1 - weight_2) + * ((1 - weight_1) * corner(0, 0, 1) + weight_1 * corner(1, 0, 1)) + + weight_2 + * ((1 - weight_1) * corner(0, 1, 1) + weight_1 * corner(1, 1, 1)) + ) + interpolated_value = (1 - weight_3) * face_low_3 + weight_3 * face_high_3 + return interpolated_value, below_grid + + +def _carbon_monoxide_cooling( + abundances, + temperature_kelvin, + carbon_monoxide_index, + hydrogen_index, + molecular_hydrogen_index, + cooling_table, + cooling_bounds, +): + """Tabulated CO rotational-line cooling (Neufeld & Kaufman 1993). + + Ported from KROME's ``cooling_CO``: interpolate the cooling coefficient on a + uniform 3D log grid of (temperature, H+H2 density, CO column density), then + scale by the collider density and the CO abundance. The CO column density + uses KROME's default local Jeans estimate ``N_CO = 1.87e21 (n_CO/1e3)^(2/3)``. + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + carbon_monoxide_index: Index of CO (or -1). + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + cooling_table: The 3D log-cooling table, shape (n_T, n_n, n_NCO). + cooling_bounds: The six uniform-log-grid limits. + + Returns: + The CO cooling rate [erg cm^-3 s^-1], or ``0.0`` if CO is absent. + """ + if carbon_monoxide_index < 0: + return 0.0 + + carbon_monoxide_density = abundances[carbon_monoxide_index] + collider_density = _species_density(abundances, hydrogen_index) + _species_density( + abundances, molecular_hydrogen_index + ) + + # Local Jeans column density of CO (KROME num2col default). + column_density = CO_COLUMN_COEFFICIENT * jnp.maximum( + carbon_monoxide_density * 1e-3, 1e-40 + ) ** (2.0 / 3.0) + + log_temperature = jnp.log10(temperature_kelvin) + log_collider_density = jnp.log10(jnp.maximum(collider_density, 1e-40)) + log_column_density = jnp.log10(column_density) + + log_cooling_coefficient, below_grid = _trilinear_table_lookup( + cooling_table, + cooling_bounds, + log_temperature, + log_collider_density, + log_column_density, + ) + + cooling = ( + 10.0**log_cooling_coefficient * collider_density * carbon_monoxide_density + ) + # Outside the fitted grid (too cold / too diffuse) KROME returns zero. + return jnp.where(below_grid, 0.0, cooling) + + +def cooling_rate( + abundances, + temperature_kelvin, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + atomic_oxygen_index, + ionized_hydrogen_index, + helium_index, + ionized_carbon_index, + co_cooling, + carbon_monoxide_index, + co_cooling_table, + co_cooling_bounds, +): + """Total gas cooling rate [erg cm^-3 s^-1]. + + Sums Lyman-alpha (high-T atomic), [C II] 158 um and [O I] 63 um + fine-structure lines (the dominant cold-neutral coolants), Glover & Abel + multi-collider H2 line cooling, and — when enabled — tabulated CO rotational + cooling. + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + electron_index: Index of the electron (or -1). + atomic_oxygen_index: Index of atomic O (or -1). + ionized_hydrogen_index: Index of H+ (or -1). + helium_index: Index of He (or -1). + ionized_carbon_index: Index of C+ (or -1). + co_cooling: Whether to add tabulated CO rotational cooling. + carbon_monoxide_index: Index of CO (or -1). + co_cooling_table: The 3D CO cooling table (unused when co_cooling False). + co_cooling_bounds: The CO table grid limits. + + Returns: + The total cooling rate. + """ + hydrogen_density = _species_density(abundances, hydrogen_index) + electron_density = _species_density(abundances, electron_index) + + lyman_alpha_cooling = ( + 7.3e-19 + * hydrogen_density + * electron_density + * jnp.exp(-118400.0 / temperature_kelvin) + ) + carbon_ion_cooling = _ionized_carbon_cooling( + abundances, + temperature_kelvin, + ionized_carbon_index, + electron_index, + hydrogen_index, + molecular_hydrogen_index, + ) + oxygen_cooling = _atomic_oxygen_cooling( + abundances, + temperature_kelvin, + atomic_oxygen_index, + hydrogen_index, + molecular_hydrogen_index, + ) + h2_cooling = _glover_abel_h2_cooling( + abundances, + temperature_kelvin, + hydrogen_index, + molecular_hydrogen_index, + ionized_hydrogen_index, + helium_index, + electron_index, + ) + + total = lyman_alpha_cooling + carbon_ion_cooling + oxygen_cooling + h2_cooling + + # The CO table is only present (and only makes sense) when requested; the + # ``co_cooling`` flag is static, so this branch is resolved at trace time. + if co_cooling: + total = total + _carbon_monoxide_cooling( + abundances, + temperature_kelvin, + carbon_monoxide_index, + hydrogen_index, + molecular_hydrogen_index, + co_cooling_table, + co_cooling_bounds, + ) + + return total + + +def temperature_derivative( + abundances, + temperature_kelvin, + adiabatic_index, + cosmic_ray_rate, + fuv_field, + dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + atomic_oxygen_index, + ionized_hydrogen_index, + helium_index, + ionized_carbon_index, + co_cooling, + carbon_monoxide_index, + co_cooling_table, + co_cooling_bounds, +): + """Gas temperature derivative dT/dt [K s^-1] from net heating minus cooling. + + ``dT/dt = (gamma - 1) * (heating - cooling) / (k_B * n_total)`` with + ``n_total`` the total particle number density (sum of abundances). + + Args: + abundances: The per-cell abundances [cm^-3]. + temperature_kelvin: The gas temperature [K]. + adiabatic_index: The adiabatic index gamma. + cosmic_ray_rate: The cosmic-ray ionization rate [s^-1]. + fuv_field: The FUV radiation field (Draine units). + dust_to_gas_ratio: The dust-to-gas mass ratio. + hydrogen_molecule_formation_rate_coefficient: The H + H -> H2 grain + formation rate coefficient [cm^3 s^-1]. + hydrogen_index: Index of atomic H (or -1). + molecular_hydrogen_index: Index of H2 (or -1). + electron_index: Index of the electron (or -1). + atomic_oxygen_index: Index of atomic O (or -1). + ionized_hydrogen_index: Index of H+ (or -1). + helium_index: Index of He (or -1). + ionized_carbon_index: Index of C+ (or -1). + co_cooling: Whether to add tabulated CO rotational cooling. + carbon_monoxide_index: Index of CO (or -1). + co_cooling_table: The 3D CO cooling table. + co_cooling_bounds: The CO table grid limits. + + Returns: + The temperature derivative dT/dt. + """ + net_heating = heating_rate( + abundances, + temperature_kelvin, + cosmic_ray_rate, + fuv_field, + dust_to_gas_ratio, + hydrogen_molecule_formation_rate_coefficient, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + ) - cooling_rate( + abundances, + temperature_kelvin, + hydrogen_index, + molecular_hydrogen_index, + electron_index, + atomic_oxygen_index, + ionized_hydrogen_index, + helium_index, + ionized_carbon_index, + co_cooling, + carbon_monoxide_index, + co_cooling_table, + co_cooling_bounds, + ) + + total_number_density = jnp.sum(abundances) + return ( + (adiabatic_index - 1.0) + * net_heating + / (BOLTZMANN_CONSTANT_CGS * total_number_density) + ) diff --git a/astronomix/_modules/_chemistry/chemistry_options.py b/astronomix/_modules/_chemistry/chemistry_options.py new file mode 100644 index 00000000..840cd800 --- /dev/null +++ b/astronomix/_modules/_chemistry/chemistry_options.py @@ -0,0 +1,157 @@ +""" +Configuration and parameter containers for the chemistry module. + +Defines the integer tags selecting the stiff ODE solver and the NamedTuples that +carry the reaction network (split into a static structure and dynamic array +leaves, following the neural-network cooling pattern), the unit-conversion +factors between astronomix code units and the CGS units carbox expects, and the +per-network physical parameters. +""" + +# typing +from typing import NamedTuple, Tuple, Union +from types import NoneType +from jaxtyping import PyTree + +# jax +import jax.numpy as jnp + +# Stiff-solver tags (select which Diffrax solver integrates the network). Only +# Kvaerno5 is wired up for now; the others leave room to switch without changing +# the container layout. +KVAERNO5 = 1 +DOPRI5 = 2 +TSIT5 = 3 + + +class ChemistryConfig(NamedTuple): + """Top-level chemistry configuration. + + Every field must be hashable because the simulation configuration is passed + as a static argument to the jitted update. The number of species and their + order are needed at registry-build time to allocate the contiguous species + block in the state array. The reaction network itself is deliberately *not* + stored here: a carbox ``JNetwork`` holds a Python list of reactions and is + therefore unhashable, so it lives in ``ChemistryParams`` (a dynamic argument) + instead. + + Attributes: + chemistry: Master switch for the chemistry module. + number_of_chemical_species: Length of the species block in the state. + species_names: Species ordering, matching both the state block and the + carbox network's ``species`` list. + number_of_reactions: Number of reactions in the network (length of the + rate-modifier vectors). + solver: Stiff-solver tag (see the module-level constants). + max_steps: Maximum internal Diffrax steps per cell per hydro step. + thermochemistry: When True, the temperature is evolved together with the + abundances (heating/cooling feedback) and written back into the + pressure field. When False the species react at a fixed temperature + and the energy field is left untouched. + hydrogen_index: Position of atomic H in the species block, or -1 if the + network has no such species. Resolved once from ``species_names``; + the thermochemistry terms that reference a missing species contribute + zero. + molecular_hydrogen_index: Position of H2, or -1. + electron_index: Position of the electron (E), or -1. + atomic_oxygen_index: Position of atomic O, or -1. + ionized_hydrogen_index: Position of H+, or -1 (H2-H+ cooling collider). + helium_index: Position of He, or -1 (H2-He cooling collider). + ionized_carbon_index: Position of C+, or -1 ([C II] 158 um cooling). + co_cooling: When True, add tabulated CO rotational-line cooling + (Neufeld & Kaufman 1993). Requires the table in ``ChemistryParams``. + carbon_monoxide_index: Position of CO, or -1 (CO cooling). + """ + + chemistry: bool = False + number_of_chemical_species: int = 0 + species_names: Tuple[str, ...] = () + number_of_reactions: int = 0 + solver: int = KVAERNO5 + max_steps: int = 4096 + + # thermochemistry (chemistry-driven heating/cooling of the energy field) + thermochemistry: bool = False + hydrogen_index: int = -1 + molecular_hydrogen_index: int = -1 + electron_index: int = -1 + atomic_oxygen_index: int = -1 + ionized_hydrogen_index: int = -1 + helium_index: int = -1 + ionized_carbon_index: int = -1 + co_cooling: bool = False + carbon_monoxide_index: int = -1 + + +class ChemistryParams(NamedTuple): + """Runtime chemistry parameters (differentiable leaves). + + Bundles the network's array leaves, the code-unit <-> CGS conversion factors, + the physical parameters the rate expressions read, and the per-reaction rate + modifiers. + + Attributes: + network: The carbox reaction network (``JNetwork``). It is an Equinox + module — a JAX pytree of arrays plus structure — so it rides here as + a dynamic argument and stays differentiable in its rate constants. + hydrogen_mass_fraction: Hydrogen mass fraction X, used for the mean + molecular weight in the pressure-to-temperature conversion. + metal_mass_fraction: Metal mass fraction Z, used likewise. + number_density_unit_cgs: Number density [cm^-3] per code number-density + unit; converts species and hydrogen densities into CGS. + temperature_unit_kelvin: Kelvin per code (rescaled) temperature unit. + time_unit_seconds: Seconds per code time unit. + cosmic_ray_rate: Cosmic-ray ionization rate [s^-1]. + fuv_field: FUV radiation field strength (Draine units). + visual_extinction: Visual extinction Av [mag]. + dust_to_gas_ratio: Dust-to-gas mass ratio, used by the grain + photoelectric heating term. + floor_temperature: Lower bound [K] on the post-reaction temperature + before it is written back to the pressure field. + hydrogen_molecule_formation_rate_coefficient: The H + H -> H2 grain + formation rate coefficient [cm^3 s^-1], used by the H2 formation + heating term. + atol: Absolute tolerance of the stiff solver. + rtol: Relative tolerance of the stiff solver. + rate_modifier_a: Per-reaction multiplicative rate modifier (ones = no-op). + rate_modifier_b: Per-reaction additive rate modifier (zeros = no-op). + co_cooling_table: The 3D Neufeld & Kaufman CO cooling table, shape + ``(n_logT, n_logn, n_logNCO)`` of ``log10`` cooling coefficients. + Empty when CO cooling is inactive. + co_cooling_bounds: The six uniform-log-grid limits + ``[logT_min, logT_max, logn_min, logn_max, logNCO_min, logNCO_max]``. + """ + + network: Union[PyTree, NoneType] = None + + # composition used to convert pressure to temperature (mean molecular weight) + hydrogen_mass_fraction: float = 0.76 + metal_mass_fraction: float = 0.02 + + # code-unit <-> CGS conversion factors (the chemistry boundary) + number_density_unit_cgs: float = 1.0 + temperature_unit_kelvin: float = 1.0 + time_unit_seconds: float = 1.0 + + # physical parameters entering the rate expressions + cosmic_ray_rate: float = 1e-17 + fuv_field: float = 1.0 + visual_extinction: float = 2.0 + + # thermochemistry parameters (heating/cooling) + dust_to_gas_ratio: float = 1e-2 + floor_temperature: float = 1e1 + # H + H -> H2 grain formation rate coefficient [cm^3 s^-1] (formation heating) + hydrogen_molecule_formation_rate_coefficient: float = 3e-17 + + # stiff-solver tolerances + atol: float = 1e-18 + rtol: float = 1e-10 + + # per-reaction rate modifiers, rate -> rate * a + b + rate_modifier_a: jnp.ndarray = jnp.array([]) + rate_modifier_b: jnp.ndarray = jnp.array([]) + + # tabulated CO rotational cooling (Neufeld & Kaufman 1993) + co_cooling_table: jnp.ndarray = jnp.array([]) + co_cooling_bounds: jnp.ndarray = jnp.array([]) diff --git a/astronomix/_modules/_iteration_level_updates.py b/astronomix/_modules/_iteration_level_updates.py index ed1fa5ba..56a48ab1 100644 --- a/astronomix/_modules/_iteration_level_updates.py +++ b/astronomix/_modules/_iteration_level_updates.py @@ -36,6 +36,7 @@ from astronomix.variable_registry.registered_variables import RegisteredVariables # astronomix functions +from astronomix._modules._chemistry._chemistry import update_chemistry from astronomix._modules._cnn_mhd_corrector._cnn_mhd_corrector import _cnn_mhd_corrector from astronomix._modules._cooling._cooling import update_pressure_by_cooling from astronomix._modules._cosmic_rays.cr_injection import inject_crs_at_strongest_shock @@ -145,6 +146,17 @@ def _iteration_level_updates( dt, ) + # Astrochemistry. The species advect with the flow (handled by the solver); + # here they are reacted per cell. Finite-volume only for now. + if config.chemistry_config.chemistry and config.solver_mode == FINITE_VOLUME: + primitive_state = update_chemistry( + primitive_state, + registered_variables, + config.chemistry_config, + params, + dt, + ) + # Neural-network body force. if config.neural_net_force_config.neural_net_force: primitive_state = _neural_net_force( diff --git a/astronomix/option_classes/simulation_config.py b/astronomix/option_classes/simulation_config.py index c6a99b40..7aa46bbf 100644 --- a/astronomix/option_classes/simulation_config.py +++ b/astronomix/option_classes/simulation_config.py @@ -25,6 +25,7 @@ from astronomix._modules._cnn_mhd_corrector._cnn_mhd_corrector_options import ( CNNMHDconfig, ) +from astronomix._modules._chemistry.chemistry_options import ChemistryConfig from astronomix._modules._cooling.cooling_options import CoolingConfig from astronomix._modules._cosmic_rays.cosmic_ray_options import CosmicRayConfig from astronomix._modules._neural_net_force._neural_net_force_options import ( @@ -595,6 +596,9 @@ class SimulationConfig(NamedTuple): #: The configuration for the cooling module. cooling_config: CoolingConfig = CoolingConfig() + #: The configuration for the chemistry module. + chemistry_config: ChemistryConfig = ChemistryConfig() + #: Frame tracking in z-direction #: shifting the frame to follow a #: turbulent radiative mixing layer diff --git a/astronomix/option_classes/simulation_params.py b/astronomix/option_classes/simulation_params.py index 6f10b6f1..13578897 100644 --- a/astronomix/option_classes/simulation_params.py +++ b/astronomix/option_classes/simulation_params.py @@ -15,6 +15,7 @@ # astronomix containers from astronomix._modules._cnn_mhd_corrector._cnn_mhd_corrector_options import CNNMHDconfig +from astronomix._modules._chemistry.chemistry_options import ChemistryParams from astronomix._modules._cooling.cooling_options import CoolingParams from astronomix._modules._cosmic_rays.cosmic_ray_options import CosmicRayParams from astronomix._modules._neural_net_force._neural_net_force_options import NeuralNetForceParams @@ -128,6 +129,9 @@ class SimulationParams(NamedTuple): #: The parameters of the cooling module. cooling_params: CoolingParams = CoolingParams() + #: The parameters of the chemistry module. + chemistry_params: ChemistryParams = ChemistryParams() + #: The parameters of the neural network force module. neural_net_force_params: NeuralNetForceParams = NeuralNetForceParams() diff --git a/astronomix/setup_helpers/__init__.py b/astronomix/setup_helpers/__init__.py index 033370d7..80619dc0 100644 --- a/astronomix/setup_helpers/__init__.py +++ b/astronomix/setup_helpers/__init__.py @@ -1,6 +1,7 @@ """Setup / orchestration helpers for astronomix simulations.""" # astronomix functions +from astronomix.setup_helpers.chemistry_setup import build_chemistry_from_network_file from astronomix.setup_helpers.restart import ( latest_checkpoint_step, restart_from_latest_checkpoint, @@ -9,4 +10,5 @@ __all__ = [ "restart_from_latest_checkpoint", "latest_checkpoint_step", + "build_chemistry_from_network_file", ] diff --git a/astronomix/setup_helpers/chemistry_setup.py b/astronomix/setup_helpers/chemistry_setup.py new file mode 100644 index 00000000..dbce27d2 --- /dev/null +++ b/astronomix/setup_helpers/chemistry_setup.py @@ -0,0 +1,211 @@ +""" +Build the chemistry configuration from a carbox reaction network file. + +Turns a carbox network CSV into the pair of containers the chemistry module +consumes: a static ``ChemistryConfig`` (species count and ordering, solver, and +the network's structural half) and a dynamic ``ChemistryParams`` (the network's +array half, unit-conversion factors, physical parameters and rate modifiers). + +The reaction network (a carbox ``JNetwork``) is an Equinox module, i.e. a JAX +pytree of arrays plus structure. Because it holds a Python list of reactions it +is not hashable, so it is carried in the dynamic ``ChemistryParams`` rather than +the static ``ChemistryConfig``; it stays differentiable in its rate constants. + +WARNING: This helper imports carbox. It is an optional dependency; the import is +kept local so that ``import astronomix`` does not require carbox unless this +helper is actually called. +""" + +# typing +from typing import Tuple + +# numerics +import numpy as np + +# jax +import jax.numpy as jnp + +# astronomix constants +from astronomix._modules._chemistry.chemistry_options import KVAERNO5 + +# astronomix containers +from astronomix._modules._chemistry.chemistry_options import ( + ChemistryConfig, + ChemistryParams, +) + +# Default location of KROME's Neufeld & Kaufman CO cooling table. +DEFAULT_CO_COOLING_TABLE_PATH = "/export/scratch/lbranca/krome/data/coolCO.dat" + + +def _load_co_cooling_table(table_path): + """Load KROME's ``coolCO.dat`` into a dense 3D table and its grid bounds. + + The file lists, per row, the integer grid indices (i, j, k) followed by + ``log10(T)``, ``log10(n_H + n_H2)``, ``log10(N_CO)`` and ``log10`` of the + cooling coefficient. The grid is uniform in each log axis. + + Args: + table_path: Path to ``coolCO.dat``. + + Returns: + A tuple ``(table, bounds)`` with the table shaped + ``(n_logT, n_logn, n_logNCO)`` and the six grid limits + ``[logT_min, logT_max, logn_min, logn_max, logNCO_min, logNCO_max]``. + """ + rows = np.loadtxt(table_path, comments="#") + index_temperature = rows[:, 0].astype(int) - 1 + index_density = rows[:, 1].astype(int) - 1 + index_column = rows[:, 2].astype(int) - 1 + log_temperature = rows[:, 3] + log_density = rows[:, 4] + log_column = rows[:, 5] + log_cooling = rows[:, 6] + + table = np.zeros( + ( + index_temperature.max() + 1, + index_density.max() + 1, + index_column.max() + 1, + ) + ) + table[index_temperature, index_density, index_column] = log_cooling + bounds = np.array( + [ + log_temperature.min(), + log_temperature.max(), + log_density.min(), + log_density.max(), + log_column.min(), + log_column.max(), + ] + ) + return jnp.asarray(table), jnp.asarray(bounds) + + +def build_chemistry_from_network_file( + network_csv_path: str, + network_format: str, + number_density_unit_cgs: float, + temperature_unit_kelvin: float, + time_unit_seconds: float, + cosmic_ray_rate: float = 1e-17, + fuv_field: float = 1.0, + visual_extinction: float = 2.0, + absolute_tolerance: float = 1e-18, + relative_tolerance: float = 1e-10, + solver: int = KVAERNO5, + max_steps: int = 4096, + thermochemistry: bool = False, + dust_to_gas_ratio: float = 1e-2, + floor_temperature: float = 1e1, + hydrogen_molecule_formation_rate_coefficient: float = 3e-17, + co_cooling: bool = False, + co_cooling_table_path: str = DEFAULT_CO_COOLING_TABLE_PATH, +) -> Tuple[ChemistryConfig, ChemistryParams, Tuple[str, ...]]: + """Assemble the chemistry containers from a carbox network file. + + Args: + network_csv_path: Path to the carbox reaction-network CSV. + network_format: carbox parser format, e.g. ``"latent_tgas"``. + number_density_unit_cgs: Number density [cm^-3] per code number-density + unit (converts species densities into CGS for the network). + temperature_unit_kelvin: Kelvin per code (rescaled) temperature unit. + time_unit_seconds: Seconds per code time unit. + cosmic_ray_rate: Cosmic-ray ionization rate [s^-1]. + fuv_field: FUV radiation field strength (Draine units). + visual_extinction: Visual extinction Av [mag]. + absolute_tolerance: Absolute tolerance of the stiff solver. + relative_tolerance: Relative tolerance of the stiff solver. + solver: Stiff-solver tag (see ``chemistry_options``). + max_steps: Maximum internal Diffrax steps per cell per hydro step. + thermochemistry: When True, evolve the temperature with the abundances + (heating/cooling) and write the result back into the pressure field. + dust_to_gas_ratio: Dust-to-gas mass ratio (grain photoelectric heating). + floor_temperature: Lower bound [K] on the post-reaction temperature. + hydrogen_molecule_formation_rate_coefficient: H + H -> H2 grain formation + rate coefficient [cm^3 s^-1] (H2 formation heating). + co_cooling: When True, add tabulated CO rotational cooling (loads the + Neufeld & Kaufman table from ``co_cooling_table_path``). + co_cooling_table_path: Path to KROME's ``coolCO.dat``. + + Returns: + A tuple ``(chemistry_config, chemistry_params, species_names)``. The + species ordering is returned so the caller can place initial abundances + into the matching state slots. + """ + + # carbox is an optional dependency; import locally so the package imports + # without it when chemistry is unused. + from carbox.network import Network + from carbox.parsers import parse_chemical_network + + # A dense stoichiometry matrix keeps the per-cell right-hand side a plain + # matmul, which vmaps cleanly across the grid (the default sparse BCOO does + # not). + parsed_network = parse_chemical_network(network_csv_path, network_format) + dense_network = Network( + parsed_network.species, + parsed_network.reactions, + use_sparse=False, + ) + reaction_network = dense_network.get_ode() + + species_names = tuple(species.name for species in dense_network.species) + number_of_species = len(species_names) + number_of_reactions = reaction_network.reactions_number + + # Resolve the species the thermochemistry terms reference. A species absent + # from the network maps to -1, and its heating/cooling contribution drops out. + def species_index(name): + return species_names.index(name) if name in species_names else -1 + + # Load the CO cooling table only when requested. + co_cooling_table = jnp.array([]) + co_cooling_bounds = jnp.array([]) + if co_cooling: + co_cooling_table, co_cooling_bounds = _load_co_cooling_table( + co_cooling_table_path + ) + + chemistry_config = ChemistryConfig( + chemistry=True, + number_of_chemical_species=number_of_species, + species_names=species_names, + number_of_reactions=number_of_reactions, + solver=solver, + max_steps=max_steps, + thermochemistry=thermochemistry, + hydrogen_index=species_index("H"), + molecular_hydrogen_index=species_index("H2"), + electron_index=species_index("E"), + atomic_oxygen_index=species_index("O"), + ionized_hydrogen_index=species_index("H+"), + helium_index=species_index("He"), + ionized_carbon_index=species_index("C+"), + co_cooling=co_cooling, + carbon_monoxide_index=species_index("CO"), + ) + + chemistry_params = ChemistryParams( + network=reaction_network, + number_density_unit_cgs=number_density_unit_cgs, + temperature_unit_kelvin=temperature_unit_kelvin, + time_unit_seconds=time_unit_seconds, + cosmic_ray_rate=cosmic_ray_rate, + fuv_field=fuv_field, + visual_extinction=visual_extinction, + dust_to_gas_ratio=dust_to_gas_ratio, + floor_temperature=floor_temperature, + hydrogen_molecule_formation_rate_coefficient=( + hydrogen_molecule_formation_rate_coefficient + ), + atol=absolute_tolerance, + rtol=relative_tolerance, + rate_modifier_a=jnp.ones(number_of_reactions), + rate_modifier_b=jnp.zeros(number_of_reactions), + co_cooling_table=co_cooling_table, + co_cooling_bounds=co_cooling_bounds, + ) + + return chemistry_config, chemistry_params, species_names diff --git a/astronomix/variable_registry/registered_variables.py b/astronomix/variable_registry/registered_variables.py index d0c87ace..6b27d6fd 100644 --- a/astronomix/variable_registry/registered_variables.py +++ b/astronomix/variable_registry/registered_variables.py @@ -118,6 +118,16 @@ class RegisteredVariables(NamedTuple): cosmic_ray_n_index: int = -1 cosmic_ray_n_active: bool = False + #: chemical species abundances + # The chemistry module carries the species number densities in a contiguous + # block of the state array, starting at ``chemistry_species_index`` and + # spanning ``num_chemical_species`` slots (ordered as the reaction network's + # species list). They are advected as passive volumetric densities by the + # finite-volume solver, exactly like the cosmic-ray and wind-density tracers. + chemistry_species_index: int = -1 + num_chemical_species: int = 0 + chemistry_species_active: bool = False + # here you can add more variables @@ -213,6 +223,24 @@ def get_registered_variables(config: SimulationConfig) -> RegisteredVariables: ) registered_variables = registered_variables._replace(cosmic_ray_n_active=True) + # NOTE: CURRENTLY ONLY IMPLEMENTED FOR FINITE VOLUME MODE + # The chemical species occupy one contiguous block; the base index marks + # its start and the block spans ``number_of_chemical_species`` slots. + if config.chemistry_config.chemistry: + registered_variables = registered_variables._replace( + chemistry_species_index=registered_variables.num_vars + ) + registered_variables = registered_variables._replace( + num_chemical_species=config.chemistry_config.number_of_chemical_species + ) + registered_variables = registered_variables._replace( + num_vars=registered_variables.num_vars + + config.chemistry_config.number_of_chemical_species + ) + registered_variables = registered_variables._replace( + chemistry_species_active=True + ) + if config.solver_mode == FINITE_DIFFERENCE: From 37307159033bdfeb6c160fb22048cd277149de0b Mon Sep 17 00:00:00 2001 From: Lorenzo Branca Date: Wed, 15 Jul 2026 11:51:17 +0200 Subject: [PATCH 2/3] Make full-chemistry branch PR-ready: GPL marking + remove hardcoded path - Add LICENSE.md marking the _chemistry subtree (and chemistry_setup) as GPL-3.0 (KROME-derived thermochemistry + carbox dependency), distinct from MIT core; note it in the package and _thermochemistry docstrings. - Remove the machine-specific default CO-table path; require the caller to supply co_cooling_table_path when co_cooling is enabled (no KROME data bundled). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01UhwF9zJfmnzuLeABHu75QW --- astronomix/_modules/_chemistry/LICENSE.md | 27 +++++++++++++++++++ astronomix/_modules/_chemistry/__init__.py | 4 +++ .../_modules/_chemistry/_thermochemistry.py | 3 +++ astronomix/setup_helpers/chemistry_setup.py | 14 +++++++--- 4 files changed, 45 insertions(+), 3 deletions(-) create mode 100644 astronomix/_modules/_chemistry/LICENSE.md diff --git a/astronomix/_modules/_chemistry/LICENSE.md b/astronomix/_modules/_chemistry/LICENSE.md new file mode 100644 index 00000000..6762b6c3 --- /dev/null +++ b/astronomix/_modules/_chemistry/LICENSE.md @@ -0,0 +1,27 @@ +# License notice — `astronomix/_modules/_chemistry/` (and `setup_helpers/chemistry_setup.py`) + +**These chemistry files are licensed under the GNU General Public License v3.0 +(GPL-3.0), NOT the MIT license that covers the rest of astronomix.** + +Reason: the thermochemistry (`_thermochemistry.py`) is a derivative work of +**KROME** (Grassi et al. 2014, GPL-3.0) — the Glover & Abel (2008) multi-collider +H2 cooling, the `wCool`/`sigmoid` smoothing, the Neufeld & Kaufman (1993) CO +trilinear lookup, and the `heatingChem` critical-density partition were +transcribed from KROME source. The driver (`_chemistry.py`) and setup +(`chemistry_setup.py`) also depend at runtime on **carbox** (GPL-3.0-or-later). + +Implications a maintainer must weigh before merging: +- GPL is copyleft: distributing astronomix with these files combined in generally + subjects the combined/distributed work to GPL-3.0 for the parts that include or + link this code. Keeping astronomix MIT while carrying this subtree is a + deliberate mixed-license choice, and its reach should be confirmed. +- No KROME data is bundled. CO cooling requires the caller to supply their own + `coolCO.dat` (KROME/Omukai, GPL) via `co_cooling_table_path`. +- carbox must be installed for the feature to run (it is imported lazily, so + `import astronomix` itself is unaffected). + +Alternatives that avoid the GPL entanglement (see the project discussion): +1. keep only the MIT-clean species/source-term *mechanism* in astronomix and ship + this engine as a separate GPL companion package; or +2. reimplement the cooling from the primary papers (not KROME source) so it can be + MIT, leaving carbox as an optional dependency. diff --git a/astronomix/_modules/_chemistry/__init__.py b/astronomix/_modules/_chemistry/__init__.py index b6c517be..98aab8b1 100644 --- a/astronomix/_modules/_chemistry/__init__.py +++ b/astronomix/_modules/_chemistry/__init__.py @@ -5,4 +5,8 @@ step, reacts them per cell using a carbox reaction network integrated with a stiff Diffrax solver. See ``_chemistry.py`` for the driver and ``chemistry_options.py`` for the configuration / parameter containers. + +LICENSE: unlike the rest of astronomix (MIT), this package is GPL-3.0 — the +thermochemistry is derived from KROME (GPL) and the coupling depends on carbox +(GPL). See ``LICENSE.md`` in this directory. """ diff --git a/astronomix/_modules/_chemistry/_thermochemistry.py b/astronomix/_modules/_chemistry/_thermochemistry.py index 789bbd71..9c5d9e7e 100644 --- a/astronomix/_modules/_chemistry/_thermochemistry.py +++ b/astronomix/_modules/_chemistry/_thermochemistry.py @@ -1,6 +1,9 @@ """ Heating and cooling for the chemistry module (thermochemistry). +LICENSE: GPL-3.0 (NOT astronomix's MIT). This file is a derivative work of KROME +(Grassi et al. 2014, GPL-3.0); see ``LICENSE.md`` in this directory. + Provides the net gas heating minus cooling rate and the corresponding temperature derivative used to evolve the gas temperature alongside the chemical abundances. Everything works in CGS with the temperature in Kelvin, matching the diff --git a/astronomix/setup_helpers/chemistry_setup.py b/astronomix/setup_helpers/chemistry_setup.py index dbce27d2..f41cdfe1 100644 --- a/astronomix/setup_helpers/chemistry_setup.py +++ b/astronomix/setup_helpers/chemistry_setup.py @@ -34,8 +34,10 @@ ChemistryParams, ) -# Default location of KROME's Neufeld & Kaufman CO cooling table. -DEFAULT_CO_COOLING_TABLE_PATH = "/export/scratch/lbranca/krome/data/coolCO.dat" +# The Neufeld & Kaufman CO cooling table is KROME/Omukai data (GPL) and is NOT +# bundled here; the caller must supply the path to their own ``coolCO.dat`` when +# CO cooling is requested. There is deliberately no default path. +DEFAULT_CO_COOLING_TABLE_PATH = None def _load_co_cooling_table(table_path): @@ -101,7 +103,7 @@ def build_chemistry_from_network_file( floor_temperature: float = 1e1, hydrogen_molecule_formation_rate_coefficient: float = 3e-17, co_cooling: bool = False, - co_cooling_table_path: str = DEFAULT_CO_COOLING_TABLE_PATH, + co_cooling_table_path=DEFAULT_CO_COOLING_TABLE_PATH, ) -> Tuple[ChemistryConfig, ChemistryParams, Tuple[str, ...]]: """Assemble the chemistry containers from a carbox network file. @@ -164,6 +166,12 @@ def species_index(name): co_cooling_table = jnp.array([]) co_cooling_bounds = jnp.array([]) if co_cooling: + if co_cooling_table_path is None: + raise ValueError( + "co_cooling=True requires co_cooling_table_path pointing to a " + "Neufeld & Kaufman CO cooling table (e.g. KROME's coolCO.dat); " + "no table is bundled." + ) co_cooling_table, co_cooling_bounds = _load_co_cooling_table( co_cooling_table_path ) From 982206466ad24285c45a42656c07a4a4d886baae Mon Sep 17 00:00:00 2001 From: Lorenzo Branca Date: Wed, 15 Jul 2026 16:47:30 +0200 Subject: [PATCH 3/3] Fix FV MHD field slicing so it composes with trailing tracers The finite-volume MHD state evolution split the magnetic field off the primitive state with `state[:-3]` / `state[-3:]`, assuming B occupies the last three state indices. That holds for pure MHD, but the variable registry appends advected tracers (chemistry species, cosmic-ray number density, wind density) *after* the magnetic block, so with any of those enabled `[-3:]` sliced tracer rows instead of B -- silently corrupting the magnetic update (no error raised). Slice B out of its actual (middle) slot via `magnetic_index` instead, and reinsert it there on return. `magnetic_update` is unchanged: it already takes the gas sub-state and B as separate arguments and only reads density / velocity / pressure from the front of the gas state. For the no-tracer case the new slicing is bit-identical to the old one, so pure MHD is unaffected. Verified with a 32^3 FV MHD box run plain vs. with the chemistry module (thermochemistry off, so species are passive): the magnetic field is bit-identical between the two runs while the species react, whereas before this fix the fields diverged at O(1). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01UhwF9zJfmnzuLeABHu75QW --- .../_state_evolution/evolve_state.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/astronomix/_finite_volume/_state_evolution/evolve_state.py b/astronomix/_finite_volume/_state_evolution/evolve_state.py index 3555aa28..9bfd8193 100644 --- a/astronomix/_finite_volume/_state_evolution/evolve_state.py +++ b/astronomix/_finite_volume/_state_evolution/evolve_state.py @@ -576,15 +576,26 @@ def _evolve_state_fv( if config.mhd: if config.dimensionality > 1: - # WARNING: this relies on the last three state indices being the - # magnetic-field components, so that stripping them off yields the - # pure gas sub-state and a matching gas variable registry. + # The magnetic field occupies three contiguous slots starting at + # ``magnetic_index.x`` (right after the pressure) — NOT necessarily the + # last three, because advected tracers (chemistry species, cosmic rays, + # wind density) are registered after the magnetic block. Strip the + # magnetic rows out of the middle so the gas sub-state keeps density / + # velocity / pressure at the front with the tracers trailing (they + # advect with the flow), then reinsert the field into its slot below. + magnetic_start = registered_variables.magnetic_index.x registered_variables_gas = registered_variables._replace( num_vars=registered_variables.num_vars - 3 ) - gas_state = primitive_state[:-3, ...] - magnetic_field = primitive_state[-3:, ...] + gas_state = jnp.concatenate( + [ + primitive_state[:magnetic_start], + primitive_state[magnetic_start + 3 :], + ], + axis=0, + ) + magnetic_field = primitive_state[magnetic_start : magnetic_start + 3, ...] if config.split == UNSPLIT: evolved_gas = _evolve_gas_state_unsplit( @@ -638,7 +649,15 @@ def _evolve_state_fv( registered_variables_gas, ) - return jnp.concatenate((evolved_gas, magnetic_field), axis=0) + # reinsert the magnetic field into its original (middle) slot + return jnp.concatenate( + [ + evolved_gas[:magnetic_start], + magnetic_field, + evolved_gas[magnetic_start:], + ], + axis=0, + ) else: raise ValueError("MHD currently not supported in 1D.")