From 1aeb427e7e5fbc8130422ff6de63f8b4e65a3c31 Mon Sep 17 00:00:00 2001
From: Corentin Cadiou <corentin.cadiou@cphyc.me>
Date: Tue, 5 Mar 2024 13:42:22 +0100
Subject: [PATCH] Add support for self-shielding in heating/cooling rates
 calculation

---
 yt/frontends/ramses/data_structures.py | 57 ++++++++++++++++++--------
 yt/frontends/ramses/fields.py          | 11 +++--
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/yt/frontends/ramses/data_structures.py b/yt/frontends/ramses/data_structures.py
index d037fd772a5..43576b63341 100644
--- a/yt/frontends/ramses/data_structures.py
+++ b/yt/frontends/ramses/data_structures.py
@@ -795,6 +795,12 @@ class RAMSESDataset(Dataset):
     _field_info_class = RAMSESFieldInfo
     gamma = 1.4  # This will get replaced on hydro_fn open
 
+    # RAMSES-specific parameters
+    force_cosmological: Optional[bool]
+    _force_max_level: tuple[int, str]
+    _bbox: Optional[list[list[float]]]
+    self_shielding: Optional[bool]
+
     def __init__(
         self,
         filename,
@@ -809,6 +815,7 @@ def __init__(
         max_level=None,
         max_level_convention=None,
         default_species_fields=None,
+        self_shielding=None,
     ):
         # Here we want to initiate a traceback, if the reader is not built.
         if isinstance(fields, str):
@@ -880,6 +887,7 @@ def __init__(
                 self.fluid_types += (FH.ftype,)
 
         self.storage_filename = storage_filename
+        self.self_shielding = self_shielding
 
     @staticmethod
     def _sanitize_max_level(max_level, max_level_convention):
@@ -1111,27 +1119,42 @@ def caster(val):
             self.group_size = rheader["ncpu"] // self.num_groups
 
         # Read namelist.txt file (if any)
-        self.read_namelist()
+        has_namelist = self.read_namelist()
+
+        if self.self_shielding is None and has_namelist:
+            nml = self.ds.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.self_shielding = physics_params.get("self_shielding", False)
+            self.self_shielding |= cooling_params.get("self_shielding", False)
 
-    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 e:
+            nml = f"An error occurred when reading the namelist: {str(e)}"
+            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:
diff --git a/yt/frontends/ramses/fields.py b/yt/frontends/ramses/fields.py
index 2aee2c1719f..86c5a25b5da 100644
--- a/yt/frontends/ramses/fields.py
+++ b/yt/frontends/ramses/fields.py
@@ -387,10 +387,15 @@ def create_cooling_fields(self):
         # Function to create the cooling fields
         def _create_field(name, interp_object, unit):
             def _func(field, data):
-                shape = data[("gas", "temperature")].shape
+                shape = data["gas", "temperature"].shape
+                nH = _X * data["gas", "density"] / mh
+                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(),
-                    "logT": np.log10(data[("gas", "temperature")]).ravel(),
+                    "lognH": np.log10(nH * boost).ravel(),
+                    "logT": np.log10(data["gas", "temperature"]).ravel(),
                 }
                 rv = interp_object(d).reshape(shape)
                 if name[-1] != "mu":