Update the genie weighter

src/simweights/_genie_weighter.py

@@ -10,7 +10,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector
-from ._utils import Column, Const, get_column, get_table
+from ._utils import Column, Const, get_column, get_table, has_column
 from ._weighter import Weighter
 
 
@@ -36,7 +36,9 @@ def genie_surface(table: Iterable[Mapping[str, float]]) -> GenerationSurface:
         global_probability_scale = get_column(table, "global_probability_scale")[i]
         surfaces.append(
             nevents
-            * generation_surface(pdgid, Const(1 / spatial.etendue / global_probability_scale), Column("wght"), spectrum),
+            * generation_surface(
+                pdgid, Const(1 / spatial.etendue / global_probability_scale), Column("wght"), Column("volscale"), spectrum
+            ),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
@@ -49,12 +51,21 @@ def GenieWeighter(file_obj: Any) -> Weighter:  # noqa: N802
 
     Reads ``I3GenieInfo`` from S-Frames and ``I3GenieResult`` from Q-Frames.
     """
-    weight_table = get_table(file_obj, "I3GenieInfo")
-    surface = genie_surface(weight_table)
+    info_table = get_table(file_obj, "I3GenieInfo")
+    result_table = get_table(file_obj, "I3GenieResult")
+    surface = genie_surface(info_table)
 
     weighter = Weighter([file_obj], surface)
-    weighter.add_weight_column("pdgid", weighter.get_column("I3GenieResult", "neu").astype(np.int32))
-    weighter.add_weight_column("energy", weighter.get_column("I3GenieResult", "Ev"))
-    weighter.add_weight_column("wght", weighter.get_column("I3GenieResult", "wght"))
+    weighter.add_weight_column("pdgid", get_column(result_table, "neu").astype(np.int32))
+    weighter.add_weight_column("energy", get_column(result_table, "Ev"))
+    weighter.add_weight_column("cos_zen", get_column(result_table, "pzv"))
+    weighter.add_weight_column("wght", get_column(result_table, "wght"))
+
+    # Include the effect of the muon scaling introduced in icecube/icetray#3607, if present.
+    if has_column(result_table, "volscale"):
+        volscale = get_column(result_table, "volscale")
+    else:
+        volscale = np.ones_like(get_column(result_table, "wght"))
+    weighter.add_weight_column("volscale", volscale)
 
     return weighter

tests/test_genie_datasets.py

@@ -61,7 +61,7 @@ def test_dataset(fname):
 
         assert w.get_weight_column("wght") == approx(genie_weight)
 
-        power_law = next(iter(w.surface.spectra.values()))[0].dists[2]
+        power_law = next(iter(w.surface.spectra.values()))[0].dists[-1]
         energy_term = 1 / power_law.pdf(w.get_weight_column("energy"))
         assert energy_term == approx(energy_factor)
 

tests/test_genie_weighter.py

@@ -22,51 +22,60 @@
     ("power_law_index", np.float64),
 ]
 
-result_dtype = [("neu", np.int32), ("Ev", np.float64), ("wght", np.float64)]
+result_dtype = [("neu", np.int32), ("pzv", np.float64), ("Ev", np.float64), ("wght", np.float64)]
 
 
-class TestCorsikaWeighter(unittest.TestCase):
-    def test_triggered_corsika(self):
+class TestGenieWeighter(unittest.TestCase):
+    def test_genie(self):
         nevents = 10000
+        coszen = 0.7
         pdgid = 12
         c1 = simweights.CircleInjector(300, 0, 1)
         p1 = simweights.PowerLaw(0, 1e3, 1e4)
 
-        weight = np.zeros(nevents, dtype=result_dtype)
-        weight["neu"] = pdgid
-        weight["Ev"] = p1.ppf(np.linspace(0, 1, nevents))
-
         for event_weight in [1e-6, 1e-3, 1]:
-            weight["wght"] = event_weight
-
             for nfiles in [1, 5, 50]:
-                rows = nfiles * [
-                    (
-                        pdgid,
-                        nevents,
-                        1,
-                        c1.radius,
-                        np.arccos(c1.cos_zen_max),
-                        np.arccos(c1.cos_zen_min),
-                        p1.a,
-                        p1.b,
-                        p1.g,
-                    ),
-                ]
-                info = np.array(rows, dtype=info_dtype)
-                d = {"I3GenieResult": weight, "I3GenieInfo": info}
-
-                for flux in [0.1, 1, 10]:
-                    wobj = simweights.GenieWeighter(d)
-                    w = wobj.get_weights(flux)
-                    np.testing.assert_allclose(
-                        w.sum(),
-                        flux * event_weight * c1.etendue * p1.integral / nfiles,
-                    )
-                    E = d["I3GenieResult"]["Ev"]
-                    y, x = np.histogram(E, weights=w, bins=51, range=[p1.a, p1.b])
-                    Ewidth = np.ediff1d(x)
-                    np.testing.assert_allclose(y, flux * event_weight * Ewidth * c1.etendue / nfiles, 5e-3)
+                for include_volscale in [True, False]:
+                    result_dtype = [("neu", np.int32), ("pzv", np.float64), ("Ev", np.float64), ("wght", np.float64)]
+                    if include_volscale:
+                        result_dtype.append(("volscale", np.float64))
+
+                    weight = np.zeros(nevents, dtype=result_dtype)
+                    weight["neu"] = pdgid
+                    weight["pzv"] = coszen
+                    weight["Ev"] = p1.ppf(np.linspace(0, 1, nevents))
+                    weight["wght"] = event_weight
+
+                    if include_volscale:
+                        weight["volscale"] = 1
+
+                    rows = nfiles * [
+                        (
+                            pdgid,
+                            nevents,
+                            1,
+                            c1.radius,
+                            np.arccos(c1.cos_zen_max),
+                            np.arccos(c1.cos_zen_min),
+                            p1.a,
+                            p1.b,
+                            p1.g,
+                        ),
+                    ]
+                    info = np.array(rows, dtype=info_dtype)
+                    d = {"I3GenieResult": weight, "I3GenieInfo": info}
+
+                    for flux in [0.1, 1, 10]:
+                        wobj = simweights.GenieWeighter(d)
+                        w = wobj.get_weights(flux)
+                        np.testing.assert_allclose(
+                            w.sum(),
+                            flux * event_weight * c1.etendue * p1.integral / nfiles,
+                        )
+                        E = d["I3GenieResult"]["Ev"]
+                        y, x = np.histogram(E, weights=w, bins=51, range=[p1.a, p1.b])
+                        Ewidth = np.ediff1d(x)
+                        np.testing.assert_allclose(y, flux * event_weight * Ewidth * c1.etendue / nfiles, 5e-3)
 
         with self.assertRaises(TypeError):
             simweights.GenieWeighter(d, nfiles=10)