diff --git a/src/simweights/_weighter.py b/src/simweights/_weighter.py index c3c7380..de49465 100644 --- a/src/simweights/_weighter.py +++ b/src/simweights/_weighter.py @@ -142,20 +142,25 @@ def effective_area( 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. + *, + return_stddev: bool = False, + ) -> NDArray[np.float64] | tuple[NDArray[np.float64], NDArray[np.float64]]: + r"""Calculate the effective area for the given energy and zenith bins. - This is accomplished by histogramming the generation surface the simulation sample + This is accomplished by histogramming the generation surface of the simulation sample in energy and zenith bins and dividing by the size of the energy and solid angle of each bin. - If mask is passed as a parameter, only events which are included in the mask are used. - Effective areas are given units of :math:`\mathrm{m}^2` + If a mask is passed as a parameter, only events which are included in the mask are used. + Effective areas are given in units of :math:`\mathrm{m}^2` .. Note :: If the sample contains more than one type of primary particle, then the result will be averaged over the number of particles. This is usually what you want. However, it can cause strange behavior if there is a small number of one type. In this case, the mask - should be used to select the particle types individually. + should be used to select the particle types individually. Alternatively, a flux model + can be specified to apply a weighting to the individual primary particle types. If + return_stddev is set to True, a tuple of effective area histograms and their standard + deviation are returned. Args: @@ -169,11 +174,16 @@ def effective_area( flux: Any A model describing the flux of the primaries to weight against. For possible types, see get_weights() + return_stddev: bool + boolean to specify if only effective area (default) or a tuple of + effective area and its standard deviation is returned Returns: - array_like + array_like | tuple(array_like, 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. + If return_stddev, a tuple of two NxM arrays (effective area and its + standard deviation) """ cm2_to_m2 = 1e-4 @@ -200,6 +210,13 @@ def effective_area( bins=[cos_zenith_bins, energy_bins], ) + hist_val_variance, _, _ = np.histogram2d( + cos_zen[maska], + energy[maska], + weights=(weights[maska]) ** 2, + bins=[cos_zenith_bins, energy_bins], + ) + assert np.array_equal(enbin, energy_bins) assert np.array_equal(czbin, cos_zenith_bins) pdgids = np.unique(self.get_weight_column("pdgid")[maska]) @@ -217,7 +234,15 @@ def effective_area( 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) + output: NDArray[np.float64] | tuple[NDArray[np.float64], NDArray[np.float64]] + if return_stddev: + output = ( + np.asarray(cm2_to_m2 * hist_val / (e_int * 2 * np.pi * z_int), dtype=np.float64), + np.asarray(cm2_to_m2 * np.sqrt(hist_val_variance) / (e_int * 2 * np.pi * z_int), dtype=np.float64), + ) + else: + output = np.asarray(cm2_to_m2 * hist_val / (e_int * 2 * np.pi * z_int), dtype=np.float64) + return output def __add__(self: Weighter, other: Weighter | int) -> Weighter: if other == 0: diff --git a/tests/test_weighter.py b/tests/test_weighter.py index c3327e5..897f966 100755 --- a/tests/test_weighter.py +++ b/tests/test_weighter.py @@ -222,6 +222,13 @@ def test_effective_area(self): 149998.7936752823, 6, ) + self.assertAlmostEqual( + self.weighter1.effective_area([5e5, 5e6], [0, 1], flux=lambda energy, pdgid: energy ** (-2.7), return_stddev=True)[ + 1 + ][0][0], + 102170.85106127625, + 6, + ) with self.assertRaises(ValueError): self.weighter1.effective_area([5e5, 5e6], [0, 1], flux="flux")