Fix yad_grids.py

extras/heavy_em_n3lo/yad_grids.py

@@ -21,8 +21,7 @@
 
 order = int(sys.argv[5])
 
-hs_version = sys.argv[6]
-if hs_version not in ["klmv", "abmp", "gm"]:
+if sys.argv[6] not in ["klmv", "abmp", "gm"]:
     raise ValueError("Set hs_version to 'klmv', 'abmp' or 'gm'")
 # klmv = Kawamura, Lo Presti, Moch, Vogt: approximation from [arXiv:1205.5727]
 # abmp = Alekhin, Blumlein, Moch, Placakyte: approximation from [arXiv:1701.05838]
@@ -32,7 +31,7 @@
 mufrac = 1.0
 verbose = True
 
-hs_version = "exact" if channel == "q" else "gm"
+hs_version = "exact" if channel == "q" else sys.argv[6]
 if order > 1:
     massive = adani.ApproximateCoefficientFunction(
         order, kind, channel, True, hs_version
@@ -72,19 +71,19 @@ def run(n_threads, eta_grid, xi_grid):
 if __name__ == "__main__":
     output_files = {}
     for variation in range(-1, 1 + 1):
-        output_files[variation] = f"C{channel}_nf{nf}_var{variation}.npy"
+        output_files[variation] = f"C{kind}{channel}_nf{nf}_var{variation}.npy"
     etafname = here / "eta.npy"
     eta_grid = np.load(etafname)
     xifname = here / "xi.npy"
     xi_grid = np.load(xifname)
 
     if verbose:
         print(
-            f"Computation of the grid for the coefficient function C{channel} for nf = {nf}, and µ/Q = {mufrac}, variation = {variation}"
+            f"Computation of the grid for the coefficient function C{kind}{channel} for nf = {nf}, and µ/Q = {mufrac}"
         )
         print(f"Size of the grid (eta,xi) = ({len(eta_grid)},{len(xi_grid)})")
         print(
-            "This may take a while (depending on the number of threads you choose). In order to spend this time, I would suggest you this interesting view:"
+            "This may take a while (depending on the number of threads you chose). In order to spend this time, I would suggest you this interesting view:"
         )
         print("https://www.youtube.com/watch?v=53pG68KCUMI")
 
@@ -94,9 +93,10 @@ def run(n_threads, eta_grid, xi_grid):
         print("total running time: ", time.perf_counter() - start)
 
     res_mat = res_vec.reshape(len(xi_grid), len(eta_grid), 3)
+
     for variation in range(-1, 1 + 1):
         if order == 1 and variation in [-1, 1]:
             continue
         if verbose:
             print(f"Saving {variation} grid in ", here / output_files[variation])
-        np.save(here / output_files[variation], res_mat[variation + 1])
+        np.save(here / output_files[variation], res_mat[:, :, variation + 1])