overload type updates, and added c-parameter #165

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[:]))",

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