test_target.py

import os
from math import pi
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

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['sp_model'] = 'blackbody'
        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,
            'sp_model': 'hybrid_planet'
        }
        old_planet = TargetOjbect(d_planet, cstar = cstar)
        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)
        self.assertEqual(old_planet.sp_model, 'bcc_planet')
        
    
    def test_bcc_albedo_spectrum(self):
        spectrum = AlbedoCat(90, 1, 0)
        self.assertIsInstance(spectrum, S.spectrum.SpectralElement)
        
        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.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_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_target_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)