Only sort T's by anti or not anti. add tests.

feynamp/form/color.py

@@ -149,7 +149,7 @@ def get_full_color_correlation_matrix(fds, legs, model, s2):
     swap = []
     ids = []
     for i in range(len(legs)):
-        swap += [is_swapped_color_vector(legs[i], s2)]
+        swap += [is_swapped_color_vector(fds[0], legs[i], model, s2)]
         vec += [get_color_vector(fds[0], legs[i], model)]
         ids += [color_vector_to_id(vec[i])]
         mom += [get_leg_momentum(legs[i])]

feynamp/leg.py

@@ -42,12 +42,25 @@ def color_vector_to_index(color_vector):
     return None
 
 
-def is_swapped_color_vector(leg, s2):
+def is_swapped_color_vector(fd, leg, model, s2):
     """
     For colorcorrelations the chains of T's must be properly ordered so we check if they must be swapped
 
     TODO: check this for gluons, maybe sign is wrong...
     """
+    p = find_leg_in_model(fd, leg, model)
+    if leg.is_incoming():
+        if p.color == 3:
+            return False
+        if p.color == -3:
+            return True
+    elif leg.is_outgoing():
+        if p.color == 3:
+            return True
+        if p.color == -3:
+            return False
+    return False
+
     if re.search(r"T\(Color" + leg.id + r",.*?,.*?\)", s2):
         return True
     elif re.search(r"T\(.*?,Color" + leg.id + r",.*?\)", s2):
@@ -59,7 +72,8 @@ def is_swapped_color_vector(leg, s2):
     elif re.search(r"f\(.*?,.*?,Glu" + leg.id + r"\)", s2):
         warning("leg color third in f, check colorcorrelations")
         return True
-    raise ValueError(f"Color vector for {leg} not found in squared amplitude")
+    return False
+    # raise ValueError(f"Color vector for {leg} not found in squared amplitude")
 
 
 def get_color_vector(fd, leg, model):

tests/form/qcd/2_to_2/colorcorrelation/test_colorcorrelation_ee_qq.py

@@ -0,0 +1,37 @@
+import logging
+
+import equation_database.isbn_9780511628788 as ref
+import sympy
+from feynml.interface.qgraf import style
+from feynmodel.interface.qgraf import feynmodel_to_qgraf
+from feynmodel.interface.ufo import load_ufo_model
+from pyfeyn2.feynmandiagram import FeynML
+from pyqgraf import qgraf
+from xsdata.formats.dataclass.parsers import XmlParser
+
+from feynamp.form import compute_squared
+
+logger = logging.getLogger("feynamp")
+logger.setLevel(logging.DEBUG)
+
+
+def test_colorcorrelation_ee_qq():
+    fm = load_ufo_model("ufo_sm")
+    qfm = feynmodel_to_qgraf(fm, True, False)
+
+    qgraf.install()
+    xml_string = qgraf.run(
+        "e_minus[p1], e_plus[p2]",
+        "u[p3], u_bar[p4]",
+        loops=0,
+        loop_momentum="l",
+        model=qfm,
+        style=style,
+    )
+
+    parser = XmlParser()
+    fml = parser.from_string(xml_string, FeynML)
+    fds = [fml.diagrams[0]]
+    born = compute_squared(fds, fm, colorcorrelated=False)
+    cc = compute_squared(fds, fm, colorcorrelated=True)
+    assert (cc / born).simplify().equals(sympy.parse_expr("Cf*colorcorrelation(p3,p4)"))

tests/form/qcd/2_to_2/colorcorrelation/test_colorcorrelation_qq_qq.py

@@ -0,0 +1,71 @@
+import logging
+
+import equation_database.isbn_9780511628788 as ref
+import sympy
+from feynml.interface.qgraf import style
+from feynmodel.interface.qgraf import feynmodel_to_qgraf
+from feynmodel.interface.ufo import load_ufo_model
+from pyfeyn2.feynmandiagram import FeynML
+from pyqgraf import qgraf
+from xsdata.formats.dataclass.parsers import XmlParser
+
+from feynamp.form import compute_squared
+
+logger = logging.getLogger("feynamp")
+logger.setLevel(logging.DEBUG)
+
+
+def test_colorcorrelation_qq_photon_qq():
+    fm = load_ufo_model("ufo_sm")
+    qfm = feynmodel_to_qgraf(fm, True, False)
+
+    qgraf.install()
+    xml_string = qgraf.run(
+        "u[p1], u_bar[p2]",
+        "u[p3], u_bar[p4]",
+        loops=0,
+        loop_momentum="l",
+        model=qfm,
+        style=style,
+    )
+
+    parser = XmlParser()
+    fml = parser.from_string(xml_string, FeynML)
+    fds = [fml.diagrams[0]]
+    born = compute_squared(fds, fm, colorcorrelated=False)
+    cc = compute_squared(fds, fm, colorcorrelated=True)
+    assert (
+        (cc / born)
+        .simplify()
+        .equals("Cf*(colorcorrelation(p1, p2) + colorcorrelation(p3, p4))")
+    )
+
+
+def test_colorcorrelation_qq_gluon_qq():
+    fm = load_ufo_model("ufo_sm")
+    qfm = feynmodel_to_qgraf(fm, True, False)
+
+    qgraf.install()
+    xml_string = qgraf.run(
+        "u[p1], u_bar[p2]",
+        "u[p3], u_bar[p4]",
+        loops=0,
+        loop_momentum="l",
+        model=qfm,
+        style=style,
+    )
+
+    parser = XmlParser()
+    fml = parser.from_string(xml_string, FeynML)
+    fds = [fml.diagrams[2]]
+    born = compute_squared(fds, fm, colorcorrelated=False)
+    cc = compute_squared(fds, fm, colorcorrelated=True)
+    assert (
+        (cc / born)
+        .simplify()
+        .equals(
+            sympy.parse_expr(
+                "Cf*colorcorrelation(p1, p2) + Cf*colorcorrelation(p3, p4) - Nc*colorcorrelation(p1, p2)/2 + Nc*colorcorrelation(p1, p3)/4 - Nc*colorcorrelation(p1, p4)/4 - Nc*colorcorrelation(p2, p3)/4 + Nc*colorcorrelation(p2, p4)/4 - Nc*colorcorrelation(p3, p4)/2 + Nc**2*colorcorrelation(p1, p3)/(8*Cf) + Nc**2*colorcorrelation(p1, p4)/(8*Cf) + Nc**2*colorcorrelation(p2, p3)/(8*Cf) + Nc**2*colorcorrelation(p2, p4)/(8*Cf) - 5*colorcorrelation(p1, p3)/(8*Cf) - 5*colorcorrelation(p1, p4)/(8*Cf) - 5*colorcorrelation(p2, p3)/(8*Cf) - 5*colorcorrelation(p2, p4)/(8*Cf) + colorcorrelation(p1, p3)/(2*Cf*Nc**2) + colorcorrelation(p1, p4)/(2*Cf*Nc**2) + colorcorrelation(p2, p3)/(2*Cf*Nc**2) + colorcorrelation(p2, p4)/(2*Cf*Nc**2)"
+            )
+        )
+    )