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

astronomix/_modules/_chemistry/__init__.py

@@ -0,0 +1,11 @@
+"""
+Advected chemical species with a pluggable reaction source term.
+
+astronomix can carry chemical-species number densities as extra scalar fields in
+the state array; in finite-volume mode they advect with the flow for free (like
+the cosmic-ray / wind-density tracers). Once per hydro step an optional,
+user-supplied source term is applied to them (and, if it chooses, the energy
+field) as an operator-split update. astronomix provides only the mechanism —
+registration, advection and the hook — so no particular chemistry engine is
+baked into the core.
+"""

astronomix/_modules/_chemistry/chemistry_options.py

@@ -0,0 +1,54 @@
+"""
+Configuration and parameter containers for advected chemical species.
+
+Carries the number and ordering of the species tracked in the state array and an
+optional reaction ``source_term`` — a user-supplied operator-split update applied
+once per hydro step. astronomix does not implement any chemistry itself; it only
+registers the species (so they advect) and calls the source term if one is given.
+This keeps the core free of any specific chemistry engine or its dependencies.
+"""
+
+# typing
+from typing import Callable, NamedTuple, Tuple, Union
+from types import NoneType
+from jaxtyping import PyTree
+
+
+class ChemistryConfig(NamedTuple):
+    """Static configuration for advected chemical species.
+
+    Every field is hashable, since the configuration is passed as a static
+    argument to the jitted update; the ``source_term`` is a plain callable.
+
+    Attributes:
+        chemistry: Master switch for carrying/advecting species and applying the
+            source term.
+        number_of_chemical_species: Number of species carried in the state array
+            (a contiguous, advected scalar block).
+        species_names: Species ordering — labels for the state block.
+        source_term: Optional operator-split update applied once per hydro step,
+            with signature
+            ``source_term(primitive_state, registered_variables, chemistry_config,
+            chemistry_params, dt) -> primitive_state``.
+            ``None`` leaves the species advected but chemically inert. The update
+            is entirely user-defined (e.g. a reaction network plus heating /
+            cooling); astronomix only invokes it.
+    """
+
+    chemistry: bool = False
+    number_of_chemical_species: int = 0
+    species_names: Tuple[str, ...] = ()
+    source_term: Union[Callable, NoneType] = None
+
+
+class ChemistryParams(NamedTuple):
+    """Runtime parameters forwarded, uninterpreted, to the reaction source term.
+
+    Attributes:
+        source_term_params: An opaque pytree the source term needs (e.g. rate
+            data, unit-conversion factors, a network object). astronomix does not
+            read it — it is threaded straight through to ``source_term`` so the
+            update stays differentiable in whatever it contains.
+    """
+
+    source_term_params: Union[PyTree, NoneType] = None

astronomix/_modules/_iteration_level_updates.py

@@ -145,6 +145,24 @@ def _iteration_level_updates(
             dt,
         )
 
+    # Chemical-species reaction source term (operator split). The species advect
+    # with the flow (they live in the state array); this hook applies a
+    # user-supplied per-cell update to them — e.g. a reaction network and its
+    # heating/cooling. astronomix only invokes the callable; the physics (and any
+    # extra dependencies) live entirely in the user's ``source_term``. FV only.
+    if (
+        config.chemistry_config.chemistry
+        and config.solver_mode == FINITE_VOLUME
+        and config.chemistry_config.source_term is not None
+    ):
+        primitive_state = config.chemistry_config.source_term(
+            primitive_state,
+            registered_variables,
+            config.chemistry_config,
+            params.chemistry_params,
+            dt,
+        )
+
     # Neural-network body force.
     if config.neural_net_force_config.neural_net_force:
         primitive_state = _neural_net_force(

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()
 
+    #: Configuration for advected chemical species and their reaction source term.
+    chemistry_config: ChemistryConfig = ChemistryConfig()
+
     #: Frame tracking in z-direction
     #: shifting the frame to follow a
     #: turbulent radiative mixing layer

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()
 
+    #: Runtime parameters forwarded to the chemistry reaction source term.
+    chemistry_params: ChemistryParams = ChemistryParams()
+
     #: The parameters of the neural network force module.
     neural_net_force_params: NeuralNetForceParams = NeuralNetForceParams()
 

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 ``species_names``).
+    # 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:
 

pytests/chemistry/test_species_hook.py

@@ -0,0 +1,155 @@
+"""
+Tests for the advected chemical-species block and the reaction source-term hook.
+
+These exercise the mechanism only (no chemistry engine): that species register
+as a contiguous state block, advect with the finite-volume flow, and that a
+user-supplied ``source_term`` is invoked once per step and can modify them.
+"""
+
+# ==== GPU selection (repo convention) ====
+from autocvd import autocvd
+
+# least_used=True picks the least-loaded GPU immediately rather than blocking
+# until one is fully idle, so the test does not hang on a shared machine.
+autocvd(num_gpus=1, least_used=True)
+# ruff: noqa: E402
+# =========================================
+
+import jax
+
+# astronomix integrates in double precision; enable it up front so the initial
+# state and the time-loop carry share a dtype.
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp
+
+from astronomix.data_classes.simulation_helper_data import get_helper_data
+from astronomix.initial_condition_generation.construct_primitive_state import (
+    construct_primitive_state,
+)
+from astronomix.option_classes.simulation_config import (
+    FINITE_VOLUME,
+    NATIVE_JAX,
+    PERIODIC_BOUNDARY,
+    BoundarySettings1D,
+    SimulationConfig,
+    finalize_config,
+)
+from astronomix.option_classes.simulation_params import SimulationParams
+from astronomix.time_stepping.time_integration import time_integration
+from astronomix.variable_registry.registered_variables import get_registered_variables
+from astronomix._modules._chemistry.chemistry_options import ChemistryConfig
+
+NUM_CELLS = 64
+ADVECTION_VELOCITY = 0.5
+END_TIME = 0.2
+
+
+def _run(source_term):
+    """Run a 1-D periodic box that advects a single species bump at a uniform
+    velocity, with an optional reaction source term. Returns (positions, initial
+    species, final species)."""
+    config = SimulationConfig(
+        solver_mode=FINITE_VOLUME,
+        backend=NATIVE_JAX,
+        dimensionality=1,
+        num_cells=NUM_CELLS,
+        box_size=1.0,
+        boundary_settings=BoundarySettings1D(PERIODIC_BOUNDARY, PERIODIC_BOUNDARY),
+        progress_bar=False,
+        chemistry_config=ChemistryConfig(
+            chemistry=True,
+            number_of_chemical_species=1,
+            species_names=("tracer",),
+            source_term=source_term,
+        ),
+    )
+    params = SimulationParams(t_end=END_TIME, C_cfl=0.4)
+    registered_variables = get_registered_variables(config)
+
+    position = get_helper_data(config).geometric_centers
+    density = jnp.ones_like(position)
+    velocity = jnp.full_like(position, ADVECTION_VELOCITY)
+    pressure = jnp.ones_like(position)
+    state = construct_primitive_state(
+        config=config,
+        registered_variables=registered_variables,
+        density=density,
+        velocity_x=velocity,
+        gas_pressure=pressure,
+    )
+
+    # A smooth interior bump in the (single) species field.
+    bump = jnp.exp(-(((position - 0.3) / 0.08) ** 2)) + 1e-6
+    species_index = registered_variables.chemistry_species_index
+    state = state.at[species_index].set(bump)
+
+    config = finalize_config(config, state.shape)
+    final_state = time_integration(state, config, params, registered_variables)
+    return position, state[species_index], final_state[species_index]
+
+
+def test_species_block_registration():
+    """The species block is appended to the state only when chemistry is on."""
+    base = SimulationConfig(solver_mode=FINITE_VOLUME, dimensionality=1, num_cells=16)
+    registered_off = get_registered_variables(base)
+    core_vars = registered_off.num_vars
+
+    assert not registered_off.chemistry_species_active
+
+    with_chemistry = base._replace(
+        chemistry_config=ChemistryConfig(
+            chemistry=True,
+            number_of_chemical_species=4,
+            species_names=("a", "b", "c", "d"),
+        )
+    )
+    registered_on = get_registered_variables(with_chemistry)
+
+    assert registered_on.chemistry_species_active
+    assert registered_on.num_chemical_species == 4
+    # the block is contiguous and appended after the core variables
+    assert registered_on.chemistry_species_index == core_vars
+    assert registered_on.num_vars == core_vars + 4
+
+
+def test_species_advects_with_the_flow():
+    """With no source term, the species bump is transported at the flow speed and
+    its total mass is conserved."""
+    position, initial_species, final_species = _run(source_term=None)
+
+    def center_of_mass(field):
+        return jnp.sum(position * field) / jnp.sum(field)
+
+    shift = float(center_of_mass(final_species) - center_of_mass(initial_species))
+    expected = ADVECTION_VELOCITY * END_TIME
+
+    assert abs(shift - expected) < 0.03, f"bump advected by {shift}, expected {expected}"
+
+    initial_mass = float(jnp.sum(initial_species))
+    final_mass = float(jnp.sum(final_species))
+    assert abs(final_mass - initial_mass) / initial_mass < 0.02
+
+    assert bool(jnp.all(final_species >= 0.0))
+    assert bool(jnp.all(jnp.isfinite(final_species)))
+
+
+def test_source_term_hook_is_invoked():
+    """A source term that scales the species down each step must leave the final
+    mass well below the advection-only (no-source) case, proving the hook fires."""
+
+    def scaling_source(primitive_state, registered_variables, chemistry_config, chemistry_params, dt):
+        start = registered_variables.chemistry_species_index
+        count = registered_variables.num_chemical_species
+        return primitive_state.at[start : start + count].multiply(0.5)
+
+    _, _, final_no_source = _run(source_term=None)
+    _, _, final_with_source = _run(source_term=scaling_source)
+
+    mass_no_source = float(jnp.sum(final_no_source))
+    mass_with_source = float(jnp.sum(final_with_source))
+
+    assert mass_with_source < 0.9 * mass_no_source, (
+        f"source term had no effect: {mass_with_source} vs {mass_no_source}"
+    )
+    assert bool(jnp.all(jnp.isfinite(final_with_source)))