Merge pull request yt-project#4739 from jzuhone/gamer_cr

Support for cosmic ray fields in GAMER

yt/frontends/gamer/cfields.pyx

@@ -7,7 +7,7 @@ cimport numpy as np
 import numpy as np
 
 
-cdef np.float64_t h_eos4(np.float64_t kT, np.float64_t g) noexcept nogil:
+cdef np.float64_t h_eos_tb(np.float64_t kT, np.float64_t g) noexcept nogil:
     cdef np.float64_t x
     x = 2.25 * kT * kT
     return 2.5 * kT + x / (1.0 + math.sqrt(x + 1.0))
@@ -18,16 +18,16 @@ cdef np.float64_t h_eos(np.float64_t kT, np.float64_t g) noexcept nogil:
 cdef np.float64_t gamma_eos(np.float64_t kT, np.float64_t g) noexcept nogil:
     return g
 
-cdef np.float64_t gamma_eos4(np.float64_t kT, np.float64_t g) noexcept nogil:
+cdef np.float64_t gamma_eos_tb(np.float64_t kT, np.float64_t g) noexcept nogil:
     cdef np.float64_t x, c_p, c_v
     x = 2.25 * kT / math.sqrt(2.25 * kT * kT + 1.0)
     c_p = 2.5 + x
     c_v = 1.5 + x
     return c_p / c_v
 
-cdef np.float64_t cs_eos4(np.float64_t kT, np.float64_t c, np.float64_t g) noexcept nogil:
+cdef np.float64_t cs_eos_tb(np.float64_t kT, np.float64_t c, np.float64_t g) noexcept nogil:
     cdef np.float64_t hp, cs2
-    hp = h_eos4(kT, 0.0) + 1.0
+    hp = h_eos_tb(kT, 0.0) + 1.0
     cs2 = kT / (3.0 * hp)
     cs2 *= (5.0 * hp - 8.0 * kT) / (hp - kT)
     return c * math.sqrt(cs2)
@@ -52,14 +52,14 @@ cdef class SRHDFields:
         self._c = clight
         self._c2 = clight * clight
         # Select aux functions based on eos no.
-        if (eos == 4):
-            self.h = h_eos4
-            self.gamma = gamma_eos4
-            self.cs = cs_eos4
-        else:
+        if eos == 1:
             self.h = h_eos
             self.gamma = gamma_eos
             self.cs = cs_eos
+        else:
+            self.h = h_eos_tb
+            self.gamma = gamma_eos_tb
+            self.cs = cs_eos_tb
 
     cdef np.float64_t _lorentz_factor(
             self,

yt/frontends/gamer/data_structures.py

@@ -356,6 +356,7 @@ def _parse_parameter_file(self):
         # make aliases to some frequently used variables
         if parameters["Model"] == "Hydro":
             self.gamma = parameters["Gamma"]
+            self.gamma_cr = self.parameters.get("CR_Gamma", None)
             self.eos = parameters.get("EoS", 1)  # Assume gamma-law by default
             # default to 0.6 for old data format
             self.mu = parameters.get(

yt/frontends/gamer/fields.py

@@ -23,6 +23,7 @@ class GAMERFieldInfo(FieldInfoContainer):
         ("MomY", (mom_units, ["momentum_density_y"], None)),
         ("MomZ", (mom_units, ["momentum_density_z"], None)),
         ("Engy", (erg_units, ["total_energy_density"], None)),
+        ("CRay", (erg_units, ["cosmic_ray_energy_density"], None)),
         ("Pote", (pot_units, ["gravitational_potential"], None)),
         # MHD fields on disk (CC=cell-centered)
         ("CCMagX", (b_units, [], "B_x")),
@@ -61,10 +62,12 @@ def setup_fluid_fields(self):
         if self.ds.srhd:
             c2 = pc.clight * pc.clight
             c = pc.clight.in_units("code_length / code_time")
-            if self.ds.eos == 4:
-                fgen = SRHDFields(self.ds.eos, 0.0, c.d)
-            else:
+            if self.ds.eos == 1:
+                # gamma-law EOS
                 fgen = SRHDFields(self.ds.eos, self.ds.gamma, c.d)
+            else:
+                # Taub-Mathews EOS
+                fgen = SRHDFields(self.ds.eos, 0.0, c.d)
 
             def _sound_speed(field, data):
                 out = fgen.sound_speed(data["gamer", "Temp"].d)
@@ -106,15 +109,22 @@ def _four_velocity(field, data):
                 )
 
             # lorentz factor
-            def _lorentz_factor(field, data):
-                out = fgen.lorentz_factor(
-                    data["gamer", "Dens"].d,
-                    data["gamer", "MomX"].d,
-                    data["gamer", "MomY"].d,
-                    data["gamer", "MomZ"].d,
-                    data["gamer", "Temp"].d,
-                )
-                return data.ds.arr(out, "dimensionless")
+            if ("gamer", "Lrtz") in self.field_list:
+
+                def _lorentz_factor(field, data):
+                    return data["gamer", "Lrtz"]
+
+            else:
+
+                def _lorentz_factor(field, data):
+                    out = fgen.lorentz_factor(
+                        data["gamer", "Dens"].d,
+                        data["gamer", "MomX"].d,
+                        data["gamer", "MomY"].d,
+                        data["gamer", "MomZ"].d,
+                        data["gamer", "Temp"].d,
+                    )
+                    return data.ds.arr(out, "dimensionless")
 
             self.add_field(
                 ("gas", "lorentz_factor"),
@@ -125,16 +135,27 @@ def _lorentz_factor(field, data):
 
             # velocity
             def velocity_xyz(v):
-                def _velocity(field, data):
-                    out = fgen.velocity_xyz(
-                        data["gamer", "Dens"].d,
-                        data["gamer", "MomX"].d,
-                        data["gamer", "MomY"].d,
-                        data["gamer", "MomZ"].d,
-                        data["gamer", "Temp"].d,
-                        data["gamer", f"Mom{v.upper()}"].d,
-                    )
-                    return data.ds.arr(out, "code_velocity").to(unit_system["velocity"])
+                if ("gamer", f"Vel{v.upper()}") in self.field_list:
+
+                    def _velocity(field, data):
+                        return data.ds.arr(data["gamer", f"Vel{v.upper()}"].d, "c").to(
+                            unit_system["velocity"]
+                        )
+
+                else:
+
+                    def _velocity(field, data):
+                        out = fgen.velocity_xyz(
+                            data["gamer", "Dens"].d,
+                            data["gamer", "MomX"].d,
+                            data["gamer", "MomY"].d,
+                            data["gamer", "MomZ"].d,
+                            data["gamer", "Temp"].d,
+                            data["gamer", f"Mom{v.upper()}"].d,
+                        )
+                        return data.ds.arr(out, "code_velocity").to(
+                            unit_system["velocity"]
+                        )
 
                 return _velocity
 
@@ -289,6 +310,9 @@ def et(data):
                 if self.ds.mhd:
                     # magnetic_energy is a yt internal field
                     Et -= data["gas", "magnetic_energy_density"]
+                if getattr(self.ds, "gamma_cr", None):
+                    # cosmic rays are included in this dataset
+                    Et -= data["gas", "cosmic_ray_energy_density"]
                 return Et
 
             # thermal energy per mass (i.e., specific)
@@ -334,6 +358,20 @@ def _thermal_energy_density(field, data):
             units=unit_system["pressure"],
         )
 
+        if getattr(self.ds, "gamma_cr", None):
+
+            def _cr_pressure(field, data):
+                return (data.ds.gamma_cr - 1.0) * data[
+                    "gas", "cosmic_ray_energy_density"
+                ]
+
+            self.add_field(
+                ("gas", "cosmic_ray_pressure"),
+                _cr_pressure,
+                sampling_type="cell",
+                units=self.ds.unit_system["pressure"],
+            )
+
         # mean molecular weight
         if hasattr(self.ds, "mu"):
 

yt/frontends/gamer/tests/test_outputs.py

@@ -125,3 +125,20 @@ def test_stress_energy():
                 else:
                     Tmunu += p
                 assert_array_almost_equal(sp[f"T{mu}{nu}"], Tmunu)
+
+
+cr_shock = "CRShockTube/Data_000005"
+
+
+@requires_ds(cr_shock)
+def test_cosmic_rays():
+    ds = data_dir_load(cr_shock)
+    assert_array_almost_equal(ds.gamma_cr, 4.0 / 3.0)
+    ad = ds.all_data()
+    p_cr = ad["gas", "cosmic_ray_pressure"]
+    e_cr = ad["gas", "cosmic_ray_energy_density"]
+    assert_array_almost_equal(p_cr, e_cr / 3.0)
+    e_kin = ad["gas", "kinetic_energy_density"]
+    e_int = ad["gas", "thermal_energy_density"]
+    e_tot = ad["gas", "total_energy_density"]
+    assert_array_almost_equal(e_tot, e_kin + e_int + e_cr)