diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 79b86e0..dbd95ca 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,46 +7,42 @@ ci: autoupdate_schedule: quarterly repos: - repo: https://github.com/pre-commit/mirrors-prettier - rev: v3.0.3 + rev: v3.1.0 hooks: - id: prettier - repo: https://github.com/fsfe/reuse-tool rev: v2.1.0 hooks: - id: reuse - - repo: https://github.com/psf/black - rev: 23.9.1 + - repo: https://github.com/codespell-project/codespell + rev: v2.2.6 hooks: - - id: black + - id: codespell - repo: https://github.com/asottile/blacken-docs rev: "1.16.0" hooks: - id: blacken-docs args: [-l 100] - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.292 + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.7.0 hooks: - - id: ruff - args: [--fix, --show-fixes] + - id: mypy + files: simweights + additional_dependencies: [numpy] + exclude: ^contrib/ - repo: https://github.com/pycqa/pylint - rev: v3.0.0 + rev: v3.0.1 hooks: - id: pylint files: simweights exclude: ^contrib/ additional_dependencies: [numpy, pandas] - - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.5.1 - hooks: - - id: mypy - files: simweights - additional_dependencies: [numpy] - exclude: ^contrib/ - - repo: https://github.com/codespell-project/codespell - rev: v2.2.6 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.6 hooks: - - id: codespell - args: [-L, livetime] + - id: ruff + args: [--fix, --show-fixes] + - id: ruff-format - repo: https://github.com/PyCQA/doc8 rev: "v1.1.1" hooks: @@ -67,10 +63,9 @@ repos: - id: forbid-crlf - id: forbid-tabs - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 + rev: v4.5.0 hooks: - id: check-added-large-files - - id: check-ast - id: check-builtin-literals - id: check-case-conflict - id: check-docstring-first diff --git a/contrib/book_simweights_testdata.py b/contrib/book_simweights_testdata.py index 1925fa5..d41445a 100755 --- a/contrib/book_simweights_testdata.py +++ b/contrib/book_simweights_testdata.py @@ -4,6 +4,8 @@ # # SPDX-License-Identifier: BSD-2-Clause +"""Script to generate the test data used by simweights testing.""" + import os.path import sys import tarfile @@ -15,6 +17,7 @@ def fake_event_header(frame: dict) -> None: + """Create a fake event header so the splitter can do its thing.""" header = dataclasses.I3EventHeader() header.run_id = 0 header.event_id = fake_event_header.event_id diff --git a/contrib/print_flux.py b/contrib/print_flux.py index 0afb330..cec6d46 100755 --- a/contrib/print_flux.py +++ b/contrib/print_flux.py @@ -4,6 +4,8 @@ # # SPDX-License-Identifier: BSD-2-Clause +"""Script to print out values of the specified flux.""" + from sys import argv import numpy as np diff --git a/docs/conf.py b/docs/conf.py index 75bf846..d469fe0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,7 +2,7 @@ # # SPDX-License-Identifier: BSD-2-Clause -# Configuration file for the Sphinx documentation builder. +"""Configuration file for the Sphinx documentation builder.""" # # This file only contains a selection of the most common options. For a full # list see the documentation: diff --git a/examples/tutorial_corsika.py b/examples/tutorial_corsika.py index 03d1a65..56ba0f5 100755 --- a/examples/tutorial_corsika.py +++ b/examples/tutorial_corsika.py @@ -13,8 +13,7 @@ CORSIKA_DATASET_DIR = Path("/data/sim/IceCube/2016/filtered/level2/CORSIKA-in-ice/20789/") corsika_filelist = sorted( - str(f) - for f in CORSIKA_DATASET_DIR.glob("0000000-0000999/Level2_IC86.2016_corsika.020789.00000*.i3.zst") + str(f) for f in CORSIKA_DATASET_DIR.glob("0000000-0000999/Level2_IC86.2016_corsika.020789.00000*.i3.zst") ) weight_keys = [ diff --git a/examples/tutorial_nugen.py b/examples/tutorial_nugen.py index a726ceb..8183459 100755 --- a/examples/tutorial_nugen.py +++ b/examples/tutorial_nugen.py @@ -18,10 +18,7 @@ def get_most_energetic_muon(mmclist: simclasses.I3MMCTrackList) -> float: emax = 0 for muon in list(mmclist): particle = muon.particle - if ( - particle.type in (dataclasses.I3Particle.MuMinus, dataclasses.I3Particle.MuPlus) - and particle.total_energy > emax - ): + if particle.type in (dataclasses.I3Particle.MuMinus, dataclasses.I3Particle.MuPlus) and particle.total_energy > emax: emax = particle.total_energy return emax @@ -43,9 +40,7 @@ def get_most_energetic_muon(mmclist: simclasses.I3MMCTrackList) -> float: ] DATASET_DIR = Path("/data/sim/IceCube/2016/filtered/level2/neutrino-generator/21217") -filelist = sorted( - str(f) for f in DATASET_DIR.glob("0000000-0000999/Level2_IC86.2016_NuMu.021217.00000*.i3.zst") -) +filelist = sorted(str(f) for f in DATASET_DIR.glob("0000000-0000999/Level2_IC86.2016_NuMu.021217.00000*.i3.zst")) MCmuonEnergy_nugen = np.array([]) I3MCWeightDict: dict = {k: [] for k in weight_keys} diff --git a/pyproject.toml b/pyproject.toml index 3ad93ab..1f63fa3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,15 +44,26 @@ Source = "https://github.com/icecube/simweights" Documentation = "https://docs.icecube.aq/simweights/main" Collaboration = "https://icecube.wisc.edu" -[tool.black] -line-length = 108 -target-version = ['py38'] +[tool.pytest.ini_options] +minversion = 7.0 +testpaths = ["tests"] +log_cli_level = "INFO" +xfail_strict = true +filterwarnings = ["error"] +addopts = ["-ra", "--strict-config", "--strict-markers", "--cov=simweights", "-W ignore"] [tool.coverage.run] source = ["simweights"] command_line = "-m pytest" omit = ["*/simweights/cmdline.py"] +[tool.coverage.report] +exclude_also = ["from numpy.typing import"] + +[tool.codespell] +skip = '_build' +ignore-words-list = 'livetime' + [tool.mypy] plugins = "numpy.typing.mypy_plugin" show_error_codes = true @@ -60,45 +71,30 @@ warn_unreachable = true enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"] [tool.pylint.format] -max-line-length = "108" +max-line-length = "128" [tool.pylint.messages_control] disable = "C0114" -[tool.pytest.ini_options] -minversion = 7.0 -testpaths = ["tests"] -log_cli_level = "INFO" -xfail_strict = true -filterwarnings = ["error"] -addopts = ["-ra", "--strict-config", "--strict-markers", "--cov=simweights", "-W ignore"] - -[tool.doc8] -max-line-length=108 - [tool.ruff] select = ["ALL"] fixable = ["I"] ignore = [ - "ANN101", # missing-type-self "ANN401", # any-type "FBT", # flake8-boolean-trap - "INP", # flake8-no-pep420 - "T20", # flake8-print - "TCH", # flake8-type-checking "S101", # assert-used - "N801", - "PYI024", + "COM812", #confilcts with ruff formatter + "ISC001", #confilcts with ruff formatter ] -line-length = 108 +line-length = 128 target-version = "py38" +namespace-packages = ["examples", "contrib", "docs"] [tool.ruff.pydocstyle] convention = "google" [tool.ruff.per-file-ignores] -"contrib/*" = ["D"] -"docs/*" = ["D"] +"contrib/*" = ["T201"] "tests/*" = [ "D", # pydocstyle "N", # pep8-naming @@ -113,4 +109,8 @@ convention = "google" "examples/*" = [ "D", # pydocstyle "F401", # unused-import + "T201", # flake8-print ] + +[tool.doc8] +max-line-length=128 diff --git a/src/simweights/_fluxes.py b/src/simweights/_fluxes.py index 62042e6..6ff85b5 100644 --- a/src/simweights/_fluxes.py +++ b/src/simweights/_fluxes.py @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: BSD-2-Clause # mypy: disable-error-code="no-any-return" +# ruff: noqa: N801 """A collection of cosmic ray flux parametrizations. @@ -18,15 +19,18 @@ from __future__ import annotations from pathlib import Path -from typing import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Callable, Mapping, Sequence from numpy import asfarray, bool_, broadcast_arrays, exp, float64, genfromtxt, int32, piecewise, sqrt from numpy import sum as nsum -from numpy.typing import ArrayLike, NDArray from scipy.interpolate import CubicSpline # pylint: disable=import-error from ._pdgcode import PDGCode +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + + # pylint: disable=too-few-public-methods # flake8: noqa: N803 @@ -44,14 +48,14 @@ class CosmicRayFlux: _funcs: Sequence[float | Callable[[float], float]] = () def _condition( - self, + self: CosmicRayFlux, energy: NDArray[float64], # noqa: ARG002 pdgid: NDArray[int32], ) -> list[NDArray[bool_]]: # pylint: disable=unused-argument return [pdgid == p for p in self.pdgids] - def __call__(self, energy: ArrayLike, pdgid: ArrayLike) -> NDArray[float64]: + def __call__(self: CosmicRayFlux, energy: ArrayLike, pdgid: ArrayLike) -> NDArray[float64]: energy_arr, pdgid_arr = broadcast_arrays(energy, pdgid) pcond = self._condition(energy_arr, pdgid_arr) return piecewise(energy, pcond, self._funcs) @@ -124,6 +128,7 @@ class Hoerandel5(CosmicRayFlux): After Becherini et al.\ [#Becherini]_ (These are the same as used by Arne Schoenwald's version\ [#Schoenwald]_). """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -158,6 +163,7 @@ class GaisserHillas(CosmicRayFlux): From an internal report\ [#Gaisser1]_ and in Astropart. Phys\ [#Gaisser2]_ by Tom Gaisser. """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -185,6 +191,7 @@ class GaisserH3a(CosmicRayFlux): instead just the extra-galactic accelerators reaching their highest energy. """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -218,6 +225,7 @@ class GaisserH4a(CosmicRayFlux): In the model H4a, on the other hand, the extra-galactic component is assumed to be all protons. """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -243,6 +251,7 @@ class GaisserH4a_IT(CosmicRayFlux): This is the flux used as an "a priori" estimate of mass-composition to produce the IceTop-only flux\ [#Aartsen]_. """ + # pylint: disable=invalid-name pdgids = (PDGCode.PPlus, PDGCode.He4Nucleus, PDGCode.O16Nucleus, PDGCode.Fe56Nucleus) _funcs = ( @@ -266,6 +275,7 @@ class Honda2004(CosmicRayFlux): Note: the E_k notation means energy per nucleon! """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -282,7 +292,7 @@ class Honda2004(CosmicRayFlux): lambda E: (0.000445 / 56.26) * (E / 56.26 + 3.07 * exp(-0.41 * sqrt(E / 56.26))) ** (-2.68), ) - def _condition(self, energy: NDArray[float64], pdgid: NDArray[int32]) -> list[NDArray[bool_]]: + def _condition(self: Honda2004, energy: NDArray[float64], pdgid: NDArray[int32]) -> list[NDArray[bool_]]: energy_break = 100 return [ (pdgid == PDGCode.PPlus) * (energy < energy_break), @@ -301,10 +311,11 @@ class TIG1996(CosmicRayFlux): The parameterization was taken directly from an earlier paper by Thunman et al\ [#Thunman]_. Only the nucleon flux was given, so for simplicity we treat it as a proton-only flux. """ + pdgids = (PDGCode.PPlus,) _funcs = (lambda E: 1.70 * E**-2.7, lambda E: 1.74e2 * E**-3.0, 0) - def _condition(self, energy: NDArray[float64], pdgid: NDArray[int32]) -> list[NDArray[bool_]]: + def _condition(self: TIG1996, energy: NDArray[float64], pdgid: NDArray[int32]) -> list[NDArray[bool_]]: energy_break = 5e6 return [ (pdgid == PDGCode.PPlus) * (energy < energy_break), @@ -314,6 +325,7 @@ def _condition(self, energy: NDArray[float64], pdgid: NDArray[int32]) -> list[ND class GlobalFitGST(CosmicRayFlux): r"""Spectral fits by Gaisser, Stanev and Tilav\ [#GaisserStanevTilav]_.""" + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -322,9 +334,7 @@ class GlobalFitGST(CosmicRayFlux): PDGCode.Fe56Nucleus, ) _funcs = ( - lambda E: 0.7 * E**-2.66 * exp(-E / 1.2e5) - + 0.015 * E**-2.4 * exp(-E / 4e6) - + 0.0014 * E**-2.4 * exp(-E / 1.3e9), + lambda E: 0.7 * E**-2.66 * exp(-E / 1.2e5) + 0.015 * E**-2.4 * exp(-E / 4e6) + 0.0014 * E**-2.4 * exp(-E / 1.3e9), lambda E: 0.32 * E**-2.58 * exp(-E / 1.2e5 / 2) + 0.0065 * E**-2.3 * exp(-E / 4e6 / 2), lambda E: 0.01 * E**-2.40 * exp(-E / 1.2e5 / 7) + 0.0006 * E**-2.3 * exp(-E / 4e6 / 7), lambda E: 0.013 * E**-2.40 * exp(-E / 1.2e5 / 13) + 0.0007 * E**-2.3 * exp(-E / 4e6 / 13), @@ -336,6 +346,7 @@ class GlobalFitGST(CosmicRayFlux): class GlobalSplineFit(CosmicRayFlux): r"""Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski].""" + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -367,7 +378,7 @@ class GlobalSplineFit(CosmicRayFlux): PDGCode.Ni59Nucleus, ) - def __init__(self) -> None: + def __init__(self: GlobalSplineFit) -> None: data = genfromtxt(Path(__file__).parent / "gsf_data_table.txt") energy = data.T[0] elements = data.T[1:] @@ -379,6 +390,7 @@ class GlobalSplineFit5Comp(CosmicRayFlux): GSF is a Data-driven spline fit of the cosmic ray spectrum by Dembinski et. al. \ [#GSFDembinski]. """ + pdgids = ( PDGCode.PPlus, PDGCode.He4Nucleus, @@ -389,7 +401,7 @@ class GlobalSplineFit5Comp(CosmicRayFlux): groups = ((1, 1), (2, 5), (6, 11), (12, 15), (16, 27)) - def __init__(self) -> None: + def __init__(self: GlobalSplineFit5Comp) -> None: data = genfromtxt(Path(__file__).parent / "gsf_data_table.txt") energy = data.T[0] elements = data.T[1:] @@ -409,7 +421,7 @@ class FixedFractionFlux(CosmicRayFlux): """ def __init__( - self, + self: FixedFractionFlux, fractions: Mapping[PDGCode, float], basis: CosmicRayFlux | None = None, normalized: bool = True, @@ -432,7 +444,7 @@ def __init__( if normalized: assert sum(self.fracs) == 1.0 # noqa: PLR2004 - def __call__(self, energy: ArrayLike, pdgid: ArrayLike) -> NDArray[float64]: + def __call__(self: FixedFractionFlux, energy: ArrayLike, pdgid: ArrayLike) -> NDArray[float64]: energy_arr, pdgid_arr = broadcast_arrays(energy, pdgid) fluxsum = sum(self.flux(energy_arr, p) for p in self.pdgids) cond = self._condition(energy_arr, pdgid_arr) diff --git a/src/simweights/_generation_surface.py b/src/simweights/_generation_surface.py index 2888203..7ffa71d 100644 --- a/src/simweights/_generation_surface.py +++ b/src/simweights/_generation_surface.py @@ -4,17 +4,27 @@ from __future__ import annotations -from collections import namedtuple from copy import deepcopy +from typing import TYPE_CHECKING, NamedTuple import numpy as np -from numpy.typing import ArrayLike, NDArray from ._pdgcode import PDGCode -from ._powerlaw import PowerLaw -from ._spatial import SpatialDist -SurfaceTuple = namedtuple("SurfaceTuple", ["pdgid", "nevents", "energy_dist", "spatial_dist"]) +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + + from ._powerlaw import PowerLaw # pragma: no cover + from ._spatial import SpatialDist # pragma: no cover + + +class SurfaceTuple(NamedTuple): + """Container for the components Generation Surfaces.""" + + pdgid: int | PDGCode + nevents: float + energy_dist: PowerLaw + spatial_dist: SpatialDist class GenerationSurface: @@ -145,12 +155,7 @@ def __eq__(self: GenerationSurface, other: object) -> bool: return True def __repr__(self: GenerationSurface) -> str: - return ( - self.__class__.__name__ - + "(" - + ",".join(repr(y) for x in self.spectra.values() for y in x) - + ")" - ) + return self.__class__.__name__ + "(" + ",".join(repr(y) for x in self.spectra.values() for y in x) + ")" def __str__(self: GenerationSurface) -> str: outstrs = [] @@ -160,9 +165,7 @@ def __str__(self: GenerationSurface) -> str: except ValueError: ptype = str(pdgid) - collections = [ - f"N={subspec.nevents} {subspec.energy_dist} {subspec.spatial_dist}" for subspec in specs - ] + collections = [f"N={subspec.nevents} {subspec.energy_dist} {subspec.spatial_dist}" for subspec in specs] outstrs.append(f" {ptype:>11} : " + "\n ".join(collections)) return "< " + self.__class__.__name__ + "\n" + "\n".join(outstrs) + "\n>" diff --git a/src/simweights/_icetop_weighter.py b/src/simweights/_icetop_weighter.py index 730fa92..e95fb5c 100644 --- a/src/simweights/_icetop_weighter.py +++ b/src/simweights/_icetop_weighter.py @@ -32,8 +32,7 @@ def sframe_icetop_surface(table: Any) -> GenerationSurface: np.cos(get_column(table, "min_zenith")[i]), ) surfaces.append( - get_column(table, "n_events")[i] - * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial), + get_column(table, "n_events")[i] * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial), ) retval = sum(surfaces) assert isinstance(retval, GenerationSurface) diff --git a/src/simweights/_nugen_weighter.py b/src/simweights/_nugen_weighter.py index ec83cba..29f0f0f 100644 --- a/src/simweights/_nugen_weighter.py +++ b/src/simweights/_nugen_weighter.py @@ -27,10 +27,7 @@ def nugen_spatial(table: Any) -> SpatialDist: # vector of the primary, this can be determined by checking `InjectionSurfaceR`. It will # be > 0 for circle injection and -1 for surface injection. In new versions >V6-00-00 it is not even # present indicating surface mode - if has_column(table, "InjectionSurfaceR"): - injection_radius = constcol(table, "InjectionSurfaceR") - else: - injection_radius = -1 + injection_radius = constcol(table, "InjectionSurfaceR") if has_column(table, "InjectionSurfaceR") else -1 if injection_radius > 0: return CircleInjector(injection_radius, min_cos, max_cos) diff --git a/src/simweights/_powerlaw.py b/src/simweights/_powerlaw.py index e492219..827c9a2 100644 --- a/src/simweights/_powerlaw.py +++ b/src/simweights/_powerlaw.py @@ -2,14 +2,17 @@ # # SPDX-License-Identifier: BSD-2-Clause +from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np -from numpy.typing import ArrayLike, NDArray from ._utils import SeedType, check_random_state +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + class PowerLaw: r"""A power-law continuous probability distribution. @@ -35,7 +38,7 @@ class PowerLaw: # pylint: disable=invalid-name - def __init__(self, g: float, a: float, b: float) -> None: + def __init__(self: PowerLaw, g: float, a: float, b: float) -> None: assert b > a self.g = float(g) self.a = float(a) @@ -48,20 +51,20 @@ def __init__(self, g: float, a: float, b: float) -> None: self.span = b - a - def _pdf(self, x: NDArray[np.float64]) -> NDArray[np.float64]: + def _pdf(self: PowerLaw, x: NDArray[np.float64]) -> NDArray[np.float64]: return np.asfarray(x**self.g / self.integral) - def _cdf(self, x: NDArray[np.float64]) -> NDArray[np.float64]: + def _cdf(self: PowerLaw, x: NDArray[np.float64]) -> NDArray[np.float64]: if self.G == 0: return np.asfarray(np.log(x / self.a) / self.integral) return np.asfarray((x**self.G - self.a**self.G) / self.G / self.integral) - def _ppf(self, q: NDArray[np.float64]) -> NDArray[np.float64]: + def _ppf(self: PowerLaw, q: NDArray[np.float64]) -> NDArray[np.float64]: if self.G == 0: return np.asfarray(self.a * np.exp(q * self.integral)) return np.asfarray((q * self.G * self.integral + self.a**self.G) ** (1 / self.G)) - def pdf(self, x: ArrayLike) -> NDArray[np.float64]: + def pdf(self: PowerLaw, x: ArrayLike) -> NDArray[np.float64]: r"""Probability density function. Args: @@ -74,7 +77,7 @@ def pdf(self, x: ArrayLike) -> NDArray[np.float64]: xa = np.asfarray(x) return np.piecewise(xa, [(xa >= self.a) & (xa <= self.b)], [self._pdf]) - def cdf(self, x: ArrayLike) -> NDArray[np.float64]: + def cdf(self: PowerLaw, x: ArrayLike) -> NDArray[np.float64]: r"""Cumulative distribution function. Args: @@ -86,7 +89,7 @@ def cdf(self, x: ArrayLike) -> NDArray[np.float64]: qa = np.asfarray(x) return np.piecewise(qa, [qa < self.a, qa > self.b], [0, 1, self._cdf]) - def ppf(self, q: ArrayLike) -> NDArray[np.float64]: + def ppf(self: PowerLaw, q: ArrayLike) -> NDArray[np.float64]: """Percent point function (inverse of `cdf`) at `q`. Args: @@ -98,7 +101,7 @@ def ppf(self, q: ArrayLike) -> NDArray[np.float64]: qa = np.asfarray(q) return np.piecewise(qa, [(qa >= 0) & (qa <= 1)], [self._ppf, np.nan]) - def rvs(self, size: Any = None, random_state: SeedType = None) -> NDArray[np.float64]: + def rvs(self: PowerLaw, size: Any = None, random_state: SeedType = None) -> NDArray[np.float64]: """Random variates. Args: @@ -115,10 +118,10 @@ def rvs(self, size: Any = None, random_state: SeedType = None) -> NDArray[np.flo rand_state = check_random_state(random_state) return self._ppf(np.asfarray(rand_state.uniform(0, 1, size))) - def __repr__(self) -> str: + def __repr__(self: PowerLaw) -> str: return f"{self.__class__.__name__}({self.g} ,{self.a}, {self.b})" - def __eq__(self, other: object) -> bool: + def __eq__(self: PowerLaw, other: object) -> bool: if not isinstance(other, PowerLaw): mesg = f"{self} cannot be compared to {other}" raise TypeError(mesg) diff --git a/src/simweights/_spatial.py b/src/simweights/_spatial.py index 1fa4862..966c0c1 100644 --- a/src/simweights/_spatial.py +++ b/src/simweights/_spatial.py @@ -1,18 +1,21 @@ # SPDX-FileCopyrightText: © 2022 the SimWeights contributors # # SPDX-License-Identifier: BSD-2-Clause +from __future__ import annotations -from typing import Union +from typing import TYPE_CHECKING, Union import numpy as np -from numpy.typing import ArrayLike, NDArray + +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray class CylinderBase: """Abstract base class for cylinder pdf classes.""" def __init__( - self, + self: CylinderBase, length: float, radius: float, cos_zen_min: float, @@ -32,7 +35,7 @@ def __init__( self._cap = 1e4 * np.pi * self.radius**2 self.etendue = float(self._diff_etendue(self.cos_zen_max) - self._diff_etendue(self.cos_zen_min)) - def projected_area(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def projected_area(self: CylinderBase, cos_zen: ArrayLike) -> NDArray[np.float64]: """Cross sectional area of a cylinder in cm^2. As seen from the angle described by cos_zen. @@ -42,29 +45,22 @@ def projected_area(self, cos_zen: ArrayLike) -> NDArray[np.float64]: assert np.all(cosz <= +1) return np.asfarray(self._cap * np.abs(cosz) + self._side * np.sqrt(1 - cosz**2)) - def _diff_etendue(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def _diff_etendue(self: CylinderBase, cos_zen: ArrayLike) -> NDArray[np.float64]: cosz = np.asfarray(cos_zen) assert np.all(cosz >= -1) assert np.all(cosz <= +1) return np.asfarray( - np.pi - * ( - self._cap * cosz * np.abs(cosz) - + self._side * (cosz * np.sqrt(1 - cosz**2) - np.arccos(cosz)) - ), + np.pi * (self._cap * cosz * np.abs(cosz) + self._side * (cosz * np.sqrt(1 - cosz**2) - np.arccos(cosz))), ) - def pdf(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def pdf(self: CylinderBase, cos_zen: ArrayLike) -> NDArray[np.float64]: """The probability density function for the given zenith angle.""" raise NotImplementedError - def __repr__(self) -> str: - return ( - f"{self.__class__.__name__}" - f"({self.length}, {self.radius}, {self.cos_zen_min}, {self.cos_zen_max})" - ) + def __repr__(self: CylinderBase) -> str: + return f"{self.__class__.__name__}" f"({self.length}, {self.radius}, {self.cos_zen_min}, {self.cos_zen_max})" - def __eq__(self, other: object) -> bool: + def __eq__(self: CylinderBase, other: object) -> bool: return ( isinstance(other, type(self)) and self.length == other.length @@ -89,10 +85,10 @@ class UniformSolidAngleCylinder(CylinderBase): """ - def _pdf(self, cos_zen: NDArray[np.float64]) -> NDArray[np.float64]: + def _pdf(self: UniformSolidAngleCylinder, cos_zen: NDArray[np.float64]) -> NDArray[np.float64]: return 1 / (2 * np.pi * (self.cos_zen_max - self.cos_zen_min) * self.projected_area(cos_zen)) - def pdf(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def pdf(self: UniformSolidAngleCylinder, cos_zen: ArrayLike) -> NDArray[np.float64]: cosz = np.asfarray(cos_zen) return np.piecewise(cosz, [(cosz >= self.cos_zen_min) & (cosz <= self.cos_zen_max)], [self._pdf]) @@ -113,11 +109,11 @@ class NaturalRateCylinder(CylinderBase): I \propto \pi\cdot r^2\cdot\sin\theta\cdot(\cos\theta+2/\pi\cdot l/r\cdot\sin\theta) """ - def __init__(self, length: float, radius: float, cos_zen_min: float, cos_zen_max: float) -> None: + def __init__(self: NaturalRateCylinder, length: float, radius: float, cos_zen_min: float, cos_zen_max: float) -> None: super().__init__(length, radius, cos_zen_min, cos_zen_max) self._normalization = 1 / self.etendue - def pdf(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def pdf(self: NaturalRateCylinder, cos_zen: ArrayLike) -> NDArray[np.float64]: cosz = np.asfarray(cos_zen) return np.piecewise( cosz, @@ -134,7 +130,7 @@ class CircleInjector: The etendue is just the area of the circle times the solid angle. """ - def __init__(self, radius: float, cos_zen_min: float, cos_zen_max: float) -> None: + def __init__(self: CircleInjector, radius: float, cos_zen_min: float, cos_zen_max: float) -> None: self.radius = radius self.cos_zen_min = cos_zen_min self.cos_zen_max = cos_zen_max @@ -142,12 +138,12 @@ def __init__(self, radius: float, cos_zen_min: float, cos_zen_max: float) -> Non self.etendue = 2 * np.pi * (self.cos_zen_max - self.cos_zen_min) * self._cap self._normalization = 1 / self.etendue - def projected_area(self, cos_zen: float) -> float: # noqa: ARG002 + def projected_area(self: CircleInjector, cos_zen: float) -> float: # noqa: ARG002 """Returns the cross sectional area of the injection area in cm^2.""" # pylint: disable=unused-argument return self._cap - def pdf(self, cos_zen: ArrayLike) -> NDArray[np.float64]: + def pdf(self: CircleInjector, cos_zen: ArrayLike) -> NDArray[np.float64]: """The probability density function for the given zenith angle.""" cosz = np.asfarray(cos_zen) return np.piecewise( @@ -156,10 +152,10 @@ def pdf(self, cos_zen: ArrayLike) -> NDArray[np.float64]: [self._normalization], ) - def __repr__(self) -> str: + def __repr__(self: CircleInjector) -> str: return f"CircleInjector({self.radius}, {self.cos_zen_min}, {self.cos_zen_max})" - def __eq__(self, other: object) -> bool: + def __eq__(self: CircleInjector, other: object) -> bool: return ( isinstance(other, CircleInjector) and self.radius == other.radius diff --git a/src/simweights/_utils.py b/src/simweights/_utils.py index 4afac58..6aebccb 100644 --- a/src/simweights/_utils.py +++ b/src/simweights/_utils.py @@ -5,14 +5,17 @@ from __future__ import annotations import numbers -from typing import Any, Union +from typing import TYPE_CHECKING, Any, Union import numpy as np from numpy.random import Generator, RandomState -from numpy.typing import ArrayLike, NDArray from ._pdgcode import PDGCode +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + + IntNumber = Union[int, np.integer] GeneratorType = Union[Generator, RandomState] SeedType = Union[GeneratorType, IntNumber, None] @@ -106,7 +109,7 @@ def check_run_counts(table: Any, nfiles: float) -> bool: # pragma: no cover else: status = "Fail" ret = False - print(f"Claimed Runs = {nfiles}, Found Runs = {len(runs)}, {status}") + print(f"Claimed Runs = {nfiles}, Found Runs = {len(runs)}, {status}") # noqa: T201 return ret diff --git a/src/simweights/_weighter.py b/src/simweights/_weighter.py index 18a105a..d7734cd 100644 --- a/src/simweights/_weighter.py +++ b/src/simweights/_weighter.py @@ -6,14 +6,17 @@ import inspect import warnings -from typing import Any, Callable, Iterable +from typing import TYPE_CHECKING, Any, Callable, Iterable import numpy as np -from numpy.typing import ArrayLike, NDArray -from ._generation_surface import GenerationSurface from ._utils import get_column, get_table +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + + from ._generation_surface import GenerationSurface # pragma: nocover + class Weighter: """Abstract base class from which all weighers derive. @@ -24,21 +27,21 @@ class Weighter: added together to form samples with different simulation parameters """ - def __init__(self, data: Iterable[Any], surface: GenerationSurface) -> None: + def __init__(self: Weighter, data: Iterable[Any], surface: GenerationSurface) -> None: self.data = list(data) self.surface = surface self.weight_cols: dict[str, NDArray[np.float64]] = {} self.colnames = sorted(self.weight_cols.keys()) self.size: int | None = None - def get_column(self, table: str, column: str) -> NDArray[np.float64]: + def get_column(self: Weighter, table: str, column: str) -> NDArray[np.float64]: """Helper function to get a specific column from the file.""" retval: NDArray[np.float64] = np.array([]) for datafile in self.data: retval = np.append(retval, get_column(get_table(datafile, table), column)) return retval - def add_weight_column(self, name: str, column: ArrayLike) -> None: + def add_weight_column(self: Weighter, name: str, column: ArrayLike) -> None: """Add a new column to be passed as parameters to flux models.""" col = np.array(column) if self.size: @@ -51,11 +54,11 @@ def add_weight_column(self, name: str, column: ArrayLike) -> None: self.weight_cols[name] = col self.colnames = sorted(self.weight_cols.keys()) - def get_weight_column(self, name: str) -> NDArray[np.float64]: + def get_weight_column(self: Weighter, name: str) -> NDArray[np.float64]: """Helper function to get a column needed in the weight calculation.""" return self.weight_cols[name] - def get_weights(self, flux: Any) -> NDArray[np.float64]: + def get_weights(self: Weighter, flux: Any) -> NDArray[np.float64]: """Calculate the weights for the sample for the given flux. Args: @@ -128,7 +131,7 @@ def get_weights(self, flux: Any) -> NDArray[np.float64]: return weights def effective_area( - self, + self: Weighter, energy_bins: ArrayLike, cos_zenith_bins: ArrayLike, mask: ArrayLike | None = None, @@ -193,7 +196,7 @@ def effective_area( e_width, z_width = np.meshgrid(np.ediff1d(enbin), np.ediff1d(czbin)) return np.asfarray(hist_val / (e_width * 2 * np.pi * z_width * nspecies)) - def __add__(self, other: Weighter | int) -> Weighter: + def __add__(self: Weighter, other: Weighter | int) -> Weighter: if other == 0: return self if not isinstance(other, Weighter): @@ -207,10 +210,10 @@ def __add__(self, other: Weighter | int) -> Weighter: weighter.add_weight_column(colname, np.append(column, other.weight_cols[colname])) return weighter - def __radd__(self, other: Weighter | int) -> Weighter: + def __radd__(self: Weighter, other: Weighter | int) -> Weighter: return self + other - def tostring(self, flux: None | object | Callable[[Any], ArrayLike] | ArrayLike = None) -> str: + def tostring(self: Weighter, flux: None | object | Callable[[Any], ArrayLike] | ArrayLike = None) -> str: """Creates a string with important information about this weighting object. Generation surface, event map, number of events, and effective area. @@ -232,5 +235,5 @@ def tostring(self, flux: None | object | Callable[[Any], ArrayLike] | ArrayLike output += f"Livetime : {weights.sum() / (weights ** 2).sum():8.6g} s\n" return output - def __str__(self) -> str: + def __str__(self: Weighter) -> str: return self.tostring() diff --git a/src/simweights/cmdline.py b/src/simweights/cmdline.py index 5fca6c9..156cfa4 100644 --- a/src/simweights/cmdline.py +++ b/src/simweights/cmdline.py @@ -63,5 +63,5 @@ def main() -> int: else: flux_model = None - print(wobj.tostring(flux_model)) + print(wobj.tostring(flux_model)) # noqa: T201 return 0 diff --git a/tests/test_corsika_datasets.py b/tests/test_corsika_datasets.py index 5b9444e..820b813 100755 --- a/tests/test_corsika_datasets.py +++ b/tests/test_corsika_datasets.py @@ -38,14 +38,7 @@ def untriggered_weights(self, f): ) / (cwm["PrimarySpectralIndex"] + 1) energy_weight = cwm["PrimaryEnergy"] ** cwm["PrimarySpectralIndex"] - return ( - 1e4 - * pflux - * energy_integral - / energy_weight - * cwm["AreaSum"] - / (cwm["NEvents"] * cwm["OverSampling"]) - ) + return 1e4 * pflux * energy_integral / energy_weight * cwm["AreaSum"] / (cwm["NEvents"] * cwm["OverSampling"]) def triggered_weights(self, f): i3cw = f["I3CorsikaWeight"] @@ -70,12 +63,8 @@ def triggered_weights(self, f): cap = 1e4 * np.pi * cylinder_radius**2 cos_minz = np.cos(min_zenith) cos_maxz = np.cos(max_zenith) - ET1 = cap * cos_minz * np.abs(cos_minz) + side * ( - cos_minz * np.sqrt(1 - cos_minz**2) - min_zenith - ) - ET2 = cap * cos_maxz * np.abs(cos_maxz) + side * ( - cos_maxz * np.sqrt(1 - cos_maxz**2) - max_zenith - ) + ET1 = cap * cos_minz * np.abs(cos_minz) + side * (cos_minz * np.sqrt(1 - cos_minz**2) - min_zenith) + ET2 = cap * cos_maxz * np.abs(cos_maxz) + side * (cos_maxz * np.sqrt(1 - cos_maxz**2) - max_zenith) etendue = np.pi * (ET1 - ET2) mask = ptype == i3cw["type"] diff --git a/tests/test_generation_surface.py b/tests/test_generation_surface.py index 935cbec..942247b 100755 --- a/tests/test_generation_surface.py +++ b/tests/test_generation_surface.py @@ -196,9 +196,7 @@ def test_powerlaw(self): cz = np.linspace(self.c1.cos_zen_min, self.c1.cos_zen_max, N) w = 1 / self.s0.get_epdf(2212, E, cz) - area = (self.p1.b - self.p1.a) * ( - 2e4 * self.c1.radius * np.pi**2 * (self.c1.radius + self.c1.length) - ) + area = (self.p1.b - self.p1.a) * (2e4 * self.c1.radius * np.pi**2 * (self.c1.radius + self.c1.length)) self.assertAlmostEqual( area, @@ -357,7 +355,6 @@ def test_equal_gsc(self): def test_repr_gsc(self): PPlus = PDGCode.PPlus # noqa: F841 - print(repr(self.gsc1)) self.assertEqual(self.gsc1, eval(repr(self.gsc1))) self.assertEqual(self.gsc2, eval(repr(self.gsc2))) self.assertEqual(self.gsc3, eval(repr(self.gsc3))) diff --git a/tests/test_genie_datasets.py b/tests/test_genie_datasets.py index 85359da..2ca69a5 100755 --- a/tests/test_genie_datasets.py +++ b/tests/test_genie_datasets.py @@ -35,9 +35,7 @@ def cmp_dataset(self, fname): total_prob = wd["TotalInteractionProbabilityWeight"] pli = -wd["PowerLawIndex"][0] - energy_integral = ( - (10 ** wd["MaxEnergyLog"][0]) ** (pli + 1) - (10 ** wd["MinEnergyLog"][0]) ** (pli + 1) - ) / (pli + 1) + energy_integral = ((10 ** wd["MaxEnergyLog"][0]) ** (pli + 1) - (10 ** wd["MinEnergyLog"][0]) ** (pli + 1)) / (pli + 1) energy_factor = 1 / (wd["PrimaryNeutrinoEnergy"] ** pli / energy_integral) one_weight = total_prob * energy_factor * solid_angle * injection_area np.testing.assert_allclose(one_weight, wd["OneWeight"]) diff --git a/tests/test_nugen_datasets.py b/tests/test_nugen_datasets.py index 0e1436b..957e267 100755 --- a/tests/test_nugen_datasets.py +++ b/tests/test_nugen_datasets.py @@ -71,9 +71,7 @@ def cmp_dataset(self, fname): power_law = w.surface.spectra[pdgid][0].energy_dist energy_factor = 1 / power_law.pdf(w.get_weight_column("energy")) - one_weight = ( - w.get_weight_column("event_weight") * energy_factor * solid_angle * injection_area - ) + one_weight = w.get_weight_column("event_weight") * energy_factor * solid_angle * injection_area np.testing.assert_allclose(one_weight, wd["OneWeight"]) np.testing.assert_allclose(w.get_weights(1), w0, 1e-5) diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 98cbcf3..cdf6736 100755 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -33,18 +33,18 @@ def check_diff_etendue(self, c, le, r): with self.assertRaises(AssertionError): c.projected_area(1.01) - self.assertAlmostEqual(c._diff_etendue(-1), -np.pi**2 * r * (r + 2 * le), 4) + self.assertAlmostEqual(c._diff_etendue(-1), -(np.pi**2) * r * (r + 2 * le), 4) self.assertAlmostEqual( c._diff_etendue(-0.5), - -np.pi**2 / 4 * r * (r + 2 / 3 * le * (3**1.5 / np.pi + 8)), + -(np.pi**2) / 4 * r * (r + 2 / 3 * le * (3**1.5 / np.pi + 8)), 4, ) self.assertAlmostEqual( c._diff_etendue(-(0.5**0.5)), - -np.pi**2 / 2 * r * (r + le * (2 / np.pi + 3)), + -(np.pi**2) / 2 * r * (r + le * (2 / np.pi + 3)), 4, ) - self.assertAlmostEqual(c._diff_etendue(0), -np.pi**2 * le * r) + self.assertAlmostEqual(c._diff_etendue(0), -(np.pi**2) * le * r) self.assertAlmostEqual( c._diff_etendue(0.5**0.5), np.pi**2 / 2 * r * (r + le * (2 / np.pi - 1)),