Fixing the 4-component version of GST

src/simweights/__init__.py

@@ -28,6 +28,7 @@
     "GenieSurface",
     "GenieWeighter",
     "GlobalFitGST",
+    "GlobalFitGST4Comp",
     "GlobalFitGST_IT",
     "GlobalSplineFit",
     "GlobalSplineFit5Comp",
@@ -57,6 +58,7 @@
     GaisserH4a_IT,
     GaisserHillas,
     GlobalFitGST,
+    GlobalFitGST4Comp,
     GlobalFitGST_IT,
     GlobalSplineFit,
     GlobalSplineFit5Comp,

src/simweights/_fluxes.py

@@ -87,7 +87,7 @@ def _condition(
     def __call__(self: CosmicRayFlux, energy: ArrayLike, pdgid: ArrayLike) -> NDArray[float64]:
         energy_arr, pdgid_arr = broadcast_arrays(energy, pdgid)
         pcond = self._condition(energy_arr, pdgid_arr)
-        return piecewise(energy, pcond, self._funcs)
+        return piecewise(energy, pcond, self._funcs)  # type: ignore[no-any-return,call-overload]
 
 
 class Hoerandel(CosmicRayFlux):
@@ -348,6 +348,39 @@ class GlobalFitGST_IT(CosmicRayFlux):  # pylint: disable=invalid-name
     )
 
 
+class GlobalFitGST4Comp(CosmicRayFlux):  # pylint: disable=invalid-name
+    r"""GlobalFitGST4Comp (4 populations) for four components [p, He, O, Fe].
+
+    The Oxygen group is the sum of Carbon and Oxygen groups.
+    The Iron group is the sum of Iron, Tellurium and Mercury groups.
+    Provided by J. Saffer.
+    """
+
+    pdgids = PDGID_4COMP
+    _funcs = (
+        lambda E: (
+            0.7000 * E**-2.66 * exp(-E / 1.2e5)
+            + 0.0150 * E**-2.4 * exp(-E / 4e6)
+            + 0.0012 * E**-2.4 * exp(-E / 1.5e9)
+            + 0.00012 * E**-2.4 * exp(-E / 40e9)
+        ),
+        lambda E: 0.3200 * E**-2.58 * exp(-E / 1.2e5 / 2) + 0.0065 * E**-2.3 * exp(-E / 4e6 / 2),
+        lambda E: (
+            0.0100 * E**-2.40 * exp(-E / 1.2e5 / 6)
+            + 0.0006 * E**-2.3 * exp(-E / 4e6 / 6)
+            + 0.0130 * E**-2.40 * exp(-E / 1.2e5 / 8)
+            + 0.0007 * E**-2.3 * exp(-E / 4e6 / 8)
+        ),
+        lambda E: (
+            0.0060 * E**-2.30 * exp(-E / 1.2e5 / 26)
+            + 0.00021 * E**-2.2 * exp(-E / 4e6 / 26)
+            + 0.0000011 * E**-2.2 * exp(-E / 1.5e9 / 26)
+            + 0.00001 * E**-2.2 * exp(-E / 4e6 / 52)
+            + 0.000053 * E**-2.2 * exp(-E / 4e6 / 80)
+        ),
+    )
+
+
 class GlobalSplineFitBase(CosmicRayFlux):
     r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]_.
 

tests/test_fluxes.py

@@ -29,6 +29,7 @@
     simweights.TIG1996(),
     simweights.GlobalFitGST(),
     simweights.GlobalFitGST_IT(),
+    simweights.GlobalFitGST4Comp(),
     simweights.GlobalSplineFit(),
     simweights.GlobalSplineFit5Comp(),
     simweights.GlobalSplineFit_IT(),
@@ -39,12 +40,6 @@
 
 @pytest.mark.parametrize("flux", flux_models, ids=[x.__class__.__name__ for x in flux_models])
 def test_flux_model(flux, ndarrays_regression):
-    # this is the old regression test it can stick around for a bit but will be deleted at a certain point
-    for pdgid in flux.pdgids:
-        v1 = flux(E, pdgid)
-        v2 = np.array(flux_values[flux.__class__.__name__][str(int(pdgid))]) / 1e4
-    assert v1 == pytest.approx(v2, rel=1e-13)
-
     ndarrays_regression.check({pdgid.name: flux(E, pdgid) for pdgid in flux.pdgids}, default_tolerance={"rtol": 1e-13})
     # make sure you get zero for non CR primaries
     assert flux(E, 22) == pytest.approx(0)