Make vector package optional

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