Merge pull request #18 from arjunsavel/chunking_tests

Chunking tests

src/cortecs/opac/chunking.py

@@ -24,6 +24,7 @@
 """
 
 import os
+from glob import glob
 
 import numpy as np
 from tqdm import tqdm
@@ -244,6 +245,7 @@ def add_previous(num_to_add, file, previous_file):
         f2.write(x)
         f2.close()
         if ticker == num_to_add:
+            pdb.set_trace()
             return
 
 
@@ -270,10 +272,14 @@ def add_overlap(filename, v_max=11463.5):
     Side effects:
         Modifies every 'filename*.dat' file.
     """
-    for i in tqdm(
-        range(len(os.listdir()[:-1])), position=0, leave=True
+    print("for sure adding overlap")
+    files = glob(filename + "*.dat")
+    for i, file in tqdm(
+        enumerate(files), total=len(files), position=0, leave=True, desc="eeeeee"
     ):  # don't include the last file
-        file = filename + str(i) + ".dat"
+        if file == filename + ".dat":
+            continue  # only add over lap to the chunked file
+        print("for sure adding overlapppp")
 
         next_file = filename + str(i + 1) + ".dat"
 
@@ -295,8 +301,11 @@ def add_overlap(filename, v_max=11463.5):
         delta_lam = 2 * max_curr_lam * v_max / c  # delta_lambda/lambda = v/c
 
         # add another 20 indices to be safe!
-        max_lam_to_add_ind = (
-            np.argmin(np.abs(next_lams - (max_curr_lam + delta_lam))) + 20
-        )
+        # max_lam_to_add_ind = (
+        #     np.argmin(np.abs(next_lams - (max_curr_lam + delta_lam))) + 20
+        # )
+        max_lam_to_add_ind = np.argmin(np.abs(next_lams - (max_curr_lam + delta_lam)))
+        print("max lam to add")
+        print(max_lam_to_add_ind)
 
         add_lams(max_lam_to_add_ind, file, next_file)

src/cortecs/opac/io.py

@@ -370,8 +370,10 @@ def get_lams_and_opacities(self, file):
         f.close()
         # pdb.set_trace()
         del f1
-
-        return np.array(wavelengths), np.array(opacities)
+        try:
+            return np.array(wavelengths), np.array(opacities)
+        except:
+            pdb.set_trace()
 
     def load(self, filename):
         """

src/cortecs/tests/test_chunking.py

@@ -4,7 +4,7 @@
 import unittest
 import os
 import numpy as np
-from cortecs.opac.chunking import chunk_wavelengths
+from cortecs.opac.chunking import chunk_wavelengths, add_overlap
 from cortecs.opac.opac import Opac
 
 
@@ -14,6 +14,7 @@ class TestIntegration(unittest.TestCase):
     second_file = (
         os.path.abspath(".") + "/src/cortecs/tests/" + "opacCH4_narrow_wl1.dat"
     )
+    file_base = os.path.abspath(".") + "/src/cortecs/tests/" + "opacCH4_narrow_wl"
 
     def test_chunking_two_files(self):
         """
@@ -80,3 +81,41 @@ def test_vals_of_each_created_file(self):
         )
 
         # self.assertEqual(opac_obj_ref.wl), len(opac_obj0.wl) + len(opac_obj1.wl))
+
+    def test_add_overlap_wl_increase_or_same(self):
+        """
+        Test that the overlap is added to the end of the first file.
+        """
+        # clean up the files already made!!
+        # os.remove(self.first_file)
+        # os.remove(self.second_file)
+
+        chunk_wavelengths(self.opacity_file, wav_per_chunk=2)
+        opac_obj0_orig = Opac(
+            self.first_file,
+            loader="exotransmit",
+        )
+        opac_obj1_orig = Opac(
+            self.second_file,
+            loader="exotransmit",
+        )
+
+        add_overlap(self.file_base, v_max=0.0)
+
+        # now get the wavelengths of each file
+        opac_obj_ref = Opac(self.opacity_file, loader="exotransmit")
+        opac_obj0 = Opac(
+            self.first_file,
+            loader="exotransmit",
+        )
+        opac_obj1 = Opac(
+            self.second_file,
+            loader="exotransmit",
+        )
+        # pdb.set_trace()
+        self.assertTrue(
+            len(opac_obj0.wl) >= len(opac_obj0_orig.wl)
+            and len(opac_obj1.wl) >= len(opac_obj1_orig.wl)
+        )
+
+        # self.assertEqual(opac_obj_ref.wl), len(opac_obj0.wl) + len(opac_obj1.wl))