diff --git a/graphicle/calculate.py b/graphicle/calculate.py index ed38e95..3105fc7 100644 --- a/graphicle/calculate.py +++ b/graphicle/calculate.py @@ -805,6 +805,10 @@ def thrust( The thrust of the event. axis : ndarray[float64], optional The x, y, and z components of the thrust axis, respectively. + + See Also + -------- + spherocity : Computes the spherocity from final state momenta. """ half_pi = 0.5 * math.pi domain = (-half_pi, half_pi) @@ -825,10 +829,32 @@ def thrust( @nb.njit(nb.types.Tuple((nb.float64, nb.float64[:]))(nb.float64[:], PMU_DTYPE)) -def spherocity_with_grad( +def _spherocity_with_grad( axis_angles: base.DoubleVector, momenta: base.VoidVector ) -> ty.Tuple[float, base.DoubleVector]: - """ """ + """Computes the spherocity for a unit vector given by a pair of + azimuth and inclination angles, as well as the gradient of the + spherocity with respect to these angles. + + Parameters + ---------- + axis_angles : ndarray[float64] + Array of shape (2,), containing azimuth and inclination angles + of the unit vector. + momenta : ndarray[void] + Structured array, containing fields labeled 'x', 'y', 'z', + of 64-bit floating point numbers for the components of a + momentum point cloud. + + Returns + ------- + spherocity_val : float + The value of spherocity, with the given unit vector. + spherocity_grad : ndarray[float64] + Array of shape (2,) containing the gradient of the spherocity + with respect to azimuth and inclination angles of the unit + vector, respectively. + """ phi, theta = axis_angles[0], axis_angles[1] # pre-compute trig ratios and spherocity axis cos_phi, sin_phi = math.cos(phi), math.sin(phi) @@ -876,6 +902,77 @@ def spherocity_with_grad( return spherocity, np.array([grad_phi, grad_theta], dtype=np.float64) +@ty.overload +def spherocity( + pmu: "MomentumArray", + return_axis: ty.Literal[False], + rng_seed: ty.Optional[int], +) -> float: + ... + + +@ty.overload +def spherocity( + pmu: "MomentumArray", + return_axis: ty.Literal[True], + rng_seed: ty.Optional[int], +) -> ty.Tuple[float, base.DoubleVector]: + ... + + +def spherocity( + pmu: "MomentumArray", + return_axis: bool = False, + rng_seed: ty.Optional[int] = None, +): + """Computes the spherocity of an event, from the final-state momenta. + + :group: calculate + + .. versionadded:: 0.3.8 + + Parameters + ---------- + pmu : MomentumArray + Momentum of hadronised particles in the final state of the event + record. + return_axis : bool + If ``True``, will return a tuple with the spherocity, and the + axis unit vector which was found to minimize spherocity. + rng_seed : int, optional + Initial guess for the axis unit vector is sampled from a uniform + random distribution, over the surface of a sphere. If passed, + will initialise the random number generator with the provided + seed, enabling reproducible results. + + Returns + ------- + spherocity : float + The spherocity of the event. + axis : ndarray[float64], optional + The x, y, and z components of the spherocity axis, respectively. + + See Also + -------- + thrust : Computes the thrust from final state momenta. + """ + half_pi = 0.5 * math.pi + domain = (-half_pi, half_pi) + rng = np.random.default_rng(seed=rng_seed) + guess = rng.uniform(*domain, size=2) + optim = spo.minimize( + fun=_spherocity_with_grad, + x0=guess, + bounds=(domain, domain), + jac=True, + args=(pmu.data,), + ) + sph_val = optim.fun + if return_axis: + return sph_val, _angles_to_axis(optim.x) + return sph_val + + @nb.njit( [ "float64(bool_[:], bool_[:], Optional(float64[:]))",