diff --git a/src/cortecs/opac/chunking.py b/src/cortecs/opac/chunking.py index 2c25a73..a6bc8f9 100644 --- a/src/cortecs/opac/chunking.py +++ b/src/cortecs/opac/chunking.py @@ -28,6 +28,8 @@ import numpy as np from tqdm import tqdm +from cortecs.opac.opac import * + def chunk_wavelengths(file, nchunks=None, wav_per_chunk=None, adjust_wavelengths=False): """ @@ -103,7 +105,7 @@ def chunk_wavelengths(file, nchunks=None, wav_per_chunk=None, adjust_wavelengths return -def write_to_file(line, file, file_suffix): +def write_to_file(line, file, file_suffix, numfiles=None): """ Writes (appends, really) a line to a file. @@ -118,12 +120,29 @@ def write_to_file(line, file, file_suffix): Side effects: Writes a line to a file! """ + if type(numfiles) != type(None): + file_suffix = (file_suffix) % numfiles + true_filename = f"{file[:-4] + str(file_suffix) + '.dat'}" f = open(true_filename, "a") f.write(line) f.close() +def get_temp_grid(file, progress=False, filetype="opacity"): + """ + gets just the temperature grid of a file. + :param file: + :return: + """ + if filetype == "cia": + raise NotImplementedError("cia files not yet supported!") + opac = Opac(file, loader="exotransmit") + temp = opac.T.copy() + del opac + return temp + + def get_header(file): """ Gets the header of a file. @@ -168,36 +187,50 @@ def count_wavelengths(file): return ticker -def get_lams(file): +def get_lams(file, progress=False, filetype="opacity"): + """ + Returns the wavelength grid used in an opacity file. + Inputs + ------- + :file: (str) path to opacity file with the wavelength grid of interest. e.g., + 'opacFe/opacFe.dat' + :progress: (bool) whether or not to include a progress bar (if tqdm is installed). + Useful for the biggest opacity file! + Outputs + ------- + :wav_grid: (np.array) wavelength values for which the opacity file had been + computed. """ - Takes in an opacity file and returns an array of all wavelengths within the file. - - Inputs: - :file: (str) path to opacity file. - Outputs: - :wavelengths: (numpy.array) individual wavelength points within the opacity file [m] + wav_grid = [] - todo: refactor so this isn't largely replicated in io. - """ f = open(file) f1 = f.readlines() + f.close() - ticker = 0 - wavelengths = [] + if progress: + iterator = tqdm(f1[2:], desc="Grabbing wavelength grid") + else: + iterator = f1[2:] - # read through all lines in the opacity file - for x in f1: - # check if blank line + # read through all lines in the opacity file; first few lines are header! + for x in iterator: + # skip blank lines if not x: continue + commad = ",".join(x.split()) # check if a wavelength line - commad = x.replace(" ", ",") - if len(np.array([eval(commad)]).flatten()) == 1: - wavelengths += [eval(x[:-1])] - f.close() - return np.array(wavelengths) + if filetype == "CIA": + constraint = len(np.array([eval(commad)]).flatten()) > 1 + elif filetype == "opacity": + constraint = len(np.array([eval(commad)]).flatten()) == 1 + else: + raise ValueError("filetype not recognized!") + + if constraint: + wav_grid += [np.array([eval(commad)]).flatten()[0]] + return wav_grid def add_lams(max_lam_to_add_ind, file, next_file): diff --git a/src/cortecs/opac/interpolate_cia.py b/src/cortecs/opac/interpolate_cia.py index d3bc785..5ff8034 100644 --- a/src/cortecs/opac/interpolate_cia.py +++ b/src/cortecs/opac/interpolate_cia.py @@ -26,27 +26,119 @@ ######################### Pt. 1: Interpolation #################################### -def interpolate_CIA(CIA_file, reference_file): +def get_n_species(CIA_file, verbose=False): """ - Interpolates a CIA file to a higher resolution, using the wavelength grid - of a reference file. Note: This function assumes that the CIA file has - Hels, HeHs, CH4CH4s, H2Hes, H2CH4s, H2Hs, H2H2s, and CO2CO2s. + Returns the number of species in a CIA file. Inputs - ------ - :CIA_file: (str) path to CIA file to be interpolated. e.g., - 'opacCIA/opacCIA.dat' - :reference_file: (str) path to opacity file with the wavelength grid of interest. e.g., - 'opacFe/opacFe.dat' + ------- + :CIA_file: (str) path to CIA file. e.g., 'opacCIA/opacCIA.dat' Outputs ------- - None - Side effects - ------------- - Creates a file with "hires" attached to the end of CIA_file that has been interpolated - to higher resolution. + :n_species: (int) number of species in the CIA file. + """ + + f = open(CIA_file) + f1 = f.readlines() + f.close() + # the third line fshould be the first normal one. + + line = f1[3] + n_species = len(line.split()[1:]) + + if verbose: + print(f"Number of species in CIA file {CIA_file}: {n_species}") + if n_species == 8: + print( + "Species are likely: H-el, He-H, CH4-CH4, H2-He, H2-CH4, H2-H, H2-H2, CO2-CO2" + ) + elif n_species == 14: + print( + "Species are likely: H2-H2, H2-He, H2-H, H2-CH4, CH4-Ar, CH4-CH4, CO2-CO2, He-H, N2-CH4, N2-H2, N2-N2, O2-CO2, O2-N2, O2-O2)" + ) + return n_species + + +def get_empty_species_dict(CIA_file, verbose=False): """ + returns a species dictioanry given a CIA file + :param CIA_file: + :param verbose: + :return: + """ + n_species = get_n_species(CIA_file, verbose=verbose) + if n_species == 8: + # todo: refactor. has to be a cleaner way to do this! infer the columns, etc. + Hels = [] + HeHs = [] + CH4CH4s = [] + H2Hes = [] + H2CH4s = [] + H2Hs = [] + H2H2s = [] + CO2CO2s = [] + + # python > 3.6 has ordered dictionaries! + species_dict = { + "Hels": Hels, + "HeHs": HeHs, + "CH4CH4s": CH4CH4s, + "H2Hes": H2Hes, + "H2CH4s": H2CH4s, + "H2Hs": H2Hs, + "H2H2s": H2H2s, + "CO2CO2s": CO2CO2s, + } - real_wavelength_grid = get_wav_grid(reference_file, progress=True) + elif n_species == 14: + H2H2s = [] + H2Hes = [] + H2Hs = [] + H2CH4s = [] + CH4Ar = [] + CH4CH4s = [] + CO2CO2s = [] + HeHs = [] + N2CH4s = [] + N2H2s = [] + N2N2s = [] + O2CO2s = [] + O2N2s = [] + O2O2s = [] + species_dict = { + "H2H2s": H2H2s, + "H2Hes": H2Hes, + "H2Hs": H2Hs, + "H2CH4s": H2CH4s, + "CH4Ar": CH4Ar, + "CH4CH4s": CH4CH4s, + "CO2CO2s": CO2CO2s, + "HeHs": HeHs, + "N2CH4s": N2CH4s, + "N2H2s": N2H2s, + "N2N2s": N2N2s, + "O2CO2s": O2CO2s, + "O2N2s": O2N2s, + "O2O2s": O2O2s, + } + else: + print("Number of species in CIA file not recognized. Check the file!") + species_dict = {} + species_keys = np.arange(n_species) + for species_key in species_keys: + species_dict[species_key] = [] + return species_dict + + +def read_cia(CIA_file, verbose=False): + """ + Reads in a CIA file. + Inputs + ------ + :CIA_file: (str) path to CIA file. e.g., 'opacCIA/opacCIA.dat' + Outputs + ------- + :df: (pd.DataFrame) dataframe with the CIA data. + """ f = open(CIA_file) f1 = f.readlines() @@ -54,15 +146,8 @@ def interpolate_CIA(CIA_file, reference_file): temperatures = [] wavelengths = [] - # todo: refactor. has to be a cleaner way to do this! infer the columns, etc. - Hels = [] - HeHs = [] - CH4CH4s = [] - H2Hes = [] - H2CH4s = [] - H2Hs = [] - H2H2s = [] - CO2CO2s = [] + + species_dict = get_empty_species_dict(CIA_file, verbose=verbose) # read through all lines in the CIA file for line in tqdm(f1[1:], desc="Reading CIA file"): @@ -74,62 +159,100 @@ def interpolate_CIA(CIA_file, reference_file): values = [eval(value) for value in line.split(" ")[::3][:-1]] temperatures += [temperature] wavelengths += [values[0]] - Hels += [values[1]] - HeHs += [values[2]] - CH4CH4s += [values[3]] - H2Hes += [values[4]] - H2CH4s += [values[5]] - H2Hs += [values[6]] - H2H2s += [values[7]] - CO2CO2s += [values[8]] - - # easier to slice with a dataframe later on + + for species_ind, species in enumerate(species_dict.keys()): + species_dict[species] += [values[species_ind + 1]] + + # hm. how to do this with different numbers of species? df = pd.DataFrame( { "temp": temperatures, "wav": wavelengths, - "Hel": Hels, - "HeH": HeHs, - "CH4CH4": CH4CH4s, - "H2He": H2Hes, - "H2CH4": H2CH4s, - "H2H": H2Hs, - "H2H2": H2H2s, - "CO2CO2": CO2CO2s, } ) + for species in species_dict.keys(): + df[species] = species_dict[species] - # perform interpolation + return df + + +def check_temp_grid(df, real_temperature_grid, CIA_file): + """ + Checks that the temperature grid of the CIA file is the same as the reference file. + Inputs + ------ + :df: (pd.DataFrame) dataframe with the CIA data. + :real_temperature_grid: (np.array) temperature values for which the opacity file had been + computed. + Outputs + ------- + None + Side effects + ------------- + Raises a ValueError if the temperature grid is not the same. + """ + + for temp in real_temperature_grid: + if temp not in df.temp.unique(): + print( + f"Temperature {temp} not in CIA file {CIA_file}! Cannot interpolate in temperature yet. Will set these values to 0." + ) + return + + +def interpolate_CIA(CIA_file, reference_file, outfile=None): + """ + Interpolates a CIA file to a higher resolution, using the wavelength grid + of a reference file. Note: This function assumes that the CIA file has + Hels, HeHs, CH4CH4s, H2Hes, H2CH4s, H2Hs, H2H2s, and CO2CO2s. - interped_Hels = [] - interped_HeHs = [] - interped_CH4CH4s = [] - interped_H2Hes = [] - interped_H2CH4s = [] - interped_H2Hs = [] - interped_H2H2s = [] - interped_CO2CO2s = [] + TODO: add 2d interpolation. + + Inputs + ------ + :CIA_file: (str) path to CIA file to be interpolated. e.g., + 'opacCIA/opacCIA.dat' + :reference_file: (str) path to opacity file with the wavelength grid of interest. e.g., + 'opacFe/opacFe.dat' + Outputs + ------- + None + Side effects + ------------- + Creates a file with "hires" attached to the end of CIA_file that has been interpolated + to higher resolution. + """ + + real_wavelength_grid = get_lams(reference_file, progress=True, filetype="opacity") + # need to put it on the right temperature grid, too! + real_temperature_grid = get_temp_grid( + reference_file, progress=True, filetype="opacity" + ) + + df = read_cia(CIA_file) + + check_temp_grid(df, real_temperature_grid, CIA_file) + + species_dict_interped = get_empty_species_dict(CIA_file) interped_wavelengths = [] interped_temps = [] - for unique_temp in tqdm(df.temp.unique()): - sub_df = df[df.temp == unique_temp] - - # convert to lists so that they're neatly nested lists - interped_Hels += list(np.interp(real_wavelength_grid, sub_df.wav, sub_df.Hel)) - interped_HeHs += list(np.interp(real_wavelength_grid, sub_df.wav, sub_df.HeH)) - interped_CH4CH4s += list( - np.interp(real_wavelength_grid, sub_df.wav, sub_df.CH4CH4) - ) - interped_H2Hes += list(np.interp(real_wavelength_grid, sub_df.wav, sub_df.H2He)) - interped_H2CH4s += list( - np.interp(real_wavelength_grid, sub_df.wav, sub_df.H2CH4) - ) - interped_H2Hs += list(np.interp(real_wavelength_grid, sub_df.wav, sub_df.H2H)) - interped_H2H2s += list(np.interp(real_wavelength_grid, sub_df.wav, sub_df.H2H2)) - interped_CO2CO2s += list( - np.interp(real_wavelength_grid, sub_df.wav, sub_df.CO2CO2) - ) + # perform interpolation + for unique_temp in tqdm(real_temperature_grid, desc="Interpolating"): + if unique_temp not in df.temp.unique(): + for species_key in species_dict_interped.keys(): + species_dict_interped[species_key] += list( + np.ones_like(real_wavelength_grid) * 0.0 + ) + + else: + sub_df = df[df.temp == unique_temp] + + # convert to lists so that they're neatly nested lists + for species_key in species_dict_interped.keys(): + species_dict_interped[species_key] += list( + np.interp(real_wavelength_grid, sub_df.wav, sub_df[species_key]) + ) interped_wavelengths += real_wavelength_grid interped_temps += list(np.ones_like(real_wavelength_grid) * unique_temp) @@ -138,100 +261,46 @@ def interpolate_CIA(CIA_file, reference_file): new_string = [] + # this is pretty gross below : ( + reference_species = species_dict_interped[list(species_dict_interped.keys())[0]] + + # todo: include the different temperature range on which to interpolate. + buffer = " " # there's a set of spaces between each string! - for i in tqdm(range(len(interped_Hels))): + for i in tqdm(range(len(reference_species))): # the first line gets different treatment! if i == 0: - temp = 500 + temp = np.min( + interped_temps + ) # add the LOWEST temperature in the temperature grid! new_string += ["{:.12e}".format(temp) + "\n"] if interped_temps[i] != temp: temp = interped_temps[i] new_string += ["{:.12e}".format(temp) + "\n"] - wavelength = "{:.12e}".format(interped_wavelengths[i]) - Hel_string = "{:.12e}".format(interped_Hels[i]) - HeH_string = "{:.12e}".format(interped_HeHs[i]) - CH4CH4_string = "{:.12e}".format(interped_CH4CH4s[i]) - H2He_string = "{:.12e}".format(interped_H2Hes[i]) - H2CH4_string = "{:.12e}".format(interped_H2CH4s[i]) - H2H_string = "{:.12e}".format(interped_H2Hs[i]) - H2H2_string = "{:.12e}".format(interped_H2H2s[i]) - CO2CO2_string = "{:.12e}".format(interped_CO2CO2s[i]) - new_string += [ - wavelength - + buffer - + Hel_string - + buffer - + HeH_string - + buffer - + CH4CH4_string - + buffer - + H2He_string - + buffer - + H2CH4_string - + buffer - + H2H_string - + buffer - + H2H2_string - + buffer - + CO2CO2_string - + buffer - + "\n" - ] + wavelength_string = "{:.12e}".format(interped_wavelengths[i]) + + line_string = wavelength_string + buffer + + for species_key in species_dict_interped.keys(): + line_string += ( + "{:.12e}".format(species_dict_interped[species_key][i]) + buffer + ) + + new_string += [line_string + "\n"] # todo: check the insert. and can pull wavelength grid. - new_string.insert( - 0, - "500.000 600.000 700.000 800.000 900.000 1000.000 1100.000 1200.000 1300.000 1400.000 1500.000 1600.000 1700.000 1800.000 1900.000 2000.000 2100.000 2200.000 2300.000 2400.000 2500.000 2600.000 2700.000 2800.000 2900.000 3000.000 3100.000 3200.000 3300.000 3400.000 3500.000 3600.000 3700.000 3800.000 3900.000 4000.000 4100.000 4200.000 4300.000 4400.000 4500.000 4600.000 4700.000 4800.000 4900.000 5000.000 \n", - ) - new_file = CIA_file.split(".dat")[0] - new_file += "_highres.dat" - f2 = open(new_file, "w") + temp_string = " ".join(str(temp) for temp in real_temperature_grid) + " \n" + new_string.insert(0, temp_string) + if type(outfile) == type(None): + outfile = CIA_file.split(".dat")[0] + outfile += "_highres.dat" + f2 = open(outfile, "w") f2.writelines(new_string) f2.close() return -def get_wav_grid(file, progress=False): - """ - Returns the wavelength grid used in an opacity file. - Inputs - ------- - :file: (str) path to opacity file with the wavelength grid of interest. e.g., - 'opacFe/opacFe.dat' - :progress: (bool) whether or not to include a progress bar (if tqdm is installed). - Useful for the biggest opacity file! - Outputs - ------- - :wav_grid: (np.array) wavelength values for which the opacity file had been - computed. - """ - - wav_grid = [] - - f = open(file) - f1 = f.readlines() - f.close() - - if progress: - iterator = tqdm(f1[2:], desc="Grabbing wavelength grid") - else: - iterator = f1[2:] - - # read through all lines in the opacity file; first few lines are header! - for x in iterator: - # skip blank lines - if not x: - continue - commad = x.replace(" ", ",") - - # check if a wavelength line - if len(np.array([eval(commad)]).flatten()) == 1: - wav_grid += [np.array([eval(commad)]).flatten()[0]] - - return wav_grid - - ######################### Pt. 2: Chunking #################################### @@ -288,8 +357,8 @@ def chunk_wavelengths_CIA(file, ref_file_base, numfiles): true_file_suffix = (file_suffix) % numfiles wav_per_chunk = chunk_list[true_file_suffix] - write_to_file(temperature, file, file_suffix, ntemps, numfiles) - write_to_file(line, file, file_suffix, ntemps, numfiles) # just writes line + write_to_file(temperature, file, file_suffix, numfiles), + write_to_file(line, file, file_suffix, numfiles) # just writes line true_header = header[0] @@ -314,7 +383,7 @@ def get_wav_per_chunk(file_suffix, ref_file_base): file = ref_file_base + str(file_suffix) + ".dat" - real_wavelength_grid = get_wav_grid(file) + real_wavelength_grid = get_lams(file, filetype="opacity") len_grid = len(real_wavelength_grid) @@ -341,26 +410,3 @@ def prepend_line(file_name, line): os.remove(file_name) # Rename dummy file as the original file os.rename(dummy_file, file_name) - - -def write_to_file(line, file, file_suffix, ntemps, numfiles): - """ - Writes a line to a file. - Inputs - ------ - :line: (str) line to be written. - :file: (str) path to file being chunked. e.g., opacCIA_highres.dat - :file_suffix: (int) number of chunk. - :ntemps: (int) number of temperatures in grid. - Outputs - -------- - None - Side effects - ------------- - Writes a line to a file! - """ - true_file_suffix = (file_suffix) % numfiles - true_filename = f"{file[:-4] + str(true_file_suffix) + '.dat'}" - f = open(true_filename, "a") - f.write(line) - f.close() diff --git a/src/cortecs/opac/io.py b/src/cortecs/opac/io.py index 0191a20..bcf5f67 100644 --- a/src/cortecs/opac/io.py +++ b/src/cortecs/opac/io.py @@ -197,6 +197,42 @@ def load(self, cross_sec_filename, T_filename, wl_filename, species_name): return wl, T, cross_section +class loader_exotransmit_cia(loader_base): + """ + loads in opacity data that are used with the PLATON code's collision-induced absorption.. + """ + + def load(self, cross_sec_filename, T_filename, wl_filename, species_name): + """ + loads in opacity data that's built for PLATON. To be passed on to Opac object. + + The temperature grid, pressure grid, and wavelength grid are saved as separate files for PLATON. + + Inputs + ------ + filename : str + name of file to load + cross_sec_filename : str + name of cross section file + T_filename : str + name of temperature file + wl_filename : str + name of wavelength file + species_name : tuples + name of two colliding species. E.g., ('H2', 'CH4'). todo: all at once? + """ + # todo: check wl units. They're in meters here. + # temperatures are in K. + # pressures are in Pa, I believe. + wl = np.load(wl_filename) + T = np.load(T_filename) + data = pickle.load( + open(cross_sec_filename, "rb"), encoding="latin1" + ) # packaged as T x P x wl. todo: check packing + cross_section = data[species_name] + return wl, T, cross_section + + class loader_exotransmit(loader_base): """ loads in opacity data in the exo-transmit format. This format is also used by PLATON, I believe. @@ -271,3 +307,6 @@ def load(self, filename): T, P = self.get_t_p(filename) return wl, T, P, opacities + + +# todo: write a writer for each of these. diff --git a/src/cortecs/opac/opac.py b/src/cortecs/opac/opac.py index e6f257f..908105e 100644 --- a/src/cortecs/opac/opac.py +++ b/src/cortecs/opac/opac.py @@ -84,9 +84,10 @@ class Opac_cia(Opac): method_dict = { "platon_cia": loader_platon_cia, + "exotransmit_cia": loader_exotransmit_cia, } - def __init__(self, filename, loader="chimera"): + def __init__(self, filename, loader="platon_cia"): """ wraps around the loaders. @@ -104,6 +105,7 @@ def __init__(self, filename, loader="chimera"): self.filename = filename load_obj = loader_base() self.wl, self.T, self.cross_section = load_obj.load(filename, loader=loader) + # at least for the exotransmit case, we have...wl x temp. self.n_wav, self.n_t = len(self.wl), len(self.T)