fix: align Lorentz/planar coordinate conventions across signatures

src/vector/_compute/lorentz/Et.py

@@ -34,7 +34,7 @@
 
 
 def xy_z_t(lib, x, y, z, t):
-    return lib.sqrt(Et2.xy_z_t(lib, x, y, z, t))
+    return lib.copysign(lib.sqrt(Et2.xy_z_t(lib, x, y, z, t)), t)
 
 
 def xy_z_tau(lib, x, y, z, tau):

src/vector/_compute/lorentz/Mt.py

@@ -34,51 +34,63 @@
 
 
 def xy_z_t(lib, x, y, z, t):
-    return lib.sqrt(Mt2.xy_z_t(lib, x, y, z, t))
+    mt2 = Mt2.xy_z_t(lib, x, y, z, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def xy_z_tau(lib, x, y, z, tau):
-    return lib.sqrt(Mt2.xy_z_tau(lib, x, y, z, tau))
+    mt2 = Mt2.xy_z_tau(lib, x, y, z, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def xy_theta_t(lib, x, y, theta, t):
-    return lib.sqrt(Mt2.xy_theta_t(lib, x, y, theta, t))
+    mt2 = Mt2.xy_theta_t(lib, x, y, theta, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def xy_theta_tau(lib, x, y, theta, tau):
-    return lib.sqrt(Mt2.xy_theta_tau(lib, x, y, theta, tau))
+    mt2 = Mt2.xy_theta_tau(lib, x, y, theta, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def xy_eta_t(lib, x, y, eta, t):
-    return lib.sqrt(Mt2.xy_eta_t(lib, x, y, eta, t))
+    mt2 = Mt2.xy_eta_t(lib, x, y, eta, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def xy_eta_tau(lib, x, y, eta, tau):
-    return lib.sqrt(Mt2.xy_eta_tau(lib, x, y, eta, tau))
+    mt2 = Mt2.xy_eta_tau(lib, x, y, eta, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_z_t(lib, rho, phi, z, t):
-    return lib.sqrt(Mt2.rhophi_z_t(lib, rho, phi, z, t))
+    mt2 = Mt2.rhophi_z_t(lib, rho, phi, z, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_z_tau(lib, rho, phi, z, tau):
-    return lib.sqrt(Mt2.rhophi_z_tau(lib, rho, phi, z, tau))
+    mt2 = Mt2.rhophi_z_tau(lib, rho, phi, z, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_theta_t(lib, rho, phi, theta, t):
-    return lib.sqrt(Mt2.rhophi_theta_t(lib, rho, phi, theta, t))
+    mt2 = Mt2.rhophi_theta_t(lib, rho, phi, theta, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_theta_tau(lib, rho, phi, theta, tau):
-    return lib.sqrt(Mt2.rhophi_theta_tau(lib, rho, phi, theta, tau))
+    mt2 = Mt2.rhophi_theta_tau(lib, rho, phi, theta, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_eta_t(lib, rho, phi, eta, t):
-    return lib.sqrt(Mt2.rhophi_eta_t(lib, rho, phi, eta, t))
+    mt2 = Mt2.rhophi_eta_t(lib, rho, phi, eta, t)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 def rhophi_eta_tau(lib, rho, phi, eta, tau):
-    return lib.sqrt(Mt2.rhophi_eta_tau(lib, rho, phi, eta, tau))
+    mt2 = Mt2.rhophi_eta_tau(lib, rho, phi, eta, tau)
+    return lib.copysign(lib.sqrt(lib.absolute(mt2)), mt2)
 
 
 dispatch_map = {

src/vector/_compute/lorentz/Mt2.py

@@ -39,47 +39,47 @@ def xy_z_t(lib, x, y, z, t):
 
 
 def xy_z_tau(lib, x, y, z, tau):
-    return lib.maximum(tau2.xy_z_tau(lib, x, y, z, tau) + x**2 + y**2, 0)
+    return tau2.xy_z_tau(lib, x, y, z, tau) + x**2 + y**2
 
 
 def xy_theta_t(lib, x, y, theta, t):
     return xy_z_t(lib, x, y, z.xy_theta(lib, x, y, theta), t)
 
 
 def xy_theta_tau(lib, x, y, theta, tau):
-    return lib.maximum(tau2.xy_theta_tau(lib, x, y, theta, tau) + x**2 + y**2, 0)
+    return tau2.xy_theta_tau(lib, x, y, theta, tau) + x**2 + y**2
 
 
 def xy_eta_t(lib, x, y, eta, t):
     return xy_z_t(lib, x, y, z.xy_eta(lib, x, y, eta), t)
 
 
 def xy_eta_tau(lib, x, y, eta, tau):
-    return lib.maximum(tau2.xy_eta_tau(lib, x, y, eta, tau) + x**2 + y**2, 0)
+    return tau2.xy_eta_tau(lib, x, y, eta, tau) + x**2 + y**2
 
 
 def rhophi_z_t(lib, rho, phi, z, t):
     return t**2 - z**2
 
 
 def rhophi_z_tau(lib, rho, phi, z, tau):
-    return lib.maximum(tau2.rhophi_z_tau(lib, rho, phi, z, tau) + rho**2, 0)
+    return tau2.rhophi_z_tau(lib, rho, phi, z, tau) + rho**2
 
 
 def rhophi_theta_t(lib, rho, phi, theta, t):
     return rhophi_z_t(lib, rho, phi, z.rhophi_theta(lib, rho, phi, theta), t)
 
 
 def rhophi_theta_tau(lib, rho, phi, theta, tau):
-    return lib.maximum(tau2.rhophi_theta_tau(lib, rho, phi, theta, tau) + rho**2, 0)
+    return tau2.rhophi_theta_tau(lib, rho, phi, theta, tau) + rho**2
 
 
 def rhophi_eta_t(lib, rho, phi, eta, t):
     return rhophi_z_t(lib, rho, phi, z.rhophi_eta(lib, rho, phi, eta), t)
 
 
 def rhophi_eta_tau(lib, rho, phi, eta, tau):
-    return lib.maximum(tau2.rhophi_eta_tau(lib, rho, phi, eta, tau) + rho**2, 0)
+    return tau2.rhophi_eta_tau(lib, rho, phi, eta, tau) + rho**2
 
 
 dispatch_map = {

src/vector/_compute/lorentz/scale.py

@@ -39,7 +39,7 @@ def xy_z_t(lib, factor, x, y, z, t):
 
 def xy_z_tau(lib, factor, x, y, z, tau):
     sx, sy, sz = scale3d.xy_z(lib, factor, x, y, z)
-    return (sx, sy, sz, tau * factor)
+    return (sx, sy, sz, tau * lib.absolute(factor))
 
 
 def xy_theta_t(lib, factor, x, y, theta, t):
@@ -49,7 +49,7 @@ def xy_theta_t(lib, factor, x, y, theta, t):
 
 def xy_theta_tau(lib, factor, x, y, theta, tau):
     sx, sy, stheta = scale3d.xy_theta(lib, factor, x, y, theta)
-    return (sx, sy, stheta, tau * factor)
+    return (sx, sy, stheta, tau * lib.absolute(factor))
 
 
 def xy_eta_t(lib, factor, x, y, eta, t):
@@ -59,7 +59,7 @@ def xy_eta_t(lib, factor, x, y, eta, t):
 
 def xy_eta_tau(lib, factor, x, y, eta, tau):
     sx, sy, seta = scale3d.xy_eta(lib, factor, x, y, eta)
-    return (sx, sy, seta, tau * factor)
+    return (sx, sy, seta, tau * lib.absolute(factor))
 
 
 def rhophi_z_t(lib, factor, rho, phi, z, t):
@@ -69,7 +69,7 @@ def rhophi_z_t(lib, factor, rho, phi, z, t):
 
 def rhophi_z_tau(lib, factor, rho, phi, z, tau):
     srho, sphi, sz = scale3d.rhophi_z(lib, factor, rho, phi, z)
-    return (srho, sphi, sz, tau * factor)
+    return (srho, sphi, sz, tau * lib.absolute(factor))
 
 
 def rhophi_theta_t(lib, factor, rho, phi, theta, t):
@@ -79,7 +79,7 @@ def rhophi_theta_t(lib, factor, rho, phi, theta, t):
 
 def rhophi_theta_tau(lib, factor, rho, phi, theta, tau):
     srho, sphi, stheta = scale3d.rhophi_theta(lib, factor, rho, phi, theta)
-    return (srho, sphi, stheta, tau * factor)
+    return (srho, sphi, stheta, tau * lib.absolute(factor))
 
 
 def rhophi_eta_t(lib, factor, rho, phi, eta, t):
@@ -89,7 +89,7 @@ def rhophi_eta_t(lib, factor, rho, phi, eta, t):
 
 def rhophi_eta_tau(lib, factor, rho, phi, eta, tau):
     srho, sphi, seta = scale3d.rhophi_eta(lib, factor, rho, phi, eta)
-    return (srho, sphi, seta, tau * factor)
+    return (srho, sphi, seta, tau * lib.absolute(factor))
 
 
 dispatch_map = {

src/vector/_compute/planar/unit.py

@@ -35,10 +35,8 @@ def xy(lib, x, y):
 
 
 def rhophi(lib, rho, phi):
-    return (1, phi)
-
-
-rhophi.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+    norm = lib.sqrt(rho**2)
+    return (lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi)
 
 
 dispatch_map = {

tests/compute/lorentz/test_Et.py

@@ -118,3 +118,28 @@ def test_rhophi_eta_tau():
         vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.Et == pytest.approx(math.sqrt(80))
+
+
+def test_negative_energy_sign_consistency():
+    # ROOT's TLorentzVector::Et() is sign-preserving in the energy:
+    # E < 0 ? -sqrt(Et2) : sqrt(Et2). All coordinate-system representations
+    # of the same vector must therefore agree on the (negative) sign.
+    xy_z = vector.obj(px=3.0, py=4.0, pz=0.0, E=-13.0)
+    rhophi_z = vector.obj(pt=5.0, phi=0.0, pz=0.0, E=-13.0)
+    assert xy_z.Et == pytest.approx(-13.0)
+    assert rhophi_z.Et == pytest.approx(-13.0)
+    assert xy_z.Et == pytest.approx(rhophi_z.Et)
+
+    # non-axial pz, all four az/long combinations should still agree in sign
+    px, py, pz, E = 3.0, 4.0, 10.0, -20.0
+    expected = -math.sqrt(80)
+    base = vector.obj(px=px, py=py, pz=pz, E=E)
+    for vec in (
+        base,
+        vector.obj(pt=base.pt, phi=base.phi, pz=pz, E=E),
+        vector.obj(px=px, py=py, theta=base.theta, E=E),
+        vector.obj(px=px, py=py, eta=base.eta, E=E),
+        vector.obj(pt=base.pt, phi=base.phi, theta=base.theta, E=E),
+        vector.obj(pt=base.pt, phi=base.phi, eta=base.eta, E=E),
+    ):
+        assert vec.Et == pytest.approx(expected)

tests/compute/lorentz/test_Mt.py

@@ -118,3 +118,37 @@ def test_rhophi_eta_tau():
         vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.Mt == pytest.approx(math.sqrt(300))
+
+
+def test_transverse_spacelike_consistency():
+    # For a transverse-spacelike vector (t**2 - z**2 < 0), ROOT's Mt() is
+    # sign-preserving: Mt = copysign(sqrt(|Mt2|), Mt2). The T- and
+    # Tau-coordinate representations of the same vector must agree.
+    base = vector.obj(px=3.0, py=4.0, pz=10.0, E=2.0)
+    expected = -math.sqrt(96)  # Mt2 = 4 - 100 = -96
+    assert base.Mt2 == pytest.approx(-96)
+    assert base.Mt == pytest.approx(expected)
+
+    tau_vec = vector.obj(px=3.0, py=4.0, pz=10.0, tau=base.tau)
+    assert tau_vec.Mt2 == pytest.approx(-96)
+    assert tau_vec.Mt == pytest.approx(expected)
+    assert base.Mt == pytest.approx(tau_vec.Mt)
+
+    # all azimuthal/longitudinal representations agree (both T and Tau)
+    for vec in (
+        vector.obj(pt=base.pt, phi=base.phi, pz=10.0, E=2.0),
+        vector.obj(px=3.0, py=4.0, theta=base.theta, E=2.0),
+        vector.obj(px=3.0, py=4.0, eta=base.eta, E=2.0),
+        vector.obj(pt=base.pt, phi=base.phi, theta=base.theta, tau=base.tau),
+        vector.obj(px=3.0, py=4.0, eta=base.eta, tau=base.tau),
+    ):
+        assert vec.Mt2 == pytest.approx(-96)
+        assert vec.Mt == pytest.approx(expected)
+
+
+def test_timelike_consistency():
+    # Sanity: ordinary timelike vector still matches between T and Tau coords.
+    base = vector.obj(px=3.0, py=4.0, pz=0.0, E=10.0)
+    tau_vec = vector.obj(px=3.0, py=4.0, pz=0.0, tau=base.tau)
+    assert base.Mt == pytest.approx(10.0)
+    assert tau_vec.Mt == pytest.approx(10.0)

tests/compute/lorentz/test_Mt2.py

@@ -116,3 +116,13 @@ def test_rhophi_eta_tau():
         vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.Mt2 == pytest.approx(300)
+
+
+def test_transverse_spacelike_consistency():
+    # Mt2 must be the signed t**2 - z**2 in both T- and Tau-coordinates
+    # (no clamp to zero), so the two representations agree.
+    base = vector.obj(px=3.0, py=4.0, pz=10.0, E=2.0)
+    tau_vec = vector.obj(px=3.0, py=4.0, pz=10.0, tau=base.tau)
+    assert base.Mt2 == pytest.approx(-96)
+    assert tau_vec.Mt2 == pytest.approx(-96)
+    assert base.Mt2 == pytest.approx(tau_vec.Mt2)

tests/compute/spatial/test_eta.py

@@ -56,3 +56,21 @@ def test_rhophi_eta():
         longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
     )
     assert vec.eta == pytest.approx(1.4436354751788103)
+
+
+def test_theta_out_of_range_consistency():
+    # For theta outside (0, pi), -log(tan(theta/2)) is nan; both the xy_theta
+    # and rhophi_theta variants must apply the same nan->0 guard so the two
+    # representations of the same direction agree.
+    for theta in (-0.5, 2.0 * 3.141592653589793):
+        xy = vector.backends.object.VectorObject3D(
+            azimuthal=vector.backends.object.AzimuthalObjectXY(1.0, 1.0),
+            longitudinal=vector.backends.object.LongitudinalObjectTheta(theta),
+        )
+        rhophi = vector.backends.object.VectorObject3D(
+            azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(1.0, 0.0),
+            longitudinal=vector.backends.object.LongitudinalObjectTheta(theta),
+        )
+        assert xy.eta == pytest.approx(0.0)
+        assert rhophi.eta == pytest.approx(0.0)
+        assert xy.eta == pytest.approx(rhophi.eta)

tests/compute/sympy/lorentz/test_Mt.py

@@ -25,7 +25,7 @@ def test_xy_z_t():
         vector.backends.sympy.LongitudinalSympyZ(z),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))
 
 
@@ -51,7 +51,7 @@ def test_xy_theta_t():
         ),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt.simplify() == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))
 
 
@@ -79,7 +79,7 @@ def test_xy_eta_t():
         ),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt.simplify() == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))
 
 
@@ -105,7 +105,7 @@ def test_rhophi_z_t():
         vector.backends.sympy.LongitudinalSympyZ(z),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt.simplify() == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))
 
 
@@ -129,7 +129,7 @@ def test_rhophi_theta_t():
         ),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt.simplify() == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))
 
 
@@ -153,7 +153,7 @@ def test_rhophi_eta_t():
         vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)),
         vector.backends.sympy.TemporalSympyT(t),
     )
-    assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2)
+    assert vec.Mt.simplify() == sympy.sqrt(sympy.Abs(t**2 - z**2))
     assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300))