Add compute_sqaured_correlated function

feynamp/form/__init__.py

@@ -5,6 +5,7 @@
 import form
 import sympy
 from feynml.feynmandiagram import FeynmanDiagram
+from feynml.leg import Leg
 from feynmodel.feyn_model import FeynModel
 
 import feynamp.amplitude as amplitude
@@ -19,13 +20,30 @@
     get_mandelstamm,
     get_onshell,
 )
+from feynamp.leg import get_leg_momentum
 from feynamp.log import debug
 
 # TODO compute squared  functino which coutns legs!!"!!!" and picks right mandelstamm,s
 
 
 def compute_squared(fds: List[FeynmanDiagram], fm: FeynModel, tag=False):
+    return compute_squared_correlated(fds, fm, leg1=None, leg2=None, tag=tag)
+
+
+def compute_squared_correlated(
+    fds: List[FeynmanDiagram],
+    fm: FeynModel,
+    color_leg1: Leg,
+    color_leg2: Leg,
+    tag=False,
+):
     assert len(fds) > 0, "No FeynmanDiagrams to compute"
+    assert not (
+        color_leg1 is not None and color_leg2 is None
+    ), "If you want to color correlate legs, you need to provide both"
+    assert not (
+        color_leg1 is None and color_leg2 is not None
+    ), "If you want to color correlate legs, you need to provide both"
     dims = fds[0].get_externals_size()
     for fd in fds:
         assert (
@@ -34,7 +52,9 @@ def compute_squared(fds: List[FeynmanDiagram], fm: FeynModel, tag=False):
     s2 = amplitude.square_parallel(fds, fm, tag=tag)
     debug(f"{s2=}")
 
-    s2 = apply_color_parallel(s2)
+    s2 = apply_color_parallel(
+        s2, get_leg_momentum(color_leg1), get_leg_momentum(color_leg2)
+    )
 
     fs = ""
     fs += get_metrics()

feynamp/form/color.py

@@ -163,9 +163,11 @@ def apply_color_ids(string_expr):
     return run(init + colorh_init + f"Local TMP = {s};" + get_color_ids())
 
 
-def apply_color_parallel(string_exprs: List[str]):
+def apply_color_parallel(string_exprs: List[str], mom1=None, mom2=None):
     return run_parallel(
-        init + colorh_init, get_color(), [string_to_form(a) for a in string_exprs]
+        init + colorh_init,
+        get_color(mom1, mom2),
+        [string_to_form(a) for a in string_exprs],
     )
 
 

feynamp/leg.py

@@ -6,12 +6,19 @@ def get_leg_math_string(leg, fd, model):
     return get_leg_math(leg, fd, model)
 
 
+def get_leg_momentum(leg):
+    if leg is None:
+        return None
+    if leg.momentum is None or leg.momentum.name is None:
+        raise ValueError("Momentum not set for particle")
+    mom = insert_momentum(leg.momentum.name)
+    return mom
+
+
 def get_leg_math(fd, leg, model):  # epsilons or u/v optionally also barred
     p = find_leg_in_model(fd, leg, model)
-    if p.particle.momentum is None or p.particle.momentum.name is None:
-        raise ValueError("Momentum not set for particle")
+    mom = get_leg_momentum(leg)
     ret = ""
-    mom = insert_momentum(p.particle.momentum.name)
     # give particles color vectors to sum over them in the end (or better average)
     # TODO this could be also done as incoming vs outcoming
     if p.color == 8: