Merge pull request #23 from arjunsavel/diff_fit_axis

pass the fitting args to eval!

src/cortecs/eval/eval.py

@@ -78,6 +78,7 @@ def eval(self, temperature, pressure, wavelength):
             self.wl,
             self.fitter_results,
             **self.eval_kwargs,
+            **self.fit_kwargs,
         )
         # we need to recast the units if PLATON was used.
         if isinstance(self.load_obj, loader_platon):

src/cortecs/eval/eval_neural_net.py

@@ -69,6 +69,7 @@ def eval_neural_net(
     n_layers=None,
     weights=None,
     biases=None,
+    **kwargs
 ):
     """
     evaluates the neural network at a given temperature and pressure.

src/cortecs/eval/eval_pca.py

@@ -7,7 +7,7 @@
 
 
 @jax.jit
-def eval_pca_ind_wav(temperature_ind, pressure_ind, vectors, pca_coeffs):
+def eval_pca_ind_wav(first_ind, second_ind, vectors, pca_coeffs):
     """
     Evaluates the PCA fit at a given temperature and pressure.
 
@@ -17,10 +17,10 @@ def eval_pca_ind_wav(temperature_ind, pressure_ind, vectors, pca_coeffs):
 
     Inputs
     ------
-    temperature_ind: int
-        The index of the temperature to evaluate at.
-    pressure_ind: int
-        The index of the pressure to evaluate at.
+    first_ind: int
+        The index of the first axis quantity (default temperature) to evaluate at.
+    second_ind: int
+        The index of the second axis quantity (default pressure) to evaluate at.
     vectors: array
         The PCA vectors.
     pca_coeffs: array
@@ -34,14 +34,20 @@ def eval_pca_ind_wav(temperature_ind, pressure_ind, vectors, pca_coeffs):
     xsec_val = 0.0
     n_components = vectors.shape[1]
     for component in range(n_components):
-        xsec_val += (
-            vectors[temperature_ind, component] * pca_coeffs[component, pressure_ind]
-        )
+        xsec_val += vectors[first_ind, component] * pca_coeffs[component, second_ind]
     return xsec_val
 
 
 def eval_pca(
-    temperature, pressure, wavelength, T, P, wl, fitter_results, fit_axis="pressure"
+    temperature,
+    pressure,
+    wavelength,
+    T,
+    P,
+    wl,
+    fitter_results,
+    fit_axis="pressure",
+    **kwargs
 ):
     """
     Evaluates the PCA fit at a given temperature, pressure, and wavelength.
@@ -67,22 +73,25 @@ def eval_pca(
     pca_coeffs = pca_coeffs_all_wl[wavelength_ind, :, :]
 
     # todo: figure out how to order the pressure and temperature inds!
+    # pdb.set_trace()
     if fit_axis == "pressure":
-        first_arg = temperature_ind
-        second_arg = pressure_ind
-    elif fit_axis == "temperature":
         first_arg = pressure_ind
         second_arg = temperature_ind
+    elif fit_axis == "temperature":
+        first_arg = temperature_ind
+        second_arg = pressure_ind
+
     elif fit_axis == "best":
         T_length = len(T)
         P_length = len(P)
 
         # todo: what if equal?
         if T_length > P_length:
-            first_arg = pressure_ind
-            second_arg = temperature_ind
+            first_arg = temperature_ind
+            second_arg = pressure_ind
         else:
             first_arg = pressure_ind
             second_arg = temperature_ind
-
+    # print("first_arg, second_arg", first_arg, second_arg)
+    # print("shapes:", pca_vectors.shape, pca_coeffs.shape)
     return eval_pca_ind_wav(first_arg, second_arg, pca_vectors, pca_coeffs)

src/cortecs/eval/eval_polynomial.py

@@ -11,6 +11,7 @@ def eval_polynomial(
     pressures,
     wavelengths,
     fitter_results,
+    **kwargs,
 ):
     """
     evaluates the polynomial at a given temperature and pressure.

src/cortecs/fit/fit_pca.py

@@ -119,14 +119,14 @@ def fit_pca(cross_section, P, T, xMat, fit_axis="pressure", **kwargs):
     -------
         :beta: (nc x pixels) PCA coefficients
     """
-    print("shapes for everything:", cross_section.shape, P.shape, T.shape, xMat.shape)
-    cross_section = move_cross_section_axis(cross_section, fit_axis)
+    # print("shapes for everything:", cross_section.shape, P.shape, T.shape, xMat.shape)
 
+    cross_section = move_cross_section_axis(cross_section, fit_axis)
     beta = fit_mlr(cross_section, xMat)
     return beta
 
 
-def move_cross_section_axis(cross_section, fit_axis):
+def move_cross_section_axis(cross_section, fit_axis, dim=2):
     """
     todo: add docstring
     :param cross_section:
@@ -140,7 +140,10 @@ def move_cross_section_axis(cross_section, fit_axis):
     # the current shape prefers pressure. let's auto-check, though
     if fit_axis == "best":
         # actually want SECOND longest.
-        longest_axis = np.argmax(cross_section[:, :, 0].shape)
+        if dim == 3:
+            longest_axis = np.argmax(cross_section[:, :, 0].shape)
+        else:
+            longest_axis = np.argmax(cross_section.shape)
 
         # now move that longest axis to 0.
         cross_section = np.moveaxis(cross_section, longest_axis, 1)
@@ -180,8 +183,7 @@ def prep_pca(
     -------
         :xMat: (n_exp x nc) PCA components
     """
-    cross_section = move_cross_section_axis(cross_section, fit_axis)
-
+    cross_section = move_cross_section_axis(cross_section, fit_axis, dim=3)
     single_pres_single_temp = cross_section[:, :, wav_ind]
     if (
         np.all(single_pres_single_temp == single_pres_single_temp[0, 0])

src/cortecs/tests/test_integration.py

@@ -208,7 +208,7 @@ def test_optimize(self):
             }
         )
 
-    def test_pca_different_axis(self):
+    def test_pca_temperature_axis(self):
         """
         test that the PCA fit works well enough even when fit along a different axis.
         """
@@ -231,5 +231,29 @@ def test_pca_different_axis(self):
         )
 
         # check that the median absolute deviation is less than 10%
-        median_err = np.median(np.abs(percent_diffs))
-        self.assertTrue(median_err < 10)
+        median_err = np.median(np.abs(abs_diffs))
+        self.assertTrue(median_err < 1)  # a reasonable fit
+
+    def test_pca_best_axis(self):
+        """
+        test that the PCA fit works well enough even when fit along a different axis.
+        """
+        load_kwargs = {
+            "T_filename": self.T_filename,
+            "P_filename": self.P_filename,
+            "wl_filename": self.wl_filename,
+        }
+        opac_obj = Opac(
+            self.cross_sec_filename, loader="platon", load_kwargs=load_kwargs
+        )
+        # fit axis is temperature now
+        fitter = Fitter(opac_obj, method="pca", wav_ind=-2, nc=3, fit_axis="best")
+        fitter.fit()
+        # run the metrics to see what the median absolute deviation is
+        vals, orig_vals, abs_diffs, percent_diffs = calc_metrics(
+            fitter, tp_undersample_factor=1, wl_under_sample_factor=8, plot=False
+        )
+
+        # check that the median absolute deviation is less than 10%
+        median_err = np.median(np.abs(abs_diffs))
+        self.assertTrue(median_err < 1)  # a reasonable fit