diff --git a/astronomix/_modules/_chemistry/__init__.py b/astronomix/_modules/_chemistry/__init__.py new file mode 100644 index 00000000..9f23731c --- /dev/null +++ b/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. +""" diff --git a/astronomix/_modules/_chemistry/chemistry_options.py b/astronomix/_modules/_chemistry/chemistry_options.py new file mode 100644 index 00000000..091eba38 --- /dev/null +++ b/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 diff --git a/astronomix/_modules/_iteration_level_updates.py b/astronomix/_modules/_iteration_level_updates.py index ed1fa5ba..ac35d3cc 100644 --- a/astronomix/_modules/_iteration_level_updates.py +++ b/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( diff --git a/astronomix/option_classes/simulation_config.py b/astronomix/option_classes/simulation_config.py index c6a99b40..b001ab3a 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() + #: 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 diff --git a/astronomix/option_classes/simulation_params.py b/astronomix/option_classes/simulation_params.py index 6f10b6f1..1e90ad42 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() + #: 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() diff --git a/astronomix/variable_registry/registered_variables.py b/astronomix/variable_registry/registered_variables.py index d0c87ace..ba5cb34a 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 ``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: diff --git a/pytests/chemistry/test_species_hook.py b/pytests/chemistry/test_species_hook.py new file mode 100644 index 00000000..6a8ebbe6 --- /dev/null +++ b/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)))