Skip to content
test_target.py 8.59 KiB
Newer Older
GZhao's avatar
GZhao committed
import os
Chen Yili's avatar
Chen Yili committed
from math import pi
GZhao's avatar
GZhao committed
import numpy as np
import unittest
from csst_cpic_sim.target import _sptype2num, hybrid_albedo_spectrum, star_photlam
from csst_cpic_sim.target import spectrum_generator, target_file_load, detect_template_path
from csst_cpic_sim.target import AlbedoCat, bcc_spectrum, planet_contrast, extract_target_x_y, TargetOjbect
from csst_cpic_sim.config import S
Chen Yili's avatar
Chen Yili committed

GZhao's avatar
GZhao committed
tests_folder = os.path.dirname(os.path.abspath(__file__))

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

GZhao's avatar
GZhao committed
        d_cstar['sp_model'] = 'blackbody'
GZhao's avatar
GZhao committed
        cstar = TargetOjbect(d_cstar)
GZhao's avatar
GZhao committed
        self.assertEqual(cstar.sp_model, 'blackbody')
GZhao's avatar
GZhao committed

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

        d_planet['sp_model'] = 'bcc_planet'
        d_planet['coe_cloud'] = 1
        d_planet['coe_metal'] = 0
GZhao's avatar
GZhao committed
        old_planet = TargetOjbect(d_planet, cstar = cstar)
        self.assertEqual(old_planet.sp_model, 'bcc_planet')
GZhao's avatar
GZhao committed
        
    
    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()

Chen Yili's avatar
Chen Yili committed
    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'))
GZhao's avatar
GZhao committed
        
    def test_detect_template_path(self):
        for f in os.listdir():
            if os.path.isfile(f):
                break
        self_file = detect_template_path(f)
        self.assertEqual(os.path.basename(self_file), f)
        self.assertEqual(os.path.dirname(self_file), os.getcwd())

        target = detect_template_path('demo_5_25.yaml')
        self.assertEqual(os.path.basename(target), 'demo_5_25.yaml')

        self.assertRaises(FileExistsError, detect_template_path, 'demo_5_35.yaml')


    def test_template_file_load(self):
        t0 = target_file_load({0:0})
        self.assertEqual(t0[0], 0)

        
        t_error = target_file_load(['1'])
        self.assertEqual(t_error, {})

        t1 = target_file_load('')
        self.assertEqual(t1, {})


        t1 = target_file_load(' ')
        self.assertEqual(t1, {})

        t2 = target_file_load('*5')
        self.assertEqual(t2['cstar']['magnitude'], 5)
        self.assertEqual(t2['cstar']['sp_model'], 'reference')

        t3 = target_file_load('*5/25(0.7,-0.7)')
        self.assertEqual(t3['cstar']['magnitude'], 5)
        self.assertEqual(t3['cstar']['sp_model'], 'reference')
        self.assertEqual(len(t3['objects']), 1)
        obj = t3['objects'][0]
        self.assertEqual(obj['sp_model'], 'reference')
        self.assertEqual(obj['magnitude'], 25)
        self.assertEqual(obj['separation'], np.sqrt(0.7**2*2))

        t3 = target_file_load('*5/25(0.7;-0.7)')
        self.assertEqual(t3['cstar']['magnitude'], 5)
        self.assertEqual(t3['cstar']['sp_model'], 'reference')
        self.assertEqual(len(t3['objects']), 0)


        t4 = target_file_load('demo_0_20')
        self.assertEqual(t4['cstar']['magnitude'], 0)
        self.assertEqual(t4['cstar']['sp_model'], 'reference')
        self.assertEqual(len(t4['objects']), 1)
        obj = t4['objects'][0]
        self.assertEqual(obj['sp_model'], 'reference')
        self.assertEqual(obj['magnitude'], 20)

        t4 = target_file_load('demo_0_20.yaml')
        self.assertEqual(t4['cstar']['magnitude'], 0)
        self.assertEqual(t4['cstar']['sp_model'], 'reference')
        self.assertEqual(len(t4['objects']), 1)
        obj = t4['objects'][0]
        self.assertEqual(obj['sp_model'], 'reference')
        self.assertEqual(obj['magnitude'], 20)

        t6 = target_file_load('file not exist')
        self.assertEqual(t6, {})

    def test_example_targets(self):
        target1 = target_file_load('demo_0_20.yaml')
        spectrums = spectrum_generator(target1)
        self.assertEqual(len(spectrums), 2)
        for sp in spectrums:
            self.assertEqual(f"{sp[2].waveunits}", 'angstrom')
            self.assertEqual(f"{sp[2].fluxunits}", 'photlam')

        bcc = target_file_load('demo_bcc')
        spectrums = spectrum_generator(bcc)
        self.assertEqual(len(spectrums), 2)

        del bcc['cstar']['distance']
        self.assertRaises(ValueError, spectrum_generator, bcc)

        tp = target_file_load('demo_planet_temp')
        spectrums = spectrum_generator(tp)
        self.assertEqual(len(spectrums), 2)

        disk = target_file_load('demo_disk')
        contrast = disk['objects'][0]['contrast']
        spectrums = spectrum_generator(disk)
        self.assertEqual(len(spectrums), 2)
        self.assertIsNotNone(spectrums[1][3])
        diff = spectrums[0][2] * contrast - spectrums[1][2]
        self.assertTrue(diff.flux.mean() < 1)

        template_star = target_file_load('demo_std_star')
        spectrums = spectrum_generator(template_star)
        self.assertEqual(len(spectrums), 1)
        
Chen Yili's avatar
Chen Yili committed




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