Add advected chemical species with a pluggable reaction source-term hook

[lorenzobranca]

Add advected chemical species with a pluggable reaction source-term hook

Summary

Adds the mechanism for carrying chemical species in astronomix and reacting
them once per hydro step, without baking any chemistry engine into the core.
Concretely, a simulation can now:

1. carry an arbitrary number of chemical-species number densities as a contiguous
   block of extra scalar fields in the state array, and
2. apply an optional, user-supplied operator-split source term to those
   fields (and, if it chooses, the energy field) once per step.

This is the minimal infrastructure needed to couple a chemical network to the
fluid; the network itself is provided by the user, so astronomix stays free of
any specific chemistry dependency.

Motivation

Astrochemistry-in-hydro (advect species with the flow, react them per cell,
optionally let heating/cooling feed back into the energy) is a common need, and
astronomix already advects extra state slots for free (cosmic-ray n,
wind density). This PR generalises that into a first-class, documented extension
point so a reaction network can be plugged in cleanly — the differentiable JAX
core makes such a coupling attractive for gradient-based / solver-in-the-loop
work.

What this adds

• Species registration (variable_registry/registered_variables.py): a
  contiguous chemistry_species_index … +num_chemical_species block, allocated
  in finite-volume mode exactly like the existing CR/wind tracers. They advect
  with the flow through the existing Euler-flux path — no new advection code.
• ChemistryConfig / ChemistryParams (_modules/_chemistry/): the species
  count/names, an optional source_term callable, and an opaque source_term_params
  pytree threaded straight through to it (uninterpreted, so it stays
  differentiable in whatever it holds).
• A gated hook in _iteration_level_updates that calls the source term once
  per step (finite-volume only), alongside the existing cooling/wind/etc. updates.
• Registration of chemistry_config / chemistry_params on SimulationConfig /
  SimulationParams.

What this deliberately does NOT add

No reaction network, no thermochemistry, no third-party dependency. The
source_term is entirely user-defined:

def my_source_term(primitive_state, registered_variables,
                   chemistry_config, chemistry_params, dt):
    # advance the species block (and optionally the pressure) however you like,
    # e.g. integrate a reaction network per cell; return the updated state.
    return primitive_state

config = SimulationConfig(
    solver_mode=FINITE_VOLUME,
    chemistry_config=ChemistryConfig(
        chemistry=True,
        number_of_chemical_species=16,
        species_names=(...),
        source_term=my_source_term,
    ),
)
params = SimulationParams(
    chemistry_params=ChemistryParams(source_term_params=my_rate_data),
)

A working source_term built on the carbox reaction network + a KROME-grade
thermochemistry (heating/cooling with feedback into pressure) exists as a
separate companion package — kept out of core deliberately, because it (and
its carbox/KROME-derived physics) is GPL, whereas astronomix is MIT. This PR
keeps the MIT core dependency-free and the GPL engine external.

Verification

• import astronomix unaffected (no new imports at module load; the hook is a
  no-op unless chemistry=True and a source_term is supplied).
• New tests in pytests/chemistry/test_species_hook.py (all passing, ~8 s),
  mechanism-only (no chemistry engine):
  • registration — the species block is appended to the state iff chemistry
    is on, contiguous, with num_vars growing by the species count;
  • advection — with no source term a species bump is transported at the flow
    speed (centre-of-mass shift ≈ u·t) and its total mass is conserved;
  • source-term hook — a scaling source term drives the final species mass
    well below the advection-only case, confirming the hook fires each step.
• The diff is small (+274 lines, 7 files incl. tests).

Open items / for discussion

• Finite-volume only for now (species allocation is FV-gated, matching the
  CR/wind tracers); the finite-difference registry would need explicit handling.
• The species advect as passive volumetric densities via _euler_flux; that
  matches the existing tracers, but let me know if you'd prefer a different
  normalisation convention.
• The tests use autocvd(num_gpus=1, least_used=True) so they pick a GPU
  immediately rather than blocking on a fully-idle one — happy to match the
  plain autocvd(num_gpus=1) convention if you prefer.

🤖 Generated with Claude Code

https://claude.ai/code/session_01UhwF9zJfmnzuLeABHu75QW