Merge pull request #37 from arjunsavel/fix_exotransmit_readwrite

faster exotransmit io. plus! better.

src/cortecs/opac/chunking.py

@@ -38,6 +38,7 @@ def chunk_wavelengths(
     wav_per_chunk=None,
     adjust_wavelengths=False,
     loader="exotransmit",
+    load_kwargs=None,
 ):
     """
     Performs wavelength-chunking.
@@ -74,7 +75,7 @@ def chunk_wavelengths(
         )
 
     if not wav_per_chunk:
-        opac = Opac(file, loader=loader)
+        opac = Opac(file, loader=loader, load_kwargs=load_kwargs)
         num_wavelengths = len(opac.wl)
         del opac  # clean it up
         wav_per_chunk = round(num_wavelengths / nchunks)
@@ -205,7 +206,7 @@ def add_lams(max_lam_to_add_ind, file, next_file):
             return
 
 
-def add_overlap(filename, v_max=11463.5):
+def add_overlap(filename, v_max=11463.5, load_kwargs=None):
     """
     Adds overlap from file n+1 to file n. The last file has nothing added to it. This
     step is necessary for the Doppler-on version of the RT code.
@@ -240,12 +241,12 @@ def add_overlap(filename, v_max=11463.5):
         next_file = filename + str(i + 1) + ".dat"
 
         # go into file n to determine the delta lambda
-        opac = Opac(file, loader="exotransmit")
+        opac = Opac(file, loader="exotransmit", load_kwargs=load_kwargs)
         curr_lams = opac.wl
         del opac
 
         try:
-            opac = Opac(next_file, loader="exotransmit")
+            opac = Opac(next_file, loader="exotransmit", load_kwargs=load_kwargs)
             next_lams = opac.wl
             del opac
         except FileNotFoundError:

src/cortecs/opac/io.py

@@ -3,6 +3,7 @@
 
 author: @arjunsavel
 """
+
 import pickle
 
 import h5py
@@ -80,6 +81,7 @@ def load(self, filename):
 
         if self.wl_style == "wno":
             wl = 1e4 / wl
+            wl *= 1e-6  # now in meters
 
         if np.all(np.diff(wl)) <= 0:
             wl = wl[::-1]
@@ -353,29 +355,40 @@ def get_lams_and_opacities(self, file):
 
         # read through all lines in the opacity file
         # skip through the header!
-        for x in tqdm(f1[2:], desc="reading wavelengths"):
-            # check if blank line
+        for x in tqdm(f1[2:], desc="reading exotransmit wavelengths"):
+            # # check if blank line
             if not x:
                 continue
             # check if a wavelength line
-            commad = x.replace(" ", ",")
-            if len(np.array([eval(commad)]).flatten()) == 1:
-                wavelengths += [eval(x[:-1])]
+            # commad = x.replace(" ", ",")
+            # if len(np.array([eval(commad)]).flatten()) == 1:
+            #     wavelengths.append(eval(x[:-1]))
+            # else:
+            #     # the first entry in each opacity line is the pressure
+            #     opacity_string = x.split()[1:]
+            #     opacity_vals = [eval(opacity) for opacity in opacity_string]
+            #     opacities.append(opacity_vals)
+            if not x.strip():
+                continue
+                # check if a wavelength line
+            line_values = x.split()
+            if len(line_values) == 1:
+                wavelengths.append(float(line_values[0]))
             else:
                 # the first entry in each opacity line is the pressure
-                opacity_string = x.split()[1:]
-                opacity_vals = np.array([eval(opacity) for opacity in opacity_string])
-                opacities += [opacity_vals]
+                opacity_vals = [float(opacity) for opacity in line_values[1:]]
+                opacities.append(opacity_vals)
 
         f.close()
         # pdb.set_trace()
+        # oh this isn't reshaped haha
         del f1
         try:
             return np.array(wavelengths), np.array(opacities)
         except:
             pdb.set_trace()
 
-    def load(self, filename):
+    def load(self, filename, fullfile=True):
         """
         Loads file.
 
@@ -394,10 +407,19 @@ def load(self, filename):
             pressures
         """
 
-        wl, opacities = self.get_lams_and_opacities(filename)
+        wl, cross_section = self.get_lams_and_opacities(filename)
         T, P = self.get_t_p(filename)
+        P /= 1e5  # internally in bars
 
-        return wl, T, P, opacities
+        cross_section[np.less(cross_section, 1e-104)] = 1e-104
+
+        # and now make it log10
+        cross_section = np.log10(cross_section)
+
+        if fullfile:  # only reshape if it's not a "test" file.
+            cross_section = cross_section.reshape(len(T), len(P), len(wl))
+
+        return wl, T, P, cross_section
 
 
 # todo: put these in the opac class?
@@ -516,10 +538,22 @@ def get_n_species(CIA_file, verbose=False):
 
 
 class writer_base(object):
-    def __init__(self):
+    def __init__(
+        self,
+        wl_key="wno",
+        T_key="T",
+        P_key="P",
+        cross_section_key="xsec",
+        wl_style="wno",
+    ):
         """
         nothing to do here
         """
+        self.wl_key = wl_key
+        self.T_key = T_key
+        self.P_key = P_key
+        self.cross_section_key = cross_section_key
+        self.wl_style = wl_style
         pass
 
 
@@ -542,7 +576,7 @@ def write(self, opac, outfile, verbose=False):
         """
 
         new_string = []
-        print('optimized time')
+        print("optimized time")
         # don't want to write temp and wav
         columns = list(opac.cross_section.columns)
         if "temp" in columns:
@@ -588,30 +622,30 @@ def append_line_string(
         temp,
     ):
         # the first line gets different treatment!
-        #if i == 0:
+        # if i == 0:
         #    temp = np.min(
         #        self.interped_temps
         #    )  # add the LOWEST temperature in the temperature grid!
         #    new_string += ["{:.12e}".format(temp) + "\n"]
-        #if self.interped_temps[i] != temp:
+        # if self.interped_temps[i] != temp:
         #    temp = self.interped_temps[i]
         #    new_string += ["{:.12e}".format(temp) + "\n"]
-        #wavelength_string = "{:.12e}".format(self.interped_wavelengths[i])######
+        # wavelength_string = "{:.12e}".format(self.interped_wavelengths[i])######
 
-        #line_string = wavelength_string + self.buffer
+        # line_string = wavelength_string + self.buffer
 
-        #for species_key in self.species_dict_interped.keys():
-            # again, this works because python dicts are ordered in 3.6+
+        # for species_key in self.species_dict_interped.keys():
+        # again, this works because python dicts are ordered in 3.6+
         #    line_string += (
         #        "{:.12e}".format(
         #            list(self.species_dict_interped[species_key].values())[i]
         #        )
         #        + self.buffer
         #    )
 
-        #new_string += [line_string + "\n"]
+        # new_string += [line_string + "\n"]
 
-        #return new_string, temp
+        # return new_string, temp
 
         if i == 0:
             temp = np.min(self.interped_temps)
@@ -623,9 +657,46 @@ def append_line_string(
         wavelength_string = "{:.12e}".format(self.interped_wavelengths[i])
         line_string = wavelength_string + self.buffer
 
-        species_values = [species_dict_value[i] for species_dict_value in self.species_dict_interped.values()]
-        line_string += self.buffer.join("{:.12e}".format(value) for value in species_values)
+        species_values = [
+            species_dict_value[i]
+            for species_dict_value in self.species_dict_interped.values()
+        ]
+        line_string += self.buffer.join(
+            "{:.12e}".format(value) for value in species_values
+        )
 
         new_string.append(line_string + "\n")
         return new_string, temp
         # todo: maybe the loader objects should also take an opac object. for parallel structure : )
+
+
+class writer_chimera(writer_base):
+    """
+    does writing for the exotransmit object, takng in an opac object.
+    """
+
+    def write(self, opac, outfile, verbose=False):
+        """
+        loads in opacity data that's built for exo-transmit. To be passed on to Opac object.
+        """
+
+        # below is the write
+        wl = 1e6 * opac.wl  # now in microns
+        wno = 1e4 / wl  # now in cm^-1
+        wno = wno[::-1]
+        cross_section = opac.cross_section[:, ::-1]
+        cross_section = np.moveaxis(cross_section, 0, -1)
+        # want temperature index 0, pressure to 1, wavelength to 2 for standard usage.
+
+        hf = h5py.File(outfile, "w")
+        hf.create_dataset(self.wl_key, data=wno)
+        hf.create_dataset(self.T_key, data=opac.T)
+        hf.create_dataset(self.P_key, data=opac.P)
+        hf.create_dataset(self.cross_section_key, data=cross_section)
+
+        # ah man darn the xsec shape is wrong.
+
+        hf.close()
+
+        # todo: once this is done, check that everything is in the correct units.
+        return

src/cortecs/tests/test_chunking.py

@@ -1,6 +1,7 @@
 """
 just do a few tests for chunking...
 """
+
 import unittest
 import os
 import numpy as np
@@ -9,6 +10,8 @@
 
 
 class TestIntegration(unittest.TestCase):
+    load_kwargs = {"fullfile": False}
+
     opacity_file = os.path.abspath(".") + "/src/cortecs/tests/opacCH4_narrow_wl.dat"
     first_file = os.path.abspath(".") + "/src/cortecs/tests/" + "opacCH4_narrow_wl0.dat"
     second_file = (
@@ -41,14 +44,14 @@ def test_wls_of_each_created_file(self):
         chunk_wavelengths(self.opacity_file, wav_per_chunk=2)
 
         # now get the wavelengths of each file
-        opac_obj_ref = Opac(self.opacity_file, loader="exotransmit")
+        opac_obj_ref = Opac(
+            self.opacity_file, loader="exotransmit", load_kwargs=self.load_kwargs
+        )
         opac_obj0 = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1 = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         np.testing.assert_array_equal(
             opac_obj_ref.wl, np.concatenate((opac_obj0.wl, opac_obj1.wl))
@@ -65,14 +68,14 @@ def test_vals_of_each_created_file(self):
         chunk_wavelengths(self.opacity_file, wav_per_chunk=2)
 
         # now get the wavelengths of each file
-        opac_obj_ref = Opac(self.opacity_file, loader="exotransmit")
+        opac_obj_ref = Opac(
+            self.opacity_file, loader="exotransmit", load_kwargs=self.load_kwargs
+        )
         opac_obj0 = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1 = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         # pdb.set_trace()
         np.testing.assert_array_equal(
@@ -92,25 +95,23 @@ def test_add_overlap_wl_increase_or_same(self):
 
         chunk_wavelengths(self.opacity_file, wav_per_chunk=2)
         opac_obj0_orig = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1_orig = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
 
-        add_overlap(self.file_base, v_max=0.0)
+        add_overlap(self.file_base, v_max=0.0, load_kwargs=self.load_kwargs)
 
         # now get the wavelengths of each file
-        opac_obj_ref = Opac(self.opacity_file, loader="exotransmit")
+        opac_obj_ref = Opac(
+            self.opacity_file, loader="exotransmit", load_kwargs=self.load_kwargs
+        )
         opac_obj0 = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1 = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         # pdb.set_trace()
         self.assertTrue(
@@ -128,30 +129,28 @@ def add_overlap_with_single_overlap_point(self):
 
         chunk_wavelengths(self.opacity_file, wav_per_chunk=2)
         opac_obj0_orig = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1_orig = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
 
         # calculate the vmax so that one point is changed
         # velocity is C * delta lambda over lambda
         c = 3e8
         max_curr_lam = np.max(opac_obj0_orig.wl)
         v_max = c * (opac_obj1_orig.wl[0] - max_curr_lam) / max_curr_lam
-        add_overlap(self.file_base, v_max=v_max)
+        add_overlap(self.file_base, v_max=v_max, load_kwargs=self.load_kwargs)
 
         # now get the wavelengths of each file
-        opac_obj_ref = Opac(self.opacity_file, loader="exotransmit")
+        opac_obj_ref = Opac(
+            self.opacity_file, loader="exotransmit", load_kwargs=self.load_kwargs
+        )
         opac_obj0 = Opac(
-            self.first_file,
-            loader="exotransmit",
+            self.first_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         opac_obj1 = Opac(
-            self.second_file,
-            loader="exotransmit",
+            self.second_file, loader="exotransmit", load_kwargs=self.load_kwargs
         )
         self.assertTrue(
             len(opac_obj1.wl.min) == len(opac_obj0.wl.max())