feat: add aligned_four_vectors function and tests (#70)

Author: mmikhasenko

.cspell/project.txt

@@ -40,6 +40,7 @@ refζ
 refζs
 Registrator
 servedocs
+sinθ
 subchannel
 sumrules
 typeof

src/ThreeBodyDecays.jl

@@ -21,6 +21,7 @@ export cosθ12, cosθ23, cosθ31
 export σiofk, σjofk
 export σ1of2, σ2of3, σ3of1, σ1of3, σ2of1, σ3of2
 export breakup, breakup_Rk, breakup_ij
+export aligned_four_vectors
 include("kinematics.jl")
 
 export minusone

src/kinematics.jl

@@ -7,7 +7,8 @@ Kibble(σs, msq) = Kallen(
     Kallen(msq[4], msq[2], σs[2]),
     Kallen(msq[4], msq[3], σs[3]),
 )
-#
+
+
 """
 	σjofk(z,σi,msq; k::Int)
 
@@ -92,3 +93,52 @@ function breakup_ij(ms; k)
     (i, j) = ij_from_k(k)
     return σ -> breakup(sqrt(σ), ms[i], ms[j])
 end
+
+"""
+    aligned_four_vectors(σs,ms; k::Int)
+
+Computes the four-momenta of the three particles in the center of momentum frame aligning the k-th particle with the -z-axis.
+
+## Arguments
+- `σs`: Tuple of mandelstam variables,
+- `ms`: Tuple of masses in the order m1, m2, m3, m0.
+- `k`: Index of the particle to be aligned with the -z-axis.
+
+## Returns
+
+A tuple of three four-momenta in the form of (px, py, pz, E).
+
+## Example
+````julia
+ms = ThreeBodyMasses(1.0, 1.0, 1.0; m0=4.0)
+σs = x2σs([0.5, 0.5], ms; k=2)
+p1, p2, p3 = aligned_four_vectors(σs, ms; k=1)
+````
+"""
+function aligned_four_vectors(σs, ms; k::Int)
+    #
+    i, j = ij_from_k(k)
+    m0 = ms[4]
+    s = m0^2
+    #
+    mi, mj, mk = ms[i], ms[j], ms[k]
+    σi, σj, σk = σs[i], σs[j], σs[k]
+    misq, mjsq, mksq = (mi, mj, mk) .^ 2
+    #
+    Ei = (s + misq - σi) / (2m0)
+    Ej = (s + mjsq - σj) / (2m0)
+    Ek = (s + mksq - σk) / (2m0)
+    #
+    p3i = sqrt(Kallen(s, misq, σi)) / (2m0)
+    p3j = sqrt(Kallen(s, mjsq, σj)) / (2m0)
+    p3k = sqrt(Kallen(s, mksq, σk)) / (2m0)
+    #
+    cosθ = -cosζ(wr(k, i, 0), σs, ms^2) # pi-angle(1,2) in rest of 0
+    sinθ = sqrt(1 - cosθ^2)
+    #
+    pk = (0, 0, -p3k, Ek)
+    pi = (p3i * sinθ, 0, p3i * cosθ, Ei)
+    pj = (-p3i * sinθ, 0, p3k - p3i * cosθ, Ej)
+    #
+    return circleorigin(-k, (pi, pj, pk))
+end

test/test_invariants.jl

@@ -89,3 +89,46 @@ rp = randomPoint(tbsS)
     @test Kibble(rp.σs, tbsS.ms^2) < 0
     @test Tuple(rp.two_λs) ∈ collect(itr(tbsS.two_js))
 end
+
+
+@testset "Aligned four vectors" begin
+    # Test 1: Check four-momentum conservation
+    p1, p2, p3 = aligned_four_vectors(σs, ms; k = 1)
+    @test all(p1 .+ p2 .+ p3 .≈ (0, 0, 0, ms.m0))
+    #
+    # aligned vector k, px=0, is on kth place
+    @test aligned_four_vectors(σs, ms; k = 1)[1][1] == 0.0
+    @test aligned_four_vectors(σs, ms; k = 2)[2][1] == 0.0
+    @test aligned_four_vectors(σs, ms; k = 3)[3][1] == 0.0
+
+    metric = (-1, -1, -1, 1)
+    # Test 2: Check mass shell conditions
+    @test sum(p1 .^ 2 .* metric) ≈ ms.m1^2
+    @test sum(p2 .^ 2 .* metric) ≈ ms.m2^2
+    @test sum(p3 .^ 2 .* metric) ≈ ms.m3^2
+
+    # Test 3: Verify invariant masses
+    @test sum((p2 .+ p3) .^ 2 .* metric) ≈ σs.σ1
+    @test sum((p3 .+ p1) .^ 2 .* metric) ≈ σs.σ2
+    @test sum((p1 .+ p2) .^ 2 .* metric) ≈ σs.σ3
+end
+
+@testset "Aligned four vectors are cyclic" begin
+    p1_1, p2_1, p3_1 = aligned_four_vectors(σs, ms; k = 1)
+    #
+    p1_2, p2_2, p3_2 = aligned_four_vectors(
+        MandelstamTuple{Float64}((σs[3], σs[1], σs[2])),
+        ThreeBodyMasses(ms[3], ms[1], ms[2]; ms.m0);
+        k = 2,
+    )
+    #
+    p1_3, p2_3, p3_3 = aligned_four_vectors(
+        MandelstamTuple{Float64}((σs[2], σs[3], σs[1])),
+        ThreeBodyMasses(ms[2], ms[3], ms[1]; ms.m0);
+        k = 3,
+    )
+    #
+    @test all(p1_1 .≈ p2_2 .≈ p3_3)
+    @test all(p2_1 .≈ p3_2 .≈ p1_3)
+    @test all(p3_1 .≈ p1_2 .≈ p2_3)
+end