add test for hdf5 and flat field normalization and get sinogram

tomoscan/esrf/test/test_hdf5scan.py

@@ -217,16 +217,99 @@ class TestHDF5Scan(HDF5TestBaseClass):
     def testCompactedProjs(self):
         projs_compacted = self.scan.projections_compacted
         self.assertTrue(projs_compacted.keys() == self.scan.projections.keys())
-        # ~ print(set(map(str, list(projs_compacted.values()))))
         for i in range(22, 1520 + 1):
             self.assertTrue(projs_compacted[i].data_slice() == slice(22, 1521, None))
         for i in range(1542, 1543):
             self.assertTrue(projs_compacted[i].data_slice() == slice(1542, 1543, None))
 
 
+class TestFlatFieldCorrection(HDF5TestBaseClass):
+    """Test the flat field correction"""
+
+    def setUp(self) -> None:
+        super(TestFlatFieldCorrection, self).setUp()
+        self.dataset_file = self.get_dataset("frm_edftomomill_twoentries.nx")
+        self.scan = HDF5TomoScan(scan=self.dataset_file)
+
+    def testFlatAndDarksSet(self):
+        self.scan.set_normed_flats(
+            {
+                21: numpy.random.random(20 * 20).reshape((20, 20)),
+                1520: numpy.random.random(20 * 20).reshape((20, 20)),
+            }
+        )
+        self.scan.set_normed_darks({0: numpy.random.random(20 * 20).reshape((20, 20))})
+
+        projs = []
+        proj_indexes = []
+        for proj_index, proj in self.scan.projections.items():
+            projs.append(proj)
+            proj_indexes.append(proj_index)
+        normed_proj = self.scan.flat_field_correction(
+            projs=projs, proj_indexes=proj_indexes
+        )
+        self.assertEqual(len(normed_proj), len(self.scan.projections))
+        raw_data = get_data(projs[50])
+        self.assertFalse(numpy.array_equal(raw_data, normed_proj[50]))
+
+    def testNoFlatOrDarkSet(self):
+        projs = []
+        proj_indexes = []
+        for proj_index, proj in self.scan.projections.items():
+            projs.append(proj)
+            proj_indexes.append(proj_index)
+        with self.assertLogs("tomoscan", level="ERROR"):
+            normed_proj = self.scan.flat_field_correction(
+                projs=projs, proj_indexes=proj_indexes
+            )
+        self.assertEqual(len(normed_proj), len(self.scan.projections))
+        raw_data = get_data(projs[50])
+        self.assertTrue(numpy.array_equal(raw_data, normed_proj[50]))
+
+
+class TestGetSinogram(HDF5TestBaseClass):
+    """Test the get_sinogram function"""
+
+    def setUp(self) -> None:
+        super(TestGetSinogram, self).setUp()
+        self.dataset_file = self.get_dataset("frm_edftomomill_twoentries.nx")
+        self.scan = HDF5TomoScan(scan=self.dataset_file)
+        # set some random dark and flats
+        self.scan.set_normed_flats(
+            {
+                21: numpy.random.random(20 * 20).reshape((20, 20)),
+                1520: numpy.random.random(20 * 20).reshape((20, 20)),
+            }
+        )
+        self.scan.set_normed_darks({0: numpy.random.random(20 * 20).reshape((20, 20))})
+
+    def testGetSinogram1(self):
+        sinogram = self.scan.get_sinogram(line=12, subsampling=1)
+        self.assertEqual(sinogram.shape, (1500, 20))
+
+    def testGetSinogram2(self):
+        """Test if subsampling is negative"""
+        with self.assertRaises(ValueError):
+            self.scan.get_sinogram(line=0, subsampling=-1)
+
+    def testGetSinogram3(self):
+        sinogram = self.scan.get_sinogram(line=0, subsampling=3)
+        self.assertEqual(sinogram.shape, (500, 20))
+
+    def testGetSinogram4(self):
+        """Test if line is not in the projection"""
+        with self.assertRaises(ValueError):
+            self.scan.get_sinogram(line=-1, subsampling=1)
+
+    def testGetSinogram5(self):
+        """Test if line is not in the projection"""
+        with self.assertRaises(ValueError):
+            self.scan.get_sinogram(line=25, subsampling=1)
+
+
 def suite():
     test_suite = unittest.TestSuite()
-    for ui in (TestHDF5Scan,):
+    for ui in (TestHDF5Scan, TestFlatFieldCorrection, TestGetSinogram):
         test_suite.addTest(unittest.defaultTestLoader.loadTestsFromTestCase(ui))
     return test_suite