Merge pull request #14 from arjunsavel/diff_fit_axis

fit diff axis!

src/cortecs/eval/eval.py

@@ -28,6 +28,7 @@ def __init__(self, opac, fitter, **eval_kwargs):
         self.wl = opac.wl
         self.P = opac.P
         self.T = opac.T
+        self.fit_kwargs = fitter.fitter_kwargs
 
         if not hasattr(fitter, "fitter_results"):
             raise ValueError("Fitter must be run before being passed to an Evaluator.")

src/cortecs/eval/eval_neural_net.py

@@ -61,7 +61,14 @@ def feed_forward(x, n_layers, weights, biases):
 
 # @jax.jit
 def eval_neural_net(
-    T, P, temperatures, pressures, wavelengths, n_layers=None, weights=None, biases=None
+    T,
+    P,
+    temperatures=None,
+    pressures=None,
+    wavelengths=None,
+    n_layers=None,
+    weights=None,
+    biases=None,
 ):
     """
     evaluates the neural network at a given temperature and pressure.

src/cortecs/eval/eval_pca.py

@@ -40,7 +40,9 @@ def eval_pca_ind_wav(temperature_ind, pressure_ind, vectors, pca_coeffs):
     return xsec_val
 
 
-def eval_pca(temperature, pressure, wavelength, T, P, wl, fitter_results):
+def eval_pca(
+    temperature, pressure, wavelength, T, P, wl, fitter_results, fit_axis="pressure"
+):
     """
     Evaluates the PCA fit at a given temperature, pressure, and wavelength.
 
@@ -64,4 +66,23 @@ def eval_pca(temperature, pressure, wavelength, T, P, wl, fitter_results):
     pca_vectors, pca_coeffs_all_wl = fitter_results
     pca_coeffs = pca_coeffs_all_wl[wavelength_ind, :, :]
 
-    return eval_pca_ind_wav(temperature_ind, pressure_ind, pca_vectors, pca_coeffs)
+    # todo: figure out how to order the pressure and temperature inds!
+    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 == "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
+        else:
+            first_arg = pressure_ind
+            second_arg = temperature_ind
+
+    return eval_pca_ind_wav(first_arg, second_arg, pca_vectors, pca_coeffs)

src/cortecs/fit/fit.py

@@ -54,7 +54,8 @@ def __init__(self, opac, method="pca", **fitter_kwargs):
             'polynomial', 'pca', and 'neural_net'. The more complex the model, the larger the model size (i.e., potentially
             the lower the compression factor), and the more likely it is to fit well.
         fitter_kwargs: dict
-            kwargs that are passed to the fitter.
+            kwargs that are passed to the fitter. one kwarg, for instance, is the fit_axis: for PCA,
+            this determines what axis is fit against.
         """
         self.opac = opac
         self.fitter_kwargs = fitter_kwargs
@@ -70,6 +71,7 @@ def __init__(self, opac, method="pca", **fitter_kwargs):
         self.P = self.opac.P
         self.T = self.opac.T
 
+        # todo: figure out how to change the fitting...based on the fit axis?
         return
 
     def fit(self, parallel=False):

src/cortecs/fit/fit_pca.py

@@ -104,7 +104,7 @@ def do_pca(cube, nc=3):
     return xMat, standardized_cube, s, vh, u
 
 
-def fit_pca(cross_section, P, T, xMat, **kwargs):
+def fit_pca(cross_section, P, T, xMat, fit_axis="pressure", **kwargs):
     """
     Fits the PCA to the opacity data.
 
@@ -119,12 +119,41 @@ def fit_pca(cross_section, P, T, xMat, **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)
 
     beta = fit_mlr(cross_section, xMat)
     return beta
 
 
-def prep_pca(cross_section, wav_ind=-1, nc=2, force_fit_constant=False):
+def move_cross_section_axis(cross_section, fit_axis):
+    """
+    todo: add docstring
+    :param cross_section:
+    :param fit_axis:
+    :return:
+    """
+    fit_axis_options = ["best", "temperature", "pressure"]
+    if fit_axis not in fit_axis_options:
+        raise ValueError(f"fit_axis param must be one of: {fit_axis_options}")
+
+    # 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)
+
+        # now move that longest axis to 0.
+        cross_section = np.moveaxis(cross_section, longest_axis, 1)
+
+    elif fit_axis == "temperature":
+        cross_section = np.moveaxis(cross_section, 0, 1)
+
+    return cross_section
+
+
+def prep_pca(
+    cross_section, wav_ind=-1, nc=2, force_fit_constant=False, fit_axis="pressure"
+):
     """
     Prepares the opacity data for PCA. That is, it calculates the PCA components to be fit along the entire
     dataset by fitting the PCA to a single wavelength.
@@ -144,11 +173,15 @@ def prep_pca(cross_section, wav_ind=-1, nc=2, force_fit_constant=False):
         :force_fit_constant: (bool) if True, will allow the PCA to fit an opacity function without temperature and pressure
         dependence. This usually isn't recommended, if these PCA vectors are to e used to fit other wavelengths that
         *do* have temperature and pressure dependence.
+        :fit_axis: (str) the axis to fit against. determines the shape of the final vectors and components.
+        if "best", chooses the largest axis. otherwise, can select "temperature" or "pressure".
 
     Returns
     -------
         :xMat: (n_exp x nc) PCA components
     """
+    cross_section = move_cross_section_axis(cross_section, fit_axis)
+
     single_pres_single_temp = cross_section[:, :, wav_ind]
     if (
         np.all(single_pres_single_temp == single_pres_single_temp[0, 0])

src/cortecs/opac/io.py

@@ -1,8 +1,6 @@
 """
 Reads opacity data from various sources.
 
-
-
 author: @arjunsavel
 """
 import pickle
@@ -18,19 +16,32 @@
 class loader_base(object):
     """
     loads in opacity data from various sources. To be passed on to Opac object.
-    """
 
-    wl_key = "wno"
-    T_key = "T"
-    P_key = "P"
-    cross_section_key = "xsec"
-    wl_style = "wno"
+    todo: tutorial on how to use this?
+    """
 
-    def __init__(self):
+    def __init__(
+        self,
+        wl_key="wno",
+        T_key="T",
+        P_key="P",
+        cross_section_key="xsec",
+        wl_style="wno",
+        temperature_axis=0,
+        pressure_axis=1,
+        wavelength_axis=2,
+    ):
         """
-        nothing to do here
+        sets the keys for the loader object.
         """
-        pass
+        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
+        self.temperature_axis = temperature_axis
+        self.pressure_axis = pressure_axis
+        self.wavelength_axis = wavelength_axis
 
     def load(self, filename):
         """
@@ -60,11 +71,29 @@ def load(self, filename):
         cross_section = np.array(hf[self.cross_section_key], dtype=np.float64)
         hf.close()
 
+        # want temperature index 0, pressure to 1, wavelength to 2 for standard usage.
+        cross_section = np.moveaxis(
+            cross_section,
+            [self.temperature_axis, self.pressure_axis, self.wavelength_axis],
+            [0, 1, 2],
+        )
+
         if self.wl_style == "wno":
             wl = 1e4 / wl
+
+        if np.all(np.diff(wl)) < 0:
+            wl = wl[::-1]
+            cross_section = cross_section[:, :, ::-1]
+            # reverse!
         return wl, T, P, cross_section
 
 
+class loader_chimera(loader_base):
+    """
+    loads in opacity data that are produced with the CHIMERA code.
+    """
+
+
 class loader_helios(loader_base):
     """
     loads in opacity data that are produced with the HELIOS ktable function.

src/cortecs/opac/opac.py

@@ -18,7 +18,7 @@ class Opac(object):
     """
 
     method_dict = {
-        "chimera": loader_base,
+        "chimera": loader_chimera,
         "helios": loader_helios,
         "platon": loader_platon,
         "exotransmit": loader_exotransmit,
@@ -28,7 +28,12 @@ class Opac(object):
     T = None
     P = None
 
-    def __init__(self, filename, loader="chimera", load_kwargs={}):
+    def __init__(
+        self,
+        filename,
+        loader="chimera",
+        load_kwargs={},
+    ):
         """
         wraps around the loaders.