From 7fd277c45901bf2a5f078adeb8cabfc800f14a07 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Tue, 14 Nov 2023 14:46:56 +0000 Subject: [PATCH] overload type updates, and added c-parameter #165 --- graphicle/calculate.py | 91 +++++++++++++++++++++++++++++++++++++++++ tests/test_calculate.py | 76 ++++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 tests/test_calculate.py diff --git a/graphicle/calculate.py b/graphicle/calculate.py index ae4f2e6..f57ca6e 100644 --- a/graphicle/calculate.py +++ b/graphicle/calculate.py @@ -767,6 +767,21 @@ def thrust( ... +@ty.overload +def thrust(pmu: "MomentumArray", return_axis: ty.Literal[False]) -> float: + ... + + +@ty.overload +def thrust(pmu: "MomentumArray", rng_seed: ty.Optional[int]) -> float: + ... + + +@ty.overload +def thrust(pmu: "MomentumArray") -> float: + ... + + @ty.overload def thrust( pmu: "MomentumArray", @@ -776,6 +791,14 @@ def thrust( ... +@ty.overload +def thrust( + pmu: "MomentumArray", + return_axis: ty.Literal[True], +) -> ty.Tuple[float, base.DoubleVector]: + ... + + def thrust( pmu: "MomentumArray", return_axis: bool = False, @@ -811,6 +834,7 @@ def thrust( See Also -------- spherocity : Computes the spherocity from final state momenta. + c_parameter : Computes the C-parameter from the final state momenta. """ half_pi = 0.5 * math.pi domain = (-half_pi, half_pi) @@ -913,6 +937,21 @@ def spherocity( ... +@ty.overload +def spherocity(pmu: "MomentumArray", return_axis: ty.Literal[False]) -> float: + ... + + +@ty.overload +def spherocity(pmu: "MomentumArray", rng_seed: ty.Optional[int]) -> float: + ... + + +@ty.overload +def spherocity(pmu: "MomentumArray") -> float: + ... + + @ty.overload def spherocity( pmu: "MomentumArray", @@ -922,6 +961,14 @@ def spherocity( ... +@ty.overload +def spherocity( + pmu: "MomentumArray", + return_axis: ty.Literal[True], +) -> ty.Tuple[float, base.DoubleVector]: + ... + + def spherocity( pmu: "MomentumArray", return_axis: bool = False, @@ -957,6 +1004,7 @@ def spherocity( See Also -------- thrust : Computes the thrust from final state momenta. + c_parameter : Computes the C-parameter from the final state momenta. """ half_pi = 0.5 * math.pi domain = (-half_pi, half_pi) @@ -975,6 +1023,49 @@ def spherocity( return sph_val +@nb.njit(nb.float64(PMU_DTYPE)) +def _c_parameter(momenta: base.VoidVector) -> float: + output = norm_sum = 0.0 + for idx_i, pmu_i in enumerate(momenta): + x_i, y_i, z_i = pmu_i["x"], pmu_i["y"], pmu_i["z"] + norm_i = _three_norm(x_i, y_i, z_i) + norm_sum += norm_i + for pmu_j in momenta[idx_i:]: + x_j, y_j, z_j = pmu_j["x"], pmu_j["y"], pmu_j["z"] + norm_j = _three_norm(x_j, y_j, z_j) + norm_prod = norm_i * norm_j + i_dot_j = (x_i * x_j) + (y_i * y_j) + (z_i * z_j) + output += 2.0 * (norm_prod - (i_dot_j * i_dot_j / norm_prod)) + return 1.5 * output / (norm_sum * norm_sum) + + +def c_parameter(pmu: "MomentumArray") -> float: + """Computes the C-parameter for the 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 C-parameter of the event. + + See Also + -------- + thrust : Computes the thrust from final state momenta. + spherocity : Computes the spherocity from final state momenta. + """ + return _c_parameter(pmu.data) + + @nb.njit( [ "float64(bool_[:], bool_[:], Optional(float64[:]))", diff --git a/tests/test_calculate.py b/tests/test_calculate.py new file mode 100644 index 0000000..72ad892 --- /dev/null +++ b/tests/test_calculate.py @@ -0,0 +1,76 @@ +import math + +import numba as nb +import numpy as np + +import graphicle as gcl + + +@nb.njit("float64[:, :](int32, float64)") +def _fibonacci_sphere_momenta( + num_points: int, + energy: float, +) -> gcl.base.DoubleVector: + pmu = np.empty((num_points, 4), dtype=np.float64) + phi = math.pi * (3.0 - math.sqrt(5.0)) # golden angle in radians + for i in range(num_points): + y = 1.0 - (i / (num_points - 1.0)) * 2.0 # y goes from 1 to -1 + radius = math.sqrt(1.0 - y * y) # radius at y + theta = phi * i # golden angle increment + pmu[i, 0] = math.cos(theta) * radius + pmu[i, 1] = y + pmu[i, 2] = math.sin(theta) * radius + pmu[i, 3] = energy + return pmu + + +def spherical_momenta( + num_points: int, energy: float = 2.0 +) -> gcl.MomentumArray: + return gcl.MomentumArray(_fibonacci_sphere_momenta(num_points, energy)) + + +def pencil_momenta() -> gcl.MomentumArray: + return gcl.MomentumArray([[1.0, 0.0, 0.0, 2.0], [-1.0, 0.0, 0.0, 2.0]]) + + +def test_thrust_spherical() -> None: + pmu = spherical_momenta(10_000) + thrust = gcl.calculate.thrust(pmu) + msg = "Thrust for spherical event is not 0.5." + assert math.isclose(thrust, 0.5, rel_tol=1.0e-3), msg + + +def test_thrust_pencil() -> None: + pmu = pencil_momenta() + thrust = gcl.calculate.thrust(pmu) + msg = "Thrust for pencil-like event is not 1.0." + assert math.isclose(thrust, 1.0, abs_tol=1.0e-6), msg + + +def test_spherocity_spherical() -> None: + pmu = spherical_momenta(10_000) + spherocity = gcl.calculate.spherocity(pmu) + msg = "Spherocity for spherical event is not 1.0." + assert math.isclose(spherocity, 1.0, rel_tol=1.0e-3), msg + + +def test_spherocity_pencil() -> None: + pmu = pencil_momenta() + spherocity = gcl.calculate.spherocity(pmu) + msg = "Spherocity for pencil-like event is not 0.0." + assert math.isclose(spherocity, 0.0, abs_tol=1.0e-6), msg + + +def test_c_parameter_spherical() -> None: + pmu = spherical_momenta(10_000) + c_param = gcl.calculate.c_parameter(pmu) + msg = "C-parameter for spherical event is not 1.0." + assert math.isclose(c_param, 1.0, rel_tol=1.0e-3), msg + + +def test_c_parameter_pencil() -> None: + pmu = pencil_momenta() + c_param = gcl.calculate.c_parameter(pmu) + msg = "C-parameter for pencil-like event is not 0.0." + assert math.isclose(c_param, 0.0, abs_tol=1.0e-6), msg