Upload New File

tests/test_optics.py

+import unittest
+from CpicImgSim.target import star_photlam
+from CpicImgSim.optics import make_focus_image, focal_mask, filter_throughput
+from CpicImgSim.config import which_focalplane, S
+import numpy as np
+
+
+class TestOptics(unittest.TestCase):
+    def test_filter_throughtput(self):
+        bandpass = filter_throughput('f661')
+        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)
+
+        bandpass = filter_throughput('F1265')
+        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)
+
+        bandpass = filter_throughput('deFault')
+        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)
+
+        bandpass = filter_throughput('none')
+        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)
+
+    def test_which_focalpalne(self):
+        self.assertEqual(which_focalplane('f565'), 'vis')
+        self.assertEqual(which_focalplane('F661'), 'vis')
+        self.assertEqual(which_focalplane('f743'), 'vis')
+        self.assertEqual(which_focalplane('f883'), 'vis')
+        self.assertEqual(which_focalplane('F940'), 'nir')
+        self.assertEqual(which_focalplane('f1265'), 'nir')
+        self.assertEqual(which_focalplane('F1425'), 'nir')
+        self.assertEqual(which_focalplane('f1542'), 'nir')
+        self.assertEqual(which_focalplane('wfs'), 'wfs')
+        self.assertRaises(ValueError, which_focalplane, 'what')
+
+    def test_make_focus_image(self):
+        # test fuction to generate psf
+        def gaussian_psf(band, spectrum, shape, error=0.1):
+            psf_shape = [shape, shape]
+
+            xx, yy = np.mgrid[0:psf_shape[0], 0:psf_shape[1]]
+            center = np.array([(psf_shape[0]-1)/2, (psf_shape[1]-1)/2])
+            sigma = 10
+            psf = np.exp(-((xx-center[0])**2 +
+                         (yy-center[1])**2) / (2*sigma**2))
+            psf = psf / psf.sum()
+
+            filter = filter_throughput(band)
+            return psf * (spectrum * filter).integrate()
+
+        # test targets
+        cstar = star_photlam(0, 'F2V', is_blackbody=True)
+        targets = [
+            [0, 0, cstar],
+            [1, 1, star_photlam(10-5, 'B2V', is_blackbody=True)],
+            [-2, 2, star_photlam(11-5, 'A2V', is_blackbody=True)],
+            [3, -3, star_photlam(12-5, 'G2V', is_blackbody=True)],
+            [100, 100, star_photlam(12, 'K2V', is_blackbody=True)],
+            [100, 100, star_photlam(12, 'K2V', is_blackbody=True)],
+        ]
+
+        focus_image = make_focus_image(
+            'f661',
+            targets,
+            gaussian_psf,
+            init_shifts=[1, 1],
+            rotation=45,
+            platesize=[1024, 1024]
+        )
+        self.assertIsNotNone(focus_image)
+
+        focus_image = make_focus_image(
+            'f661',
+            [],
+            gaussian_psf,
+            init_shifts=[1, 1],
+            rotation=45,
+            platesize=[1024, 1024]
+        )
+
+        self.assertIsNotNone(focus_image)
+
+    def test_focal_mask(self):
+        image = np.zeros((100, 100)) + 1
+        image_out = focal_mask(image, 1, 0.1, throughtput=0)
+        self.assertEqual((image - image_out).sum(), 2000+2000-400)
+
+
+if __name__ == '__main__':
+    unittest.main()