test_target.py

import os
import unittest
from CpicImgSim.target import _sptype2num, hybrid_albedo_spectrum, star_photlam, spectrum_generator
from CpicImgSim.target import AlbedoCat, bcc_spectrum, planet_contrast, extract_target_x_y, TargetOjbect
from CpicImgSim.config import S
from math import pi

tests_folder = os.path.dirname(os.path.abspath(__file__))


class TestTarget(unittest.TestCase):
    def test_target_object(self):
        d_cstar= {
            'magnitude': 5,
            'ra': '120d',
            'dec': '40d',
            'distance': 10,
            'sptype': 'F0III',
        }
        cstar = TargetOjbect(d_cstar)
        self.assertEqual(cstar.sp_model == 'star')

        d_cstar['is_blackbody'] = True
        cstar = TargetOjbect(d_cstar)
        self.assertEqual(cstar.sp_model == 'blackbody')

        d_planet = {
            'radius': 2,
            'pangle': 60,
            'coe_b': 0.3,
            'coe_r': 0.7,
            'separation': 0.5,
            'phase_angle': 90,
        }
        old_planet = TargetOjbect(d_planet, cstar = cstar, sp_model='planet')
        self.assertEqual(old_planet.sp_model == 'hybrid_planet')

        d_planet['sp_model'] = 'bcc_planet'
        d_planet['coe_cloud'] = 1
        d_planet['coe_metal'] = 0
        old_planet = TargetOjbect(d_planet, cstar = cstar, sp_model='planet')
        self.assertEqual(old_planet.sp_model == 'bcc_planet')

        self.assertRaises(Warning, TargetOjbect, d_cstar, cstar = cstar, sp_model='planet')

    def test_image_input(self):
        

    
    def test_bcc_class(self):
        spectrum = AlbedoCat(90, 1, 0)
        self.assertIsInstance(spectrum, S.spectrum.SpectralElement)
        
        # import matplotlib.pyplot as plt
        # plt.plot(spectrum.wave, spectrum.throughput)
        # plt.show()

    def test_bcc_func(self):
        spectrum = bcc_spectrum(0.5, 0.5)
        self.assertIsInstance(spectrum, S.spectrum.SpectralElement)
        self.assertEqual(spectrum.waveunits.name, 'angstrom')

        # import matplotlib.pyplot as plt
        # plt.plot(spectrum.wave, spectrum.throughput)
        # plt.xlabel(spectrum.waveunits)
        # plt.show()

    def test_hybrid_albedo_spectrum(self):
        planet = hybrid_albedo_spectrum(0.5, 1)
        self.assertIsInstance(planet, S.spectrum.SpectralElement)

    def test_sptype(self):
        self.assertEqual(_sptype2num('M1V'), (6.1, 5))
        self.assertEqual(_sptype2num('O5IV'), (0.5, 4))
        self.assertEqual(_sptype2num('F3V'), (3.3, 5))
        self.assertEqual(_sptype2num('K4.5II'), (5.45, 2))
        self.assertEqual(_sptype2num('M1'), (6.1, 5))
        self.assertEqual(_sptype2num('M'), (6, 5))
        self.assertEqual(_sptype2num('K3Vvar'), (5.3, 5))
        self.assertEqual(_sptype2num('F6Vbwvar'), (3.6, 5))
        self.assertEqual(_sptype2num('K0IV SB'), (5.0, 4))
        self.assertEqual(_sptype2num('F5V+'), (3.5, 5))
        self.assertEqual(_sptype2num('G5IV/V +K1IV'), (4.5, 4))
        self.assertEqual(_sptype2num('F7IV-V'), (3.7, 4))
        self.assertEqual(_sptype2num('O4/O5IV'), (0.4, 5))
        self.assertEqual(_sptype2num('G3/G5V'), (4.3, 5))
        self.assertRaises(ValueError, _sptype2num, 'errorinput')

    def test_star_photlam(self):
        photlam = star_photlam(5, 'F0III', is_blackbody=True)
        self.assertIsInstance(photlam, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F2II', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F2', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        self.assertRaises(ValueError, star_photlam, 5, 'T0III', is_blackbody=1)

    def test_planet_contrast(self):
        contrast = planet_contrast(1, 0, 90, 1)
        contrast_cal = (69_911 / 149_597_870.7)**2 / pi
        self.assertAlmostEqual(contrast, contrast_cal)
        self.assertRaises(ValueError, planet_contrast, 1, 0, 0, 1)

    def test_extract_target_x_y(self):
        target = {
            'ra': '120d',
            'dec': '40d',
            'distance': 10,
        }
        x, y = extract_target_x_y(target, '120d', '40d')
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 0)

        x, y = extract_target_x_y(target, '120d', '41d')
        self.assertAlmostEqual(y, -3600)

        x, y = extract_target_x_y(target, '8h', '40d')
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 0)

        x, y = extract_target_x_y(dict(pangle=0, separation=1))
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 1)

        x, y = extract_target_x_y(dict(pangle=90, separation=1))
        self.assertAlmostEqual(x, 1)
        self.assertAlmostEqual(y, 0)

        self.assertRaises(ValueError, extract_target_x_y, dict(pangle=0))
        self.assertRaises(ValueError, extract_target_x_y, dict(ra='120d'))
        self.assertRaises(ValueError, extract_target_x_y,
                          dict(ra='120d', dec='40d'))

    def test_spectrum_generator(self):
        target_example = {
            'cstar': {
                'magnitude': 5,
                'ra': '120d',
                'dec': '40d',
                'distance': 10,
                'sptype': 'F0III',
            },
            'stars': [
                {
                    'magnitude': 5,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0III',
                    'isblackbody': True
                },
                {
                    'magnitude': 2,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0II',
                    'isblackbody': True
                },
                {
                    'magnitude': 2,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0',
                    'isblackbody': True
                },
            ],
            'planets': [
                {
                    'radius': 2,
                    'pangle': 60,
                    'separation': 0.5,
                    'phase_angle': 90,
                },
                {
                    'radius': 2,
                    'pangle': 60,
                    'separation': 0.5,
                    'phase_angle': 90,
                    'model': 'bcc_planet',
                    'coe_c': 1,
                    'coe_m': 0, 
                }
            ]
        }

        sps = spectrum_generator(target_example)
        num_ouput = len([sp for sp in sps if sp is not None])
        num_input = len(target_example['stars']) + \
            len(target_example['planets']) + 1
        self.assertEqual(num_ouput, num_input)


if __name__ == '__main__':
    unittest.main()