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

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.")
 

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.

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.
+"""