
(*         Copyright (C) 1990  Eran Yehudai, Alexander C. K. Hsieh          *)

BeginPackage["RelativisticKinematics`"]

Clear[PrepareMomentum, Mandelstam, SetOnShell, SetMass, SetDotProduct,
DotProductRules, SetMandelstam, MandelstamRules, SetReal, DotProduct,
Mass, Boost, BoostAmount, Decay, PerpendicularMomentum, Energy,
Momentum, SpaceDirection, Eps, SpaceTimeDimension, ZAxis]

SetReal::usage = "SetReal[a] sets Conjugate[a] = a.  SetReal[a, b,...]
applies SetReal to a, b,... in turn."

SpaceTimeDimension::usage = "SpaceTimeDimension is what you'd expect."

DotProduct::usage = "DotProduct[p, q] is the Lorentz invariant dot
product.  It is evaluated explicitly (using the metric) if both p and
q are lists of length SpaceTimeDimension."

Mass::usage = "Mass[p] is the mass of a four-vector p.  It is
evaluated explicitly if p is a list of length SpaceTimeDimension.  "

Eps::usage = "Eps[a, b, c, ...] is the completely anti-symmetric
object.  Its arguments are automatically sorted alphabetically.
Eps[..., p, ...] is the same as Eps[..., mu, ...] p[mu]."

Boost::usage = "Boost[fv, rap, dir] gives a four-vector obtained by
boosting fv by rapidity rap in the direction specified by dir.  If dir
is of the form {cth, phi}, cth is the Cosine of the angle between the
direction and the z axis and phi is the angle between the direction's
projection on the x-y plane and the x axis.  If dir is of the form {x,
y, z}, the direction is taken to be the direction of the vector {x, y,
z}. Boost[fv, rap] boosts fv by amount rap in the z direction."
Options[Boost] = {ZAxis->{0,0,1}}

BoostAmount::usage = "BoostAmount[fv] gives a vector {rap, {x, y, z}}
such that Boost[{0, 0, 0, Mass[fv]}, rap, {x, y, z}] is fv."
Options[BoostAmount] = {ZAxis -> {0,0,1}}

Decay::usage = "Decay[fv, dir, {m1, m2}] gives two four-vectors {fv1,
fv2} such that fv1+fv2 is fv, Mass[fv1] is m1, Mass[fv2] is m2 and the
direction of fv1 in the fv center-of-mass frame is given by dir.  If
dir is of the form {cth, phi}, cth is the Cosine of the angle between
the direction and the z axis and phi is the angle between the
direction's projection on the x-y plane and the x axis.  If dir is of
the form {x, y, z}, the direction is taken to be the direction of the
vector {x, y, z}.  Decay[fv, dir] is the same as Decay[fv, dir, {0,
0}]."
Options[Decay] = {ZAxis->{0,0,1}}

ZAxis::usage = "ZAxis is an option used with functions receiving a
direction as an argument. It specifies the axis with respect to which
the direction is given."

PerpendicularMomentum::usage = "PerpendicularMomentum[fv] gives the
component of fv perpendicular to the z-axis."
Options[PerpendicularMomentum] = {ZAxis->{0,0,1}}

Energy::usage = "Energy[fv] is the energy of fv."

Momentum::usage = "Momentum[fv] is the three-component momentum of fv."

SpaceDirection::usage = "SpaceDirection[fv] gives the pair {cth, phi}
describing the direction of the momentum of fv."
Options[SpaceDirection] = {ZAxis->{0,0,1}}

SetMass::usage = "SetMass[{p, m}] sets p to be a four-vector with
invariant mass m.  SetMass[{p1, p2,..., m}] sets p1, p2,.. to be
four-vector with invariant mass m.  SetMass[{p1, m1}, {p2, m2},...]
applies SetMass to {p1, m1}, {p2, m2},.. in turn."

SetDotProduct::usage = "SetDotProduct[{p, q, u}] sets
DotProduct[p, q] to be u.  SetDotProduct[{p1, q1, u1}, {p2, q2,
u2},...] sets DotProduct[p1, q1] to be u1, DotProduct[p2, q2] to be
u2, etc."

DotProductRules::usage = "DotProductRules[{p, q, u}] returns
{toRule, fromRule}.  toRule takes DotProduct[p, q] to u and
fromRule takes u to DotProduct[p, q].  DotProductRules[{p1, q1, u1},
{p2, q2, u2},...] returns {toRules, fromRules} where toRules are the
rules for converting DotProducts to u's and fromRules are the rules
for converting u's to DotProducts.  DotProductRules[{p, q, u},
Mandelstam -> {s}] will define fromRules in terms of s."

Options[DotProductRules] = {Mandelstam -> {}}

Mandelstam::usage = "Mandelstam is an option for specifying a list of
Mandelstam-type variables."

SetMandelstam::usage = "SetMandelstam[{p1, p2, p3, p4}, {m1, m2, m3,
m4}, s, t, u] sets p1, p2, p3, p4 to be on-shell with masses m1, m2, m3, m4
respectively and sets the DotProducts of p1, p2, p3 and p4 in terms of
the Mandelstam variables s, t and u and the masses."

MandelstamRules::usage = "MandelstamRules[{p1, p2, p3, p4}, {m1, m2,
m3, m4}, s, t, u] returns {toRules, fromRules}.  toRules are rules for
converting DotProducts of p1, p2, p3 and p4 to expressions containing
s, t and u and the masses.  fromRules are the rules for converting s,
t and u to expressions containing the DotProducts of p1, p2, p3 and p4
and the masses.  It sets p1, p2, p3, p4 to be on-shell with masses m1,
m2, m3, m4 respectively."

LongitudinalPolarization::usage = "LongitudinalPolarization[p] is a 
four-vector longitudinal polarization of the four-vector p."

Begin["`Private`"]

hh[x_] := Block[{h = x}, While[Head[h] =!= Symbol, h = Head[h]]; h]

SetReal[a__] :=
  (Function[
    	x,
    	(Release[hh[x]]/: Conjugate[x] = x); 
	(Release[hh[x]]/: Conjugate[b_x] := Conjugate /@ b)
  ] /@ {a};)

setOnShell[a__] :=  Function[x, x/: DotProduct[x, x] = Mass[x]^2] /@ {a};

(* functional programming *)

protected = Unprotect[Plus, Times]

(f_ +  g_) [x___] := f[x] + g[x]

(f_ *  g_) [x___] := f[x] * g[x]

(f_ *  g_?NumberQ) [x___] := f[x] * g

Protect[Release[protected]]

(* Eps *)

Literal[Eps[a___, b_, c_, d___]] := -Eps[a, c, b, d] /; !OrderedQ[{b, c}]

Literal[Eps[___, a_, ___, a_, ___]] := 0

Literal[Eps[a___, b_ + c_, d___]] := Eps[a, b, d] + Eps[a, c, d]

Literal[Eps[a___, b_?NumberQ c_, d___]] := b Eps[a, c, d]

(* DotProduct *)

SetAttributes[DotProduct, Orderless]

SetReal[DotProduct]

Literal[DotProduct[{px_, py_, pz_, pt_}, {kx_, ky_, kz_, kt_}]] :=
  kt pt - kz pz - ky py - kx px

Literal[DotProduct[p_, c_?NumberQ k_]] := c DotProduct[p, k]
Literal[DotProduct[p_, a_Plus]] := DotProduct[p, #]& /@ a

SetReal[Mass]

Mass[p:{_, _, _, _}] := Sqrt[DotProduct[p, p]]

(* SpaceDirection, etc. *)

unit[{x_, y_, z_}] := Block[{len = Sqrt[x^2 + y^2 + z^2]},
			{x/len, y/len, z/len}]

unit[{ct_, p_}] := Block[{st = Sqrt[1-ct^2]}, {st Cos[p], st Sin[p], ct}]

SpaceDirection[fv:{_, _, _, _}] := SpaceDirection[Momentum[fv]]

SpaceDirection[{px_, py_, pz_}] := 
	{pz / Sqrt[px^2 + py^2 + pz^2], ArcTan[px, py]}

newDirection[{ct_, p_}, {0, 0, _}] := 
  Block[{st = Sqrt[1-ct^2]}, {st Cos[p], st Sin[p], ct}]

newDirection[{ct_, p_}, {1, 0}] := 
  Block[{st = Sqrt[1-ct^2]}, {st Cos[p], st Sin[p], ct}]

newDirection[v:{_, _, _}, {0, 0, _}] := unit[v]

newDirection[v:{_, _, _}, {1, 0}] := unit[v]

newDirection[{ct1_, p1_}, {ct_, p_}] := 
  {Sqrt[1 - ct^2] ct1 Cos[p] + Sqrt[1 - ct1^2] Cos[p1] - 
    Sqrt[1 - ct1^2] Cos[p]^2 Cos[p1] + 
    ct Sqrt[1 - ct1^2] Cos[p]^2 Cos[p1] - 
    Sqrt[1 - ct1^2] Cos[p] Sin[p] Sin[p1] + 
    ct Sqrt[1 - ct1^2] Cos[p] Sin[p] Sin[p1], 
   Sqrt[1 - ct^2] ct1 Sin[p] - Sqrt[1 - ct1^2] Cos[p] Cos[p1] Sin[p] + 
    ct Sqrt[1 - ct1^2] Cos[p] Cos[p1] Sin[p] + 
    ct Sqrt[1 - ct1^2] Sin[p1] + Sqrt[1 - ct1^2] Cos[p]^2 Sin[p1] - 
    ct Sqrt[1 - ct1^2] Cos[p]^2 Sin[p1], 
   ct ct1 - Sqrt[1 - ct^2] Sqrt[1 - ct1^2] Cos[p] Cos[p1] - 
    Sqrt[1 - ct^2] Sqrt[1 - ct1^2] Sin[p] Sin[p1]}

newDirection[od:{_, _, _}, {ct_, p_}] :=
  {{ct Cos[p]^2 + Sin[p]^2, -Cos[p] Sin[p] + ct Cos[p] Sin[p], 
    Sqrt[1 - ct^2] Cos[p]},
    {-Cos[p] Sin[p] + ct Cos[p] Sin[p], Cos[p]^2 + ct Sin[p]^2, 
    Sqrt[1 - ct^2] Sin[p]}, 
   {-Sqrt[1 - ct^2] Cos[p], -Sqrt[1 - ct^2] Sin[p], ct}} . od

newDirection[od:{_, _, _}, {px_, py_, pz_}] :=
  {{py^2/(px^2 + py^2) + (px^2 pz)/((px^2 + py^2) Sqrt[px^2 + py^2 + pz^2]), 
    -((px py)/(px^2 + py^2)) + (px py pz)/
      ((px^2 + py^2) Sqrt[px^2 + py^2 + pz^2]), 
    px/Sqrt[px^2 + py^2 + pz^2]}, 
   {-((px py)/(px^2 + py^2)) + (px py pz)/
      ((px^2 + py^2) Sqrt[px^2 + py^2 + pz^2]), 
    px^2/(px^2 + py^2) + (py^2 pz)/((px^2 + py^2) Sqrt[px^2 + py^2 + pz^2]), 
    py/Sqrt[px^2 + py^2 + pz^2]}, 
   {-(px/Sqrt[px^2 + py^2 + pz^2]), -(py/Sqrt[px^2 + py^2 + pz^2]), 
    pz/Sqrt[px^2 + py^2 + pz^2]}} . od

newDirection[od:{_, _}, za:{_, _, _}] := newDirection[unit[od], za]

(*  if ct = 1, st = 0 then it's irrelevant what cp and sp are  *)
dirToctstcpsp[{x_, y_, _}] := {1, 0, 0, 0} /; Chop[x] == 0 && Chop[y] == 0

dirToctstcpsp[{x_, y_, z_}] := {z/Sqrt[x^2+y^2+z^2],
Sqrt[x^2+y^2]/Sqrt[x^2+y^2+z^2], x/Sqrt[x^2+y^2], y/Sqrt[x^2+y^2]}

Boost[{px_, py_, pz_, ee_}, rap_, dir_List, opts___] :=
  Block[{cp, sp, st, ct, cr = Cosh[rap], sr = Sinh[rap]},
  {ct, st, cp, sp} = dirToctstcpsp[newDirection[dir, 
        ZAxis /. Join[{opts}, Options[Boost]]]];
  {px - cp ct pz st + cp cr ct pz st + cp ee sr st - cp^2 px st^2 + 
    cp^2 cr px st^2 - cp py sp st^2 + cp cr py sp st^2, 
   py - ct pz sp st + cr ct pz sp st + ee sp sr st - cp px sp st^2 + 
    cp cr px sp st^2 - py sp^2 st^2 + cr py sp^2 st^2, 
   pz - ct^2 pz + cr ct^2 pz + ct ee sr - cp ct px st + cp cr ct px st - 
    ct py sp st + cr ct py sp st, 
   cr ee + ct pz sr + cp px sr st + py sp sr st}
]

Boost[fv:{_, _, _, _}, rap_, opts___] := 
  Boost[fv, rap, ZAxis /. Join[{opts}, Options[Boost]]]

opposite[{x_, y_, z_}] := {-x, -y, z}

opposite[{cth_, phi_}] := {cth, phi+Pi}

BoostAmount[fv:{px_, py_, pz_, e_}, opts___] := 
    {ArcCosh[e / Mass[fv]], 
	newDirection[{px, py, pz}, 
		opposite[ZAxis /. Join[{opts}, Options[BoostAmount]]]]}

phaseSpaceFactor[_, 0, 0] = 1

phaseSpaceFactor[m_, m1_, 0] := 1 - m1^2/m^2

phaseSpaceFactor[m_, 0, m1_] := phaseSpaceFactor[m, m1, 0]

phaseSpaceFactor[m_, m1_, m1_] := Sqrt[m^4-4 m^2 m1^2] / m^2

phaseSpaceFactor[m_, m1_, m2_] := 
         Sqrt[m^4+m1^4+m2^4-2m^2 m1^2-2m^2 m2^2-2m1^2 m2^2]/m^2

Decay[fv:{px_, py_, pz_, pt_}, dir_, {m1_, m2_}, opts___] :=
  Block[{n = newDirection[dir, ZAxis /. Join[{opts}, Options[Decay]]],
	m = ExpandAll[Mass[fv]], pp = Drop[fv, -1], dot, p1, p,
	e1 = (m^2+m1^2-m2^2)/2/m},
	p = m phaseSpaceFactor[m, m1, m2] / 2; dot = n.pp;
	p1 = Append[p (n + dot pp/m/(m+pt)) + e1 pp/m, pt e1/m + p dot/m];
	Expand[{p1, fv-p1}]
 ]

Decay[fv:{px_, py_, pz_, pt_}, dir_, opts___] :=
  Block[{n = newDirection[dir, ZAxis /. Join[{opts}, Options[Decay]]],
	m = ExpandAll[Mass[fv]], pp = Drop[fv, -1], dot, p1},
	dot = n.pp;
	p1 = Append[m (n + dot pp/m/(m+pt)) + pp, pt + dot] / 2;
	Expand[{p1, fv-p1}]
 ]

Decay[fv_, dir_, {m_}, opts___] := Decay[fv, dir, {m, 0}, opts]

PerpendicularMomentum[fv:{px_, py_, pz_, _}, opts__] :=
  Block[{n = unit[ZAxis /. Join[{opts}, Options[BoostAmount]]],
        pt, p = {px, py, pz}},
	pt = p - (n.p) n;
	Sqrt[pt.pt]]
	
Energy[fv:{_, _, _, e_}] := e

Momentum[fv:{px_, py_, pz_, _}] := {px, py, pz}

Eps[fv1:{_, _, _, _}, fv2:{_, _, _, _}, fv3:{_, _, _, _}, fv4:{_, _, _, _}] :=
  Det[{fv1, fv2, fv3, fv4}]

(* Mandelstam Variables *)

SetMandelstam[{p1_, p2_, p3_, p4_},
	{m1_, m2_, m3_, m4_}, s_, t_, u_] :=
	(Mass[p1] ^= m1;
	Mass[p2] ^= m2;
	Mass[p3] ^= m3;
	Mass[p4] ^= m4;
	setOnShell[p1,p2,p3,p4];
	SetReal[u, t, s, p1, p2, p3, p4];
	p1/: DotProduct[p1, p1] = m1^2;
	p2/: DotProduct[p2, p2] = m2^2;
	p3/: DotProduct[p3, p3] = m3^2;
	p4/: DotProduct[p4, p4] = m4^2;
	p1/: DotProduct[p1, p2] = (s - m1^2 - m2^2) / 2;
	p1/: DotProduct[p1, p3] = (-t + m1^2 + m3^2) / 2;
	p1/: DotProduct[p1, p4] = (-u + m1^2 + m4^2) / 2;
	p2/: DotProduct[p2, p3] = (-u + m2^2 + m3^2) / 2;
	p2/: DotProduct[p2, p4] = (-t + m2^2 + m4^2) / 2;
	p3/: DotProduct[p3, p4] = (s - m3^2 - m4^2) / 2;
	)

MandelstamRules[{p1_, p2_, p3_, p4_},
	{m1_, m2_, m3_, m4_}, s_, t_, u_] := (
	Mass[p1] ^= m1;
	Mass[p2] ^= m2;
	Mass[p3] ^= m3;
	Mass[p4] ^= m4;
	setOnShell[p1,p2,p3,p4];
	SetReal[p1,p2,p3,p4,s,t,u];
	If[NumberQ[#], Null, SetReal[#]]& /@ {m1,m2,m3,m4};
	{{DotProduct[p1, p2] -> (s - m1^2 - m2^2) / 2,
	DotProduct[p1, p3] -> (-t + m1^2 + m3^2) / 2,
	DotProduct[p1, p4] -> (-u + m1^2 + m4^2) / 2,
	DotProduct[p2, p3] -> (-u + m2^2 + m3^2) / 2,
	DotProduct[p2, p4] -> (-t + m2^2 + m4^2) / 2,
	DotProduct[p3, p4] -> (s - m3^2 - m4^2) / 2},
	{s -> m1^2 + m2^2 + 2 DotProduct[p1,p2],
	t -> m1^2 + m3^2 - 2 DotProduct[p1,p3],
	u -> m1^2 + m4^2 - 2 DotProduct[p1,p4]}}
	)

SetMass[a:{__}..] := Scan[setmasses, {a}]

SetDotProduct[a:{_, _, _}..] := 
	Scan[(DotProduct[#[[1]], #[[2]]] ^= #[[3]])&, {a}]

DotProductRules[a__] :=
	Block[{list, opts, man, backrules},
	list = Cases[{a}, {_, _, _}];
	opts = Cases[{a}, _Rule];
	man = Mandelstam /. opts /. Options[DotProductRules];
	If[man === {}, backrules = backwards /@ list,
	backrules = {DotProduct[#[[1]], #[[2]]], #[[3]]}& /@ list;
	If[backrules === {}, Null,
	backrules = Flatten[{Reduce[Apply[Equal, backrules, 1], man] /.
				And -> List /. Or -> List}];
	backrules = ToRules[
			Union[Select[backrules, MemberQ[man, First[#]]&]] /.
				List -> And]]];
	{forwards /@ list, backrules}]
	

setmasses[{p_, m_}] := (setOnShell[p]; SetReal[p]; Mass[p] ^= m)

setmasses[l_List] := Block[{mass = Last[l]},
			(setOnShell[#]; SetReal[#]; Mass[#] ^= mass;)& /@
				Drop[l, -1]]

forwards[{p1_, p2_, dp_}] := DotProduct[p1,p2] -> dp

backwards[{p1_, p2_, dp_}] := dp -> DotProduct[p1,p2]

LongitudinalPolarization/: 
  Literal[DotProduct[LongitudinalPolarization[p_], p_]] = 0

LongitudinalPolarization/:
  DotProduct[LongitudinalPolarization[p_], LongitudinalPolarization[p_]] = -1

LongitudinalPolarization[{px_, py_, pz_, e_}] :=
  Block[{pp = Sqrt[px^2+py^2+pz^2]}, {e px/pp, e py/pp, e pz/pp, pp}]

End[]

EndPackage[];
