Make PDFs more explicit

One of the criticistms of simweights is that it is hard to find the
weighting functions. This was because the weighting was a list of
objects and the weights were calculating by taking the product of them.
This PR refactors it so that each type of simulation has a class that
is derived from GenerationSurface that overrides get_epdf(). This is a
clear function that is 1/weight. However since it makes use of Spatial
and PowerLaw it is not as explicit as it could be but DRY. As a side
effecth Spatial disturcitions and PowerLaw no longer have column names
as a member variable which i thought was just weird.

Author: kjmeagher

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
       - id: blacken-docs
         args: [-l 100]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.16.1
+    rev: v1.17.1
     hooks:
       - id: mypy
         files: simweights
@@ -42,7 +42,7 @@ repos:
         exclude: ^contrib/
         additional_dependencies: [numpy, pandas]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.2
+    rev: v0.12.7
     hooks:
       - id: ruff
         args: [--fix, --show-fixes]

pyproject.toml

@@ -70,7 +70,6 @@ max-line-length = 128
 
 [tool.mypy]
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
-plugins = "numpy.typing.mypy_plugin"
 strict = true
 warn_unreachable = true
 
@@ -84,7 +83,7 @@ disable = "C0114,R0902,R0913,R0917,R0914,R0911"
 addopts = ["-ra", "--strict-config", "--strict-markers", "--cov=simweights", "-W ignore"]
 filterwarnings = ["error"]
 log_cli_level = "INFO"
-minversion = 7.0
+minversion = "7.0"
 testpaths = ["tests"]
 xfail_strict = true
 
@@ -93,7 +92,7 @@ line-length = 128
 namespace-packages = ["examples", "contrib", "docs"]
 
 [tool.ruff.lint]
-fixable = ["I"]
+fixable = ["I", "Q"]
 ignore = [
   "ANN401",  # any-type
   "S101",  # assert-used

src/simweights/__init__.py

@@ -24,6 +24,7 @@
     "GaisserH4a_IT",
     "GaisserHillas",
     "GenerationSurface",
+    "GenieSurface",
     "GenieWeighter",
     "GlobalFitGST",
     "GlobalFitGST_IT",
@@ -34,15 +35,16 @@
     "Hoerandel5",
     "Hoerandel_IT",
     "Honda2004",
+    "IceTopSurface",
     "IceTopWeighter",
     "NaturalRateCylinder",
+    "NuGenSurface",
     "NuGenWeighter",
     "PDGCode",
     "PowerLaw",
     "UniformSolidAngleCylinder",
     "Weighter",
     "corsika_to_pdg",
-    "generation_surface",
 ]
 
 from ._corsika_weighter import CorsikaWeighter
@@ -63,10 +65,10 @@
     Hoerandel_IT,
     Honda2004,
 )
-from ._generation_surface import GenerationSurface, generation_surface
-from ._genie_weighter import GenieWeighter
-from ._icetop_weighter import IceTopWeighter
-from ._nugen_weighter import NuGenWeighter
+from ._generation_surface import GenerationSurface
+from ._genie_weighter import GenieSurface, GenieWeighter
+from ._icetop_weighter import IceTopSurface, IceTopWeighter
+from ._nugen_weighter import NuGenSurface, NuGenWeighter
 from ._pdgcode import PDGCode
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector, NaturalRateCylinder, UniformSolidAngleCylinder

src/simweights/_corsika_weighter.py

@@ -6,24 +6,40 @@
 
 import numbers
 import warnings
-from typing import Any
+from typing import TYPE_CHECKING, Any, Mapping
 
 import numpy as np
 
-from ._generation_surface import GenerationSurface, generation_surface
+from ._generation_surface import CompositeSurface, GenerationSurface
 from ._powerlaw import PowerLaw
 from ._spatial import NaturalRateCylinder
-from ._utils import Column, constcol, get_column, get_table, has_column, has_table
+from ._utils import constcol, get_column, get_table, has_column, has_table
 from ._weighter import Weighter
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
 
-def sframe_corsika_surface(table: Any) -> GenerationSurface:
+
+class CorsikaSurface(GenerationSurface):
+    """Represents a surface on which CORSIKA simulation was generated on."""
+
+    def get_epdf(self: CorsikaSurface, weight_cols: Mapping[str, NDArray[np.float64]]) -> NDArray[np.float64]:
+        """Get the extended pdf of a sample of CORSIKA."""
+        return (
+            self.nevents
+            / weight_cols["event_weight"]
+            * self.power_law.pdf(weight_cols["energy"])
+            * self.spatial.pdf(weight_cols["cos_zen"])
+        )
+
+
+def sframe_corsika_surface(table: Any) -> CompositeSurface:
     """Inspect the rows of a CORSIKA S-Frame table object to generate a surface object.
 
     This function works on files generated with either triggered CORSIKA or corsika-reader because
     `I3PrimaryInjectorInfo` and `I3CorsikaInfo` use exactly the same names for quantities.
     """
-    surfaces = []
+    surfaces = CompositeSurface()
     cylinder_height = get_column(table, "cylinder_height")
     cylinder_radius = get_column(table, "cylinder_radius")
     max_zenith = get_column(table, "max_zenith")
@@ -39,32 +55,26 @@ def sframe_corsika_surface(table: Any) -> GenerationSurface:
             cylinder_radius[i],
             np.cos(max_zenith[i]),
             np.cos(min_zenith[i]),
-            "cos_zen",
         )
         spectrum = PowerLaw(
             power_law_index[i],
             min_energy[i],
             max_energy[i],
-            "energy",
         )
         oversampling_val = get_column(table, "oversampling")[i] if has_column(table, "oversampling") else 1
         pdgid = int(get_column(table, "primary_type")[i])
-        surfaces.append(
-            n_events[i] * oversampling_val * generation_surface(pdgid, Column("event_weight"), spectrum, spatial),
-        )
-    retval = sum(surfaces)
-    assert isinstance(retval, GenerationSurface)
-    return retval
+        surfaces.insert(CorsikaSurface(pdgid, n_events[i] * oversampling_val, spectrum, spatial))
+    return surfaces
 
 
-def weight_map_corsika_surface(table: Any) -> GenerationSurface:
+def weight_map_corsika_surface(table: Any) -> CompositeSurface:
     """Inspect the `CorsikaWeightMap` table object of a corsika file to generate a surface object."""
     pdgids = sorted(np.unique(get_column(table, "ParticleType").astype(int)))
 
     if len(pdgids) == 0:
         msg = "`CorsikaWeightMap` is empty. SimWeights cannot process this file"
         raise RuntimeError(msg)
-    surface: int | GenerationSurface = 0
+    surface = CompositeSurface()
     for pdgid in pdgids:
         mask = pdgid == get_column(table, "ParticleType")
 
@@ -73,7 +83,6 @@ 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)
@@ -83,11 +92,9 @@ 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)
-    assert isinstance(surface, GenerationSurface)
+        surface.insert(CorsikaSurface(pdgid, nevents, spectrum, spatial))
     return surface
 
 
@@ -145,7 +152,8 @@ def CorsikaWeighter(file_obj: Any, nfiles: float | None = None) -> Weighter:  #
             )
 
         table = get_table(file_obj, "CorsikaWeightMap")
-        surface = nfiles * weight_map_corsika_surface(table)
+        surface = weight_map_corsika_surface(table)
+        surface.scale(nfiles)
         triggered = False
 
     weighter = Weighter([file_obj], surface)