enh: add vectors for boosting and return, refactor some internals

CHANGELOG.rst

@@ -9,6 +9,8 @@ Develop
 Major Features and Improvements
 -------------------------------
 
+- integrating `vector <https://vector.readthedocs.io/en/latest/index.html>`_ support for ``generate``: ``boost_to`` can be a Momentum Lorentz vector and return the boosted particles as a vector using ``as_vectors=True``.
+
 Behavioral changes
 ------------------
 

docs/usage.rst

@@ -5,11 +5,10 @@ Usage
 The base of ``phasespace`` is the ``GenParticle`` object.
 This object, which represents a particle, either stable or decaying, has only one mandatory argument, its name.
 
-In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either
-a number or ``tf.constant``, or a function.
+In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either be array-like, or a function.
 Functions are used to specify the mass of particles such as resonances, which are not fixed but vary according to
 a broad distribution.
-These mass functions get three arguments, and must return a ``TensorFlow`` Tensor:
+These mass functions get three arguments, and must return an array-like object of the same shape as the input arguments.
 
 - The minimum mass allowed by the decay chain, which will be of shape `(n_events,)`.
 - The maximum mass available, which will be of shape `(n_events,)`.
@@ -22,24 +21,27 @@ A simple example
 --------------------------
 
 
-With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the ``GenParticle``
-class (which returns the class itself). As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which
-:math:`K^*\to K\pi` with a fixed mass, we would write:
+With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the :py:class:~`phasespace.GenParticle`
+class. As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which
+:math:`K^*\to K\pi` with a fixed mass, one would write:
 
 .. jupyter-execute::
 
-   from phasespace import GenParticle
+    from phasespace import GenParticle
+    import numpy as np
+    from particle import literals as lp
+    import vector
 
-   B0_MASS = 5279.58
-   KSTARZ_MASS = 895.81
-   PION_MASS = 139.57018
-   KAON_MASS = 493.677
+    B0_MASS = lp.B_0.mass
+    KSTARZ_MASS = lp.Kst_892_0.mass
+    PION_MASS = lp.pi_plus.mass
+    KAON_MASS = lp.K_plus.mass
 
-   pion = GenParticle('pi-', PION_MASS)
-   kaon = GenParticle('K+', KAON_MASS)
-   kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon)
-   gamma = GenParticle('gamma', 0)
-   bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma)
+    pion = GenParticle('pi-', PION_MASS)
+    kaon = GenParticle('K+', KAON_MASS)
+    kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon)
+    gamma = GenParticle('gamma', 0)
+    bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma)
 
 .. thebe-button:: Run this interactively
 
@@ -49,27 +51,42 @@ The method returns:
 
 - The normalized weights of each event, as an array of dimension (n_events,).
 - The 4-momenta of the generated particles as values of a dictionary with the particle name as key. These momenta
-  are expressed as arrays of dimension (n_events, 4).
+  are *either* expressed as arrays of dimension (n_events, 4) or :py:class:~`vector.Vector4D` objects, depending on the
+  ``as_vectors`` flag given to ``generate``.
 
 .. jupyter-execute::
 
-   N_EVENTS = 1000
+    N_EVENTS = 1000
 
-   weights, particles = bz.generate(n_events=N_EVENTS)
+    weights, particles = bz.generate(n_events=N_EVENTS, as_vectors=True)
+    # or
+    weights, particles = bz.generate(n_events=N_EVENTS)
 
-The ``generate`` method return an eager ``Tensor``: this is basically a wrapped numpy array. With ``Tensor.numpy()``,
-it can always directly be converted to a numpy array (if really needed).
+(The  array-like objects can always directly be converted to a numpy array (if really needed) through ``np.asarray(obj)``.)
 
 Boosting the particles
 --------------------------
 
 
 The particles are generated in the rest frame of the top particle.
-To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in both
-``generate`` and ~`tf.Tensor`. This latter approach can be useful if the momentum of the top particle
+To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in
+``generate``. This latter approach can be useful if the momentum of the top particle
 is generated according to some distribution, for example the kinematics of the LHC (see ``test_kstargamma_kstarnonresonant_lhc``
 and ``test_k1gamma_kstarnonresonant_lhc`` in ``tests/test_physics.py`` to see how this could be done).
 
+The ``boost_to`` argument can be a 4-momentum array of shape (n_events, 4) with (px, py, yz, energy) or a :py:class:~`vector.Vector4D` (both a momentum and a Lorentz vector).
+
+.. jupyter-execute::
+
+    N_EVENTS = 1000
+
+    # Generate the top particle with a momentum of 100 GeV
+    top_momentum = np.array([0, 0, 100, np.sqrt(100**2 + B0_MASS**2)])
+    # or
+    top_momentum = vector.array({'px': [0], 'py': [0], 'pz': [100], 'E': [np.sqrt(100**2 + B0_MASS**2)]})
+    weights, particles = bz.generate(n_events=N_EVENTS, boost_to=top_momentum)
+
+
 Weights
 --------------------------
 
@@ -83,12 +100,11 @@ and the particle dictionary.
     print(particles)
 
 
-It is worth noting that the graph generation is cached even when using ``generate``, so iterative generation
-can be performed using normal python loops without loss in performance:
+Iterative generation can be performed using normal python loops without loss in performance:
 
 .. jupyter-execute::
 
-   for i in range(5):
+    for i in range(5):
        weights, particles = bz.generate(n_events=100)
        # ...
        # (do something with weights and particles)
@@ -102,7 +118,7 @@ Resonances with variable mass
 
 To generate the mass of a resonance, we need to give a function as its mass instead of a floating number.
 This function should take as input the per-event lower mass allowed, per-event upper mass allowed and the number of
-events, and should return a ~`tf.Tensor` with the generated masses and shape (nevents,). Well suited for this task
+events, and should return an array-like object with the generated masses and shape (nevents,). Well suited for this task
 are the `TensorFlow Probability distributions <https://www.tensorflow.org/probability/api_docs/python/tfp/distributions>`_
 or, for more customized mass shapes, the
 `zfit pdfs <https://zfit.github.io/zfit/model.html#tensor-sampling>`_ (currently an
@@ -114,20 +130,20 @@ write the :math:`B^{0}\to K^{*}\gamma` decay chain as (more details can be found
 .. jupyter-execute::
     :hide-output:
 
-    import tensorflow as tf
+    from phasespace import numpy as tnp
     import tensorflow_probability as tfp
     from phasespace import GenParticle
 
     KSTARZ_MASS = 895.81
     KSTARZ_WIDTH = 47.4
 
     def kstar_mass(min_mass, max_mass, n_events):
-       min_mass = tf.cast(min_mass, tf.float64)
-       max_mass = tf.cast(max_mass, tf.float64)
-       kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64)
-       kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64)
+       min_mass = tnp.asarray(min_mass, tnp.float64)
+       max_mass = tnp.asarray(max_mass, tnp.float64)
+       kstar_width_cast = tnp.asarray(KSTARZ_WIDTH, tnp.float64)
+       kstar_mass_cast = tnp.asarray(KSTARZ_MASS, tnp.float64)
 
-       kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,))
+       kstar_mass = tnp.broadcast_to(kstar_mass_cast, shape=(n_events,))
        if KSTARZ_WIDTH > 0:
            kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass,
                                                           scale=kstar_width_cast,
@@ -157,35 +173,29 @@ If the names are not given, `top` and `p_{i}` are assigned. For example, to gene
 
 .. jupyter-execute::
 
-   import phasespace
+    import phasespace
+    from particle import literals as lp
 
-   N_EVENTS = 1000
+    N_EVENTS = 1000
 
-   B0_MASS = 5279.58
-   PION_MASS = 139.57018
-   KAON_MASS = 493.677
+    B0_MASS = lp.B_0.mass
+    PION_MASS = lp.pi_plus.mass
+    KAON_MASS = lp.K_plus.mass
 
-   decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS],
-                                  top_name="B0", names=["pi", "K"])
-   weights, particles = decay.generate(n_events=N_EVENTS)
+    decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS],
+                                   top_name="B0", names=["pi", "K"])
+    weights, particles = decay.generate(n_events=N_EVENTS)
 
 In this example, ``decay`` is simply a ``GenParticle`` with the corresponding children.
 
 
 Eager execution
 ---------------
 
-By default, `phasespace` uses TensorFlow to build a graph of the computations. This is usually more
-performant, especially if used multiple times. However, this has the disadvantage that _inside_
-`phasespac`, the actual values are not computed on Python runtime, e.g. if a breakpoint is set
-the values of a `tf.Tensor` won't be available.
-
-TensorFlow (since version 2.0) however can easily switch to so called "eager execution": in this
-mode, it behaves the same as Numpy; values are computed instantly and the Python code is not only
-executed once but every time.
+By default, `phasespace` uses JIT (*just-in-time*) compilation of TensorFlow to greatly speed up the generation of events. Simplified, this means that the first time a decay is generated, a symbolic array *without a concrete value* is used and the computation is remembered. As a user calling the function, you will not notice this, the output will be the same as if the function was executed eagerly.
+The consequence is two-fold: on one hand the initial overhead is higher with a significant speedup for subsequent generations, on the other hand, the values of the generated particles *inside the function* are not available in pure Python (e.g. for debugging basically).
 
-To switch this on or off, the global flag in TensorFlow `tf.config.run_functions_eagerly(True)` or
-the enviroment variable "PHASESPACE_EAGER" (which switches this flag) can be used.
+If you need to debug the internals, using ``tf.config.run_functions_eagerly(True)`` (or the environment variable ``"PHASESPACE_EAGER=1"``) will make everything run numpy-like.
 
 Random numbers
 --------------

phasespace/__init__.py

@@ -13,17 +13,19 @@
 
 __credits__ = ["Jonas Eschle <[email protected]>"]
 
-__all__ = ["nbody_decay", "GenParticle", "random"]
+__all__ = ["nbody_decay", "GenParticle", "random", "to_vectors", "numpy"]
 
-import tensorflow as tf
+import tensorflow.experimental.numpy as numpy
 
 from . import random
-from .phasespace import GenParticle, nbody_decay
+from .phasespace import GenParticle, nbody_decay, to_vectors
 
 
 def _set_eager_mode():
     import os
 
+    import tensorflow as tf
+
     is_eager = bool(os.environ.get("PHASESPACE_EAGER"))
     tf.config.run_functions_eagerly(is_eager)
 

phasespace/backend.py

@@ -1,21 +1,11 @@
 import tensorflow as tf
 
 RELAX_SHAPES = True
-if int(tf.__version__.split(".")[1]) < 5:  # smaller than 2.5
-    jit_compile_argname = "experimental_compile"
-else:
-    jit_compile_argname = "jit_compile"
-function = tf.function(
-    autograph=False,
-    experimental_relax_shapes=RELAX_SHAPES,
-    **{jit_compile_argname: False},
-)
+function = tf.function(autograph=False, jit_compile=False)
 function_jit = tf.function(
-    autograph=False,
-    experimental_relax_shapes=RELAX_SHAPES,
-    **{jit_compile_argname: True},
+    autograph=False, reduce_retracing=RELAX_SHAPES, jit_compile=True
 )
 
 function_jit_fixedshape = tf.function(
-    autograph=False, experimental_relax_shapes=False, **{jit_compile_argname: True}
+    autograph=False, reduce_retracing=False, jit_compile=True
 )

phasespace/fromdecay/genmultidecay.py

@@ -36,8 +36,7 @@ def from_dict(
         """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format.
 
         Args:
-            dec_dict:
-                The input dict from which the GenMultiDecay object will be created from.
+            dec_dict: The input dict from which the GenMultiDecay object will be created from.
                 A dec_dict has the same structure as the dicts used in DecayLanguage, see the examples below.
             mass_converter: A dict with mass function names and their corresponding mass functions.
                 These functions should take the particle mass and the mass width as inputs
@@ -53,6 +52,7 @@ def from_dict(
                 zfit parameter to every place where this particle appears in dec_dict. If the zfit parameter
                 is specified for a particle which is also included in particle_model_map, the zfit parameter
                 mass function name will be prioritized.
+
         Returns:
             The created GenMultiDecay object.
 
@@ -123,7 +123,7 @@ def from_dict(
 
         Notes:
             For a more in-depth tutorial, see the tutorial on GenMultiDecay in the
-            [documentation](https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html).
+            `documentation <https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html>`_.
         """
         if tolerance is None:
             tolerance = cls.MASS_WIDTH_TOLERANCE