From f1f82a6dedf2c92b00740dff26bf15f3f12657ca Mon Sep 17 00:00:00 2001
From: Felix Hekhorn <felix.hekhorn@uni-tuebingen.de>
Date: Tue, 27 Jun 2023 12:33:13 +0200
Subject: [PATCH] Make generate_missing* more symmetric

---
 src/yadism/coefficient_functions/__init__.py  |  6 +-
 .../coefficient_functions/asy/kernels.py      | 59 +++++++++----------
 .../coefficient_functions/heavy/kernels.py    |  8 +--
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/src/yadism/coefficient_functions/__init__.py b/src/yadism/coefficient_functions/__init__.py
index 0fa3480e..2d2a68c4 100644
--- a/src/yadism/coefficient_functions/__init__.py
+++ b/src/yadism/coefficient_functions/__init__.py
@@ -157,7 +157,11 @@ def heavy_components(self):
                 if self.scheme == "FFN0":
                     heavy_comps[sfh].extend(
                         asy.kernels.generate_missing_asy(
-                            self.esf, nf, ihq, self.esf.info.theory["pto_evol"]
+                            self.esf,
+                            nf,
+                            ihq,
+                            self.esf.info.theory["pto_evol"],
+                            icoupl=sfh,
                         )
                     )
                 else:
diff --git a/src/yadism/coefficient_functions/asy/kernels.py b/src/yadism/coefficient_functions/asy/kernels.py
index 26c19fc3..fb2f6449 100644
--- a/src/yadism/coefficient_functions/asy/kernels.py
+++ b/src/yadism/coefficient_functions/asy/kernels.py
@@ -18,51 +18,50 @@ def import_pc_module(kind, process, subpkg=None):
     return kernels.import_local(kind, process, subpkg)
 
 
-def generate_missing_asy(esf, nf, ihq, pto_evol):
+def generate_missing_asy(esf, nf, ihq, pto_evol, icoupl=None):
     r"""
     Collect the high-virtuality limit of missing.
 
     Parameters
     ----------
-        esf : EvaluatedStructureFunction
-            kinematic point
-        nf : int
-            number of light flavors
-        ihq : int
-            heavy quark
-        pto_evol : int
-            PTO of evolution
+    esf : EvaluatedStructureFunction
+        kinematic point
+    nf : int
+        number of light flavors
+    ihq : int
+        heavy quark
+    pto_evol : int
+        PTO of evolution
+    icoupl : None or int
+        PID of the flavor coupling (default: None)
 
     Returns
     -------
-        elems : list(yadism.kernels.Kernel)
-            list of elements
+    elems : list(yadism.kernels.Kernel)
+        list of elements
     """
-    kind = esf.info.obs_name.kind
+    # in CC there are no missing diagrams known yet
     if esf.process == "CC":
-        light_weights = kernels.cc_weights(
-            esf.info.coupling_constants,
-            esf.Q2,
-            br.quark_names[:nf],
-            nf + 1,
-            esf.info.obs_name.is_parity_violating,
-        )
-    else:
-        light_weights = light.kernels.nc_weights(
-            esf.info.coupling_constants,
-            esf.Q2,
-            nf + 1,
-            esf.info.obs_name.is_parity_violating,
-            skip_heavylight=True,
-        )
+        return ()
+    # only NC
+    weights = light.kernels.nc_weights(
+        esf.info.coupling_constants,
+        esf.Q2,
+        nf,
+        esf.info.obs_name.is_parity_violating,
+        skip_heavylight=True,
+    )
+    if icoupl is not None:
+        weights["ns"] = {k: v for k, v in weights["ns"].items() if abs(k) == icoupl}
 
-    asys = []
-    m2hq = esf.info.m2hq[ihq - 4]
+    kind = esf.info.obs_name.kind
     asy_cfs = import_pc_module(kind, esf.process)
+    m2hq = esf.info.m2hq[ihq - 4]
+    asys = []
     for res in range(pto_evol + 1):
         name = "Asy" + ("N" * res) + "LL" + "NonSinglet"
         km = kernels.Kernel(
-            light_weights["ns"], asy_cfs.__getattribute__(name)(esf, nf, m2hq=m2hq)
+            weights["ns"], asy_cfs.__getattribute__(name)(esf, nf, m2hq=m2hq)
         )
         asys.append(km)
     return asys
diff --git a/src/yadism/coefficient_functions/heavy/kernels.py b/src/yadism/coefficient_functions/heavy/kernels.py
index 894bf50b..ed4976b1 100644
--- a/src/yadism/coefficient_functions/heavy/kernels.py
+++ b/src/yadism/coefficient_functions/heavy/kernels.py
@@ -97,13 +97,10 @@ def generate_missing(esf, nf, ihq, icoupl=None):
         list of elements
 
     """
-    kind = esf.info.obs_name.kind
-    pcs = import_pc_module(kind, esf.process)
-    m2hq = esf.info.m2hq[ihq - 4]
     # in CC there are no missing diagrams known yet
     if esf.process == "CC":
         return ()
-
+    # only NC
     weights = light_nc_weights(
         esf.info.coupling_constants,
         esf.Q2,
@@ -112,6 +109,9 @@ def generate_missing(esf, nf, ihq, icoupl=None):
     )
     if icoupl is not None:
         weights["ns"] = {k: v for k, v in weights["ns"].items() if abs(k) == icoupl}
+    kind = esf.info.obs_name.kind
+    pcs = import_pc_module(kind, esf.process)
+    m2hq = esf.info.m2hq[ihq - 4]
     return (kernels.Kernel(weights["ns"], pcs.NonSinglet(esf, nf, m2hq=m2hq)),)