Sample momentum from uniform spherical distribution #178

graphicle/data.py

@@ -1333,6 +1333,84 @@ def __attrs_post_init__(self):
         self.dtype = np.dtype(list(zip("xyze", it.repeat("<f8"))))
         self._HANDLED_TYPES = (np.ndarray, nm.Number, cla.Sequence)
 
+    @classmethod
+    def from_spherical_uniform(
+        cls,
+        size: int,
+        max_energy: float,
+        massless: float = 0.0,
+        seed: ty.Optional[int] = None,
+    ) -> "MomentumArray":
+        """Returns a ``MomentumArray`` whose elements are sampled from
+        uniform distributions of energy and 3-momentum.
+
+        .. versionadded:: 0.4.0
+
+        Parameters
+        ----------
+        size : int
+            Number of elements.
+        max_energy : float
+            Upper bound for the energy of the momenta.
+        massless : float
+            Probability for any given momentum element to be massless.
+            Default is ``0.0``.
+        seed : int, optional
+            Random number generator seed. Use the same seed for
+            consistent results.
+
+        Returns
+        -------
+        MomentumArray
+            Set of momenta, sampled uniformly in the energy component,
+            and with uniform probability density over the surface of a
+            3-sphere for the spatial components.
+
+        Raises
+        ------
+        ValueError
+            If massless is not within the interval [0.0, 1.0].
+
+        Notes
+        -----
+        Some 'massless' particles may have small, but finite masses.
+        This is due to numerical errors associated with floating point
+        calculations.
+        """
+        if not (0.0 <= massless <= 1.0):
+            raise ValueError(
+                "Probability of massless particles must be in the "
+                "interval [0, 1]."
+            )
+        rng = np.random.default_rng(seed=seed)
+        energy = rng.uniform(0.0, max_energy, size)
+        p_mag = energy
+        if massless == 0.0:
+            mass = rng.uniform(0.0, energy)
+            p_mag = calculate._root_diff_two_squares(energy, mass)
+        elif massless < 1.0:
+            p_mag = p_mag.copy()
+            mask = rng.random(size) < massless
+            masked_e = energy[~mask]
+            mass = rng.uniform(0.0, masked_e)
+            p_mag[~mask] = calculate._root_diff_two_squares(masked_e, mass)
+        theta_polar = p_mag * np.exp(
+            1.0j * np.arccos(1.0 - 2.0 * rng.random(size, dtype=np.float64))
+        )
+        phi_polar = theta_polar.real * np.exp(
+            1.0j * rng.uniform(-np.pi, np.pi, size)
+        )
+        return cls(
+            np.hstack(
+                [
+                    phi_polar.real.reshape(-1, 1),
+                    phi_polar.imag.reshape(-1, 1),
+                    theta_polar.imag.reshape(-1, 1),
+                    energy.reshape(-1, 1),
+                ]
+            )
+        )
+
     @property
     def __array_interface__(self) -> ty.Dict[str, ty.Any]:
         return self._data.__array_interface__

graphicle/transform.py

@@ -178,7 +178,11 @@ def split_momentum(
     parent = _momentum_to_numpy(momentum)
     children = np.tile(parent, (2, 1))
     children[0, :] *= z
-    children[0, :3] @= rotation_matrix(angle, axis).T
+    # backwards compatibility, switched inplace syntax for explicit ufunc
+    # children[0, :3] @= rotation_matrix(angle, axis).T
+    np.matmul(
+        children[0, :3], rotation_matrix(angle, axis).T, out=children[0, :3]
+    )
     children[1, :] -= children[0, :]
     return gcl.MomentumArray(children)
 

tests/test_transform.py

@@ -0,0 +1,82 @@
+"""
+``test_transform``
+==================
+
+Unit tests for the data structures, probing their attributes and
+methods.
+"""
+import math
+
+import numpy as np
+import pytest
+
+import graphicle as gcl
+
+DoubleVector = gcl.base.DoubleVector
+
+
+def unit_vec_about_axis(vec: DoubleVector, axis: DoubleVector) -> DoubleVector:
+    """Normalises ``vec``, such that its closest distance from ``axis``
+    is unity. This enables us to consider it as a point along a 2D unit
+    circle about ``axis``, albeit embedded in 3D space.
+    """
+    if not math.isclose(np.linalg.norm(axis), 1.0):
+        raise ValueError("axis must be unit length.")
+    mag = np.linalg.norm(vec)
+    axis_angle = np.arccos(np.dot(vec, axis) / mag)
+    radius = mag * np.sin(axis_angle)
+    return vec / radius
+
+
+def angular_distance_about_axis(
+    vec_1: DoubleVector, vec_2: DoubleVector, axis: DoubleVector
+) -> float:
+    """Computes the angle subtended by the chord connecting projections
+    of ``vec_1`` and ``vec_2`` in the plane of the unit circle about
+    ``axis``.
+
+    Will always give a positive number, as direction is not considered.
+    """
+    chord = unit_vec_about_axis(vec_2, axis) - unit_vec_about_axis(vec_1, axis)
+    chord_length = np.linalg.norm(chord)
+    return abs(2.0 * np.arcsin(0.5 * chord_length).item())
+
+
+def test_momentum_split():
+    rng = np.random.default_rng()
+    for _ in range(100):
+        energy_fraction = rng.uniform(1.0e-4, 0.5)
+        angles = rng.uniform(-1.0, 1.0, 3)
+        angles[0] = np.arccos(angles[0])
+        angles[1:] *= np.pi
+        rotation_angle = angles[2].item()
+        rotation_axis = gcl.transform.SphericalAngle(*angles[:2].tolist())
+        parent = gcl.MomentumArray.from_spherical_uniform(1, 100.0)
+        children = gcl.transform.split_momentum(
+            parent, energy_fraction, rotation_angle, rotation_axis
+        )
+        child_sum = np.sum(children, axis=0)
+        assert np.allclose(parent, child_sum)
+        first_child = children[0].copy()
+        actual_angle = angular_distance_about_axis(
+            parent._data[0, :3],
+            first_child._data[0, :3],
+            np.array(gcl.transform._angles_to_axis(rotation_axis)),
+        )
+        assert math.isclose(abs(rotation_angle), actual_angle)
+        # invert the split of the first child:
+        first_child_inv = children[0].copy()
+        first_child_inv /= energy_fraction
+        np.matmul(
+            first_child_inv._data[0, :3],
+            gcl.transform.rotation_matrix(
+                -rotation_angle,
+                rotation_axis,
+            ).T,
+            out=first_child_inv._data[0, :3],
+        )
+        assert np.allclose(parent, first_child_inv)
+
+
+if __name__ == "__main__":
+    pytest.main()