Esempio n. 1
0
class CalibrationModelTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.app = QtGui.QApplication([])

    @classmethod
    def tearDownClass(cls):
        cls.app.quit()
        cls.app.deleteLater()

    def setUp(self):
        self.img_model = ImgModel()
        self.calibration_model = CalibrationModel(self.img_model)

    def tearDown(self):
        del self.img_model
        if hasattr(self.calibration_model, 'cake_geometry'):
            del self.calibration_model.cake_geometry
        del self.calibration_model.spectrum_geometry
        del self.calibration_model
        gc.collect()

    def test_loading_calibration_gives_right_pixel_size(self):
        self.calibration_model.spectrum_geometry.load(os.path.join(data_path, 'CeO2_Pilatus1M.poni'))
        self.assertEqual(self.calibration_model.spectrum_geometry.pixel1, 0.000172)

        self.calibration_model.load(os.path.join(data_path, 'LaB6_40keV_MarCCD.poni'))
        self.assertEqual(self.calibration_model.spectrum_geometry.pixel1, 0.000079)

    def test_find_peaks_automatic(self):
        self.load_pilatus_1M_and_find_peaks()
        self.assertEqual(len(self.calibration_model.points), 6)
        for points in self.calibration_model.points:
            self.assertGreater(len(points), 0)

    def test_find_peak(self):
        """
        Tests the find_peak function for several maxima and pick points

        """
        points_and_pick_points = [
            [[30, 50], [31, 49]],
            [[30, 50], [34, 46]],
            [[5, 5],  [3, 3]],
            [[298, 298], [299, 299]]
        ]

        for data in points_and_pick_points:
            self.img_model._img_data = np.zeros((300, 300))

            point = data[0]
            pick_point = data[1]
            self.img_model._img_data[point[0], point[1]] = 100

            peak_point = self.calibration_model.find_peak(pick_point[0], pick_point[1], 10, 0)
            self.assertEqual(peak_point[0][0], point[0])
            self.assertEqual(peak_point[0][1], point[1])

    def test_search_peaks_on_ring(self):
        """
        Tests to search on the first ring of the calibrant after an inital calibration
        """
        pass

    def load_pilatus_1M_and_find_peaks(self):
        self.img_model.load(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.calibration_model.find_peaks_automatic(517.664434674, 647.529865592, 0)
        self.calibration_model.find_peaks_automatic(667.380513299, 525.252854758, 1)
        self.calibration_model.find_peaks_automatic(671.110095329, 473.571503774, 2)
        self.calibration_model.find_peaks_automatic(592.788872703, 350.495296791, 3)
        self.calibration_model.find_peaks_automatic(387.395462348, 390.987901686, 4)
        self.calibration_model.find_peaks_automatic(367.94835605, 554.290314848, 5)

    def test_calibration_with_supersampling(self):
        self.load_pilatus_1M_and_find_peaks()
        self.calibration_model.set_calibrant(os.path.join(calibrant_path, 'LaB6.D'))
        self.calibration_model.calibrate()
        normal_poni1 = self.calibration_model.spectrum_geometry.poni1
        self.img_model.set_supersampling(2)
        self.calibration_model.set_supersampling(2)
        self.calibration_model.calibrate()
        self.assertAlmostEqual(normal_poni1, self.calibration_model.spectrum_geometry.poni1, places=5)

    def test_calibration1(self):
        self.img_model.load(os.path.join(data_path, 'LaB6_40keV_MarCCD.tif'))
        self.calibration_model.find_peaks_automatic(1179.6, 1129.4, 0)
        self.calibration_model.find_peaks_automatic(1268.5, 1119.8, 1)
        self.calibration_model.set_calibrant(os.path.join(calibrant_path, 'LaB6.D'))
        self.calibration_model.calibrate()

        self.assertGreater(self.calibration_model.spectrum_geometry.poni1, 0)
        self.assertAlmostEqual(self.calibration_model.spectrum_geometry.dist, 0.18, delta=0.01)
        self.assertGreater(self.calibration_model.cake_geometry.poni1, 0)

    def test_calibration2(self):
        self.img_model.load(os.path.join(data_path, 'LaB6_OffCenter_PE.tif'))
        self.calibration_model.find_peaks_automatic(1245.2, 1919.3, 0)
        self.calibration_model.find_peaks_automatic(1334.0, 1823.7, 1)
        self.calibration_model.start_values['dist'] = 500e-3
        self.calibration_model.start_values['pixel_height'] = 200e-6
        self.calibration_model.start_values['pixel_width'] = 200e-6
        self.calibration_model.set_calibrant(os.path.join(calibrant_path, 'LaB6.D'))
        self.calibration_model.calibrate()

        self.assertGreater(self.calibration_model.spectrum_geometry.poni1, 0)
        self.assertAlmostEqual(self.calibration_model.spectrum_geometry.dist, 0.500, delta=0.01)
        self.assertGreater(self.calibration_model.cake_geometry.poni1, 0)

    def test_calibration3(self):
        self.load_pilatus_1M_and_find_peaks()
        self.calibration_model.start_values['wavelength'] = 0.406626e-10
        self.calibration_model.start_values['pixel_height'] = 172e-6
        self.calibration_model.start_values['pixel_width'] = 172e-6
        self.calibration_model.set_calibrant(os.path.join(calibrant_path, 'LaB6.D'))
        self.calibration_model.calibrate()

        self.assertGreater(self.calibration_model.spectrum_geometry.poni1, 0)
        self.assertAlmostEqual(self.calibration_model.spectrum_geometry.dist, 0.100, delta=0.02)
        self.assertGreater(self.calibration_model.cake_geometry.poni1, 0)

    def test_get_pixel_ind(self):
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        self.calibration_model.load(os.path.join(data_path, 'LaB6_40keV_MarCCD.poni'))

        self.calibration_model.integrate_1d(1000)

        tth_array = self.calibration_model.spectrum_geometry.ttha
        azi_array = self.calibration_model.spectrum_geometry.chia

        for i in range(100):
            ind1 = np.random.random_integers(0, 2023)
            ind2 = np.random.random_integers(0, 2023)

            tth = tth_array[ind1, ind2]
            azi = azi_array[ind1, ind2]

            result_ind1, result_ind2 = self.calibration_model.get_pixel_ind(tth, azi)

            self.assertAlmostEqual(ind1, result_ind1, places=3)
            self.assertAlmostEqual(ind2, result_ind2, places=3)
Esempio n. 2
0
class ImgModelTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.app = QtGui.QApplication([])

    @classmethod
    def tearDownClass(cls):
        cls.app.quit()

    def setUp(self):
        self.img_model = ImgModel()
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))

    def tearDown(self):
        del self.img_model

    def perform_transformations_tests(self):
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.rotate_img_m90()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.flip_img_horizontally()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.rotate_img_p90()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.flip_img_vertically()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.reset_img_transformations()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)

    def test_load_emits_signal(self):
        callback_fcn = MagicMock()
        self.img_model.img_changed.connect(callback_fcn)
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        callback_fcn.assert_called_once_with()

    def test_flipping_images(self):
        original_image = np.copy(self.img_model._img_data)
        self.img_model.flip_img_vertically()
        self.assertTrue(np.array_equal(self.img_model._img_data, np.flipud(original_image)))

    def test_simple_background_subtraction(self):
        self.first_image = np.copy(self.img_model.get_img_data())
        self.img_model.load_next_file()
        self.second_image = np.copy(self.img_model.get_img_data())

        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        self.assertFalse(np.array_equal(self.first_image, self.img_model.get_img_data()))

        self.img_model.load_next_file()
        self.assertEqual(np.sum(self.img_model.get_img_data()), 0)

    def test_background_subtraction_with_supersampling(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        self.img_model.set_supersampling(2)
        self.img_model.get_img_data()
        self.img_model.set_supersampling(3)
        self.img_model.get_img_data()

        self.img_model.load_next_file()
        self.img_model.get_img_data()

    def test_background_subtraction_with_transformation(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        original_img = np.copy(self.img_model._img_data)
        original_background = np.copy(self.img_model._background_data)

        self.assertNotEqual(self.img_model._background_data, None)
        self.assertFalse(np.array_equal(self.img_model.img_data, self.img_model._img_data))

        original_img_background_subtracted = np.copy(self.img_model.get_img_data())
        self.assertTrue(np.array_equal(original_img_background_subtracted, original_img - original_background))

        ### now comes the main process - flipping the image
        self.img_model.flip_img_vertically()
        flipped_img = np.copy(self.img_model._img_data)
        self.assertTrue(np.array_equal(np.flipud(original_img), flipped_img))

        flipped_background = np.copy(self.img_model._background_data)
        self.assertTrue(np.array_equal(np.flipud(original_background), flipped_background))

        flipped_img_background_subtracted = np.copy(self.img_model.get_img_data())
        self.assertTrue(np.array_equal(flipped_img_background_subtracted, flipped_img - flipped_background))

        self.assertTrue(np.array_equal(np.flipud(original_img_background_subtracted),
                                       flipped_img_background_subtracted))
        self.assertEqual(np.sum(np.flipud(original_img_background_subtracted) - flipped_img_background_subtracted), 0)

        self.img_model.load(os.path.join(data_path, 'image_002.tif'))
        self.perform_transformations_tests()

    def test_background_subtraction_with_supersampling_and_image_transformation(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        self.img_model.load(os.path.join(data_path, 'image_002.tif'))

        self.img_model.set_supersampling(2)
        self.assertEqual(self.img_model.get_img_data().shape, (4096, 4096))

        self.perform_transformations_tests()

        self.img_model.set_supersampling(3)
        self.assertEqual(self.img_model.get_img_data().shape, (6144, 6144))

        self.perform_transformations_tests()

        self.img_model.load(os.path.join(data_path, 'image_002.tif'))
        self.assertEqual(self.img_model.get_img_data().shape, (6144, 6144))

        self.perform_transformations_tests()

    def test_background_scaling_and_offset(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        # assure that everything is correct before
        self.assertTrue(np.array_equal(self.img_model.get_img_data(),
                                       self.img_model._img_data - self.img_model._background_data))

        # set scaling and see difference
        self.img_model.set_background_scaling(2.4)
        self.assertTrue(np.array_equal(self.img_model.get_img_data(),
                                       self.img_model._img_data - 2.4 * self.img_model._background_data))

        # set offset and see the difference
        self.img_model.set_background_scaling(1.0)
        self.img_model.set_background_offset(100.0)
        self.assertTrue(np.array_equal(self.img_model.img_data,
                                       self.img_model._img_data - (self.img_model._background_data + 100.0)))

        # use offset and scaling combined
        self.img_model.set_background_scaling(2.3)
        self.img_model.set_background_offset(100.0)
        self.assertTrue(np.array_equal(self.img_model.img_data,
                                       self.img_model._img_data - (2.3 * self.img_model._background_data + 100)))

    def test_background_with_different_shape(self):
        self.img_model.load_background(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertEqual(self.img_model._background_data, None)

        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        self.assertTrue(self.img_model._background_data is not None)

        self.img_model.load(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertEqual(self.img_model._background_data, None)

    def test_absorption_correction_with_supersampling(self):
        original_image = np.copy(self.img_model._img_data)
        dummy_correction = DummyCorrection(self.img_model.get_img_data().shape, 0.6)

        self.img_model.add_img_correction(dummy_correction, "Dummy 1")
        self.assertAlmostEqual(np.sum(original_image / 0.6), np.sum(self.img_model.get_img_data()), places=4)

        self.img_model.set_supersampling(2)
        self.img_model.get_img_data()

    def test_absorption_correction_with_different_image_sizes(self):
        dummy_correction = DummyCorrection(self.img_model.get_img_data().shape, 0.4)
        # self.img_data.set_absorption_correction(np.ones(self.img_data._img_data.shape)*0.4)
        self.img_model.add_img_correction(dummy_correction, "Dummy 1")
        self.assertTrue(self.img_model._img_corrections.has_items())

        self.img_model.load(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertFalse(self.img_model.has_corrections())

    def test_adding_several_absorption_corrections(self):
        original_image = np.copy(self.img_model.get_img_data())
        img_shape = original_image.shape
        self.img_model.add_img_correction(DummyCorrection(img_shape, 0.4))
        self.img_model.add_img_correction(DummyCorrection(img_shape, 3))
        self.img_model.add_img_correction(DummyCorrection(img_shape, 5))

        self.assertTrue(np.sum(original_image) / (0.5 * 3 * 5), np.sum(self.img_model.get_img_data()))

        self.img_model.delete_img_correction(1)
        self.assertTrue(np.sum(original_image) / (0.5 * 5), np.sum(self.img_model.get_img_data()))

    def test_saving_data(self):
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        filename = os.path.join(data_path, 'test.tif')
        self.img_model.save(filename)
        first_img_array = np.copy(self.img_model._img_data)
        self.img_model.load(filename)
        self.assertTrue(np.array_equal(first_img_array, self.img_model._img_data))
        self.assertTrue(os.path.exists(filename))
        os.remove(filename)

    def test_negative_rotation(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

    def test_combined_rotation(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

    def test_flip_img_horizontally(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_horizontally()
        self.img_model.flip_img_horizontally()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

    def test_flip_img_vertically(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.flip_img_vertically()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

    def test_combined_rotation_and_flipping(self):
        self.img_model.flip_img_vertically()
        self.img_model.flip_img_horizontally()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.flip_img_horizontally()
        transformed_data = self.img_model.get_img_data()
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), transformed_data))

    def test_reset_img_transformation(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_p90()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_horizontally()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.flip_img_horizontally()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.flip_img_horizontally()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

    def test_loading_a_tagged_tif_file_and_retrieving_info_string(self):
        self.img_model.load(os.path.join(data_path, "attrib.tif"))
        self.assertIn("areaDetector", self.img_model.file_info)

    def test_loading_spe_file(self):
        self.img_model.load(os.path.join(spe_path, 'CeO2_PI_CCD_Mo.SPE'))
        self.assertEqual(self.img_model.img_data.shape, (1042, 1042))
Esempio n. 3
0
class ImgDataUnitTest(unittest.TestCase):
    def setUp(self):
        self.app = QtGui.QApplication([])
        self.img_model = ImgModel()
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))

    def tearDown(self):
        del self.app
        del self.img_model
        gc.collect()

    def perform_transformations_tests(self):
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.rotate_img_m90()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.flip_img_horizontally()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.rotate_img_p90()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.flip_img_vertically()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)
        self.img_model.reset_img_transformations()
        self.assertEqual(np.sum(np.absolute(self.img_model.get_img_data())), 0)

    def test_flipping_images(self):
        original_image = np.copy(self.img_model._img_data)
        self.img_model.flip_img_vertically()
        self.assertTrue(np.array_equal(self.img_model._img_data, np.flipud(original_image)))

    def test_simple_background_subtraction(self):
        self.first_image = np.copy(self.img_model.get_img_data())
        self.img_model.load_next_file()
        self.second_image = np.copy(self.img_model.get_img_data())

        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        self.assertFalse(np.array_equal(self.first_image, self.img_model.get_img_data()))

        self.img_model.load_next_file()
        self.assertEqual(np.sum(self.img_model.get_img_data()), 0)

    def test_background_subtraction_with_supersampling(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        self.img_model.set_supersampling(2)
        self.img_model.get_img_data()
        self.img_model.set_supersampling(3)
        self.img_model.get_img_data()

        self.img_model.load_next_file()
        self.img_model.get_img_data()

    def test_background_subtraction_with_transformation(self):

        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        original_img = np.copy(self.img_model._img_data)
        original_background = np.copy(self.img_model._background_data)

        self.assertNotEqual(self.img_model._background_data, None)
        self.assertFalse(np.array_equal(self.img_model.img_data,  self.img_model._img_data))

        original_img_background_subtracted = np.copy(self.img_model.get_img_data())
        self.assertTrue(np.array_equal(original_img_background_subtracted, original_img-original_background))

        ### now comes the main process - flipping the image
        self.img_model.flip_img_vertically()
        flipped_img = np.copy(self.img_model._img_data)
        self.assertTrue(np.array_equal(np.flipud(original_img), flipped_img))

        flipped_background = np.copy(self.img_model._background_data)
        self.assertTrue(np.array_equal(np.flipud(original_background), flipped_background))

        flipped_img_background_subtracted = np.copy(self.img_model.get_img_data())
        self.assertTrue(np.array_equal(flipped_img_background_subtracted, flipped_img-flipped_background))

        self.assertTrue(np.array_equal(np.flipud(original_img_background_subtracted),
                                       flipped_img_background_subtracted))
        self.assertEqual(np.sum(np.flipud(original_img_background_subtracted)-flipped_img_background_subtracted), 0)

        self.img_model.load(os.path.join(data_path, 'image_002.tif'))
        self.perform_transformations_tests()


    def test_background_subtraction_with_supersampling_and_image_transformation(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        self.img_model.load(os.path.join(data_path, 'image_002.tif'))

        self.img_model.set_supersampling(2)
        self.assertEqual(self.img_model.get_img_data().shape, (4096, 4096))

        self.perform_transformations_tests()

        self.img_model.set_supersampling(3)
        self.assertEqual(self.img_model.get_img_data().shape, (6144, 6144))

        self.perform_transformations_tests()

        self.img_model.load(os.path.join(data_path, 'image_002.tif'))
        self.assertEqual(self.img_model.get_img_data().shape, (6144, 6144))

        self.perform_transformations_tests()

    def test_background_scaling_and_offset(self):
        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))

        #assure that everything is correct before
        self.assertTrue(np.array_equal(self.img_model.get_img_data(),
                                       self.img_model._img_data-self.img_model._background_data))

        #set scaling and see difference
        self.img_model.set_background_scaling(2.4)
        self.assertTrue(np.array_equal(self.img_model.get_img_data(),
                                       self.img_model._img_data-2.4*self.img_model._background_data))

        #set offset and see the difference
        self.img_model.set_background_scaling(1.0)
        self.img_model.set_background_offset(100.0)
        self.assertTrue(np.array_equal(self.img_model.img_data,
                                       self.img_model._img_data-(self.img_model._background_data+100.0)))

        #use offset and scaling combined
        self.img_model.set_background_scaling(2.3)
        self.img_model.set_background_offset(100.0)
        self.assertTrue(np.array_equal(self.img_model.img_data,
                                       self.img_model._img_data-(2.3*self.img_model._background_data+100)))

    def test_background_with_different_shape(self):
        self.img_model.load_background(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertEqual(self.img_model._background_data, None)

        self.img_model.load_background(os.path.join(data_path, 'image_002.tif'))
        self.assertTrue(self.img_model._background_data is not None)

        self.img_model.load(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertEqual(self.img_model._background_data, None)

    def test_absorption_correction_with_supersampling(self):
        original_image = np.copy(self.img_model.get_img_data())
        dummy_correction = DummyCorrection(self.img_model.get_img_data().shape, 0.6)

        self.img_model.add_img_correction(dummy_correction, "Dummy 1")
        self.assertAlmostEqual(np.sum(original_image)/0.6, np.sum(self.img_model.get_img_data()), places=4)

        self.img_model.set_supersampling(2)
        self.img_model.get_img_data()

    def test_absorption_correction_with_different_image_sizes(self):
        dummy_correction = DummyCorrection(self.img_model.get_img_data().shape, 0.4)
        # self.img_data.set_absorption_correction(np.ones(self.img_data._img_data.shape)*0.4)
        self.img_model.add_img_correction(dummy_correction, "Dummy 1")
        self.assertTrue(self.img_model._img_corrections.has_items())


        self.img_model.load(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.assertFalse(self.img_model.has_corrections())

    def test_adding_several_absorption_corrections(self):
        original_image = np.copy(self.img_model.get_img_data())
        img_shape = original_image.shape
        self.img_model.add_img_correction(DummyCorrection(img_shape, 0.4))
        self.img_model.add_img_correction(DummyCorrection(img_shape, 3))
        self.img_model.add_img_correction(DummyCorrection(img_shape, 5))

        self.assertTrue(np.sum(original_image)/(0.5*3*5), np.sum(self.img_model.get_img_data()))

        self.img_model.delete_img_correction(1)
        self.assertTrue(np.sum(original_image)/(0.5*5), np.sum(self.img_model.get_img_data()))


    def test_saving_data(self):
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        filename = os.path.join(data_path, 'test.tif')
        self.img_model.save(filename)
        first_img_array = np.copy(self.img_model._img_data)
        self.img_model.load(filename)
        self.assertTrue(np.array_equal(first_img_array, self.img_model._img_data))
        self.assertTrue(os.path.exists(filename))
        os.remove(filename)


    def test_rotation(self):
        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))
        self.img_model.reset_img_transformations()

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))
        self.img_model.reset_img_transformations()

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_horizontally()
        self.img_model.flip_img_horizontally()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))
        self.img_model.reset_img_transformations()

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.flip_img_vertically()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))
        self.img_model.reset_img_transformations()

        self.img_model.flip_img_vertically()
        self.img_model.flip_img_horizontally()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.flip_img_horizontally()
        transformed_data = self.img_model.get_img_data()
        self.img_model.load(os.path.join(data_path, 'image_001.tif'))
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), transformed_data))
        self.img_model.reset_img_transformations()

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_m90()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.rotate_img_p90()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_horizontally()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))

        pre_transformed_data = self.img_model.get_img_data()
        self.img_model.flip_img_vertically()
        self.img_model.flip_img_horizontally()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_p90()
        self.img_model.rotate_img_m90()
        self.img_model.rotate_img_m90()
        self.img_model.flip_img_horizontally()
        self.img_model.reset_img_transformations()
        self.assertTrue(np.array_equal(self.img_model.get_img_data(), pre_transformed_data))
class CalibrationControllerTest(unittest.TestCase):
    def setUp(self):
        self.app = QtGui.QApplication(sys.argv)
        self.img_model = ImgModel()
        self.mask_model = MaskModel()
        self.calibration_model = CalibrationModel(self.img_model)
        self.calibration_model._calibrants_working_dir = os.path.join(data_path, 'calibrants')
        self.calibration_widget = CalibrationWidget()
        self.working_dir = {}
        self.calibration_controller = CalibrationController(working_dir=self.working_dir,
                                                            img_model=self.img_model,
                                                            mask_model=self.mask_model,
                                                            widget=self.calibration_widget,
                                                            calibration_model=self.calibration_model)

    def tearDown(self):
        del self.img_model
        del self.calibration_model.cake_geometry
        del self.calibration_model.spectrum_geometry
        del self.calibration_model
        del self.app
        gc.collect()

    def load_pilatus_1M_and_pick_peaks(self):
        self.calibration_controller.load_img(os.path.join(data_path,'CeO2_Pilatus1M.tif'))
        QTest.mouseClick(self.calibration_widget.automatic_peak_num_inc_cb, QtCore.Qt.LeftButton)

        self.assertFalse(self.calibration_widget.automatic_peak_num_inc_cb.isChecked())
        self.calibration_controller.search_peaks(517.664434674, 647.529865592)
        self.calibration_controller.search_peaks(667.380513299, 525.252854758)
        self.calibration_controller.search_peaks(671.110095329, 473.571503774)
        self.calibration_controller.search_peaks(592.788872703, 350.495296791)
        self.calibration_controller.search_peaks(387.395462348, 390.987901686)
        self.calibration_controller.search_peaks(367.94835605, 554.290314848)

        self.calibration_widget.sv_wavelength_txt.setText('0.406626')
        self.calibration_widget.sv_distance_txt.setText('200')
        self.calibration_widget.sv_pixel_width_txt.setText('172')
        self.calibration_widget.sv_pixel_height_txt.setText('172')
        calibrant_index = self.calibration_widget.calibrant_cb.findText('CeO2')
        self.calibration_widget.calibrant_cb.setCurrentIndex(calibrant_index)

        self.mask_model.set_dimension(self.img_model.img_data.shape)

    def test_automatic_calibration1(self):
        self.calibration_controller.load_img(os.path.join(data_path,'LaB6_40keV_MarCCD.tif'))
        self.calibration_controller.search_peaks(1179.6, 1129.4)
        self.calibration_controller.search_peaks(1268.5, 1119.8)
        self.calibration_controller.widget.sv_wavelength_txt.setText('0.31')
        self.calibration_controller.widget.sv_distance_txt.setText('200')
        self.calibration_controller.widget.sv_pixel_width_txt.setText('79')
        self.calibration_controller.widget.sv_pixel_height_txt.setText('79')
        calibrant_index = self.calibration_widget.calibrant_cb.findText('LaB6')
        self.calibration_controller.widget.calibrant_cb.setCurrentIndex(calibrant_index)
        self.calibration_controller.calibrate()
        self.calibration_controller.widget.cake_view.set_vertical_line_pos(1419.8, 653.4)

    def test_automatic_calibration2(self):
        self.calibration_controller.load_img(os.path.join(data_path,'LaB6_OffCenter_PE.tif'))
        self.calibration_controller.search_peaks(1245.2, 1919.3)
        self.calibration_controller.search_peaks(1334.0, 1823.7)
        self.calibration_controller.widget.sv_wavelength_txt.setText('0.3344')
        self.calibration_controller.widget.sv_distance_txt.setText('500')
        self.calibration_controller.widget.sv_pixel_width_txt.setText('200')
        self.calibration_controller.widget.sv_pixel_height_txt.setText('200')
        calibrant_index = self.calibration_widget.calibrant_cb.findText('LaB6')
        self.calibration_controller.widget.calibrant_cb.setCurrentIndex(calibrant_index)
        self.calibration_controller.calibrate()
        self.calibration_controller.widget.cake_view.set_vertical_line_pos(206.5, 171.6)

    def test_automatic_calibration3(self):
        self.load_pilatus_1M_and_pick_peaks()

        start_values = self.calibration_widget.get_start_values()
        self.assertAlmostEqual(start_values['wavelength'], 0.406626e-10)
        self.assertAlmostEqual(start_values['pixel_height'], 172e-6)
        self.assertAlmostEqual(start_values['pixel_width'], 172e-6)
        self.calibration_controller.load_calibrant()
        self.assertAlmostEqual(self.calibration_model.calibrant.wavelength, 0.406626e-10)

        QTest.mouseClick(self.calibration_widget.integrate_btn, QtCore.Qt.LeftButton)
        calibration_parameter = self.calibration_model.get_calibration_parameter()[0]

        self.assertAlmostEqual(calibration_parameter['dist'], .2086, places=4)

    def test_automatic_calibration_with_supersampling(self):
        self.load_pilatus_1M_and_pick_peaks()
        self.img_model.set_supersampling(2)
        self.calibration_model.set_supersampling(2)
        self.calibration_controller.calibrate()

    def test_automatic_calibration_with_supersampling_and_mask(self):
        self.load_pilatus_1M_and_pick_peaks()
        self.img_model.set_supersampling(2)
        self.mask_model.mask_below_threshold(self.img_model._img_data, 1)
        self.mask_model.set_supersampling(2)
        self.calibration_model.set_supersampling(2)
        self.calibration_controller.widget.use_mask_cb.setChecked(True)
        self.calibration_controller.calibrate()

    def test_calibrating_one_image_size_and_loading_different_image_size(self):
        self.load_pilatus_1M_and_pick_peaks()
        self.calibration_controller.calibrate()
        self.calibration_model.integrate_1d()
        self.calibration_model.integrate_2d()
        self.calibration_controller.load_img(os.path.join(data_path, 'CeO2_Pilatus1M.tif'))
        self.calibration_model.integrate_1d()
        self.calibration_model.integrate_2d()

    def test_loading_and_saving_of_calibration_files(self):
        self.calibration_controller.load_calibration(os.path.join(data_path, 'LaB6_40keV_MarCCD.poni'))
        self.calibration_controller.save_calibration(os.path.join(data_path, 'calibration.poni'))
        self.assertTrue(os.path.exists(os.path.join(data_path, 'calibration.poni')))
        os.remove(os.path.join(data_path, 'calibration.poni'))