diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 616a5173..207dd0b5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,7 +24,7 @@ repos: - id: check-ast - repo: https://github.com/PyCQA/docformatter - rev: v1.7.7 + rev: v1.7.8 hooks: - id: docformatter args: @@ -42,7 +42,7 @@ repos: - id: rst-directive-colons - repo: https://github.com/PyCQA/isort - rev: 8.0.1 + rev: 9.0.2 hooks: - id: isort @@ -81,7 +81,7 @@ repos: - id: nbstripout - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.15.9" + rev: "v0.16.10" hooks: - id: ruff types_or: [ python, pyi, jupyter ] diff --git a/benchmark/bench_phasespace.py b/benchmark/bench_phasespace.py index 66370e86..5f964fa5 100644 --- a/benchmark/bench_phasespace.py +++ b/benchmark/bench_phasespace.py @@ -11,6 +11,7 @@ from timeit import default_timer import tensorflow as tf + from phasespace import phasespace sys.path.append(os.path.dirname(__file__)) diff --git a/data/download_test_files.py b/data/download_test_files.py index 4a34edc1..a61992dc 100644 --- a/data/download_test_files.py +++ b/data/download_test_files.py @@ -18,19 +18,19 @@ def _is_silent(): FILE_URLS = [ ( "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", - "https://cernbox.cern.ch/remote.php/dav/public-files/8mN10X8U7VGfaRc/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", # noqa: E501 + "https://cernbox.cern.ch/remote.php/dav/public-files/8mN10X8U7VGfaRc/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", ), ( "B2K1Gamma_RapidSim_7TeV_Tree.root", - "https://cernbox.cern.ch/remote.php/dav/public-files/pr3aM8n2hPT4Pag/B2K1Gamma_RapidSim_7TeV_Tree.root", # noqa: E501 + "https://cernbox.cern.ch/remote.php/dav/public-files/pr3aM8n2hPT4Pag/B2K1Gamma_RapidSim_7TeV_Tree.root", ), ( "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", - "https://cernbox.cern.ch/remote.php/dav/public-files/QuP2cHeISTTSLVv/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", # noqa: E501 + "https://cernbox.cern.ch/remote.php/dav/public-files/QuP2cHeISTTSLVv/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", ), ( "B2KstGamma_RapidSim_7TeV_Tree.root", - "https://cernbox.cern.ch/remote.php/dav/public-files/EH5yrCpGko7P7Mc/B2KstGamma_RapidSim_7TeV_Tree.root", # noqa: E501 + "https://cernbox.cern.ch/remote.php/dav/public-files/EH5yrCpGko7P7Mc/B2KstGamma_RapidSim_7TeV_Tree.root", ), ] diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index d525dbe8..d3d8ece0 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -32,10 +32,10 @@ "# Import libraries\n", "from pprint import pprint\n", "\n", + "import tensorflow as tf\n", "import zfit\n", + "from decaylanguage import DecayChain, DecayChainViewer, DecayMode, DecFileParser\n", "from particle import Particle\n", - "from decaylanguage import DecFileParser, DecayChainViewer, DecayChain, DecayMode\n", - "import tensorflow as tf\n", "\n", "from phasespace.fromdecay import GenMultiDecay" ] diff --git a/docs/conf.py b/docs/conf.py index bff9557c..50354612 100755 --- a/docs/conf.py +++ b/docs/conf.py @@ -30,7 +30,7 @@ warnings.filterwarnings("ignore") # Suppress TensorFlow C++ logging to stderr -import logging # noqa: E402 +import logging logging.getLogger("tensorflow").setLevel(logging.ERROR) @@ -43,7 +43,7 @@ except ImportError: pass -import phasespace # noqa: E402 +import phasespace # -- General configuration --------------------------------------------- diff --git a/src/phasespace/__init__.py b/src/phasespace/__init__.py index 1a316ecd..f81efe2d 100644 --- a/src/phasespace/__init__.py +++ b/src/phasespace/__init__.py @@ -13,9 +13,9 @@ __credits__ = ["Jonas Eschle "] -__all__ = ["nbody_decay", "GenParticle", "random", "to_vectors", "numpy"] +__all__ = ["GenParticle", "nbody_decay", "numpy", "random", "to_vectors"] -import tensorflow.experimental.numpy as numpy +from tensorflow.experimental import numpy from . import random from .phasespace import GenParticle, nbody_decay, to_vectors diff --git a/src/phasespace/fromdecay/__init__.py b/src/phasespace/fromdecay/__init__.py index 7ef066fc..19e61c95 100644 --- a/src/phasespace/fromdecay/__init__.py +++ b/src/phasespace/fromdecay/__init__.py @@ -8,7 +8,7 @@ import sys -from .genmultidecay import GenMultiDecay # noqa: F401 +from .genmultidecay import GenMultiDecay try: import zfit # noqa: F401 diff --git a/src/phasespace/fromdecay/genmultidecay.py b/src/phasespace/fromdecay/genmultidecay.py index 39d2a544..c0053f99 100644 --- a/src/phasespace/fromdecay/genmultidecay.py +++ b/src/phasespace/fromdecay/genmultidecay.py @@ -6,6 +6,7 @@ import tensorflow as tf import tensorflow.experimental.numpy as tnp from particle import Particle + from phasespace import GenParticle from .mass_functions import DEFAULT_CONVERTER @@ -28,9 +29,9 @@ def __init__(self, gen_particles: list[tuple[float, GenParticle]]): def from_dict( cls, dec_dict: dict, - mass_converter: dict[str, Callable] = None, - tolerance: float = None, - particle_model_map: dict[str, str] = None, + mass_converter: dict[str, Callable] | None = None, + tolerance: float | None = None, + particle_model_map: dict[str, str] | None = None, ): """Create a ``GenMultiDecay`` instance from a dict in the ``DecayLanguage`` package format. @@ -215,7 +216,7 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: name += " [0]" i = 1 while name in preexisting_particles: - name = name[: name.rfind("[")] + f"[{str(i)}]" + name = name[: name.rfind("[")] + f"[{i!s}]" i += 1 preexisting_particles.add(name) return name @@ -254,7 +255,7 @@ def _recursively_traverse( mass_converter: dict[str, Callable], particle_model_map: dict[str, str], tolerance: float, - preexisting_particles: set[str] = None, + preexisting_particles: set[str] | None = None, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. diff --git a/src/phasespace/kinematics.py b/src/phasespace/kinematics.py index 1215512e..3609d26a 100644 --- a/src/phasespace/kinematics.py +++ b/src/phasespace/kinematics.py @@ -7,6 +7,7 @@ """Basic kinematics.""" import tensorflow.experimental.numpy as tnp + from phasespace.backend import function, function_jit diff --git a/src/phasespace/phasespace.py b/src/phasespace/phasespace.py index a01fcae3..a2a64530 100644 --- a/src/phasespace/phasespace.py +++ b/src/phasespace/phasespace.py @@ -97,9 +97,7 @@ class GenParticle: converted to array-like. """ - def __init__( - self, name: str, mass: Callable | int | float | np.typing.ArrayLike - ) -> None: # noqa + def __init__(self, name: str, mass: Callable | float | np.typing.ArrayLike) -> None: self.name = name self.children = [] self._mass_val = mass @@ -185,7 +183,7 @@ def get_mass( @property def has_fixed_mass(self): - """bool: Is the mass a callable function?""" + """Bool: Is the mass a callable function?""" return not callable(self._mass) def set_children(self, *children): @@ -222,12 +220,12 @@ def set_children(self, *children): @property def has_children(self): - """bool: Does the particle have children?""" + """Bool: Does the particle have children?""" return bool(self.children) @property def has_grandchildren(self): - """bool: Does the particle have grandchildren?""" + """Bool: Does the particle have grandchildren?""" if not self.children: return False return any(child.has_children for child in self.children) @@ -257,14 +255,13 @@ def _preprocess(momentum, n_events): if len(momentum.shape) not in (1, 2): raise ValueError(f"Bad shape for momentum -> {list(momentum.shape)}") # Check compatibility of inputs - if len(momentum.shape) == 2: - if n_events is not None: - momentum_shape = momentum.shape[0] - if momentum_shape is None: - momentum_shape = tf.shape(momentum)[0] - momentum_shape = tnp.asarray(momentum_shape, tnp.int64) - else: - momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) + if len(momentum.shape) == 2 and n_events is not None: + momentum_shape = momentum.shape[0] + if momentum_shape is None: + momentum_shape = tf.shape(momentum)[0] + momentum_shape = tnp.asarray(momentum_shape, tnp.int64) + else: + momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) # tf.assert_equal( # n_events, # momentum_shape, @@ -698,9 +695,8 @@ def generate( f"The number of events requested ({n_events}) doesn't match the boost_to input size " f"of {boost_to.shape}" ) - if n_events is not None: - if boost_to.shape[0] not in (n_events, 1): - raise ValueError(message) + if n_events is not None and boost_to.shape[0] not in (n_events, 1): + raise ValueError(message) if not isinstance(n_events, tf.Variable): n_events = tnp.asarray(n_events, dtype=tnp.int64) weights, weights_max, parts, _ = self._recursive_generate( @@ -754,7 +750,6 @@ def generate_tensor( ValueError: If ``n_events`` and the size of ``boost_to`` don't match. See ``GenParticle.generate_unnormalized``. """ - # Run generation raise RuntimeError( "This function is removed. Use `generate` which does not return a Tensor as well." @@ -775,7 +770,9 @@ def __init__(self): ) -def nbody_decay(mass_top: float, masses: list, top_name: str = "", names: list = None): +def nbody_decay( + mass_top: float, masses: list, top_name: str = "", names: list | None = None +): """Shortcut to build an n-body decay of a GenParticle. If the particle names are not given, the top particle is called 'top' and the @@ -831,7 +828,7 @@ def to_vectors(particles: dict[str, tf.Tensor]) -> dict[str, vector.Momentum]: newparticles = {} for name, particle in particles.items(): px, py, pz, e = np.moveaxis(particle, -1, 0) # numpy "unstack" - newparticles[name] = vector.array(dict(px=px, py=py, pz=pz, energy=e)) + newparticles[name] = vector.array({"px": px, "py": py, "pz": pz, "energy": e}) return newparticles diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 77f6e7d1..63574308 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -5,6 +5,7 @@ import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal + from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER diff --git a/tests/fromdecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py index 6ed7417e..e93e4c2f 100644 --- a/tests/fromdecay/test_mass_functions.py +++ b/tests/fromdecay/test_mass_functions.py @@ -2,12 +2,13 @@ from collections.abc import Callable -import phasespace.fromdecay.mass_functions as mf import pytest import tensorflow as tf import tensorflow_probability as tfp from particle import Particle +import phasespace.fromdecay.mass_functions as mf + _kstarz = Particle.from_evtgen_name("K*0") KSTARZ_MASS = _kstarz.mass KSTARZ_WIDTH = _kstarz.width diff --git a/tests/helpers/decays.py b/tests/helpers/decays.py index 59860f1e..6ff09099 100644 --- a/tests/helpers/decays.py +++ b/tests/helpers/decays.py @@ -8,6 +8,7 @@ import tensorflow as tf import tensorflow_probability as tfp + from phasespace import GenParticle # Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat) diff --git a/tests/test_chain.py b/tests/test_chain.py index 8a70c061..af875879 100644 --- a/tests/test_chain.py +++ b/tests/test_chain.py @@ -11,11 +11,12 @@ import numpy as np import pytest + from phasespace import GenParticle sys.path.append(os.path.dirname(__file__)) -from .helpers import decays # noqa: E402 +from .helpers import decays def test_name_clashes(): diff --git a/tests/test_generate.py b/tests/test_generate.py index 98d046dc..accc3042 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -10,12 +10,13 @@ import sys import numpy as np -import phasespace import pytest +import phasespace + sys.path.append(os.path.dirname(__file__)) -from .helpers import decays # noqa: E402 +from .helpers import decays B0_MASS = decays.B0_MASS PION_MASS = decays.PION_MASS diff --git a/tests/test_nbody_decay.py b/tests/test_nbody_decay.py index f553d050..efd11df9 100644 --- a/tests/test_nbody_decay.py +++ b/tests/test_nbody_decay.py @@ -7,6 +7,7 @@ """Test n-body decay generator.""" import pytest + from phasespace import nbody_decay from .helpers import decays diff --git a/tests/test_physics.py b/tests/test_physics.py index fb22b91b..0f77d871 100644 --- a/tests/test_physics.py +++ b/tests/test_physics.py @@ -25,12 +25,13 @@ import matplotlib.pyplot as plt import tensorflow as tf import uproot + from phasespace import phasespace sys.path.append(os.path.dirname(__file__)) -from .helpers import decays, rapidsim # noqa: E402 -from .helpers.plotting import make_norm_histo # noqa: E402 +from .helpers import decays, rapidsim +from .helpers.plotting import make_norm_histo BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) PLOT_DIR = os.path.join(BASE_PATH, "tests", "plots") @@ -191,12 +192,12 @@ def run_kstargamma(input_file, kstar_width, b_at_rest, suffix, use_vector): booster = booster.transpose() if use_vector: booster = vector.array( - dict( - px=booster[:, 0], - py=booster[:, 1], - pz=booster[:, 2], - e=booster[:, 3], - ) + { + "px": booster[:, 0], + "py": booster[:, 1], + "pz": booster[:, 2], + "e": booster[:, 3], + } ) rapidsim_getter = rapidsim.get_tree decay = decays.b0_to_kstar_gamma(kstar_width=kstar_width) diff --git a/tests/test_random.py b/tests/test_random.py index ec8bf580..a0148d29 100644 --- a/tests/test_random.py +++ b/tests/test_random.py @@ -1,8 +1,9 @@ import numpy as np -import phasespace as phsp import pytest import tensorflow as tf +import phasespace as phsp + @pytest.mark.parametrize( "seed", [lambda: 15, lambda: tf.random.Generator.from_seed(15)]