[ENH] Add support for self-shielding in heating/cooling rates calculation for RAMSES

yt/frontends/ramses/data_structures.py

@@ -790,6 +790,12 @@ class RAMSESDataset(Dataset):
     _field_info_class = RAMSESFieldInfo
     gamma = 1.4  # This will get replaced on hydro_fn open
 
+    # RAMSES-specific parameters
+    force_cosmological: bool | None
+    _force_max_level: tuple[int, str]
+    _bbox: list[list[float]] | None
+    _self_shielding: bool | None = None
+
     def __init__(
         self,
         filename,
@@ -804,6 +810,7 @@ def __init__(
         max_level=None,
         max_level_convention=None,
         default_species_fields=None,
+        self_shielding=None,
         use_conformal_time=None,
     ):
         # Here we want to initiate a traceback, if the reader is not built.
@@ -821,6 +828,10 @@ def __init__(
         cosmological:
         If set to None, automatically detect cosmological simulation.
         If a boolean, force its value.
+
+        self_shielding:
+        If set to True, assume gas is self-shielded above 0.01 mp/cm^3.
+        This affects the fields related to cooling and the mean molecular weight.
         """
 
         self._fields_in_file = fields
@@ -887,6 +898,37 @@ def __init__(
 
         self.storage_filename = storage_filename
 
+        self.self_shielding = self_shielding
+
+    @property
+    def self_shielding(self) -> bool:
+        if self._self_shielding is not None:
+            return self._self_shielding
+
+        # Read namelist.txt file (if any)
+        has_namelist = self.read_namelist()
+
+        if not has_namelist:
+            self._self_shielding = False
+            return self._self_shielding
+
+        nml = self.parameters["namelist"]
+
+        # "self_shielding" is stored in physics_params in older versions of the code
+        physics_params = nml.get("physics_params", default={})
+        # and in "cooling_params" in more recent ones
+        cooling_params = nml.get("cooling_params", default={})
+
+        self_shielding = physics_params.get("self_shielding", False)
+        self_shielding |= cooling_params.get("self_shielding", False)
+
+        self._self_shielding = self_shielding
+        return self_shielding
+
+    @self_shielding.setter
+    def self_shielding(self, value):
+        self._self_shielding = value
+
     @staticmethod
     def _sanitize_max_level(max_level, max_level_convention):
         # NOTE: the initialisation of the dataset class sets
@@ -1091,28 +1133,34 @@ def caster(val):
         if self.num_groups > 0:
             self.group_size = rheader["ncpu"] // self.num_groups
 
-        # Read namelist.txt file (if any)
         self.read_namelist()
 
-    def read_namelist(self):
+    def read_namelist(self) -> bool:
         """Read the namelist.txt file in the output folder, if present"""
         namelist_file = os.path.join(self.root_folder, "namelist.txt")
-        if os.path.exists(namelist_file):
-            try:
-                with open(namelist_file) as f:
-                    nml = f90nml.read(f)
-            except ImportError as e:
-                nml = f"An error occurred when reading the namelist: {str(e)}"
-            except (ValueError, StopIteration, AssertionError) as err:
-                # Note: f90nml may raise a StopIteration, a ValueError or an AssertionError if
-                # the namelist is not valid.
-                mylog.warning(
-                    "Could not parse `namelist.txt` file as it was malformed:",
-                    exc_info=err,
-                )
-                return
+        if not os.path.exists(namelist_file):
+            return False
+
+        try:
+            with open(namelist_file) as f:
+                nml = f90nml.read(f)
+        except ImportError as err:
+            mylog.warning(
+                "`namelist.txt` file found but missing package f90nml to read it:",
+                exc_info=err,
+            )
+            return False
+        except (ValueError, StopIteration, AssertionError) as err:
+            # Note: f90nml may raise a StopIteration, a ValueError or an AssertionError if
+            # the namelist is not valid.
+            mylog.warning(
+                "Could not parse `namelist.txt` file as it was malformed:",
+                exc_info=err,
+            )
+            return False
 
-            self.parameters["namelist"] = nml
+        self.parameters["namelist"] = nml
+        return True
 
     @classmethod
     def _is_valid(cls, filename: str, *args, **kwargs) -> bool:

yt/frontends/ramses/fields.py

@@ -436,8 +436,16 @@ def create_cooling_fields(self) -> bool:
         def _create_field(name, interp_object, unit):
             def _func(field, data):
                 shape = data["gas", "temperature_over_mu"].shape
+                # Ramses assumes a fraction X of Hydrogen within the non-metal gas.
+                # It has to be corrected by metallicity.
+                Z = data["gas", "metallicity"]
+                nH = ((1 - _Y) * (1 - Z) * data["gas", "density"] / mh).to("cm**-3")
+                if data.ds.self_shielding:
+                    boost = np.maximum(np.exp(-nH / 0.01), 1e-20)
+                else:
+                    boost = 1
                 d = {
-                    "lognH": np.log10(_X * data["gas", "density"] / mh).ravel(),
+                    "lognH": np.log10(nH / boost).ravel(),
                     "logT": np.log10(data["gas", "temperature_over_mu"]).ravel(),
                 }
                 rv = interp_object(d).reshape(shape)

yt/frontends/ramses/tests/test_outputs.py

@@ -1,4 +1,8 @@
 import os
+from contextlib import contextmanager
+from pathlib import Path
+from shutil import copytree
+from tempfile import TemporaryDirectory
 
 import numpy as np
 from numpy.testing import assert_equal, assert_raises
@@ -685,3 +689,108 @@ def _dummy_field(field, data):
     v1 = ad["gas", "test"]
 
     np.testing.assert_allclose(v0, v1)
+
+
+# Simple context manager that overrides the value of
+# the self_shielding flag in the namelist.txt file
+@contextmanager
+def override_self_shielding(root: Path, section: str, val: bool):
+    # Copy content of root in a temporary folder
+    with TemporaryDirectory() as tmp:
+        tmpdir = Path(tmp) / root.name
+        tmpdir.mkdir()
+
+        # Copy content of `root` in `tmpdir` recursively
+        copytree(root, tmpdir, dirs_exist_ok=True)
+
+        fname = Path(tmpdir) / "namelist.txt"
+
+        with open(fname) as f:
+            namelist_data = f90nml.read(f).todict()
+            sec = namelist_data.get(section, {})
+            sec["self_shielding"] = val
+            namelist_data[section] = sec
+
+        with open(fname, "w") as f:
+            new_nml = f90nml.Namelist(namelist_data)
+            new_nml.write(f)
+
+        yield tmpdir
+
+
+@requires_file(ramses_new_format)
+@requires_module("f90nml")
+def test_self_shielding_logic():
+    ##################################################
+    # Check that the self_shielding flag is correctly read
+    ds = yt.load(ramses_new_format)
+
+    assert ds.parameters["namelist"] is not None
+    assert ds.self_shielding is False
+
+    ##################################################
+    # Modify the namelist in-situ, reload and check
+    root_folder = Path(ds.root_folder)
+    for section in ("physics_params", "cooling_params"):
+        for val in (True, False):
+            with override_self_shielding(root_folder, section, val) as tmp_output:
+                # Autodetection should work
+                ds = yt.load(tmp_output)
+                assert ds.parameters["namelist"] is not None
+                assert ds.self_shielding is val
+
+                # Manually set should ignore the namelist
+                ds = yt.load(tmp_output, self_shielding=True)
+                assert ds.self_shielding is True
+
+                ds = yt.load(tmp_output, self_shielding=False)
+                assert ds.self_shielding is False
+
+    ##################################################
+    # Manually set should work, even if namelist does not contain the flag
+    ds = yt.load(ramses_new_format, self_shielding=True)
+    assert ds.self_shielding is True
+
+    ds = yt.load(ramses_new_format, self_shielding=False)
+    assert ds.self_shielding is False
+
+
+ramses_star_formation = "ramses_star_formation/output_00013"
+
+
+@requires_file(ramses_star_formation)
+@requires_module("f90nml")
+def test_self_shielding_loading():
+    ds_off = yt.load(ramses_star_formation, self_shielding=False)
+    ds_on = yt.load(ramses_star_formation, self_shielding=True)
+
+    # We do not expect any significant change at densities << 1e-2
+    reg_on = ds_on.all_data().cut_region(
+        ["obj['gas', 'density'].to('mp/cm**3') < 1e-6"]
+    )
+    reg_off = ds_off.all_data().cut_region(
+        ["obj['gas', 'density'].to('mp/cm**3') < 1e-6"]
+    )
+
+    np.testing.assert_allclose(
+        reg_on["gas", "cooling_total"],
+        reg_off["gas", "cooling_total"],
+        rtol=1e-5,
+    )
+
+    # We expect large differences from 1e-2 mp/cc
+    reg_on = ds_on.all_data().cut_region(
+        ["obj['gas', 'density'].to('mp/cm**3') > 1e-2"]
+    )
+    reg_off = ds_off.all_data().cut_region(
+        ["obj['gas', 'density'].to('mp/cm**3') > 1e-2"]
+    )
+    # Primordial cooling decreases with density (except at really high temperature)
+    # so we expect the boosted version to cool *less* efficiently
+    diff = (
+        reg_on["gas", "cooling_primordial"] - reg_off["gas", "cooling_primordial"]
+    ) / reg_off["gas", "cooling_primordial"]
+    assert (np.isclose(diff, 0, atol=1e-5) | (diff < 0)).all()
+
+    # Also make sure the difference is large for some cells
+    assert (np.abs(diff) > 0.1).any()