Add MockHDF5 from tomwer

tomoscan/esrf/mock.py

@@ -146,6 +146,197 @@ class _ScanMock:
             raise ValueError('selected scene %s is no managed' % self.scene)
 
 
+class MockHDF5(_ScanMock):
+    """
+    Mock an acquisition in a hdf5 file.
+
+    note: for now the Mock class only manage one initial ref and one final
+    """
+    _PROJ_COUNT = 1
+
+    def __init__(self, scan_path, n_proj, n_ini_proj=None, n_alignement_proj=0,
+                 scan_range=360, n_recons=0, n_pag_recons=0, recons_vol=False,
+                 dim=200, create_ini_dark=True, create_ini_ref=True,
+                 create_final_ref=False, n_refs=10):
+        """
+
+        :param scan_path: directory of the file containing the hdf5 acquisition
+        :param n_proj: number of projections (does not contain alignement proj)
+        :param n_ini_proj: number of projection do add in the constructor
+        :param n_alignement_proj: number of alignment projection
+        :param int scan_range:
+        :param n_recons:
+        :param n_pag_recons:
+        :param recons_vol:
+        :param dim: frame dim - only manage square fame for now
+        :param create_ini_dark: create one initial dark frame on construction
+        :param create_ini_ref: create the initial serie of ref (n_ref) on
+                               construction (after creation of the dark)
+        :param create_final_ref: create the final serie of ref (n_ref) on
+                               construction (after creation of the dark)
+        :param n_refs: number of refs per serie
+        """
+        self.rotation_angle = numpy.linspace(start=0, stop=scan_range,
+                                             num=n_proj + 1)
+        self.rotation_angle_return = numpy.linspace(start=scan_range, stop=0,
+                                                    num=n_alignement_proj)
+        self.scan_master_file = os.path.join(scan_path, os.path.basename((scan_path)) + '.h5')
+        self._n_refs = n_refs
+        self.scan_entry = 'entry'
+
+        super(MockHDF5, self).__init__(scan_path=scan_path, n_radio=n_proj,
+                                       n_ini_radio=n_ini_proj,
+                                       n_extra_radio=n_alignement_proj,
+                                       scan_range=scan_range,
+                                       n_recons=n_recons,
+                                       n_pag_recons=n_pag_recons,
+                                       recons_vol=recons_vol, dim=dim)
+        if create_ini_dark:
+            self.add_initial_dark()
+        if create_ini_ref:
+            self.add_initial_ref()
+        if n_ini_proj is not None:
+            for i_radio in range(n_ini_proj):
+                self.add_radio(index=i_radio)
+        if create_final_ref:
+            self.add_final_ref()
+        self.scan = HDF5TomoScan(scan=self.scan_master_file, entry='entry')
+
+    def add_initial_dark(self):
+        dark = numpy.random.random((self.det_height * self.det_width)).reshape(
+            (1, self.det_width, self.det_height))
+        self._append_frame(data_=dark,
+                           rotation_angle=self.rotation_angle[-1],
+                           image_key=ImageKey.DARK_FIELD.value,
+                           image_key_control=ImageKey.DARK_FIELD.value)
+
+    def add_initial_ref(self):
+        for i in range(self._n_refs):
+            flat = numpy.random.random((self.det_height * self.det_width)).reshape(
+                (1, self.det_width, self.det_height))
+            self._append_frame(data_=flat,
+                               rotation_angle=self.rotation_angle[0],
+                               image_key=ImageKey.FLAT_FIELD.value,
+                               image_key_control=ImageKey.FLAT_FIELD.value)
+
+    def add_final_ref(self):
+        for i in range(self._n_refs):
+            flat = numpy.random.random((self.det_height * self.det_width)).reshape(
+                (1, self.det_width, self.det_height))
+            self._append_frame(data_=flat,
+                               rotation_angle=self.rotation_angle[-1],
+                               image_key=ImageKey.FLAT_FIELD.value,
+                               image_key_control=ImageKey.FLAT_FIELD.value)
+
+    def add_radio(self, index=None):
+        radio = self._get_radio_data(index=index)
+        radio = radio.reshape((1, self.det_height, self.det_width))
+        self._append_frame(data_=radio,
+                           rotation_angle=self.rotation_angle[index],
+                           image_key=ImageKey.PROJECTION.value,
+                           image_key_control=ImageKey.PROJECTION.value)
+
+    def add_alignment_radio(self, index, angle):
+        radio = self._get_radio_data(index=index)
+        radio = radio.reshape((1, self.det_height, self.det_width))
+        self._append_frame(data_=radio,
+                           rotation_angle=angle,
+                           image_key=ImageKey.PROJECTION.value,
+                           image_key_control=ImageKey.ALIGNMENT.value)
+
+    def _append_frame(self, data_, rotation_angle, image_key,
+                      image_key_control):
+
+        with h5py.File(self.scan_master_file, 'r+') as h5_file:
+            entry_one = h5_file.require_group(self.scan_entry)
+            instrument_grp = entry_one.require_group('instrument')
+            detector_grp = instrument_grp.require_group('detector')
+            sample_grp = entry_one.require_group('sample')
+
+            # add data
+            if 'data' in detector_grp:
+                # read and remove data
+                current_dataset = detector_grp['data'][()]
+                new_dataset = numpy.append(current_dataset, data_)
+                del detector_grp['data']
+                shape = list(current_dataset.shape)
+                shape[0] += 1
+                new_dataset = new_dataset.reshape(shape)
+            else:
+                new_dataset = data_
+
+            # add rotation angle
+            if 'rotation_angle' in sample_grp:
+                new_rot_angle = sample_grp['rotation_angle'][()]
+                new_rot_angle = numpy.append(new_rot_angle, rotation_angle)
+                del sample_grp['rotation_angle']
+            else:
+                new_rot_angle = [rotation_angle, ]
+
+            # add image_key
+            if 'image_key' in detector_grp:
+                new_image_key = detector_grp['image_key'][()]
+                new_image_key = numpy.append(new_image_key,
+                                             image_key)
+                del detector_grp['image_key']
+            else:
+                new_image_key = [image_key, ]
+
+            # add image_key_control
+            if 'image_key_control' in detector_grp:
+                new_image_key_control = detector_grp['image_key_control'][()]
+                new_image_key_control = numpy.append(new_image_key_control, image_key_control)
+                del detector_grp['image_key_control']
+            else:
+                new_image_key_control = [image_key_control, ]
+
+            # add count_time
+            if 'count_time' in detector_grp:
+                new_count_time = detector_grp['count_time'][()]
+                new_count_time = numpy.append(new_count_time, self._PROJ_COUNT)
+                del detector_grp['count_time']
+            else:
+                new_count_time = [self._PROJ_COUNT, ]
+
+        with h5py.File(self.scan_master_file, 'a') as h5_file:
+            entry_one = h5_file.require_group(self.scan_entry)
+            instrument_grp = entry_one.require_group('instrument')
+            if 'NX_class' not in instrument_grp.attrs:
+                instrument_grp.attrs['NX_class'] = 'NXinstrument'
+            detector_grp = instrument_grp.require_group('detector')
+            if 'NX_class' not in detector_grp.attrs:
+                detector_grp.attrs['NX_class'] = 'NXdetector'
+            sample_grp = entry_one.require_group('sample')
+            if 'NX_class' not in sample_grp.attrs:
+                sample_grp.attrs['NX_class'] = 'NXsample'
+            # write camera information
+            detector_grp['data'] = new_dataset
+            detector_grp['image_key'] = new_image_key
+            detector_grp['image_key_control'] = new_image_key_control
+            detector_grp['count_time'] = new_count_time
+            # write sample information
+            sample_grp['rotation_angle'] = new_rot_angle
+
+    def write_metadata(self, n_radio, scan_range, ref_n, dark_n):
+        with h5py.File(self.scan_master_file, 'a') as h5_file:
+            entry_one = h5_file.require_group(self.scan_entry)
+            instrument_grp = entry_one.require_group('instrument')
+            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"
+            if 'size' not in detector_grp:
+                detector_grp['size'] = (self.det_width, self.det_height)
+            if 'x_pixel_size' not in detector_grp:
+                detector_grp['x_pixel_size'] = _ScanMock.PIXEL_SIZE
+            if 'y_pixel_size' not in detector_grp:
+                detector_grp['y_pixel_size'] = _ScanMock.PIXEL_SIZE
+
+    def end_acquisition(self):
+        # no specific operation to do
+        pass
+
 
 
 class MockEDF(_ScanMock):

tomoscan/esrf/test/test_hdf5scan.py

@@ -51,6 +51,7 @@ class HDF5TestBaseClass(unittest.TestCase):
         self.test_dir = tempfile.mkdtemp()
 
     def tearDown(self) -> None:
+        input("o_O")
         shutil.rmtree(self.test_dir)