Merge branch 'improve_mock' into 'master'

Improve mock

See merge request !57

tomoscan/esrf/mock.py

@@ -431,8 +431,8 @@ class MockHDF5(_ScanMock):
             detector_grp = instrument_grp.require_group("detector")
             sample_grp = entry_one.require_group("sample")
 
-            entry_one.attrs["NX_class"] = u"NXentry"
-            entry_one.attrs["definition"] = u"NXtomo"
+            entry_one.attrs["NX_class"] = "NXentry"
+            entry_one.attrs["definition"] = "NXtomo"
             if "size" not in detector_grp:
                 detector_grp["size"] = (self.det_width, self.det_height)
             if "x_pixel_size" not in detector_grp:
@@ -470,6 +470,10 @@ class MockEDF(_ScanMock):
 
     _PAG_RECONS_PATTERN = "_slice_pag_"
 
+    _DISTANCE = 0.25
+
+    _ENERGY = 19.0
+
     def __init__(
         self,
         scan_path,
@@ -482,6 +486,8 @@ class MockEDF(_ScanMock):
         recons_vol=False,
         dim=200,
         scene="noise",
+        dark_n=0,
+        ref_n=0,
     ):
         self._last_radio_index = -1
         super(MockEDF, self).__init__(
@@ -495,12 +501,19 @@ class MockEDF(_ScanMock):
             recons_vol=recons_vol,
             dim=dim,
             scene=scene,
+            dark_n=dark_n,
+            ref_n=ref_n,
         )
         if n_ini_radio:
             for i_radio in range(n_ini_radio):
                 self.add_radio(i_radio)
         for i_extra_radio in range(n_extra_radio):
             self.add_radio(i_extra_radio + n_ini_radio)
+
+        for i_dark in range(dark_n):
+            self.add_dark(i_dark)
+        for i_ref in range(ref_n):
+            self.add_ref(i_ref)
         for i_recons in range(n_recons):
             self.add_reconstruction(i_recons)
         for i_recons in range(n_pag_recons):
@@ -539,6 +552,8 @@ class MockEDF(_ScanMock):
                 info_file.write("Row_beg=    0" + "\n")
                 info_file.write("Row_end=    " + str(self.det_height) + "\n")
                 info_file.write("PixelSize=  " + str(_ScanMock.PIXEL_SIZE) + "\n")
+                info_file.write("Distance=  " + str(self._DISTANCE) + "\n")
+                info_file.write("Energy=  " + str(self._ENERGY) + "\n")
 
     def add_radio(self, index=None):
         if index is not None:
@@ -555,9 +570,47 @@ class MockEDF(_ScanMock):
             data = self._get_radio_data(index=index_)
             assert data is not None
             assert data.shape == (self.det_width, self.det_height)
-            edf_writer = fabio.edfimage.EdfImage(data=data, header={"tata": "toto"})
+            edf_writer = fabio.edfimage.EdfImage(
+                data=data,
+                header={
+                    "motor_pos": f"{index} 0.0 1.0 2.0",
+                    "motor_mne": "srot sx sy sz",
+                },
+            )
             edf_writer.write(f)
 
+    def add_dark(self, index):
+        file_name = "darkend{0:04d}.edf".format(index)
+        file_path = os.path.join(self.scan_path, file_name)
+        if not os.path.exists(file_path):
+            data = numpy.random.random((self.det_height * self.det_width)).reshape(
+                (self.det_width, self.det_height)
+            )
+            edf_writer = fabio.edfimage.EdfImage(
+                data=data,
+                header={
+                    "motor_pos": f"{index} 0.0 1.0 2.0",
+                    "motor_mne": "srot sx sy sz",
+                },
+            )
+            edf_writer.write(file_path)
+
+    def add_ref(self, index):
+        file_name = "refHST{0:04d}.edf".format(index)
+        file_path = os.path.join(self.scan_path, file_name)
+        if not os.path.exists(file_path):
+            data = numpy.random.random((self.det_height * self.det_width)).reshape(
+                (self.det_width, self.det_height)
+            )
+            edf_writer = fabio.edfimage.EdfImage(
+                data=data,
+                header={
+                    "motor_pos": f"{index} 0.0 1.0 2.0",
+                    "motor_mne": "srot sx sy sz",
+                },
+            )
+            edf_writer.write(file_path)
+
     @staticmethod
     def mockReconstruction(folder, nRecons=5, nPagRecons=0, volFile=False):
         """