diff --git a/pyproject.toml b/pyproject.toml index 374d6de..a6d23e6 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/simweights/_spatial.py b/src/simweights/_spatial.py index e58df28..6d4b3d3 100644 --- a/src/simweights/_spatial.py +++ b/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]: diff --git a/src/simweights/_utils.py b/src/simweights/_utils.py index b3ca7fe..1cff4fc 100644 --- a/src/simweights/_utils.py +++ b/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,7 +91,7 @@ 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 @@ -99,7 +99,7 @@ def constcol(table: Any, colname: str, mask: NDArray[np.bool_] | None = None) -> """ 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() diff --git a/src/simweights/_weighter.py b/src/simweights/_weighter.py index 2896109..df44892 100644 --- a/src/simweights/_weighter.py +++ b/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,13 +159,16 @@ 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 @@ -171,6 +176,7 @@ def effective_area( 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: diff --git a/tests/test_weighter.py b/tests/test_weighter.py index 4433188..3987313 100755 --- a/tests/test_weighter.py +++ b/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