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.") 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 new file mode 100644 index 00000000..98aab8b1 --- /dev/null +++ b/astronomix/_modules/_chemistry/__init__.py @@ -0,0 +1,12 @@ +""" +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. + +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/_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..9c5d9e7e --- /dev/null +++ b/astronomix/_modules/_chemistry/_thermochemistry.py @@ -0,0 +1,871 @@ +""" +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 +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..f41cdfe1 --- /dev/null +++ b/astronomix/setup_helpers/chemistry_setup.py @@ -0,0 +1,219 @@ +""" +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, +) + +# 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): + """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=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: + 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 + ) + + 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: