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_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_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) if __name__ == '__main__': unittest.main()