From 0d3c086af1222b84c3f5ff53b7474844b994d957 Mon Sep 17 00:00:00 2001 From: Julian Saffer <31139467+jsaffer@users.noreply.github.com> Date: Tue, 31 Dec 2024 06:06:27 +0100 Subject: [PATCH] Adding flux weighting option to the effective_area function of Weighter objects (#47) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 <11620178+kjmeagher@users.noreply.github.com> --- pyproject.toml | 2 +- src/simweights/_spatial.py | 2 +- src/simweights/_utils.py | 8 ++++---- src/simweights/_weighter.py | 32 +++++++++++++++++++++++++------- tests/test_weighter.py | 16 ++++++++++++++++ 5 files changed, 47 insertions(+), 13 deletions(-) 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