(* ::Package:: *)

(* From Feynman diagrams to the amplituhedron: a gentle review
   Shounak De, Dmitrii Pavlov, Marcus Spradlin, Anastasia Volovich
   https://arxiv.org/abs/2410.11757
   ---------------------------------------------------------------

   AMP[H] is a hard-coded known formula for the helicity amplitude A[H].

   AFD[n] computes the n-particle amplitude by summing Feynman diagrams.

   When H is a permutation of {1, 1, 1, -1, -1, -1}, AMT[H] computes the
   helicity amplitude A[H] from the amplituhedron.  Note that in contrast
   to AMP and AFD, which output analytic formulas, this one performs a
   numerical calculation and requires that Z be set to a numeric 6x4
   matrix specifying the momentum twistor variables.

   Compare[H] generates a random kinematic point Z and compares the
   (numerical) results obtained from the various approaches.

   Run CompareAll[] to scan through and compare all helicity amplitudes
   for n <= 6 particles (except for the special case when H is a
   permutation of {1, 1, -1}, for which the amplitude is nonzero on a
   kinematic branch not accessible by momentum twistors).
*)

SetAttributes[dot, Orderless];


(* --- AMP = AMPLITUDE FORMULAS FROM THE LITERATURE, FOR COMPARISON --- *)


Print["Constructing helicity amplitudes from known formulas ..."];


(* Only 3 independent NMHV helicity amplitudes need to be specified *)

AMP[{-1, -1, -1, 1, 1, 1}] = -
	braket[4,{5,6},1]^3/(bra[2,3] bra[3,4] braket[2,{3,4},5] ket[5,6] ket[6,1] S[2,3,4]) -
	braket[6,{1,2},3]^3/(bra[1,2] bra[6,1] braket[2,{3,4},5] ket[3,4] ket[4,5] S[1,2,6]);

AMP[{1, -1, 1, -1, 1, -1}] =
        ket[6,4]^4 bra[1,3]^4/(S[4,5,6] braket[1,{5,6},4] braket[3,{4,5},6] ket[4,5] ket[5,6] bra[1,2] bra[2,3]) +
        ket[6,2]^4 bra[3,5]^4/(S[6,1,2] braket[3,{1,2},6] braket[5,{6,1},2] ket[6,1] ket[1,2] bra[3,4] bra[4,5]) +
        ket[2,4]^4 bra[1,5]^4/(S[2,3,4] braket[5,{3,4},2] braket[1,{2,3},4] ket[2,3] ket[3,4] bra[5,6] bra[6,1]);

AMP[{1, 1, -1, 1, -1, -1}] =
	braket[4,{1,2},3]^4/(bra[4,5] bra[5,6] braket[4,{3,2},1] braket[6,{1,2},3] ket[1,2] ket[2,3] S[1,2,3]) +
	bra[2,4]^4 ket[5,6]^3/(bra[2,3] bra[3,4] braket[2,{3,4},5] braket[4,{3,2},1] ket[6,1] S[2,3,4]) +
	bra[1, 2]^3 ket[3,5]^4/(bra[6,1] braket[2,{3,4},5] braket[6,{5,4},3] ket[3,4] ket[4,5] S[3,4,5]);


(* We'll automatically generate everything else from these *)

(* First rewrite them in terms of spinor helicity bras [] and kets <> only *)
rewrite = {
	braket[a_, {b_, c_}, d_] -> bra[a, b] ket[b, d] + bra[a, c] ket[c, d],
	S[a__] :> (bra @@@ Subsets[{a}, {2}]) . (ket @@@ Subsets[{a}, {2}])};

AMP[{-1, -1, -1, 1, 1, 1}] = AMP[{-1, -1, -1, 1, 1, 1}]/.rewrite;
AMP[{1, -1, 1, -1, 1, -1}] = AMP[{1, -1, 1, -1, 1, -1}]/.rewrite;
AMP[{1, 1, -1, 1, -1, -1}] = AMP[{1, 1, -1, 1, -1, -1}]/.rewrite;


(* Now generate all cyclic rotations of the above *)

AMP[{-1, 1, -1, 1, -1, 1}] =
	AMP[{1, -1, 1, -1, 1, -1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+1, 6, 1], Mod[b+1, 6, 1]];

For[i = 1, i < 6, i++,

	AMP[RotateRight[{-1, -1, -1, 1, 1, 1}, i]] =
		AMP[{-1, -1, -1, 1, 1, 1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+i, 6, 1], Mod[b+i, 6, 1]];

	AMP[RotateRight[{1, 1, -1, 1, -1, -1}, i]] =
		AMP[{1, 1, -1, 1, -1, -1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+i, 6, 1], Mod[b+i, 6, 1]];

];


(* Finally the last six are obtained by parity conjugation *)

For[i = 0, i < 6, i++,

	AMP[RotateRight[{-1, -1, 1, -1, 1, 1}, i]] =
		AMP[RotateRight[{1, 1, -1, 1, -1, -1}, i]]/.{bra -> ket, ket -> bra};

];


(* Now we're done with NMHV amplitudes; here are the rest *)

(* MHV amplitudes *)

AMP[H_] := (ket @@ Flatten[Position[H, -1]])^4/Product[ket[i, Mod[i+1, Length[H], 1]], {i, Length[H]}] /; (Length[H] > 2 && Count[H, -1] == 2 && Count[H, +1] == Length[H] - 2);

(* MHV-bar amplitudes *)

AMP[H_] := (bra @@ Flatten[Position[H, +1]])^4/Product[bra[i, Mod[i+1, Length[H], 1]], {i, Length[H]}] /; (Length[H] > 2 && Count[H, +1] == 2 && Count[H, -1] == Length[H] - 2);

(* All others vanish *)

AMP[H_] := 0 /; (MemberQ[{{-1, 1}, {-1}, {1}}, Union[H]] && Length[H] < 7);

Print["... done."];


(* --- AFD = AMPLITUDES COMPUTED FROM FEYNMAN DIAGRAMS --- *)


Print["Constructing amplitudes from Feynman diagrams ..."];

(* Basic definitions and conventions *)

p[a_Integer] := - p[-a] /; a < 0;

dot[a_ + b_, c_] := dot[a, c] + dot[b, c];

dot[a_, b_ + c_] := dot[a, b] + dot[a, c];

dot[-f_[a_], b_] := - dot[f[a], b];

dot[a_, -f_[b_]] := - dot[a, f[b]];

dot[e[a_], e[a_]] := 0;

dot[e[a_], p[a_]] := 0;


(* Cubic vertex *)

V3[a_, b_, c_] := I/Sqrt[2] (
	dot[e[a], e[b]] dot[p[a] - p[b], e[c]] +
	dot[e[b], e[c]] dot[p[b] - p[c], e[a]] +
	dot[e[c], e[a]] dot[p[c] - p[a], e[b]]);


(* Quartic vertex *)

V4[a_, b_, c_, d_] :=
	I dot[e[a], e[c]] dot[e[b], e[d]] -
	I/2 dot[e[a], e[b]] dot[e[c], e[d]] -
	I/2 dot[e[b], e[c]] dot[e[d], e[a]];


(* We sum over all cyclic labelings, and divide by a symmetry factor in certain
   diagrams as needed to avoid overcounting *)

Cycle[X_, m_] := X/.{p[a_] :> p[Mod[a + 1, m, 1]], e[a_] :> e[Mod[a + 1, m, 1]]};

CycleSum[X_, m_] := Total[NestList[Cycle[#, m] &, X, m - 1]];


(* This is the differential operator associated to internal edges *)

DiffOp[X_, L_] := Module[{A, i, j},
	A = Expand[X];
	For[i = 1, i <= Length[L], i++,
		j = L[[i]];
		A = A/.dot[a_, e[j]] dot[b_, e[-j]] -> dot[a, b] (-I)/dot[p[j], p[j]];
	];
	Return[-I A];
];


(* Here we construct the amplitudes *)

AFD[3] = -I V3[1, 2, 3];


D41 = DiffOp[V3[1, 2, 5] V3[3, 4, -5], {5}]/.
	p[5] -> p[3] + p[4]/.dot[p[a_], p[a_]] -> 0;

D42 = -I V4[1, 2, 3, 4];

AFD[4] = CycleSum[1/2 D41 + 1/4 D42, 4];


D51 = DiffOp[V3[1, 2, -6] V3[4, 5, -7] V3[3, 7, 6], {6, 7}]/.
	p[6] -> p[1] + p[2]/.p[7] -> p[4] + p[5]/.dot[p[a_], p[a_]] -> 0;

D52 = DiffOp[V4[1, 2, 3, 6] V3[4, 5, -6], {6}]/.
	p[6] -> p[4] + p[5]/.dot[p[a_], p[a_]] -> 0;

AFD[5] = CycleSum[D51 + D52, 5];


D61 = DiffOp[V3[1, 2, -7] V3[7, 3, -8] V3[8, 4, -9] V3[9, 5, 6], {7, 8, 9}]/.
	p[8] -> - p[4] + p[9]/.p[9] -> - p[5] - p[6]/.p[7] -> p[1] + p[2]/.dot[p[a_], p[a_]] -> 0;

D62 = DiffOp[V3[1, 2, 7] V4[-7, 3, -8, 6] V3[8, 4, 5], {7, 8}]/.
	p[8] -> - p[4] - p[5]/.p[7] -> - p[1] - p[2]/.dot[p[a_], p[a_]] -> 0;

D63 = DiffOp[V4[1, 2, 3, 7] V4[-7, 4, 5, 6], {7}]/.
	p[7] -> - p[1] - p[2] - p[3]/.dot[p[a_], p[a_]] -> 0;

D64 = DiffOp[V4[1, 2, 3, 7] V3[-7, 4, 8] V3[-8, 5, 6], {7, 8}]/.
	p[7] -> - p[1] - p[2] - p[3]/.p[8] -> p[5] + p[6]/.dot[p[a_], p[a_]] -> 0;

D65 = DiffOp[V4[7, 8, 5, 6] V3[1, 2, -7] V3[3, 4, -8], {7, 8}]/.
	p[7] -> p[1] + p[2]/.p[8] -> p[3] + p[4]/.dot[p[a_], p[a_]] -> 0;

D66 = DiffOp[V3[1, 2, 7] V3[3, 8, -7] V3[6, -8, 9] V3[4, 5, -9], {7, 8, 9}]/.
	p[8] -> p[7] - p[3]/.p[9] -> p[4] + p[5]/.p[7] -> - p[1] - p[2]/.dot[p[a_], p[a_]] -> 0;

D67 = DiffOp[V3[1, 2, 7] V3[-7, 8, 6] V3[-8, 3, 9] V3[-9, 4, 5], {7, 8, 9}]/.
	p[8] -> p[3] + p[9]/.p[9] -> p[4] + p[5]/.p[7] -> - p[1] - p[2]/.dot[p[a_], p[a_]] -> 0;

D68 = DiffOp[V4[1, 2, 3, -7] V3[7, 8, 6] V3[4, 5, -8], {7, 8}]/.
	p[8] -> p[4] + p[5]/.p[7] -> p[1] + p[2] + p[3]/.dot[p[a_], p[a_]] -> 0;

D69 = DiffOp[V3[1, 2, -7] V3[3, 4, -8] V3[5, 6, -9] V3[7, 8, 9], {7, 8, 9}]/.
	p[9] -> - p[7] - p[8]/.p[7] -> p[1] + p[2]/.p[8] -> p[3] + p[4]/.p[9] -> p[5] + p[6]/.dot[p[a_], p[a_]] -> 0;

AFD[6] = CycleSum[D61 + 1/2 D62 + 1/2 D63 + D64 + D65 + 1/2 D66 + 1/2 D67 + D68 + 1/3 D69, 6];

Print["... done."];


(* --- AMT = EXTRACT NMHV HELICITY AMPLITUDE FROM MOMENTUM SUPERTWISTORS --- *)


(* The four-brakcet < , , , > on momentum twistors *)

det[x__] := Det[Z[[{x}]]];


(* Definition of the 5-bracket, or R-invariant *)

R[a_, b_, c_, d_, e_] := (det[b,c,d,e] chi[a] + det[c,d,e,a] chi[b] + det[d,e,a,b] chi[c] + det[e,a,b,c] chi[d] + det[a,b,c,d] chi[e])^4/(det[a,b,c,d] det[b,c,d,e] det[c,d,e,a] det[d,e,a,b] det[e,a,b,c]);


(* This function does not generate analytic formulas; it requires the momentum
   twistor matrix Z to first be set to numerical values and then returns a
   numerical result for the amplitude evaluated at that kinematic point *)

(* Here is the code to extract a desired helicity component amplitude *)

AMT[H_] := Module[{ANMHV, m1, m2, m3, eqn, sol},

	(* Here is the NMHV superamplitude, including the MHV prefactor *)

	ANMHV = (R[6, 1, 2, 3, 4] + R[6, 1, 2, 4, 5] + R[6, 2, 3, 4, 5])/(KET[1, 2] KET[2, 3] KET[3, 4] KET[4, 5] KET[5, 6] KET[6, 1]);

	(* Identify the locations of the three negative helicities *)

	m1 = Position[H, -1][[1,1]];
	m2 = Position[H, -1][[2,1]];
	m3 = Position[H, -1][[3,1]];

	(* Solve for the first two using the supermomentum conservation equations *)

	eqn = Array[eta, 6] . Array[L, 6];
	sol = Solve[Thread[eqn == 0], {eta[m1], eta[m2]}];

	(* Now we can solve for the complementary set of chi's *)

	eqn = Table[eta[i] - (KET[Mod[i-1,6,1],i] chi[Mod[i+1,6,1]] + KET[Mod[i+1,6,1],Mod[i-1,6,1]] chi[i] + KET[i,Mod[i+1,6,1]] chi[Mod[i-1,6,1]])/(KET[Mod[i-1,6,1], i] KET[i, Mod[i+1,6,1]]), {i, 1, 6}]/.sol[[1]];

	sol = Solve[Thread[eqn == 0], Delete[Array[chi, 6], {{m1}, {m2}}]];

	Return[Coefficient[Expand[KET[m1, m2]^4 ANMHV/.sol[[1]]], eta[m3]^4]];
];


(* ----- CARRY OUT A NUMERICAL COMPARISON FOR ALL HELICITY AMPLITUDES ----- *)


(* Convert from 2x2 matrix in alpha, alpha-dot space to 4-vector *)

unpack[m_] :=
        {m[[1, 1]] + m[[2, 2]], m[[1, 2]] + m[[2, 1]],
        I m[[2, 1]] - I m[[1, 2]], m[[2, 2]] - m[[1, 1]]}/2;

unpack[a_, b_] := unpack[Outer[Times, a, b]];


(* Minkowksi dot-product on 4-vectors *)

DOT[a_, b_] := a . DiagonalMatrix[{1, -1, -1, -1}] . b;


(* Choose a random numerical nx4 momentum twistor matrix Z *)

(* Care is taken to avoid a choice that might accidentally amount to
   evaluating an amplitude on a pole *)

SetKinematics[n_Integer] := Module[{R, A, vars, bad},

	bad = True;

	While[bad,

		R = Table[Random[Integer, {-9, 9}], {4}, {4}];

		bad = (Det[R] <= 0);

	];

	A = Take[{{y[1], 0, 0, 0}, {0, y[2], 0, 0}, {0, 0, y[3], 0}, {0, 0, 0, y[4]},
		{-y[5], (1 + x[3, 1])*y[5], (-1 - x[3, 1] - x[2, 1]*x[3, 1])*y[5],
		(1 + x[3, 1] + x[2, 1]*x[3, 1] + x[1, 1]*x[2, 1]*x[3, 1])*y[5]},
		{-y[6], y[6], -y[6], y[6]}}, n];

	vars = Variables[A];

	Z = A/.Thread[vars -> Table[Random[Integer, {1, 5}], {Length[vars]}]];

	Z = Z . R;

	For[i = 0, i <= n+1, i++,

		L[i] = Take[Z[[Mod[i, Length[Z], 1]]], +2];

		M[i] = Take[Z[[Mod[i, Length[Z], 1]]], -2];

	];

	bad = 1;

	For[i = 0, i <= n+1, i++, For[j = 0, j <= n+1, j++,

		KET[i, j] = L[i][[1]] L[j][[2]] - L[i][[2]] L[j][[1]];

		If[Mod[i - j, n] != 0, bad *= KET[i, j]];

	]];

	If[bad == 0, SetKinematics[n],

		For[i = 1, i <= n, i++,

			T[i] = (KET[i-1,i] M[i+1] + KET[i+1,i-1] M[i] + KET[i,i+1] M[i-1])/(KET[i-1,i] KET[i,i+1]);

			P[i] = unpack[L[i], T[i]];

		];

		bad = 1;

		For[i = 1, i <= n, i++, For[j = 1, j <= n, j++,

			BRA[i, j] = T[i][[1]] T[j][[2]] - T[i][[2]] T[j][[1]];

			If[n > 3 && Mod[i - j, n] != 0, bad *= BRA[i, j]];

		]];

		If[bad == 0, SetKinematics[n]];

	];

	Print["Setting kinematics to Z = ", Z];

] /; (n > 2);


(* Here is the main comparison engine *)

Compare[H_] := Module[{eps, i, A},

	If[Length[H] != Length[Z], SetKinematics[Length[H]]];

	For[i = 1, i <= Length[H], i++, If[H[[i]] == 1,

		eps[i] = + Sqrt[2] If[L[i][[2]] != 0, unpack[{1, 0}, T[i]]/(L[i][[2]]), unpack[{0, 1}, T[i]]/(- L[i][[1]])],
		eps[i] = - Sqrt[2] If[T[i][[1]] != 0, unpack[L[i], {0, 1}]/(T[i][[1]]), unpack[L[i], {1, 0}]/(- T[i][[2]])]

	]];

	If[Sort[H] === {-1, -1, -1, 1, 1, 1},

		A = Union[{AMT[H], AMP[H]/.bra->BRA/.ket->KET, Together[AFD[Length[H]]/.e->eps/.p->P/.dot->DOT]}],

		A = Union[{AMP[H]/.bra->BRA/.ket->KET, Together[AFD[Length[H]]/.e->eps/.p->P/.dot->DOT]}]

	];

	Print["A[", H, "] = ", A];

	Return[Length[A] == 1];

] /; (2 < Length[H] < 7 && MemberQ[{{-1, 1}, {-1}, {1}}, Union[H]] && !(Sort[H] === {-1, 1, 1}));

CompareAll[] := Module[{err, n},

	err = 0;

	For[n = 3, n <= 6, n++,

		err += Count[Compare[2 IntegerDigits[#, 2, n] - 1] & /@ Range[0, 2^n - 1], False];

	];

	Print["There were ", err, " errors."];
];

Print["Run CompareAll[] to run a full suite of numerical checks."];
