Adding flux weighting option to the effective_area function of Weighter objects

pyproject.toml

@@ -78,7 +78,7 @@ warn_unreachable = true
 max-line-length = "128"
 
 [tool.pylint.messages_control]
-disable = "C0114,R0902,R0913,R0917"
+disable = "C0114,R0902,R0913,R0917,R0914"
 
 [tool.pytest.ini_options]
 addopts = ["-ra", "--strict-config", "--strict-markers", "--cov=simweights", "-W ignore"]

src/simweights/_spatial.py

@@ -88,7 +88,7 @@ class UniformSolidAngleCylinder(CylinderBase):
 
     """
 
-    def _pdf(self: UniformSolidAngleCylinder, cos_zen: NDArray[np.float64]) -> NDArray[np.float64]:
+    def _pdf(self: UniformSolidAngleCylinder, cos_zen: NDArray[np.float64]) -> NDArray[np.floating]:
         return 1 / (2 * np.pi * (self.cos_zen_max - self.cos_zen_min) * self.projected_area(cos_zen))
 
     def pdf(self: UniformSolidAngleCylinder, cos_zen: ArrayLike) -> NDArray[np.float64]:

src/simweights/_utils.py

@@ -27,9 +27,9 @@ class Column:
     def __init__(self: Column, colname: str | None = None) -> None:
         self.columns = (colname,)
 
-    def pdf(self: Column, value: ArrayLike) -> NDArray[np.float64]:
+    def pdf(self: Column, value: ArrayLike) -> NDArray[np.floating]:
         r"""Probability density function."""
-        return 1 / np.asarray(value, dtype=np.float64)
+        return 1.0 / np.asarray(value, dtype=np.float64)
 
     def __eq__(self: Column, other: object) -> bool:
         return isinstance(other, Column) and self.columns == other.columns
@@ -91,15 +91,15 @@ def get_column(table: Any, name: str) -> NDArray[np.float64]:
     return np.asarray(column, dtype=np.float64)
 
 
-def constcol(table: Any, colname: str, mask: NDArray[np.bool_] | None = None) -> float:
+def constcol(table: Any, colname: str, mask: ArrayLike | None = None) -> float:
     """Helper function which makes sure that all of the entries in a column are exactly the same.
 
     This is necessary because CORSIKA and NuGen store generation surface parameters in every frame and we
     want to verify that they are all the same.
     """
     col = get_column(table, colname)
     if mask is not None:
-        col = col[mask]
+        col = col[np.asarray(mask, dtype=bool)]
     val = col[0]
     assert np.ndim(val) == 0
     assert (val == col).all()

src/simweights/_weighter.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Iterable
 
 import numpy as np
+from scipy.integrate import quad  # pylint: disable=import-error
 
 from ._utils import get_column, get_table
 
@@ -140,6 +141,7 @@ def effective_area(
         energy_bins: ArrayLike,
         cos_zenith_bins: ArrayLike,
         mask: ArrayLike | None = None,
+        flux: Any = 1,  # default is 1 GeV^-1 cm^-2 sr^-1 flux
     ) -> NDArray[np.float64]:
         r"""Calculate The effective area for the given energy and zenith bins.
 
@@ -157,20 +159,24 @@ def effective_area(
 
 
         Args:
-            energy_bins : array_like
+            energy_bins: array_like
                 A length N+1 array of energy bin edges
-            cos_zenith_bins : array_like
+            cos_zenith_bins: array_like
                 A length M+1 array of cos(zenith) bin edges
             mask: array_like
                 boolean array where 1 indicates to use the event in the calculation.
                 Must have the same length as the simulation sample.
+            flux: Any
+                A model describing the flux of the primaries to weight against. For
+                possible types, see get_weights()
 
         Returns:
             array_like
                 An NxM array of effective areas. Where N is the number of energy bins and M
                 is the number of cos(zenith) bins.
 
         """
+        cm2_to_m2 = 1e-4
         energy_bins = np.array(energy_bins)
         cos_zenith_bins = np.array(cos_zenith_bins)
 
@@ -182,7 +188,7 @@ def effective_area(
         energy = self.get_weight_column("energy")
         cos_zen = self.get_weight_column("cos_zen")
 
-        weights = self.get_weights(1e-4)
+        weights = self.get_weights(flux)
         maska = np.full(weights.size, 1, dtype=bool) if mask is None else np.asarray(mask, dtype=bool)
 
         assert maska.shape == weights.shape
@@ -194,12 +200,24 @@ def effective_area(
             bins=[cos_zenith_bins, energy_bins],
         )
 
-        nspecies = len(np.unique(self.get_weight_column("pdgid")[maska]))
-
         assert np.array_equal(enbin, energy_bins)
         assert np.array_equal(czbin, cos_zenith_bins)
-        e_width, z_width = np.meshgrid(np.ediff1d(enbin), np.ediff1d(czbin))
-        return np.asarray(hist_val / (e_width * 2 * np.pi * z_width * nspecies), dtype=np.float64)
+        pdgids = np.unique(self.get_weight_column("pdgid")[maska])
+
+        if np.isscalar(flux):
+            flux_integrals = len(pdgids) * np.ediff1d(enbin)
+        elif callable(flux):
+            flux_integrals = np.asarray(
+                [
+                    sum(quad(flux, energy_bins[bin_index], energy_bins[bin_index + 1], args=(p,))[0] for p in pdgids)
+                    for bin_index in range(len(energy_bins) - 1)
+                ]
+            )
+        else:
+            mesg = f"flux of type {type(flux)} is supplied but only scalar flux or cosmic ray flux models are supported so far"
+            raise TypeError(mesg)
+        e_int, z_int = np.meshgrid(flux_integrals, np.ediff1d(czbin))
+        return np.asarray(cm2_to_m2 * hist_val / (e_int * 2 * np.pi * z_int), dtype=np.float64)
 
     def __add__(self: Weighter, other: Weighter | int) -> Weighter:
         if other == 0:

tests/test_weighter.py

@@ -207,6 +207,22 @@ def test_effective_area(self):
             self.weighter1.effective_area(np.linspace(5e5, 5e6, self.N1 + 1), [0, 1]),
             [np.full(self.N1, self.c1.etendue / 2e4 / np.pi)],
         )
+        self.assertAlmostEqual(
+            self.weighter1.effective_area([5e5, 5e6], [0, 1], flux=TIG1996())[0][0],
+            149998.97822505102,
+            6,
+        )
+        self.assertAlmostEqual(
+            self.weighter1.effective_area([5e5, 5e6], [0, 1], np.ones(self.N1, dtype=bool), flux=TIG1996())[0][0],
+            149998.97822505102,
+            6,
+        )
+
+        with self.assertRaises(ValueError):
+            self.weighter1.effective_area([5e5, 5e6], [0, 1], flux="flux")
+
+        with self.assertRaises(TypeError):
+            self.weighter1.effective_area([5e5, 5e6], [0, 1], flux=list(range(self.N1)))
 
     def test_weighter_addition(self):
         weighter_sum = self.weighter1 + self.weighter1