unit tests for i3files

.github/workflows/tests.yml

@@ -32,8 +32,8 @@ jobs:
           cache-dependency-path: pyproject.toml
       - name: Install SimWeights
         run: |
-          if [ "${{matrix.python-version}}" == "3.9" ]; then
-            python3 -m pip install "numpy<2"
+          if [ "${{matrix.python-version}}" != "3.13" ]; then
+            python3 -m pip install nuflux
           fi
           python3 -m pip install flit
           python3 -m flit install --symlink --deps=production --extras=test
@@ -60,6 +60,30 @@ jobs:
         with:
           fail_ci_if_error: false
           verbose: true
+  TestsIceTray:
+    runs-on: ubuntu-latest
+    container: icecube/icetray:icetray-devel-current-ubuntu22.04-X64
+    strategy:
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install SimWeights
+        run: python3 -m pip install 'pytest>7' .[test]
+      - name: Download Test Data
+        run: |
+          curl -u icecube:${{ secrets.ICECUBE_PASSWORD }} https://convey.icecube.wisc.edu/data/ana/Software/simweights/test-data/simweights_testdata.tar.gz -O
+          tar xzvf simweights_testdata.tar.gz
+      - name: Run Unit Tests
+        env:
+          SIMWEIGHTS_TESTDATA: .
+        run: /opt/icetray/bin/icetray-shell python -m pytest --junit-xml=test-results-icetray.junit.xml
+      - name: Upload Test Results
+        uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          if-no-files-found: error
+          name: test-results-icetray.junit.xml
+          path: test-results-icetray.junit.xml
   publish-test-results:
     name: "Publish Tests Results"
     needs: Tests

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ ci:
   autoupdate_schedule: quarterly
 repos:
   - repo: https://github.com/google/yamlfmt
-    rev: v0.13.0
+    rev: v0.15.0
     hooks:
       - id: yamlfmt
   - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
@@ -16,34 +16,34 @@ repos:
       - id: pretty-format-toml
         args: [--autofix]
   - repo: https://github.com/fsfe/reuse-tool
-    rev: v4.0.3
+    rev: v5.0.2
     hooks:
       - id: reuse
   - repo: https://github.com/codespell-project/codespell
     rev: v2.3.0
     hooks:
       - id: codespell
   - repo: https://github.com/adamchainz/blacken-docs
-    rev: "1.19.0"
+    rev: "1.19.1"
     hooks:
       - id: blacken-docs
         args: [-l 100]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.11.2
+    rev: v1.14.1
     hooks:
       - id: mypy
         files: simweights
         additional_dependencies: [numpy]
         exclude: ^contrib/
   - repo: https://github.com/pycqa/pylint
-    rev: v3.3.1
+    rev: v3.3.3
     hooks:
       - id: pylint
         files: simweights
         exclude: ^contrib/
         additional_dependencies: [numpy, pandas]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.9
+    rev: v0.9.1
     hooks:
       - id: ruff
         args: [--fix, --show-fixes]

contrib/book_simweights_testdata.py

@@ -87,13 +87,6 @@ def fake_event_header(frame: dict) -> None:
     ],
     "icetop": ["I3TopInjectorInfo", "MCPrimary"],
 }
-streams = {
-    "corsika": ["InIceSplit"],
-    "nugen": ["InIceSplit", "in_ice"],
-    "genie": ["NullSplit"],
-    "icetop": ["IceTopSplit"],
-}
-
 
 if "notemp" in sys.argv:
     outdir = Path("/scratch/kmeagher/simweights/")
@@ -104,40 +97,34 @@ def fake_event_header(frame: dict) -> None:
 for simtype, filename in ((i, x) for i in filelist for x in filelist[i]):
     basename = Path(filename).name.replace(".i3.zst", "").replace(".i3.bz2", "").replace(".i3.gz", "")
     assert basename != Path(filename).name
-
-    split = simtype == "genie"
     outfile = outdir / basename
 
     print(f"Booking  : {filename}")
     print(f"  outfile: {outfile}")
     print(f"  keys   : {keys[simtype]}")
-    print(f"  streams: {streams[simtype]}")
 
     tray = icetray.I3Tray()
     tray.Add("I3Reader", FileNameList=[filename])
-
-    if split:
-        tray.Add(
-            fake_event_header,
-            Streams=[icetray.I3Frame.DAQ],
-            If=lambda f: "I3EventHeader" not in f,
-        )
-        tray.Add("I3NullSplitter", SubEventStreamName="NullSplit")
+    tray.Add(
+        fake_event_header,
+        Streams=[icetray.I3Frame.DAQ],
+        If=lambda f: "I3EventHeader" not in f,
+    )
+    tray.Add("I3NullSplitter", SubEventStreamName="weight")
     tray.Add(
         tableio.I3TableWriter,
         tableservice=[
             hdfwriter.I3HDFTableService(str(outfile) + ".hdf5"),
             rootwriter.I3ROOTTableService(str(outfile) + ".root"),
         ],
-        SubEventStreams=streams[simtype],
+        SubEventStreams=["weight"],
         keys=keys[simtype],
     )
-    tray.Add("Keep", keys=keys[simtype])
+    tray.Add("Keep", keys=["I3EventHeader"] + keys[simtype])
     tray.Add(
         "I3Writer",
         Filename=str(outfile) + ".i3.zst",
         Streams=[icetray.I3Frame.Simulation, icetray.I3Frame.DAQ],
-        DropOrphanStreams=[icetray.I3Frame.Physics],
     )
 
     tray.Execute()

docs/i3frame_support.rst

@@ -0,0 +1,33 @@
+.. SPDX-FileCopyrightText: © 2025 the SimWeights contributors
+..
+.. SPDX-License-Identifier: BSD-2-Clause
+
+I3Frame Support
+===============
+
+You can calculate the weight for a single event directly from an ``I3Frame``
+with a number of warnings. First weights are only meaningful for a sample of
+Monte Carlo events not a single event. The events calculated here will not be
+usfull to combine different samples. Second, the normalization may be off in
+subtle ways. For example in triggered CORSIKA there will be 8 S-frames for
+each primary type per file which need to be accounted for. When reading an HDF5
+file this is correctly accounted for but there is no way to account for this in
+IceTray. It is strongly encouraged that you only weight events after obtaining
+a complete sample in an HDF5 or similar file. But if you really need to
+calculate weights in IceTray you can follow the example below:
+
+.. literalinclude:: ../examples/triggered_corsika_i3file.py
+  :start-after: start-example1
+
+Note that the module keeps track of how many S-Frames there are and hence the
+factor the weight is incorrect by. But, because of the serial nature of
+IceTray, it can't retroactivly apply that correction to events that have
+already been processed. The output should look like::
+
+  PPlus       8
+  He4Nucleus  8
+  N14Nucleus  8
+  Al27Nucleus 8
+  Fe56Nucleus 8
+
+Indicating that there were 8 S-Frames of each primary type.

docs/index.rst

@@ -42,6 +42,7 @@
     reading_files
     reading_nugen
     units
+    i3frame_support
 
 ..  toctree::
     :maxdepth: 1

docs/reading_files.rst

@@ -80,7 +80,11 @@ The table below shows the what tables simweights looks for for each type of simu
 +--------------------------+---------------------------+---------------------------------------------+
 | CORSIKA without S-Frames | none                      | ``CorsikaWeightMap``, ``PolyplopiaPrimary`` |
 +--------------------------+---------------------------+---------------------------------------------+
+| IceTop CORSIKA           | ``I3TopInjectorInfo``     | ``MCPrimary``                               |
++--------------------------+---------------------------+---------------------------------------------+
 | neutrino-generator       | none                      | ``I3MCWeightDict``                          |
 +--------------------------+---------------------------+---------------------------------------------+
 | genie-reader             | ``I3GenieInfo``           | ``I3GenieResult``                           |
 +--------------------------+---------------------------+---------------------------------------------+
+| genie-icetray            | none                      | ``I3MCWeightDict``, ``I3GENIEResultDict``   |
++--------------------------+---------------------------+---------------------------------------------+

examples/triggered_corsika_i3file.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+
+# SPDX-FileCopyrightText: © 2025 the SimWeights contributors
+#
+# SPDX-License-Identifier: BSD-2-Clause
+
+# start-example1
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+from icecube import dataclasses, hdfwriter, icetray, simclasses
+
+from simweights import CorsikaWeighter, GaisserH3a
+
+FILE_DIR = Path("/data/sim/IceCube/2016/filtered/level2/CORSIKA-in-ice/21889/0000000-0000999")
+files = sorted(str(f) for f in FILE_DIR.glob("Level2_IC86.2016_corsika.021889.000000.i3.zst"))
+
+
+class Weighter(icetray.I3Module):
+    def __init__(self, context: icetray.I3Context) -> None:
+        icetray.I3Module.__init__(self, context)
+        self.weighter = None
+        self.fluxmodel = GaisserH3a()
+        self.s_frames = defaultdict(int)
+
+    def Simulation(self, frame: icetray.I3Frame) -> None:
+        pdgid = frame["I3PrimaryInjectorInfo"].primary_type
+        self.s_frames[pdgid] += 1
+
+    def DAQ(self, frame: icetray.I3Frame) -> None:
+        weighter = CorsikaWeighter(frame)
+        weight = weighter.get_weights(self.fluxmodel)[0]
+        name = f"weight_{self.fluxmodel.__class__.__name__}"
+        frame[name] = dataclasses.I3Double(weight)
+        self.PushFrame(frame)
+
+    def Finish(self) -> None:
+        print("Weights Need to be adjusted by the following factors:")
+        for pdgid, counts in self.s_frames.items():
+            print(f"{pdgid!s:>11} {counts}")
+
+
+tray = icetray.I3Tray()
+tray.Add("I3Reader", FileNameList=files)
+tray.Add(Weighter)
+tray.Execute()

pyproject.toml

@@ -78,7 +78,7 @@ warn_unreachable = true
 max-line-length = "128"
 
 [tool.pylint.messages_control]
-disable = "C0114,R0902,R0913,R0917,R0914"
+disable = "C0114,R0902,R0913,R0917,R0914,R0911"
 
 [tool.pytest.ini_options]
 addopts = ["-ra", "--strict-config", "--strict-markers", "--cov=simweights", "-W ignore"]
@@ -99,7 +99,8 @@ ignore = [
   "S101",  # assert-used
   "COM812",  # conflicts with ruff formatter
   "ISC001",  # conflicts with ruff formatter
-  "PLR0913"  # Too many arguments in function definition
+  "PLR0913",  # Too many arguments in function definition
+  "PLR0911"  # Too many return statement
 ]
 select = ["ALL"]
 
@@ -108,7 +109,9 @@ select = ["ALL"]
 "examples/*" = [
   "D",  # pydocstyle
   "F401",  # unused-import
-  "T201"  # flake8-print
+  "T201",  # flake8-print
+  "PLC0206",  # Extracting value from dictionary without calling `.items()`
+  "N802"  # Function name should be lowercase
 ]
 "tests/*" = [
   "D",  # pydocstyle

src/simweights/__init__.py

@@ -15,17 +15,16 @@
 __version__ = "0.1.3"
 
 __all__ = [
+    "TIG1996",
     "CircleInjector",
     "CorsikaWeighter",
-    "NaturalRateCylinder",
-    "UniformSolidAngleCylinder",
-    "TIG1996",
     "FixedFractionFlux",
-    "GenieWeighter",
     "GaisserH3a",
     "GaisserH4a",
     "GaisserH4a_IT",
     "GaisserHillas",
+    "GenerationSurface",
+    "GenieWeighter",
     "GlobalFitGST",
     "GlobalFitGST_IT",
     "GlobalSplineFit",
@@ -36,13 +35,14 @@
     "Hoerandel_IT",
     "Honda2004",
     "IceTopWeighter",
-    "PDGCode",
-    "corsika_to_pdg",
-    "generation_surface",
-    "GenerationSurface",
+    "NaturalRateCylinder",
     "NuGenWeighter",
+    "PDGCode",
     "PowerLaw",
+    "UniformSolidAngleCylinder",
     "Weighter",
+    "corsika_to_pdg",
+    "generation_surface",
 ]
 
 from ._corsika_weighter import CorsikaWeighter

src/simweights/_fluxes.py

@@ -331,7 +331,7 @@ class GlobalFitGST_IT(CosmicRayFlux):  # pylint: disable=invalid-name
 
 
 class GlobalSplineFitBase(CosmicRayFlux):
-    r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski].
+    r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]_.
 
     Base class all actual classes should inherit from this one.
     """
@@ -350,7 +350,7 @@ def __init__(self: GlobalSplineFitBase) -> None:
 
 
 class GlobalSplineFit(GlobalSplineFitBase):
-    r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]."""
+    r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]_."""
 
     pdgids = PDGID_ALL
 
@@ -362,7 +362,7 @@ def __init__(self: GlobalSplineFit) -> None:
 class GlobalSplineFit5Comp(GlobalSplineFitBase):
     r"""Sum of the flux of the GSF model for the standard 5 components injected by IceCube.
 
-    GSF is a Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski].
+    GSF is a Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]_.
     """
 
     pdgids = PDGID_5COMP
@@ -376,7 +376,7 @@ class GlobalSplineFit_IT(GlobalSplineFitBase):  # pylint: disable=invalid-name
     r"""Sum of the flux of the GSF model for the standard 4 components injected by IceCube.
 
     [(H),  (He),  (Li, Be, B, C, N, O, F, Ne),  (Na, Mg, Al, Si, P, S, Cl, Ar, K, Ca, Sc, Ti, V, Cr, Mn, Fe, Co, Ni)]
-    GSF is a Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski].
+    GSF is a Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]_.
     """
 
     pdgids = PDGID_4COMP

src/simweights/_genie_weighter.py

@@ -46,8 +46,8 @@ def genie_icetray_surface(
         (pid, _, _, _, _) = row
         mask = np.all(gen_schemes == row[None, :], axis=1)
 
-        spatial = nugen_spatial(mcweightdict[mask])
-        spectrum = nugen_spectrum(mcweightdict[mask])
+        spatial = nugen_spatial(mcweightdict, mask)
+        spectrum = nugen_spectrum(mcweightdict, mask)
 
         type_weight = nufraction if pid > 0 else 1 - nufraction
         n_events = type_weight * constcol(mcweightdict, "NEvents", mask)