diff --git a/graphicle/data.py b/graphicle/data.py index 2cbc334..eb7399f 100644 --- a/graphicle/data.py +++ b/graphicle/data.py @@ -1333,6 +1333,84 @@ def __attrs_post_init__(self): self.dtype = np.dtype(list(zip("xyze", it.repeat(" "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__ diff --git a/graphicle/transform.py b/graphicle/transform.py index 1b72fbf..b812f31 100644 --- a/graphicle/transform.py +++ b/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) diff --git a/tests/test_transform.py b/tests/test_transform.py new file mode 100644 index 0000000..7136aad --- /dev/null +++ b/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()