Add tests for genie-icetray files.

examples/basic_genie_icetray.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+
+# SPDX-FileCopyrightText: © 2022 the SimWeights contributors
+#
+# SPDX-License-Identifier: BSD-2-Clause
+import nuflux
+import numpy as np
+import pandas as pd
+
+import simweights
+
+simfile = pd.HDFStore("level2_genie-icetray.140000_000000.hdf")
+flux_model = nuflux.makeFlux("IPhonda2014_spl_solmax")
+weight_obj = simweights.GenieWeighter(simfile, 1)
+weights = weight_obj.get_weights(flux_model)
+print("Rate from simweights", weights.sum(), "Hz")
+simfile.close()

src/simweights/_genie_weighter.py

@@ -2,48 +2,54 @@
 #
 # SPDX-License-Identifier: BSD-2-Clause
 
+from __future__ import annotations
 
 from typing import Any, Iterable, Mapping
 
 import numpy as np
 
 from ._generation_surface import GenerationSurface, generation_surface
+from ._nugen_weighter import nugen_spatial, nugen_spectrum
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector
 from ._utils import Column, Const, constcol, get_column, get_table, has_column, has_table
-from ._nugen_weighter import nugen_spatial, nugen_spectrum
 from ._weighter import Weighter
 
-def genie_icetray_surface(mcweightdict: Iterable[Mapping[str, float]], 
+
+def genie_icetray_surface(mcweightdict: Iterable[Mapping[str, float]],
                           geniedict: Iterable[Mapping[str, float]],
                           nufraction: float = 0.7) -> GenerationSurface:
-    """Inspect the rows of a GENIE-icetray's I3MCWeightDict table object to generate a surface object.
+    """Inspect the rows of a GENIE-icetray"s I3MCWeightDict table object to generate a surface object.
 
     This is a bit of a pain: the oscillations group historically produced 4-5 energy bands with varying
     generation parameters, then merged them all into one "file". Because of this, we need to check the
-    neutrino type, volume, and spectral index to get the correct surfaces. The type weight also isn't 
-    stored in the files: this was fixed to 70/30 for oscillation-produced genie-icetray files.
+    neutrino type, volume, and spectrum to get the correct surfaces. The type weight also isn"t stored
+    in the files: this was fixed to 70/30 for oscillation-produced genie-icetray files.
     """
     pdgid = get_column(geniedict, "neu")
     volume = get_column(mcweightdict, "GeneratorVolume")
     index = get_column(mcweightdict, "PowerLawIndex")
-    unique_schemes = np.unique(np.array([pdgid, volume, index]).T, axis=0)
+    minenergylog = get_column(mcweightdict, "MinEnergyLog")
+    maxenergylog = get_column(mcweightdict, "MaxEnergyLog")
+    gen_schemes = np.array([pdgid, volume, index, minenergylog, maxenergylog]).T
+    unique_schemes = np.unique(gen_schemes, axis=0)
 
     if len(unique_schemes) == 0:
         msg = "`I3MCWeightDict` is empty. SimWeights cannot process this file"
         raise RuntimeError(msg)
 
-    spatial = nugen_spatial(mcweightdict)
-    spectrum = nugen_spectrum(mcweightdict)
-
     surfaces = []
-    for pid, _, _ in unique_schemes:
-        mask = pid == pdgid
+    for row in unique_schemes:
+        (pid, vol, idx, emin, emax) = row
+        mask = np.all(gen_schemes == row[None,:], axis=1)
+
+        spatial = nugen_spatial(mcweightdict[mask])
+        spectrum = nugen_spectrum(mcweightdict[mask])
 
         type_weight = nufraction if pid>0 else 1-nufraction
-        primary_type = int(constcol(geniedict, "neu", mask))
         n_events = type_weight * constcol(mcweightdict, "NEvents", mask)
-        surfaces.append(n_events * generation_surface(primary_type, Column("wght"), spectrum, spatial))
+
+        surfaces.append(n_events * generation_surface(pid, Column("wght"), spectrum, spatial))
     ret = sum(surfaces)
     assert isinstance(ret, GenerationSurface)
     return ret
@@ -78,8 +84,7 @@ def genie_reader_surface(table: Iterable[Mapping[str, float]]) -> GenerationSurf
     return retval
 
 
-def GenieWeighter(file_obj: Any, 
-                  nfiles:float | None = None) -> Weighter:  # noqa: N802
+def GenieWeighter(file_obj: Any, nfiles:float | None = None) -> Weighter:  # noqa: N802
     # pylint: disable=invalid-name
     """Weighter for GENIE simulation.
 
@@ -89,7 +94,7 @@ def GenieWeighter(file_obj: Any,
     if has_table(file_obj, "I3GenieInfo"):
         # Branch for newer genie-reader files
         if nfiles is not None:
-            msg = (f"GenieWeighter received an nfiles={nfiles}, but `{getattr(file_obj, 'filename', '<NONE>')}` "
+            msg = (f"GenieWeighter received an nfiles={nfiles}, but `{getattr(file_obj, "filename", "<NONE>")}` "
                    "was produced with genie-reader instead of genie-icetray. We expect to read the number of "
                    "files from the number of observed S-frames in the file, so this is unnecessary. Do not pass "
                    "in a value for nfiles for genie-reader files.")
@@ -115,25 +120,25 @@ def GenieWeighter(file_obj: Any,
     else:
         # Branch for older genie-icetray files
         if nfiles is None:
-            msg = (f"GenieWeighter received an nfiles={nfiles}, but `{getattr(file_obj, 'filename', '<NONE>')}` "
+            msg = (f"GenieWeighter received an nfiles={nfiles}, but `{getattr(file_obj, "filename", "<NONE>")}` "
                    "was produced with genie-reader instead of genie-icetray. We expect to read the number of "
                    "files from the number of observed S-frames in the file, so this is unnecessary. Do not pass "
                    "in a value for nfiles for genie-reader files.")
             raise RuntimeError(msg)
 
         weight_table = get_table(file_obj, "I3MCWeightDict")
         result_table = get_table(file_obj, "I3GENIEResultDict")
+
         surface = nfiles * genie_icetray_surface(weight_table, result_table)
 
-        momentum_vec = np.array([get_column(result_table, 'pxv'),
-                                 get_column(result_table, 'pyv'),
-                                 get_column(result_table, 'pzv')])
-        cos_zen = -1 * get_column(result_table, 'pzv') / np.sum(momentum_vec**2, axis=0)**0.5
+        momentum_vec = np.array([get_column(result_table, "pxv"),
+                                 get_column(result_table, "pyv"),
+                                 get_column(result_table, "pzv")])
+        cos_zen = -1 * get_column(result_table, "pzv") / np.sum(momentum_vec**2, axis=0)**0.5
 
         weighter = Weighter([file_obj], surface)
         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", cos_zen)
         weighter.add_weight_column("wght", get_column(result_table, "wght")*get_column(result_table, "_glbprbscale"))
-
     return weighter

tests/test_genie_icetray_datasets.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+
+# SPDX-FileCopyrightText: © 2022 the SimWeights contributors
+#
+# SPDX-License-Identifier: BSD-2-Clause
+import os
+import sys
+from pathlib import Path
+
+import h5py
+import numpy as np
+import pandas as pd
+import pytest
+import tables
+
+from simweights import GenieWeighter
+
+datasets = [
+    "genie-icetray.140000A_000000.hdf",
+    "genie-icetray.140000B_000000.hdf",
+    "genie-icetray.140000C_000000.hdf",
+    "genie-icetray.140000D_000000.hdf",
+    "level2_genie-icetray.140000_000000.hdf",
+]
+approx = pytest.approx
+datadir = os.environ.get("SIMWEIGHTS_TESTDATA", None)
+
+@pytest.mark.parametrize("fname", datasets)
+@pytest.mark.skipif(not datadir, reason="environment variable SIMWEIGHTS_TESTDATA not set")
+def test_dataset(fname):
+    filename = Path(datadir) / fname
+    reffile = h5py.File(str(filename), "r")
+
+    wd = reffile["I3MCWeightDict"]
+    grd = reffile["I3GENIEResultDict"]
+    pdgid = grd["neu"]
+    emin, emax = 10**wd["MinEnergyLog"], 10**wd["MaxEnergyLog"]
+
+    solid_angle = 2 * np.pi * (np.cos(wd["MinZenith"]) - np.cos(wd["MaxZenith"]))
+    injection_area_cm = 1e4 * np.pi * wd["InjectionSurfaceR"]**2
+    total_weight = wd["TotalInteractionProbabilityWeight"]
+
+    type_weight = np.empty_like(wd["OneWeight"])
+    type_weight[pdgid>0] = 0.7
+    type_weight[pdgid<0] = 0.3
+    w0 = wd["OneWeight"] / (wd["NEvents"] * type_weight)
+
+    fobjs = [
+        tables.open_file(str(filename), "r"),
+        pd.HDFStore(str(filename), "r"),
+    ]
+
+    for fobj in fobjs:
+        w = GenieWeighter(fobj, nfiles=1)
+
+        event_weight = w.get_weight_column("wght")
+        assert event_weight == approx(total_weight)
+
+        for particle in np.unique(pdgid):
+            for spectrum in w.surface.spectra[particle]:
+                power_min, power_max = spectrum.dists[1].a, spectrum.dists[1].b
+                event_mask = (pdgid == particle) & (emin == power_min) & (emax == power_max)
+
+                power_law = spectrum.dists[1]
+                energy_factor = 1 / power_law.pdf(w.get_weight_column("energy"))
+
+                one_weight = (w.get_weight_column("wght")[event_mask]
+                              * energy_factor[event_mask]
+                              * solid_angle[event_mask]
+                              * injection_area_cm[event_mask])
+
+                assert one_weight == approx(wd["OneWeight"][event_mask])
+
+        assert w0 == approx(w.get_weights(1), rel=1e-5)
+        fobj.close()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__, *sys.argv[1:]]))

tests/test_genie_icetray_weighter.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+
+# SPDX-FileCopyrightText: © 2022 the SimWeights contributors
+#
+# SPDX-License-Identifier: BSD-2-Clause
+
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from simweights import CircleInjector, GenieWeighter, PowerLaw
+
+mcwd_keys = [
+    "NEvents",
+    "MinZenith",
+    "MaxZenith",
+    "PowerLawIndex",
+    "MinEnergyLog",
+    "MaxEnergyLog",
+    "InjectionSurfaceR",
+    "PrimaryNeutrinoEnergy",
+    "GeneratorVolume",
+]
+
+grd_keys = [
+    "neu",
+    "pxv",
+    "pyv",
+    "pzv",
+    "Ev",
+    "wght",
+    "_glbprbscale"]
+
+
+def make_new_table(pdgid, nevents, spatial, spectrum, coszen):
+    dtype = [(k, float) for k in mcwd_keys]
+    weight = np.zeros(nevents, dtype=dtype)
+    weight["NEvents"] = nevents
+    weight["MinZenith"] = np.arccos(spatial.cos_zen_max)
+    weight["MaxZenith"] = np.arccos(spatial.cos_zen_min)
+    weight["PowerLawIndex"] = -1 * spectrum.g
+    weight["MinEnergyLog"] = np.log10(spectrum.a)
+    weight["MaxEnergyLog"] = np.log10(spectrum.b)
+    weight["InjectionSurfaceR"] = spatial.radius
+    weight["GeneratorVolume"] = 1.
+    weight["PrimaryNeutrinoEnergy"] = spectrum.ppf(np.linspace(0, 1, nevents, endpoint=False))
+
+    dtype = [(k, float) for k in grd_keys]
+    resultdict = np.zeros(nevents, dtype=dtype)
+    resultdict["neu"] = pdgid
+    resultdict["pxv"] = 1
+    resultdict["pyv"] = 1
+    resultdict["pzv"] = -coszen
+    resultdict["Ev"] = weight["PrimaryNeutrinoEnergy"]
+    resultdict["wght"] = 1.0
+    resultdict["_glbprbscale"] = 1.0
+
+    return weight, resultdict
+
+
+class TestGenieIcetrayWeighter(unittest.TestCase):
+    def test_genie_icetray(self):
+        nevents = 100000
+        coszen = 0.7
+        pdgid = 12
+        c1 = CircleInjector(300, 0, 1)
+        p1 = PowerLaw(0, 1e3, 1e4)
+
+        t1 = make_new_table(pdgid, nevents, c1, p1, coszen)
+
+        mcwd = pd.DataFrame(t1[0])
+        grd = pd.DataFrame(t1[1])
+
+        f1 = {"I3MCWeightDict": mcwd,
+              "I3GENIEResultDict": grd}
+
+        for nfiles in [1, 10, 100]:
+            wf = GenieWeighter(f1, nfiles=nfiles)
+            for flux in [1e-6, 1, 1e6]:
+                w1 = wf.get_weights(flux)
+                w2 = flux * p1.integral * c1.etendue / (0.7 * nfiles)
+                np.testing.assert_allclose(
+                    w1.sum(),
+                    w2
+                )
+                E = mcwd["PrimaryNeutrinoEnergy"]
+                y, x = np.histogram(E, weights=w1, bins=51, range=[p1.a, p1.b])
+                Ewidth = np.ediff1d(x)
+                np.testing.assert_allclose(y, flux * Ewidth * c1.etendue / (0.7 * nfiles), 6e-3)
+
+    def test_empty(self):
+        with self.assertRaises(RuntimeError):
+            x = {"I3MCWeightDict": {key:[] for key in mcwd_keys},
+                 "I3GENIEResultDict": {key:[] for key in grd_keys}}
+            GenieWeighter(x, nfiles=1)
+
+if __name__ == "__main__":
+    unittest.main()