Adding flux weighting option to the effective_area function of Weight…

…er objects (#47)

* Adding flux weighting option to the effective_area function of Weighter objects

Instead of weighting to a 1 GeV⁻¹ m⁻² sr⁻¹ flux, one can now chose to use any of the cosmic ray flux models as well in order take into account the spectral shape and varying mass composition inside energy bins. The scalar flux is still possible and the default (1e-4).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Changed order of arguments of the effective_area function

The new optional flux argument is now the last in the list of arguments. This way, confusion with the mask argument, which remains in 3rd place, is avoided.

* Addrerssing some of the warnings and suggestions from pylint and ruff

This should fix C3001 (lambda expression replaced by def function definition), R1728 (using generator for sum of flux components), D417 (added description for flux argument), E731 (lambda -> def), N803 (E -> energy), TRY003 and EM101 (assigned error message to variable), also I changed the ValueError to a TypeError which should be more appropriate

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* More fixes of pylint / ruff complaints

Adding type annotations for flux_func function, and moving the return outside of the it/elif/else condition.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean up flux weighting for effective_area

* Adding tests for weighted effective area, also covering the ValueError of invalid flux argument and the TypeError of a valid flux argument that is not supported yet

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix typing annotations

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kevin Meagher <[email protected]>

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