From 90bedf4ee7596a29e17c722f9c083327a9564837 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Sat, 4 Nov 2023 13:02:25 +0000 Subject: [PATCH] added routine for thrust #165 --- graphicle/calculate.py | 71 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/graphicle/calculate.py b/graphicle/calculate.py index 9d62e70..7cd9787 100644 --- a/graphicle/calculate.py +++ b/graphicle/calculate.py @@ -19,6 +19,7 @@ import numba as nb import numpy as np import numpy.lib.recfunctions as rfn +import scipy.optimize as spo from deprecation import deprecated from . import base @@ -36,9 +37,13 @@ "flow_trace", "aggregate_momenta", "cluster_coeff_distbn", + "thrust", ] +PMU_DTYPE = nb.typeof(np.dtype(list(zip("xyze", (" float: + max_component = max(abs(x), abs(y), abs(z)) + max_recip = 1.0 / max_component + x *= max_recip + y *= max_recip + z *= max_recip + return max_component * np.sqrt(x * x + y * y + z * z) + + +@nb.njit("UniTuple(float64, 3)(float64, float64)") +def _angles_to_axis( + azimuth: float, inclination: float +) -> ty.Tuple[float, float, float]: + polar = cmath.rect(1.0, azimuth) * cmath.rect(1.0, inclination) + real, phase = polar.real, cmath.phase(polar) + axis_x = real * math.cos(phase) + axis_y = real * math.sin(phase) + axis_z = math.sqrt(1.0 - pow(real, 2)) + return axis_x, axis_y, axis_z + + +@nb.njit(nb.float64(nb.float64[:], PMU_DTYPE)) +def _thrust_with_axis( + axis_angles: base.DoubleVector, momenta: base.VoidVector +) -> float: + axis_x, axis_y, axis_z = _angles_to_axis(axis_angles[0], axis_angles[1]) + abs_dot_sum = norm_sum = 0.0 + for momentum in momenta: + x, y, z = momentum["x"], momentum["y"], momentum["z"] + dot_prod = (axis_x * x) + (axis_y * y) + (axis_z * z) + abs_dot_sum += abs(dot_prod) + norm_sum += _three_norm(x, y, z) + return abs_dot_sum / norm_sum + + +def thrust(momenta: "MomentumArray") -> float: + """Computes the thrust 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. + + Returns + ------- + float + The thrust of the event. + """ + domain = (-math.pi, math.pi) + optim = spo.minimize( + fun=lambda n, p: -_thrust_with_axis(n, p), + x0=np.zeros(2), + bounds=(domain, domain), + args=(momenta.data,), + ) + return -optim.fun