diff --git a/phasespace/phasespace.py b/phasespace/phasespace.py index da9cbce..e73dec3 100644 --- a/phasespace/phasespace.py +++ b/phasespace/phasespace.py @@ -17,16 +17,20 @@ import inspect from collections.abc import Callable from math import pi +from typing import TYPE_CHECKING, NoReturn import numpy as np import tensorflow as tf import tensorflow.experimental.numpy as tnp -import vector from . import kinematics as kin from .backend import function, function_jit_fixedshape from .random import SeedLike, get_rng +if TYPE_CHECKING: + import vector + + RELAX_SHAPES = False @@ -669,6 +673,10 @@ def generate( """ rng = get_rng(seed) if boost_to is not None: + try: + import vector + except ImportError as error: + _raise_missing_vector_package(error) if isinstance(boost_to, vector.Vector): if not ( isinstance(boost_to, vector.Momentum) @@ -815,8 +823,18 @@ def to_vectors(particles: dict[str, tf.Tensor]) -> dict[str, vector.Momentum]: Return: dict: Dictionary of `vector.Momentum` instances with numpy arrays """ + try: + import vector + except ImportError as error: + _raise_missing_vector_package(error) 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)) return newparticles + + +def _raise_missing_vector_package(exception: ImportError) -> NoReturn: + raise ImportError( + "To use `boost_to`, the `vector` package has to be installed." + ) from exception