More versatile generation surfaces

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = ["numpy>=1.21.2", "scipy"]
 dynamic = ["version", "description"]
 
 [project.optional-dependencies]
-test = ["h5py", "tables < 3.8; python_version < '3.9'", "tables", "pandas", "uproot", "pytest-cov"]
+test = ["h5py", "tables < 3.8; python_version < '3.9'", "tables <= 3.9.1", "pandas", "uproot", "pytest-cov"]
 docs = ["sphinx","sphinx-rtd-theme","pandas"]
 dev = ["pytest","pre-commit","reuse","black","ruff","pylint","mypy"]
 
@@ -74,7 +74,7 @@ enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 max-line-length = "128"
 
 [tool.pylint.messages_control]
-disable = "C0114"
+disable = "C0114,R0902,R0913"
 
 [tool.ruff]
 select = ["ALL"]
@@ -83,8 +83,9 @@ ignore = [
     "ANN401", # any-type
     "FBT",    # flake8-boolean-trap
     "S101",   # assert-used
-    "COM812", #confilcts with ruff formatter
-    "ISC001", #confilcts with ruff formatter
+    "COM812", # confilcts with ruff formatter
+    "ISC001", # confilcts with ruff formatter
+    "PLR0913",# Too many arguments in function definition
     ]
 line-length = 128
 target-version = "py38"

src/simweights/_corsika_weighter.py

@@ -13,7 +13,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import NaturalRateCylinder
-from ._utils import constcol, get_column, get_table, has_table
+from ._utils import Column, constcol, get_column, get_table, has_table
 from ._weighter import Weighter
 
 
@@ -32,17 +32,20 @@ def sframe_corsika_surface(table: Any, oversampling: bool) -> GenerationSurface:
             get_column(table, "cylinder_radius")[i],
             np.cos(get_column(table, "max_zenith")[i]),
             np.cos(get_column(table, "min_zenith")[i]),
+            "cos_zen",
         )
         spectrum = PowerLaw(
             get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
         oversampling_val = get_column(table, "oversampling")[i] if oversampling else 1
+        pdgid = int(get_column(table, "primary_type")[i])
         surfaces.append(
             get_column(table, "n_events")[i]
             * oversampling_val
-            * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
+            * generation_surface(pdgid, Column("event_weight"), spectrum, spatial),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
@@ -61,6 +64,7 @@ def weight_map_corsika_surface(table: Any) -> GenerationSurface:
             constcol(table, "CylinderRadius", mask),
             np.cos(constcol(table, "ThetaMax", mask)),
             np.cos(constcol(table, "ThetaMin", mask)),
+            "cos_zen",
         )
 
         primary_spectral_index = round(constcol(table, "PrimarySpectralIndex", mask), 6)
@@ -70,6 +74,7 @@ def weight_map_corsika_surface(table: Any) -> GenerationSurface:
             primary_spectral_index,
             constcol(table, "EnergyPrimaryMin", mask),
             constcol(table, "EnergyPrimaryMax", mask),
+            "energy",
         )
         nevents = constcol(table, "OverSampling", mask) * constcol(table, "NEvents", mask)
         surface += nevents * generation_surface(pdgid, spectrum, spatial)

src/simweights/_generation_surface.py

@@ -5,26 +5,28 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, NamedTuple, Sequence
 
 import numpy as np
 
 from ._pdgcode import PDGCode
+from ._powerlaw import PowerLaw
+from ._spatial import CircleInjector, CylinderBase, SpatialDist
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from numpy.typing import ArrayLike, NDArray
 
-    from ._powerlaw import PowerLaw  # pragma: no cover
-    from ._spatial import SpatialDist  # pragma: no cover
+    from ._utils import Column, Const
+
+    Dist = SpatialDist | PowerLaw | Column | Const
 
 
 class SurfaceTuple(NamedTuple):
     """Container for the components Generation Surfaces."""
 
     pdgid: int | PDGCode
     nevents: float
-    energy_dist: PowerLaw
-    spatial_dist: SpatialDist
+    dists: Sequence
 
 
 class GenerationSurface:
@@ -46,7 +48,7 @@ def _insert(self: GenerationSurface, surface: SurfaceTuple) -> None:
             self.spectra[key] = []
 
         for i, spec in enumerate(self.spectra[key]):
-            if surface.energy_dist == spec.energy_dist and surface.spatial_dist == spec.spatial_dist:
+            if surface.dists == spec.dists:
                 self.spectra[key][i] = spec._replace(nevents=surface.nevents + spec.nevents)
                 break
         else:
@@ -77,30 +79,45 @@ def __mul__(self: GenerationSurface, factor: float) -> GenerationSurface:
     def __rmul__(self: GenerationSurface, factor: float) -> GenerationSurface:
         return self.__mul__(factor)
 
+    def get_keys(self: GenerationSurface) -> list[str]:
+        """Get a list of the available keys needed for weighting this surface."""
+        keys = []
+        for x in self.spectra.values():
+            for y in x:
+                for a in y.dists:
+                    keys += list(a.columns)
+        return list(set(keys))
+
     def get_epdf(
         self: GenerationSurface,
         pdgid: ArrayLike,
-        energy: ArrayLike,
-        cos_zen: ArrayLike,
+        **kwargs: ArrayLike,
     ) -> NDArray[np.float64]:
         """Get the extended pdf of an event.
 
         The pdf is the probability that an event with these parameters is generated. The pdf is multiplied
         by the number of events.
         """
-        energy = np.asarray(energy)
-        cos_zen = np.asarray(cos_zen)
-        count = np.zeros_like(energy, dtype=float)
-
+        cols = {}
+        shape = None
+        for key, value in kwargs.items():
+            cols[key] = np.asfarray(value)
+            if shape is None:
+                shape = cols[key].shape
+            else:
+                assert shape == cols[key].shape  # type: ignore[unreachable]
+        assert shape is not None
+        count = np.zeros(shape, dtype=np.float_)
+        # loop over particle type
         for ptype in np.unique(pdgid):
             mask = ptype == pdgid
-            if np.any(mask):
-                masked_energy = energy[mask]
-                masked_cos_zen = cos_zen[mask]
-                count[mask] += sum(
-                    p.nevents * p.energy_dist.pdf(masked_energy) * p.spatial_dist.pdf(masked_cos_zen)
-                    for p in self.spectra[ptype]
-                )
+            # loop over different datasets of the same particle type
+            for surface in self.spectra[ptype]:
+                result = surface.nevents
+                # loop over the different distributions in the generation surface
+                for dist in surface.dists:
+                    result *= dist.pdf(*(cols[k][mask] for k in dist.columns))
+                count[mask] += result
         return count
 
     def get_pdgids(self: GenerationSurface) -> list[int | PDGCode]:
@@ -116,8 +133,10 @@ def get_energy_range(self: GenerationSurface, pdgid: PDGCode | None) -> tuple[fl
         emax = -np.inf
         for pid in pdgids:
             for surf in self.spectra[pid]:
-                emin = min(emin, surf.energy_dist.a)
-                emax = max(emax, surf.energy_dist.b)
+                for dist in surf.dists:
+                    if isinstance(dist, PowerLaw):
+                        emin = min(emin, dist.a)
+                        emax = max(emax, dist.b)
         assert np.isfinite(emin)
         assert np.isfinite(emax)
         return emin, emax
@@ -131,8 +150,10 @@ def get_cos_zenith_range(self: GenerationSurface, pdgid: PDGCode | None) -> tupl
         czmax = -np.inf
         for pid in pdgids:
             for surf in self.spectra[pid]:
-                czmin = min(czmin, surf.spatial_dist.cos_zen_min)
-                czmax = max(czmax, surf.spatial_dist.cos_zen_max)
+                for dist in surf.dists:
+                    if isinstance(dist, (CircleInjector, CylinderBase)):
+                        czmin = min(czmin, dist.cos_zen_min)
+                        czmax = max(czmax, dist.cos_zen_max)
         assert np.isfinite(czmin)
         assert np.isfinite(czmax)
         return czmin, czmax
@@ -164,18 +185,13 @@ def __str__(self: GenerationSurface) -> str:
                 ptype = PDGCode(pdgid).name
             except ValueError:
                 ptype = str(pdgid)
-
-            collections = [f"N={subspec.nevents} {subspec.energy_dist} {subspec.spatial_dist}" for subspec in specs]
+            collections = (f"N={subspec.nevents} " + " ".join(repr(d) for d in subspec.dists) for subspec in specs)
             outstrs.append(f"     {ptype:>11} : " + "\n                   ".join(collections))
         return "< " + self.__class__.__name__ + "\n" + "\n".join(outstrs) + "\n>"
 
 
-def generation_surface(
-    pdgid: int | PDGCode,
-    energy_dist: PowerLaw,
-    spatial_dist: SpatialDist,
-) -> GenerationSurface:
+def generation_surface(pdgid: int | PDGCode, *dists: Dist) -> GenerationSurface:
     """Convenience function to generate a GenerationSurface for a single particle type."""
     return GenerationSurface(
-        SurfaceTuple(pdgid=pdgid, nevents=1.0, energy_dist=energy_dist, spatial_dist=spatial_dist),
+        SurfaceTuple(pdgid=pdgid, nevents=1.0, dists=dists),
     )

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 constcol, get_column, get_table
+from ._utils import Column, Const, get_column, get_table
 from ._weighter import Weighter
 
 
@@ -29,10 +29,14 @@ def genie_surface(table: Iterable[Mapping[str, float]]) -> GenerationSurface:
             -get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
+        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]
         surfaces.append(
-            get_column(table, "n_flux_events")[i]
-            * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
+            nevents
+            * generation_surface(pdgid, Const(1 / spatial.etendue / global_probability_scale), Column("wght"), spectrum),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
@@ -47,15 +51,10 @@ def GenieWeighter(file_obj: Any) -> Weighter:  # noqa: N802
     """
     weight_table = get_table(file_obj, "I3GenieInfo")
     surface = genie_surface(weight_table)
-    global_probability_scale = constcol(weight_table, "global_probability_scale")
 
     weighter = Weighter([file_obj], surface)
-    weighter.add_weight_column("energy", weighter.get_column("I3GenieResult", "Ev"))
     weighter.add_weight_column("pdgid", weighter.get_column("I3GenieResult", "neu").astype(np.int32))
-    weighter.add_weight_column("cos_zen", np.full(len(weighter.get_column("I3GenieResult", "Ev")), 1))
-    weighter.add_weight_column(
-        "event_weight",
-        global_probability_scale * weighter.get_column("I3GenieResult", "wght"),
-    )
+    weighter.add_weight_column("energy", weighter.get_column("I3GenieResult", "Ev"))
+    weighter.add_weight_column("wght", weighter.get_column("I3GenieResult", "wght"))
 
     return weighter

src/simweights/_icetop_weighter.py

@@ -20,20 +20,22 @@ def sframe_icetop_surface(table: Any) -> GenerationSurface:
 
     for i in range(len(get_column(table, "n_events"))):
         assert get_column(table, "power_law_index")[i] <= 0
-        spectrum = PowerLaw(
+        spectrum = PowerLaw(  # pylint: disable=duplicate-code
             get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
         spatial = NaturalRateCylinder(
             0,  # set cylinder height to 0 to get simple surface plane
             get_column(table, "sampling_radius")[i],
             np.cos(get_column(table, "max_zenith")[i]),
             np.cos(get_column(table, "min_zenith")[i]),
+            "cos_zen",
         )
-        surfaces.append(
-            get_column(table, "n_events")[i] * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
-        )
+        nevents = get_column(table, "n_events")[i]
+        pdgid = int(get_column(table, "primary_type")[i])
+        surfaces.append(nevents * generation_surface(pdgid, spectrum, spatial))
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
     return retval
@@ -49,6 +51,4 @@ def IceTopWeighter(file_obj: Any) -> Weighter:  # noqa: N802
     weighter.add_weight_column("pdgid", pdgid)
     weighter.add_weight_column("energy", weighter.get_column("MCPrimary", "energy"))
     weighter.add_weight_column("cos_zen", np.cos(weighter.get_column("MCPrimary", "zenith")))
-    weighter.add_weight_column("event_weight", np.ones_like(pdgid))
-
     return weighter

src/simweights/_nugen_weighter.py

@@ -10,7 +10,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector, SpatialDist, UniformSolidAngleCylinder
-from ._utils import constcol, get_column, get_table, has_column
+from ._utils import Column, constcol, get_column, get_table, has_column
 from ._weighter import Weighter
 
 
@@ -30,14 +30,14 @@ def nugen_spatial(table: Any) -> SpatialDist:
     injection_radius = constcol(table, "InjectionSurfaceR") if has_column(table, "InjectionSurfaceR") else -1
 
     if injection_radius > 0:
-        return CircleInjector(injection_radius, min_cos, max_cos)
+        return CircleInjector(injection_radius, min_cos, max_cos, "cos_zen")
 
     # Surface mode was added in V04-01-00 but the cylinder size was hard coded, `CylinderHeight` and
     # `CylinderRadius` were added after later V06-00-00. If they are not in the table then use the
     # hardcoded values
     cylinder_height = constcol(table, "CylinderHeight") if has_column(table, "CylinderHeight") else 1900
     cylinder_radius = constcol(table, "CylinderRadius") if has_column(table, "CylinderRadius") else 950
-    return UniformSolidAngleCylinder(cylinder_height, cylinder_radius, min_cos, max_cos)
+    return UniformSolidAngleCylinder(cylinder_height, cylinder_radius, min_cos, max_cos, "cos_zen")
 
 
 def nugen_spectrum(table: Any) -> PowerLaw:
@@ -48,7 +48,7 @@ def nugen_spectrum(table: Any) -> PowerLaw:
     # for negative slopes ie +2 means E**-2 injection spectrum
     power_law_index = -constcol(table, "PowerLawIndex")
     assert power_law_index <= 0
-    return PowerLaw(power_law_index, min_energy, max_energy)
+    return PowerLaw(power_law_index, min_energy, max_energy, "energy")
 
 
 def nugen_surface(table: Any) -> GenerationSurface:
@@ -66,9 +66,9 @@ def nugen_surface(table: Any) -> GenerationSurface:
         # newer version will explicitly put the ratio in `TypeWeight` but for older version we
         # assume it is 0.5
         type_weight = constcol(table, "TypeWeight", mask) if has_column(table, "TypeWeight") else 0.5
-        primary_type = constcol(table, "PrimaryNeutrinoType", mask)
+        primary_type = int(constcol(table, "PrimaryNeutrinoType", mask))
         n_events = type_weight * constcol(table, "NEvents", mask)
-        surfaces.append(n_events * generation_surface(int(primary_type), spectrum, spatial))
+        surfaces.append(n_events * generation_surface(primary_type, Column("event_weight"), spectrum, spatial))
     ret = sum(surfaces)
     assert isinstance(ret, GenerationSurface)
     return ret

src/simweights/_powerlaw.py

@@ -38,7 +38,7 @@ class PowerLaw:
 
     # pylint: disable=invalid-name
 
-    def __init__(self: PowerLaw, g: float, a: float, b: float) -> None:
+    def __init__(self: PowerLaw, g: float, a: float, b: float, colname: str | None = None) -> None:
         assert b > a
         self.g = float(g)
         self.a = float(a)
@@ -48,8 +48,8 @@ def __init__(self: PowerLaw, g: float, a: float, b: float) -> None:
             self.integral = np.log(self.b / self.a)
         else:
             self.integral = (self.b**self.G - self.a**self.G) / self.G
-
         self.span = b - a
+        self.columns = (colname,)
 
     def _pdf(self: PowerLaw, x: NDArray[np.float64]) -> NDArray[np.float64]:
         return np.asfarray(x**self.g / self.integral)