Chemistry–hydro coupling: advected species + carbox reactions + thermochemistry

[lorenzobranca]

Chemistry ↔ hydro coupling: advected species + carbox reactions + thermochemistry

⚠️ Licensing decision required — please read first.
This PR adds files that are GPL-3.0, not MIT: _modules/_chemistry/
(_thermochemistry.py is transcribed from KROME, GPL) and
setup_helpers/chemistry_setup.py (depends on carbox, GPL-3.0-or-later).
They are marked as GPL via _modules/_chemistry/LICENSE.md. Merging this into
MIT astronomix is a deliberate mixed-license choice with copyleft reach on
the distributed combination — your call as maintainer. Opened as a draft
precisely so we can decide this together. Alternatives are listed at the bottom.

What it adds

A full operator-split coupling of a chemical network to the fluid, differentiable
and GPU-vmapped end to end:

• Species advection — N chemical-species number densities carried as a
  contiguous state block, advected by the FV solver like the CR/wind tracers
  (registered_variables.py).
• Per-cell reaction — each hydro step, every cell integrates a carbox
  reaction network over dt with a stiff Diffrax Kvaerno5 solve, vmapped over
  the grid (_chemistry.py).
• Thermochemistry (optional) — the temperature is evolved together with the
  abundances (KROME-grade heating/cooling: cosmic-ray + grain photoelectric + H2
  formation heating; Lyman-α + [C II] 158 µm + [O I] 63 µm + Glover-Abel H2 +
  Neufeld-Kaufman CO cooling) and written back to the pressure field
  (_thermochemistry.py).
• Config/params (ChemistryConfig/ChemistryParams), a gated hook in
  _iteration_level_updates, and a builder build_chemistry_from_network_file.

Gated on finite-volume mode; a no-op unless chemistry=True. The carbox/Diffrax
imports are lazy, so import astronomix is unaffected when chemistry is unused.

Core fix (MIT code): FV MHD now composes with trailing state tracers

One commit touches MIT core code — _finite_volume/_state_evolution/evolve_state.py —
so flagging it explicitly. The FV MHD evolution split the magnetic field off the
primitive state with state[-3:], assuming B is the last three state indices. That
holds only for pure MHD: the registry appends advected tracers (chemistry species,
cosmic-ray density, wind density) after the magnetic block, so with any of them
enabled [-3:] sliced tracer rows instead of B — silently corrupting the magnetic
update (no error raised). It now slices B out of its actual slot via magnetic_index
and reinserts it there; magnetic_update is unchanged (it already takes the gas
sub-state and B as separate args). For the no-tracer case the new slicing is
bit-identical to the old, so pure MHD is unaffected. This also fixes MHD + cosmic
rays and MHD + wind, independent of chemistry.

Verified with a 32³ FV MHD box run plain vs. with the (passive) chemistry block: the
magnetic field is bit-identical between the two while the species react — before the
fix the fields diverged at O(1).

Dependencies / data

• carbox (GPL) required at runtime for the reaction network.
• CO cooling needs a Neufeld & Kaufman table (KROME's coolCO.dat, GPL); it is
  not bundled — the caller passes co_cooling_table_path, else co_cooling
  must stay off.

Validation (done during development)

• chemistry-OFF regression unchanged; chemistry-ON end-to-end on GPU (Pallas) and
  native CPU; species advect + react, non-negative and finite.
• Differentiable: finite jax.grad of final state w.r.t. cosmic_ray_rate.
• 0-D benchmark vs KROME: cooling matches to ~1 % where H2 dominates
  (verbatim Glover-Abel port); the benchmark also caught + fixed a bad
  H2-formation heating term.
• 3D driven turbulence (32³/64³): with cooling the box reaches a realistic
  cold, compressible state (~40 K, broadened density PDF) vs ~240 K without —
  the expected thermostat effect.

Known limitations / open items

• Finite-volume only (FD registry not yet handled).
• Stand-in code↔CGS unit factors; fixed cr/FUV/Av in the current helper.
• Fine-structure collision rates are representative literature fits, not KROME's
  tabulated data; collisional H2-dissociation cooling not included.
• No pytests/ for the full engine yet (the MIT mechanism has tests — see the
  companion feature/chemistry-species-hook).

Alternatives to full-GPL-in (if you'd rather keep astronomix MIT)

1. Mechanism-only in core + GPL companion — merge just the species/
   source-term hook (branch feature/chemistry-species-hook, MIT, with tests)
   and keep this engine as a separate GPL package that plugs into the hook.
2. Clean-room — reimplement the cooling from the primary papers (not KROME
   source) so it can be MIT, leaving carbox an optional dependency.

🤖 Generated with Claude Code

https://claude.ai/code/session_01UhwF9zJfmnzuLeABHu75QW