Storing photon replicas in a dictionary

Author: achiefa

validphys2/src/validphys/photon/compute.py

@@ -69,7 +69,8 @@ def __init__(
         self.save_to_disk = save_to_disk
         self.q_in = q_in
         self.luxpdfset = lux_params["luxset"].load()
-        self.luxpdfset_members = self.luxpdfset.n_members - 1  #
+        self.luxpdfset_members = self.luxpdfset.n_members - 1
+        self.path_to_eko_photon = self.theoryid.path / "eko_photon.tar"
 
         # Compute or load photon_qin
         if force_computation:
@@ -88,15 +89,14 @@ def _load_photon(self):
         log.info(f"Loading photon QED set from {path_to_photon}")
 
         # Load the needed replicas
-        photon_qin_array = []
+        photon_qin_array = {}
         for replica in self.replicas:
             # As input replica for the photon computation we take the MOD of the luxset_members to
             # avoid failing due to limited number of replicas in the luxset
             photonreplica = (replica % self.luxpdfset_members) or self.luxpdfset_members
             photon_qin = np.load(path_to_photon / f"replica_{photonreplica}.npz")["photon_qin"]
-            photon_qin_array.append(photon_qin)
+            photon_qin_array[replica] = photon_qin
 
-        photon_qin_array = np.stack(photon_qin_array, axis=0)
         return photon_qin_array
 
     def _compute_photon_set(self):
@@ -149,7 +149,7 @@ def _compute_photon_set(self):
             )
             path_to_photon.mkdir(parents=True, exist_ok=True)
 
-        photon_qin_array = []
+        photon_qin_array = {}
         for replica in replicas:
             # Avoid failing due to limited number of replicas in the luxset
             photonreplica = (replica % self.luxpdfset_members) or self.luxpdfset_members
@@ -180,9 +180,8 @@ def _compute_photon_set(self):
                 )
                 log.info(f"Saved photon replica {photonreplica} to {path_to_photon}")
 
-            photon_qin_array.append(photon_qin)
+            photon_qin_array[replica] = photon_qin
 
-        photon_qin_array = np.stack(photon_qin_array, axis=0)
         return photon_qin_array
 
     def _setup_fiatlux_runcard(self):
@@ -206,7 +205,7 @@ def _setup_fiatlux_runcard(self):
     @property
     def error_matrix(self):
         """Generate error matrix to be used in generate_errors."""
-        if "additional_errors" not in self.lux_params:
+        if not self.lux_params["additional_errors"]:
             return None
         extra_set = self.lux_params["additional_errors"].load()
         qs = [self.q_in] * len(XGRID)
@@ -241,8 +240,7 @@ def _evolve(self):
         interpolator = []
         integral = []
 
-        path_to_eko_photon = self.theoryid.path / "eko_photon.tar"
-        with EKO.read(path_to_eko_photon) as eko_photon:
+        with EKO.read(self.path_to_eko_photon) as eko_photon:
             # Check that qin mathces with the one in the EKO
             if not np.isclose(self.q_in, np.sqrt(eko_photon.mu20)):
                 log.error(
@@ -262,7 +260,7 @@ def _evolve(self):
                     pdfs_init = np.zeros_like(eko_op[0, 0])
                     for j, pid in enumerate(basis_rotation.flavor_basis_pids):
                         if pid == 22:
-                            pdfs_init[j] = photon_qin_array[replica - 1]
+                            pdfs_init[j] = photon_qin_array[replica]
                             ph_id = j
                         elif pid not in self.luxpdfset.flavors:
                             continue

validphys2/src/validphys/tests/photon/test_compute.py

@@ -37,12 +37,16 @@ def test_parameters_init():
     # we are not testing the photon here so we make it faster
     fiatlux_runcard['eps_base'] = 1e-1
 
-    photon = Photon(test_theory, fiatlux_runcard, [1, 2, 3])
+    photon = Photon(
+        test_theory, fiatlux_runcard, [1, 2, 3], save_to_disk=False, force_computation=True
+    )
 
     np.testing.assert_equal(photon.replicas, [1, 2, 3])
     np.testing.assert_equal(photon.luxpdfset._name, fiatlux_runcard["luxset"].name)
-    np.testing.assert_equal(photon.additional_errors.name, "LUXqed17_plus_PDF4LHC15_nnlo_100")
-    np.testing.assert_equal(photon.luxseed, fiatlux_runcard["luxseed"])
+    np.testing.assert_equal(
+        photon.lux_params["additional_errors"].name, "LUXqed17_plus_PDF4LHC15_nnlo_100"
+    )
+    np.testing.assert_equal(photon.lux_params["luxseed"], fiatlux_runcard["luxseed"])
     np.testing.assert_equal(photon.path_to_eko_photon, test_theory.path / "eko_photon.tar")
     np.testing.assert_equal(photon.q_in, 100.0)
 
@@ -58,7 +62,9 @@ def test_photon():
     theory = test_theory.get_description()
 
     for replica in [1, 2, 3]:
-        photon = Photon(test_theory, fiatlux_runcard, [replica])
+        photon = Photon(
+            test_theory, fiatlux_runcard, [replica], save_to_disk=False, force_computation=True
+        )
 
         # set up fiatlux
         path_to_F2 = test_theory.path / "fastkernel/FIATLUX_DIS_F2.pineappl.lz4"