add dataset to test genie volume scaling (#40)

contrib/book_simweights_testdata.py

@@ -7,7 +7,6 @@
 """Script to generate the test data used by simweights testing."""
 
 import os.path
-import shutil
 import sys
 import tarfile
 import tempfile
@@ -63,6 +62,7 @@ def fake_event_header(frame: dict) -> None:
     "genie": [
         "/data/sim/IceCubeUpgrade/genie/step3/141828/upgrade_genie_step3_141828_000000.i3.zst",
         "/data/sim/IceCube/2023/generated/GENIE/22590/0000000-0000999/GENIE_NuMu_IceCubeUpgrade_v58.22590.000000.i3.zst",
+        "/data/user/mlarson/icetray/src/genie-reader/resources/scripts/genie_numu_volume_scaling.i3.zst",
     ],
 }
 keys = {
@@ -98,10 +98,6 @@ def fake_event_header(frame: dict) -> None:
     split = simtype == "genie"
     outfile = outdir / basename
 
-    if Path(outfile.name + ".hdf5").exists():
-        print(f"Skipping {filename}: {outfile} already exists!")
-        continue
-
     print(f"Booking  : {filename}")
     print(f"  outfile: {outfile}")
     print(f"  keys   : {keys[simtype]}")
@@ -140,12 +136,9 @@ def fake_event_header(frame: dict) -> None:
 tarfilename = "/data/user/kmeagher/simweights_testdata_test.tar.gz"
 print(f"Writing tarfile {tarfilename}")
 
-shutil.copy("upgrade_genie_step3_140021_000000.root", outdir)
-shutil.copy("upgrade_genie_step3_140021_000000.hdf5", outdir)
-
 with tarfile.open(tarfilename, "w:gz") as tar:
     for f in os.listdir(outdir):
         print(f"Adding {f} to tarball")
         tar.add(outdir / f, arcname=f)
 
-print("Finished writing tarfile {tarfilename}")
+print(f"Finished writing tarfile {tarfilename}")

src/simweights/_genie_weighter.py

@@ -34,11 +34,10 @@ def genie_surface(table: Iterable[Mapping[str, float]]) -> GenerationSurface:
         pdgid = int(get_column(table, "primary_type")[i])
         nevents = get_column(table, "n_flux_events")[i]
         global_probability_scale = get_column(table, "global_probability_scale")[i]
+
+        const_factor = 1 / spatial.etendue / global_probability_scale
         surfaces.append(
-            nevents
-            * generation_surface(
-                pdgid, Const(1 / spatial.etendue / global_probability_scale), Column("wght"), Column("volscale"), spectrum
-            ),
+            nevents * generation_surface(pdgid, Const(const_factor), Column("wght"), Column("volscale"), spectrum),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)

tests/test_genie_datasets.py

@@ -16,13 +16,12 @@
 import uproot
 
 from simweights import GenieWeighter
-from simweights._utils import get_column, get_table
 
 datadir = os.environ.get("SIMWEIGHTS_TESTDATA", None)
 datasets = [
-    "upgrade_genie_step3_140021_000000",
     "upgrade_genie_step3_141828_000000",
     "GENIE_NuMu_IceCubeUpgrade_v58.22590.000000",
+    "genie_numu_volume_scaling",
 ]
 approx = pytest.approx
 
@@ -32,19 +31,44 @@
 def test_dataset(fname):
     filename = Path(datadir) / fname
     reffile = h5py.File(str(filename) + ".hdf5", "r")
+    info = reffile["I3GenieInfo"][0]
     wd = reffile["I3MCWeightDict"]
 
-    solid_angle = 2 * np.pi * (np.cos(wd["MinZenith"]) - np.cos(wd["MaxZenith"]))
-    injection_area = np.pi * (wd["InjectionSurfaceR"] * 1e2) ** 2
-    global_probability_scale = wd["GlobalProbabilityScale"]
-    genie_weight = wd["GENIEWeight"]
-
-    pli = -wd["PowerLawIndex"][0]
-    energy_integral = ((10 ** wd["MaxEnergyLog"][0]) ** (pli + 1) - (10 ** wd["MinEnergyLog"][0]) ** (pli + 1)) / (pli + 1)
-    energy_factor = 1 / (wd["PrimaryNeutrinoEnergy"] ** pli / energy_integral)
-    one_weight = global_probability_scale * genie_weight * energy_factor * solid_angle * injection_area
+    n_flux_events = info["n_flux_events"]
+    primary_type = info["primary_type"]
+    cylinder_radius = info["cylinder_radius"]
+    min_zenith = info["min_zenith"]
+    max_zenith = info["max_zenith"]
+    min_energy = info["min_energy"]
+    max_energy = info["max_energy"]
+    power_law_index = info["power_law_index"]
+    global_probability_scale = info["global_probability_scale"]
+    muon_volume_scaling = info["muon_volume_scaling"]
+
+    if "PrimaryNeutrinoType" in wd.dtype.names:
+        assert np.all(wd["PrimaryNeutrinoType"] == primary_type)
+    assert wd["InjectionSurfaceR"] == approx(cylinder_radius)
+    assert wd["MinZenith"] == approx(min_zenith)
+    assert wd["MaxZenith"] == approx(max_zenith)
+    assert 10 ** wd["MinEnergyLog"] == approx(min_energy)
+    assert 10 ** wd["MaxEnergyLog"] == approx(max_energy)
+    assert wd["PowerLawIndex"] == approx(power_law_index)
+    assert wd["GlobalProbabilityScale"] == approx(global_probability_scale)
+    if "MuonVolumeScaling" in wd.dtype.names:
+        assert wd["MuonVolumeScaling"] == approx(muon_volume_scaling)
+
+    solid_angle = 2 * np.pi * (np.cos(min_zenith) - np.cos(max_zenith))
+    injection_area = np.pi * (cylinder_radius * 1e2) ** 2
+
+    if power_law_index == 1:
+        energy_integral = np.log(max_energy / min_energy)
+    else:
+        energy_integral = (max_energy ** (1 - power_law_index) - min_energy ** (1 - power_law_index)) / (1 - power_law_index)
+    energy_factor = 1 / (wd["PrimaryNeutrinoEnergy"] ** (-power_law_index) / energy_integral)
+
+    one_weight = wd["TotalInteractionProbabilityWeight"] * energy_factor * solid_angle * injection_area
     np.testing.assert_allclose(one_weight, wd["OneWeight"])
-    final_weight = wd["OneWeight"] / (get_column(get_table(reffile, "I3GenieInfo"), "n_flux_events")[0])
+    final_weight = wd["OneWeight"] / n_flux_events
 
     fobjs = [
         reffile,
@@ -56,16 +80,17 @@ def test_dataset(fname):
     for fobj in fobjs:
         w = GenieWeighter(fobj)
 
-        pdf0 = next(iter(w.surface.spectra.values()))[0].dists[0]
-        assert 1 / pdf0.v == approx(global_probability_scale * solid_angle * injection_area, rel=1e-5)
-
-        assert w.get_weight_column("wght") == approx(genie_weight)
+        gprob, _, _, edist = next(iter(w.surface.spectra.values()))[0].dists
+        energy_term = 1 / edist.pdf(w.get_weight_column("energy"))
 
-        power_law = next(iter(w.surface.spectra.values()))[0].dists[-1]
-        energy_term = 1 / power_law.pdf(w.get_weight_column("energy"))
+        assert gprob.v == approx(1 / solid_angle / injection_area / global_probability_scale)
+        assert w.get_weight_column("wght") == approx(wd["GENIEWeight"])
         assert energy_term == approx(energy_factor)
 
-        one_weight = w.get_weight_column("wght") * energy_term * solid_angle * injection_area * global_probability_scale
+        vol_scale = w.get_weight_column("volscale")
+        one_weight = (
+            w.get_weight_column("wght") * energy_term * solid_angle * injection_area * global_probability_scale * vol_scale
+        )
         assert one_weight == approx(wd["OneWeight"], rel=1e-5)
 
         assert w.get_weights(1) == approx(final_weight, rel=1e-5)