|
| 1 | +# %load_ext autoreload |
| 2 | +# %autoreload 2 |
| 3 | + |
| 4 | +import numpy as np |
| 5 | + |
| 6 | +from pyhamimports import * |
| 7 | +from spectrum import Spectrum |
| 8 | +import glob |
| 9 | +from tqdm import tqdm |
| 10 | +from subprocess import check_output |
| 11 | + |
| 12 | +datestr = check_output(["/bin/date","+%F"]) |
| 13 | +datestr = datestr.decode().replace('\n', '') |
| 14 | + |
| 15 | +singleTemp_dir = "resources/templates/" |
| 16 | +SB2Temp_dir = "resources/templates_SB2/" |
| 17 | + |
| 18 | +singleTemp_list = np.array([os.path.basename(x) |
| 19 | + for x in glob.glob(singleTemp_dir + "*.fits")]) |
| 20 | +singleTemp_list.sort() |
| 21 | + |
| 22 | +SB2Temp_list = np.array([os.path.basename(x) |
| 23 | + for x in glob.glob(SB2Temp_dir + "*.fits")]) |
| 24 | +SB2Temp_list.sort() |
| 25 | + |
| 26 | +# 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 = O, B, A, F, G, K, M, L, C, WD |
| 27 | +single_letter_specTypes = np.array(['O', 'B', 'A', 'F', 'G', 'K', 'M', 'L', 'C', 'W']) |
| 28 | +specTypes = np.array(['O', 'B', 'A', 'F', 'G', 'K', 'M', 'L', 'C', 'WD']) |
| 29 | + |
| 30 | +new_tempLines_0 = np.empty(singleTemp_list.size, dtype=int) |
| 31 | +new_tempLines_1 = np.empty(singleTemp_list.size, dtype=int) |
| 32 | +new_tempLines_2 = np.empty(singleTemp_list.size, dtype=np.float64) |
| 33 | +new_tempLines_3 = np.ones(singleTemp_list.size, dtype=int) * 5 |
| 34 | +new_tempLines_4 = [] |
| 35 | + |
| 36 | +for ii in range(singleTemp_list.size): |
| 37 | + new_tempLines_0[ii] = np.where( |
| 38 | + single_letter_specTypes == singleTemp_list[ii][0])[0][0] |
| 39 | + if new_tempLines_0[ii] == 9: |
| 40 | + new_tempLines_1[ii] = singleTemp_list[ii][2] |
| 41 | + else: |
| 42 | + new_tempLines_1[ii] = singleTemp_list[ii][1] |
| 43 | + if len(singleTemp_list[ii].replace("_", " ").split()) == 1: |
| 44 | + new_tempLines_2[ii] = 0. |
| 45 | + else: |
| 46 | + new_tempLines_2[ii] = np.float64( |
| 47 | + singleTemp_list[ii].replace("_", " ").split()[1]) |
| 48 | + |
| 49 | +spec = Spectrum() |
| 50 | +ftype = None |
| 51 | +print("Measuring lines for single star templates:") |
| 52 | +for ii in tqdm(range(singleTemp_list.size)): |
| 53 | + message, ftype = spec.readFile(singleTemp_dir + singleTemp_list[ii], ftype) |
| 54 | + spec._lines = spec.measureLines() |
| 55 | + lines = np.array(list(spec._lines.values()))[ |
| 56 | + np.argsort(list(spec._lines.keys()))] |
| 57 | + new_tempLines_4.append(lines) |
| 58 | + |
| 59 | +SB2_index_start = new_tempLines_0.max() + 1 # 10 |
| 60 | +new_tempLines_0 = np.append(new_tempLines_0, np.arange( |
| 61 | + SB2_index_start, SB2_index_start + SB2Temp_list.size, step=1)) |
| 62 | +new_tempLines_1 = np.append(new_tempLines_1, np.zeros(SB2Temp_list.size)) |
| 63 | +new_tempLines_2 = np.append(new_tempLines_2, np.zeros(SB2Temp_list.size)) |
| 64 | +new_tempLines_3 = np.append(new_tempLines_3, np.ones(SB2Temp_list.size) * 5) |
| 65 | +# new_tempLines_4 = new_tempLines_4 |
| 66 | + |
| 67 | +spec = Spectrum() |
| 68 | +ftype = None |
| 69 | +print("Measuring lines for SB2 templates:") |
| 70 | +for ii, filename in enumerate(tqdm(SB2Temp_list)): |
| 71 | + # temp_list = [] |
| 72 | + message, ftype = spec.readFile(SB2Temp_dir + filename, ftype) |
| 73 | + measuredLines = spec.measureLines() |
| 74 | + spec._lines = measuredLines |
| 75 | + lines = np.array(list(spec._lines.values()))[ |
| 76 | + np.argsort(list(spec._lines.keys()))] |
| 77 | + linesLabels = np.array(list(spec._lines.keys()))[ |
| 78 | + np.argsort(list(spec._lines.keys()))] |
| 79 | + # temp_list.append(lines) |
| 80 | + new_tempLines_4.append(lines) |
| 81 | + |
| 82 | +new_tempLines = [new_tempLines_0, new_tempLines_1, |
| 83 | + new_tempLines_2, new_tempLines_3, new_tempLines_4] |
| 84 | + |
| 85 | +pklPath = os.path.join(spec.thisDir, 'resources', |
| 86 | + f'tempLines_{datestr}.pickle') |
| 87 | +with open(pklPath, 'wb') as pklFile: |
| 88 | + pickle.dump(new_tempLines, pklFile) |
0 commit comments