Merge pull request #3 from lauracabayol/improve_code

notebooks/NMAD.py

@@ -61,6 +61,7 @@ eval_methods=True
 # ### LOAD DATA
 
 # %%
+
 #define here the directory containing the photometric catalogues
 parent_dir = Path('/data/astro/scratch/lcabayol/insight/data/Euclid_EXT_MER_PHZ_DC2_v1.5')
 modules_dir = Path('../data/models/')
@@ -68,7 +69,6 @@ filename_calib = 'euclid_cosmos_DC2_S1_v2.1_calib_clean.fits'
 filename_valid = 'euclid_cosmos_DC2_S1_v2.1_valid_matched.fits'
 
 # %%
-filename_valid='euclid_cosmos_DC2_S1_v2.1_valid_matched.fits'
 path_file = parent_dir / filename_valid  # Creating the path to the file
 hdu_list = fits.open(path_file)
 cat = Table(hdu_list[1].data).to_pandas()
@@ -158,3 +158,5 @@ plot_photoz(df_list,
             save=False,
             samp='L15'
            )
+
+# %%

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "gradio",
     "jupytext",
     "mkdocs",
+    "typing"
 ]
 
 classifiers = [

temps/archive.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass, field
 import numpy as np
 import pandas as pd
 from astropy.io import fits
@@ -5,199 +6,263 @@ from astropy.table import Table
 from scipy.spatial import KDTree
 from matplotlib import pyplot as plt
 from matplotlib import rcParams
-from pathlib import Path  
+from pathlib import Path
 from loguru import logger
+from typing import Optional, Tuple, Union, List
 
-
+# Set matplotlib configuration
 rcParams["mathtext.fontset"] = "stix"
 rcParams["font.family"] = "STIXGeneral"
 
+@dataclass
 class Archive:
-    def __init__(self, 
-                 path_calib, 
-                 path_valid=None,
-                 drop_stars=True, 
-                 clean_photometry=True, 
-                 convert_colors=True, 
-                 extinction_corr=True, 
-                 only_zspec=True, 
-                 columns_photometry = ['FLUX_G_2','FLUX_R_2','FLUX_I_2','FLUX_Z_2','FLUX_Y_2','FLUX_J_2','FLUX_H_2'],
-                 columns_ebv = ['EB_V_corr_FLUX_G','EB_V_corr_FLUX_R','EB_V_corr_FLUX_I','EB_V_corr_FLUX_Z','EB_V_corr_FLUX_Y','EB_V_corr_FLUX_J','EB_V_corr_FLUX_H'],
-                 photoz_name="photo_z_L15",
-                 specz_name="z_spec_S15",
-                 target_test='specz', 
-                 flags_kept=[3, 3.1, 3.4, 3.5, 4]):
+    path_calib: Path
+    path_valid: Optional[Path] = None
+    drop_stars: bool = True
+    clean_photometry: bool = True
+    convert_colors: bool = True
+    extinction_corr: bool = True
+    only_zspec: bool = True
+    columns_photometry: List[str] = field(default_factory=lambda: [
+        "FLUX_G_2",
+        "FLUX_R_2",
+        "FLUX_I_2",
+        "FLUX_Z_2",
+        "FLUX_Y_2",
+        "FLUX_J_2",
+        "FLUX_H_2",
+    ])
+    columns_ebv: List[str] = field(default_factory=lambda: [
+        "EB_V_corr_FLUX_G",
+        "EB_V_corr_FLUX_R",
+        "EB_V_corr_FLUX_I",
+        "EB_V_corr_FLUX_Z",
+        "EB_V_corr_FLUX_Y",
+        "EB_V_corr_FLUX_J",
+        "EB_V_corr_FLUX_H",
+    ])
+    photoz_name: str = "photo_z_L15"
+    specz_name: str = "z_spec_S15"
+    target_test: str = "specz"
+    flags_kept: List[float] = field(default_factory=lambda: [3, 3.1, 3.4, 3.5, 4])
 
+    def __post_init__(self):
         logger.info("Starting archive")
-        self.flags_kept = flags_kept
-        self.columns_photometry=columns_photometry
-        self.columns_ebv=columns_ebv    
 
-        
-        if path_calib.suffix == ".fits":
-            with fits.open(path_calib) as hdu_list:
-                cat = Table(hdu_list[1].data).to_pandas()
-            if path_valid != None:
-                with fits.open(path_valid) as hdu_list:
-                    cat_test = Table(hdu_list[1].data).to_pandas()    
-                
-        elif path_calib.suffix == ".csv":
-            cat = pd.read_csv(path_calib)
-            if path_valid != None:
-                cat_test = pd.read_csv(path_valid)
+        # Load data based on the file format
+        if self.path_calib.suffix == ".fits":
+            with fits.open(self.path_calib) as hdu_list:
+                self.cat = Table(hdu_list[1].data).to_pandas()
+            if self.path_valid is not None:
+                with fits.open(self.path_valid) as hdu_list:
+                    self.cat_test = Table(hdu_list[1].data).to_pandas()
+
+        elif self.path_calib.suffix == ".csv":
+            self.cat = pd.read_csv(self.path_calib)
+            if self.path_valid is not None:
+                self.cat_test = pd.read_csv(self.path_valid)
         else:
             raise ValueError("Unsupported file format. Please provide a .fits or .csv file.")
 
-        cat = cat.rename(columns ={f"{specz_name}":"specz",
-                         f"{photoz_name}":"photo_z"})
-        cat_test = cat_test.rename(columns ={f"{specz_name}":"specz",
-                         f"{photoz_name}":"photo_z"})
-        
-        cat = cat[(cat['specz'] > 0) | (cat['photo_z'] > 0)]
-        
-        # Store the catalogs for later use
-        self.cat = cat
-        self.cat_test = cat_test
-        
-                
-        if drop_stars==True:
-            logger.info("dropping stars...")
-            cat = cat[cat.mu_class_L07==1]
-            cat_test = cat_test[cat_test.mu_class_L07==1]
-
-        if clean_photometry==True:
-            logger.info("cleaning stars...")
-            cat = self._clean_photometry(cat)
-            cat_test = self._clean_photometry(cat_test)
-            
-        
-        cat = self._set_combiend_target(cat)
-        cat_test = self._set_combiend_target(cat_test)
-        
-        
-        
-        cat = cat[cat.MAG_VIS<25]
-        cat_test = cat_test[cat_test.MAG_VIS<25]
-        
-        cat = cat[cat.target_z<5]
-        cat_test = cat_test[cat_test.target_z<5]
-        
-                    
-                    
-        self._set_training_data(cat, 
-                                cat_test,
-                                only_zspec=only_zspec, 
-                                extinction_corr=extinction_corr, 
-                                convert_colors=convert_colors)
-        self._set_testing_data(cat_test, 
-                               target=target_test, 
-                               extinction_corr=extinction_corr, 
-                               convert_colors=convert_colors)
-        
-            
-    def _extract_fluxes(self,catalogue):
+        self.cat = self.cat.rename(
+            columns={f"{self.specz_name}": "specz", f"{self.photoz_name}": "photo_z"}
+        )
+        self.cat_test = self.cat_test.rename(
+            columns={f"{self.specz_name}": "specz", f"{self.photoz_name}": "photo_z"}
+        )
+
+        self.cat = self.cat[(self.cat["specz"] > 0) | (self.cat["photo_z"] > 0)]
+
+        # Apply operations based on the initialized parameters
+        if self.drop_stars:
+            logger.info("Dropping stars...")
+            self.cat = self.cat[self.cat.mu_class_L07 == 1]
+            self.cat_test = self.cat_test[self.cat_test.mu_class_L07 == 1]
+
+        if self.clean_photometry:
+            logger.info("Cleaning photometry...")
+            self.cat = self._clean_photometry(catalogue=self.cat)
+            self.cat_test = self._clean_photometry(catalogue=self.cat_test)
+
+        self.cat = self._set_combined_target(self.cat)
+        self.cat_test = self._set_combined_target(self.cat_test)
+
+        # Apply magnitude and redshift cuts
+        self.cat = self.cat[self.cat.MAG_VIS < 25]
+        self.cat_test = self.cat_test[self.cat_test.MAG_VIS < 25]
+
+        self.cat = self.cat[self.cat.target_z < 5]
+        self.cat_test = self.cat_test[self.cat_test.target_z < 5]
+
+        self._set_training_data(
+            self.cat,
+            self.cat_test,
+            only_zspec=self.only_zspec,
+            extinction_corr=self.extinction_corr,
+            convert_colors=self.convert_colors,
+        )
+        self._set_testing_data(
+            self.cat_test,
+            target=self.target_test,
+            extinction_corr=self.extinction_corr,
+            convert_colors=self.convert_colors,
+        )
+
+
+    def _extract_fluxes(self, catalogue: pd.DataFrame) -> np.ndarray:
+        """Extract fluxes from the given catalogue.
+
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
+
+        Returns:
+            np.ndarray: An array of fluxes.
+        """
         f = catalogue[self.columns_photometry].values
         return f
 
-    def _to_colors(self, flux):
-        """ Convert fluxes to colors"""
-        color = flux[:,:-1] / flux[:,1:]
+    @staticmethod
+    def _to_colors(flux: np.ndarray) -> np.ndarray:
+        """Convert fluxes to colors.
+
+        Args:
+            flux (np.ndarray): The input fluxes.
+
+        Returns:
+            np.ndarray: An array of colors.
+        """
+        color = flux[:, :-1] / flux[:, 1:]
         return color
-    
-    def _set_combiend_target(self, catalogue):
-        catalogue['target_z'] = catalogue.apply(lambda row: row['specz'] 
-                                                if row['specz'] > 0 
-                                                else row['photo_z'], axis=1)
-        
+
+    @staticmethod
+    def _set_combined_target(catalogue: pd.DataFrame) -> pd.DataFrame:
+        """Set the combined target redshift based on available data.
+
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
+
+        Returns:
+            pd.DataFrame: Updated catalogue with the combined target redshift.
+        """
+        catalogue["target_z"] = catalogue.apply(
+            lambda row: row["specz"] if row["specz"] > 0 else row["photo_z"], axis=1
+        )
         return catalogue
 
-    
-    def _clean_photometry(self,catalogue):
-        """ Drops all object with FLAG_PHOT!=0"""
-        catalogue = catalogue[catalogue['FLAG_PHOT']==0]
-        
+    @staticmethod
+    def _clean_photometry(catalogue: pd.DataFrame) -> pd.DataFrame:
+        """Drops all objects with FLAG_PHOT != 0.
+
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
+
+        Returns:
+            pd.DataFrame: Cleaned catalogue.
+        """
+        catalogue = catalogue[catalogue["FLAG_PHOT"] == 0]
         return catalogue
-    
-    def _correct_extinction(self,catalogue, f, return_ext_corr=False):
-        """Corrects for extinction"""
+
+    def _correct_extinction(
+        self, catalogue: pd.DataFrame, f: np.ndarray, return_ext_corr: bool = False
+    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+        """Corrects for extinction based on the provided catalogue.
+
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
+            f (np.ndarray): The flux values to correct.
+            return_ext_corr (bool): Whether to return the extinction correction values.
+
+        Returns:
+            Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: Corrected fluxes, and optionally the extinction corrections.
+        """
         ext_correction = catalogue[self.columns_ebv].values
         f = f * ext_correction
         if return_ext_corr:
             return f, ext_correction
         else:
             return f
-    
-    def _select_only_zspec(self,catalogue,cat_flag=None):
-        """Selects only galaxies with spectroscopic redshift"""
-        if cat_flag=='Calib':
-            catalogue = catalogue[catalogue.specz>0]
-        elif cat_flag=='Valid':
-            catalogue = catalogue[catalogue.specz>0]
-        return catalogue
-    
-    def _exclude_only_zspec(self,catalogue):
-        """Selects only galaxies without spectroscopic redshift"""
-        catalogue = catalogue[(catalogue.specz<0)&(catalogue.photo_z>0)&(catalogue.photo_z<4)]
-        return catalogue
-    
-    def _select_L15_sample(self,catalogue):
-        """Selects only galaxies withoutidx spectroscopic redshift"""
-        catalogue = catalogue[(catalogue.target_z>0)]
-        catalogue = catalogue[(catalogue.target_z<4)]
 
+    @staticmethod
+    def _select_only_zspec(
+        catalogue: pd.DataFrame, cat_flag: Optional[str] = None
+    ) -> pd.DataFrame:
+        """Selects only galaxies with spectroscopic redshift.
 
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
+            cat_flag (Optional[str]): Indicates the catalogue type ('Calib' or 'Valid').
+
+        Returns:
+            pd.DataFrame: Filtered catalogue.
+        """
+        if cat_flag == "Calib":
+            catalogue = catalogue[catalogue.specz > 0]
+        elif cat_flag == "Valid":
+            catalogue = catalogue[catalogue.specz > 0]
         return catalogue
-    
-    def _take_zspec_and_photoz(self,catalogue,cat_flag=None):
+
+    @staticmethod
+    def take_zspec_and_photoz(catalogue: pd.DataFrame, cat_flag: Optional[str] = None
+    ) -> pd.DataFrame:
         """Selects only galaxies with spectroscopic redshift"""
         if cat_flag=='Calib':
             catalogue = catalogue[catalogue.target_z>0]
         elif cat_flag=='Valid':
             catalogue = catalogue[catalogue.specz>0]
         return catalogue
+    
+    @staticmethod
+    def exclude_only_zspec(catalogue: pd.DataFrame) -> pd.DataFrame:
+        """Selects only galaxies without spectroscopic redshift.
 
-    def _clean_zspec_sample(self,catalogue ,flags_kept=[3,3.1,3.4,3.5,4]):
-        #[ 2.5,  3.5,  4. ,  1.5,  1.1, 13.5,  9. ,  3. ,  2.1,  9.5,  3.1,
-        #1. ,  9.1,  2. ,  9.3,  1.4,  3.4, 11.5,  2.4, 13. , 14. , 12.1,
-        #12.5, 13.1,  9.4, 11.1]
-        
-        catalogue = catalogue[catalogue.Q_f_S15.isin(flags_kept)]
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
 
+        Returns:
+            pd.DataFrame: Filtered catalogue.
+        """
+        catalogue = catalogue[
+            (catalogue.specz < 0) & (catalogue.photo_z > 0) & (catalogue.photo_z < 4)
+        ]
         return catalogue
-        
 
+    @staticmethod
+    def _clean_zspec_sample(catalogue ,flags_kept=[3,3.1,3.4,3.5,4]):
+        catalogue = catalogue[catalogue.Q_f_S15.isin(flags_kept)]
+        return catalogue
 
-    def _match_gold_sample(self,catalogue_valid, catalogue_gold, max_distance_arcsec=2):
-        max_distance_deg = max_distance_arcsec / 3600.0 
-
-        gold_sample_radec = np.c_[catalogue_gold.RIGHT_ASCENSION,catalogue_gold.DECLINATION]
-        valid_sample_radec = np.c_[catalogue_valid['RA'],catalogue_valid['DEC']]
-
-        kdtree = KDTree(gold_sample_radec)
-        distances, indices = kdtree.query(valid_sample_radec, k=1)
-
-        specz_match_gold = catalogue_gold.FINAL_SPEC_Z.values[indices]
-
-        zs = [specz_match_gold[i] if distance < max_distance_deg else -99 for i, distance in enumerate(distances)]
-
-        catalogue_valid['z_spec_gold'] = zs
+    @staticmethod
+    def _select_L15_sample(self, catalogue: pd.DataFrame) -> pd.DataFrame:
+        """Selects only galaxies within a specific redshift range.
 
-        return catalogue_valid
+        Args:
+            catalogue (pd.DataFrame): The input catalogue.
 
-    
-    def _set_training_data(self,catalogue, catalogue_da, only_zspec=True, extinction_corr=True, convert_colors=True):
+        Returns:
+            pd.DataFrame: Filtered catalogue.
+        """
+        catalogue = catalogue[(catalogue.target_z > 0) & (catalogue.target_z < 3)]
+        return catalogue
         
-        cat_da = self._exclude_only_zspec(catalogue_da)  
+    def _set_training_data(self,
+                           catalogue: pd.DataFrame, 
+                           catalogue_da: pd.DataFrame, 
+                           only_zspec: bool = True,
+                           extinction_corr: bool = True,
+                           convert_colors: bool = True
+                          )-> None:
+                        
+        cat_da = Archive.exclude_only_zspec(catalogue_da)  
         target_z_train_DA = cat_da['photo_z'].values
   
         
         if only_zspec:
             logger.info("Selecting only galaxies with spectroscopic redshift")
-            catalogue = self._select_only_zspec(catalogue, cat_flag='Calib')
-            catalogue = self._clean_zspec_sample(catalogue, flags_kept=self.flags_kept)
+            catalogue = Archive._select_only_zspec(catalogue, cat_flag='Calib')
+            catalogue = Archive._clean_zspec_sample(catalogue, flags_kept=self.flags_kept)
         else:
             logger.info("Selecting galaxies with spectroscopic redshift and high-precision photo-z")
-            catalogue = self._take_zspec_and_photoz(catalogue, cat_flag='Calib')
+            catalogue = Archive.take_zspec_and_photoz(catalogue, cat_flag='Calib')
             
             
         self.cat_train=catalogue
@@ -230,25 +295,32 @@ class Archive:
             self.target_z_train = catalogue['target_z'].values
             
         self.VIS_mag_train = catalogue['MAG_VIS'].values
-        
-    def _set_testing_data(self,catalogue, target='specz', extinction_corr=True, convert_colors=True):
+    
+
+    def _set_testing_data(
+        self,
+        cat_test: pd.DataFrame,
+        target: str = "specz",
+        extinction_corr: bool = True,
+        convert_colors: bool = True,
+    ) -> None:
         
         if target=='specz':
-            catalogue = self._select_only_zspec(catalogue, cat_flag='Valid')
-            catalogue = self._clean_zspec_sample(catalogue)
-            self.target_z_test = catalogue['specz'].values
+            cat_test = Archive._select_only_zspec(cat_test, cat_flag='Valid')
+            cat_test = Archive._clean_zspec_sample(cat_test)
+            self.target_z_test = cat_test['specz'].values
             
         elif target=='L15':
-            catalogue = self._select_L15_sample(catalogue)
-            self.target_z_test = catalogue['target_z'].values
+            cat_test = self._select_L15_sample(cat_test)
+            self.target_z_test = cat_test['target_z'].values
                     
                         
-        self.cat_test=catalogue
+        self.cat_test=cat_test
             
-        f = self._extract_fluxes(catalogue)
+        f = self._extract_fluxes(cat_test)
         
         if extinction_corr==True:
-            f = self._correct_extinction(catalogue,f)
+            f = self._correct_extinction(cat_test,f)
             
         if convert_colors==True:
             col = self._to_colors(f)
@@ -257,9 +329,9 @@ class Archive:
             self.phot_test = f
             
         
-        self.VIS_mag_test = catalogue['MAG_VIS'].values
-            
-        
+        self.VIS_mag_test = cat_test['MAG_VIS'].values
+    
+    
     def get_training_data(self):
         return self.phot_train, self.target_z_train, self.VIS_mag_train, self.phot_train_DA, self.target_z_train_DA
 
@@ -267,4 +339,4 @@ class Archive:
         return self.phot_test, self.target_z_test, self.VIS_mag_test
 
     def get_VIS_mag(self, catalogue):
-        return catalogue[['MAG_VIS']].values
+        return catalogue[['MAG_VIS']].values       

temps/plots.py

@@ -2,127 +2,185 @@ import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
 from temps.utils import nmad
-import numpy as np
-import matplotlib.pyplot as plt
 from scipy import stats
+from typing import List, Optional, Dict
+
+
+def plot_photoz(
+    df_list: List[pd.DataFrame],
+    nbins: int,
+    xvariable: str,
+    metric: str,
+    type_bin: str = "bin",
+    label_list: Optional[List[str]] = None,
+    samp: str = "zs",
+    save: bool = False,
+) -> None:
+    """
+    Plot photo-z metrics for multiple dataframes.
+
+    Parameters:
+    - df_list (List[pd.DataFrame]): List of dataframes containing data for plotting.
+    - nbins (int): Number of bins for the histogram.
+    - xvariable (str): Variable to plot on the x-axis.
+    - metric (str): Metric to plot (e.g., 'sig68', 'bias', 'nmad', 'outliers').
+    - type_bin (str, optional): Type of binning ('bin' or 'cum'). Default is 'bin'.
+    - label_list (Optional[List[str]], optional): List of labels for each dataframe. Default is None.
+    - samp (str, optional): Sample label for saving. Default is 'zs'.
+    - save (bool, optional): If True, save the plot to a file. Default is False.
+
+    Returns:
+    None
+    """
+    # Plot properties
+    plt.rcParams["font.family"] = "serif"
+    plt.rcParams["font.size"] = 12
 
-def plot_photoz(df_list, nbins, xvariable, metric, type_bin='bin',label_list=None, samp='zs', save=False):
-    #plot properties
-    plt.rcParams['font.family'] = 'serif'
-    plt.rcParams['font.size'] = 12
-    
-    if xvariable == 'VISmag':
-        xvariable_lab = 'VIS'
-    if xvariable == 'zs':
-        xvariable_lab = r'$z_{\rm s}$'
-
-    bin_edges = stats.mstats.mquantiles(df_list[0][xvariable].values, np.linspace(0.05, 1, nbins))
-    cmap = plt.get_cmap('Dark2')  # Choose a colormap for coloring lines
-    #plt.figure(figsize=(6, 5))
-    ls = ['--',':','-']
-    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8), gridspec_kw={'height_ratios': [3, 1]})
-    
-    ydata_dict = {}
+    # Set x-axis label based on variable
+    xvariable_lab = "VIS" if xvariable == "VISmag" else r"$z_{\rm s}$"
 
+    # Calculate bin edges
+    bin_edges = stats.mstats.mquantiles(
+        df_list[0][xvariable].values, np.linspace(0.05, 1, nbins)
+    )
+    cmap = plt.get_cmap("Dark2")
+
+    # Create subplots
+    fig, (ax1, ax2) = plt.subplots(
+        2, 1, figsize=(8, 8), gridspec_kw={"height_ratios": [3, 1]}
+    )
+    ydata_dict: Dict[str, List[float]] = {}
+
+    # Loop through dataframes and calculate metrics
     for i, df in enumerate(df_list):
         ydata, xlab = [], []
-        
+
         label = label_list[i]
-        
-        if label == 'zs':
-            label_lab = r'$z_{\rm s}$'
-        if label == 'zs+L15':
-            label_lab = r'$z_{\rm s}$+L15'    
-        if label == 'TEMPS':
-            label_lab = 'TEMPS'   
-
-        for k in range(len(bin_edges)-1):
-            edge_min = bin_edges[k]
-            edge_max = bin_edges[k+1]
+        label_lab = {
+            "zs": r"$z_{\rm s}$",
+            "zs+L15": r"$z_{\rm s}$+L15",
+            "TEMPS": "TEMPS",
+        }.get(label, label)
 
+        for k in range(len(bin_edges) - 1):
+            edge_min = bin_edges[k]
+            edge_max = bin_edges[k + 1]
             mean_mag = (edge_max + edge_min) / 2
 
-            if type_bin == 'bin':
-                df_plot = df[(df[xvariable] > edge_min) & (df[xvariable] < edge_max)]
-            elif type_bin == 'cum':
-                df_plot = df[(df[xvariable] < edge_max)]
-            else:
-                raise ValueError("Only type_bin=='bin' for binned and 'cum' for cumulative are supported")
+            df_plot = (
+                df[(df[xvariable] > edge_min) & (df[xvariable] < edge_max)]
+                if type_bin == "bin"
+                else df[(df[xvariable] < edge_max)]
+            )
 
             xlab.append(mean_mag)
-            if metric == 'sig68':
+            if metric == "sig68":
                 ydata.append(sigma68(df_plot.zwerr))
-            elif metric == 'bias':
+            elif metric == "bias":
                 ydata.append(np.mean(df_plot.zwerr))
-            elif metric == 'nmad':
+            elif metric == "nmad":
                 ydata.append(nmad(df_plot.zwerr))
-            elif metric == 'outliers':
-                ydata.append(len(df_plot[np.abs(df_plot.zwerr) > 0.15]) / len(df_plot)*100)
-                
-        ydata_dict[f'{i}'] = ydata
-        color = cmap(i)  # Get a different color for each dataframe
-        ax1.plot(xlab, ydata,marker='.', lw=1, label=label_lab, color=color, ls=ls[i])
-        
-
-
-    ax1.set_ylabel(f'{metric} $[\Delta z]$', fontsize=18)
-    #ax1.set_xlabel(f'{xvariable_lab}', fontsize=16)
+            elif metric == "outliers":
+                ydata.append(
+                    len(df_plot[np.abs(df_plot.zwerr) > 0.15]) / len(df_plot) * 100
+                )
+
+        ydata_dict[f"{i}"] = ydata
+        color = cmap(i)
+        ax1.plot(
+            xlab,
+            ydata,
+            marker=".",
+            lw=1,
+            label=label_lab,
+            color=color,
+            ls=["--", ":", "-"][i],
+        )
+
+    ax1.set_ylabel(f"{metric} $[\Delta z]$", fontsize=18)
     ax1.grid(False)
     ax1.legend()
-    
-    # Plot ratios between lines in the upper panel
-    
-    ax2.plot(xlab, np.array(ydata_dict['1'])/np.array(ydata_dict['0']), marker='.', color = cmap(1))
-    ax2.plot(xlab, np.array(ydata_dict['2'])/np.array(ydata_dict['0']), marker='.', color = cmap(2))
-    ax2.set_ylabel(r'Method $X$ / $z_{\rm z}$', fontsize=14)
-   
-
-    ax2.set_xlabel(f'{xvariable_lab}', fontsize=16)
+
+    # Plot ratios
+    ax2.plot(
+        xlab,
+        np.array(ydata_dict["1"]) / np.array(ydata_dict["0"]),
+        marker=".",
+        color=cmap(1),
+    )
+    ax2.plot(
+        xlab,
+        np.array(ydata_dict["2"]) / np.array(ydata_dict["0"]),
+        marker=".",
+        color=cmap(2),
+    )
+    ax2.set_ylabel(r"Method $X$ / $z_{\rm z}$", fontsize=14)
+    ax2.set_xlabel(f"{xvariable_lab}", fontsize=16)
     ax2.grid(True)
 
-    
-    if save==True:
-        plt.savefig(f'{metric}_{xvariable}_{samp}.pdf', dpi=300, bbox_inches='tight')
+    if save:
+        plt.savefig(f"{metric}_{xvariable}_{samp}.pdf", dpi=300, bbox_inches="tight")
     plt.show()
 
 
-def plot_pz(m, pz, specz):
-    # Create a figure and axis
-    fig, ax = plt.subplots(figsize=(8, 6))
+def plot_pz(m: int, pz: np.ndarray, specz: float) -> None:
+    """
+    Plot the Probability Density Function (PDF) for a given model and compare it with the spectroscopic redshift.
 
-    # Plot the PDF with a label
-    ax.plot(np.linspace(0, 4, 1000), pz[m], label='PDF', color='navy')
+    Parameters:
+    - m (int): Index for the model.
+    - pz (np.ndarray): Probability density function values.
+    - specz (float): Spectroscopic redshift value.
 
-    # Add a vertical line for 'specz_test'
-    ax.axvline(specz[m], color='black', linestyle='--', label=r'$z_{\rm s}$')
+    Returns:
+    None
+    """
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax.plot(np.linspace(0, 4, 1000), pz[m], label="PDF", color="navy")
+    ax.axvline(specz[m], color="black", linestyle="--", label=r"$z_{\rm s}$")
+    ax.set_xlabel(r"$z$", fontsize=18)
+    ax.set_ylabel("Probability Density", fontsize=16)
+    ax.legend(fontsize=18)
+    plt.show()
 
-    # Add labels and a legend
-    ax.set_xlabel(r'$z$', fontsize = 18)
-    ax.set_ylabel('Probability Density', fontsize=16)
-    ax.legend(fontsize = 18)
 
-    # Display the plot
-    plt.show()
+def plot_zdistribution(archive, plot_test: bool = False, bins: int = 50) -> None:
+    """
+    Plot the distribution of redshifts for training and optionally test samples.
 
-    
-def plot_zdistribution(archive, plot_test=False, bins=50):
-    _,_,specz = archive.get_training_data()
-    plt.hist(specz, bins = bins, hisstype='step', color='navy', label=r'Training sample')
+    Parameters:
+    - archive: Data archive object containing the training data.
+    - plot_test (bool, optional): If True, plot test sample distribution. Default is False.
+    - bins (int, optional): Number of histogram bins. Default is 50.
 
-    if plot_test:
-        _,_,specz_test = archive.get_training_data()
-        plt.hist(specz, bins = bins, hisstype='step', color='goldenrod', label=r'Test sample',ls='--')
+    Returns:
+    None
+    """
+    _, _, specz = archive.get_training_data()
+    plt.hist(specz, bins=bins, histtype="step", color="navy", label=r"Training sample")
 
+    if plot_test:
+        _, _, specz_test = archive.get_training_data()
+        plt.hist(
+            specz_test,
+            bins=bins,
+            histtype="step",
+            color="goldenrod",
+            label=r"Test sample",
+            linestyle="--",
+        )
 
     plt.xticks(fontsize=12)
     plt.yticks(fontsize=12)
+    plt.xlabel(r"Redshift", fontsize=14)
+    plt.ylabel("Counts", fontsize=14)
+    plt.legend()
+    plt.show()
 
-    plt.xlabel(r'Redshift', fontsize=14)
-    plt.ylabel('Counts', fontsize=14)
 
-    plt.show()
-        
-def plot_som_map(som_data, plot_arg = 'z', vmin=0, vmax=1):
+def plot_som_map(
+    som_data: np.ndarray, plot_arg: str = "z", vmin: float = 0, vmax: float = 1
+) -> None:
     """
     Plot the Self-Organizing Map (SOM) data.
 
@@ -135,182 +193,159 @@ def plot_som_map(som_data, plot_arg = 'z', vmin=0, vmax=1):
     Returns:
     None
     """
-    plt.imshow(som_data, vmin=vmin, vmax=vmax, cmap='viridis')  # Choose an appropriate colormap
-    plt.colorbar(label=f'{plot_arg}')  # Add a colorbar with a label
-    plt.xlabel(r'$x$ [pixel]', fontsize=14)  # Add an appropriate X-axis label
-    plt.ylabel(r'$y$ [pixel]', fontsize=14)  # Add an appropriate Y-axis label
+    plt.imshow(som_data, vmin=vmin, vmax=vmax, cmap="viridis")
+    plt.colorbar(label=f"{plot_arg}")
+    plt.xlabel(r"$x$ [pixel]", fontsize=14)
+    plt.ylabel(r"$y$ [pixel]", fontsize=14)
     plt.show()
 
-    
-def plot_PIT(pit_list_1, pit_list_2 = None, pit_list_3=None, sample='specz', labels=None, save =True):
-    #plot properties
-    plt.rcParams['font.family'] = 'serif'
-    plt.rcParams['font.size'] = 12
-    fig, ax = plt.subplots(figsize=(8, 6))
-    kwargs=dict(bins=30, histtype='step', density=True, range=(0,1))
-    cmap = plt.get_cmap('Dark2')
-    
 
-    # Create a histogram
-    hist, bins, _ = ax.hist(pit_list_1,  color=cmap(0), ls='--', **kwargs, label=labels[0])
-    if pit_list_2!= None:
-        hist, bins, _ = ax.hist(pit_list_2,  color=cmap(1), ls=':', **kwargs, label=labels[1])
-    if pit_list_3!= None:
-        hist, bins, _ = ax.hist(pit_list_3,  color=cmap(2), ls='-', **kwargs, label=labels[2])
+def plot_PIT(
+    pit_list_1: List[float],
+    pit_list_2: Optional[List[float]] = None,
+    pit_list_3: Optional[List[float]] = None,
+    sample: str = "specz",
+    labels: Optional[List[str]] = None,
+    save: bool = True,
+) -> None:
+    """
+    Plot Probability Integral Transform (PIT) values for given lists.
 
-    
-    # Add labels and a title
-    ax.set_xlabel('PIT values', fontsize = 18)
-    ax.set_ylabel('Frequency', fontsize = 18)
+    Parameters:
+    - pit_list_1 (List[float]): First list of PIT values.
+    - pit_list_2 (Optional[List[float]], optional): Second list of PIT values. Default is None.
+    - pit_list_3 (Optional[List[float]], optional): Third list of PIT values. Default is None.
+    - sample (str, optional): Sample label for saving. Default is 'specz'.
+    - labels (Optional[List[str]], optional): List of labels for each PIT list. Default is None.
+    - save (bool, optional): If True, save the plot to a file. Default is True.
 
-    # Add grid lines
-    ax.grid(True, linestyle='--', alpha=0.7)
+    Returns:
+    None
+    """
+    plt.rcParams["font.family"] = "serif"
+    plt.rcParams["font.size"] = 12
+    fig, ax = plt.subplots(figsize=(8, 6))
+    kwargs = dict(bins=30, histtype="step", density=True, range=(0, 1))
+    cmap = plt.get_cmap("Dark2")
 
-    # Customize the x-axis
-    ax.set_xlim(0, 1)
-    #ax.set_ylim(0,3)
-    
-    plt.legend(fontsize=12)
+    # Create a histogram
+    ax.hist(pit_list_1, color=cmap(0), linestyle="--", **kwargs, label=labels[0])
+    if pit_list_2 is not None:
+        ax.hist(pit_list_2, color=cmap(1), linestyle="--", **kwargs, label=labels[1])
+    if pit_list_3 is not None:
+        ax.hist(pit_list_3, color=cmap(2), linestyle="--", **kwargs, label=labels[2])
 
-    # Make ticks larger
-    ax.tick_params(axis='both', which='major', labelsize=14)
-    if save==True:
-        plt.savefig(f'{sample}_PIT.pdf', bbox_inches='tight')
+    ax.set_xlabel("PIT values", fontsize=14)
+    ax.set_ylabel("Normalized Counts", fontsize=14)
+    ax.legend(fontsize=12)
 
-    # Show the plot
+    if save:
+        plt.savefig(f"PIT_{sample}.pdf", dpi=300, bbox_inches="tight")
     plt.show()
-    
 
 
-    
-def plot_nz(df_list, 
-            zcuts = [0.1, 0.5, 1, 1.5, 2, 3, 4],
-            save=False):
-    # Plot properties
-    plt.rcParams['font.family'] = 'serif'
-    plt.rcParams['font.size'] = 16
+def plot_outlier_ratio(
+    outliers: np.ndarray, num_samp: int = 100, plot_mean: bool = True
+) -> None:
+    """
+    Plot the outlier ratio as a function of the number of samples.
 
-    cmap = plt.get_cmap('Dark2')  # Choose a colormap for coloring lines
-    
-    # Create subplots
-    fig, axs = plt.subplots(3, 1, figsize=(20, 8), sharex=True)
+    Parameters:
+    - outliers (np.ndarray): Outlier ratio data.
+    - num_samp (int, optional): Number of samples for plotting. Default is 100.
+    - plot_mean (bool, optional): If True, plot the mean of outliers. Default is True.
 
-    for i, df in enumerate(df_list):
-        dfplot = df_list[i].copy()  # Assuming df_list contains dataframes
-        ax = axs[i]  # Selecting the appropriate subplot
-        
-        for iz in range(len(zcuts)-1):
-            dfplot_z = dfplot[(dfplot['ztarget'] > zcuts[iz]) & (dfplot['ztarget'] < zcuts[iz + 1])]
-            color = cmap(iz)  # Get a different color for each redshift
-            
-            zt_mean = np.median(dfplot_z.ztarget.values)
-            zp_mean = np.median(dfplot_z.z.values)
-
-            
-            # Plot histogram on the selected subplot
-            ax.hist(dfplot_z.z, bins=50, color=color, histtype='step', linestyle='-', density=True, range=(0, 4))
-            ax.axvline(zt_mean, color=color, linestyle='-', lw=2)
-            ax.axvline(zp_mean, color=color, linestyle='--', lw=2)
-            
-        ax.set_ylabel(f'Frequency', fontsize=14)
-        ax.grid(False)
-        ax.set_xlim(0, 3.5)
-    
-    axs[-1].set_xlabel(f'$z$', fontsize=18)
-    
-    if save:
-        plt.savefig(f'nz_hist.pdf', dpi=300, bbox_inches='tight')
-    
+    Returns:
+    None
+    """
+    plt.figure(figsize=(10, 6))
+    plt.plot(np.arange(1, num_samp + 1), outliers[:num_samp], label="Outlier Ratio")
+
+    if plot_mean:
+        plt.axhline(
+            np.mean(outliers), color="red", linestyle="--", label="Mean Outlier Ratio"
+        )
+
+    plt.xlabel("Number of Samples", fontsize=14)
+    plt.ylabel("Outlier Ratio", fontsize=14)
+    plt.legend()
+    plt.grid()
     plt.show()
 
-    
-    
 
-def plot_crps(crps_list_1, crps_list_2 = None, crps_list_3=None, labels=None,  sample='specz', save =True):
+def plot_crps(
+    crps_list_1: List[float],
+    crps_list_2: Optional[List[float]] = None,
+    crps_list_3: Optional[List[float]] = None,
+    label: Optional[List[str]] = None,
+    sample: str = "specz",
+    save: bool = True,
+) -> None:
     # Create a figure and axis
-    #plot properties
-    plt.rcParams['font.family'] = 'serif'
-    plt.rcParams['font.size'] = 12
+    # plot properties
+    plt.rcParams["font.family"] = "serif"
+    plt.rcParams["font.size"] = 12
     fig, ax = plt.subplots(figsize=(8, 6))
-    cmap = plt.get_cmap('Dark2')
+    cmap = plt.get_cmap("Dark2")
 
-    kwargs=dict(bins=50, histtype='step', density=True, range=(0,1))
+    kwargs = dict(bins=50, histtype="step", density=True, range=(0, 1))
 
     # Create a histogram
-    hist, bins, _ = ax.hist(crps_list_1,  color=cmap(0), ls='--', **kwargs, label=labels[0])
+    hist, bins, _ = ax.hist(
+        crps_list_1, color=cmap(0), ls="--", **kwargs, label=labels[0]
+    )
     if crps_list_2 is not None:
-        hist, bins, _ = ax.hist(crps_list_2,  color=cmap(1), ls=':', **kwargs, label=labels[1])
+        hist, bins, _ = ax.hist(
+            crps_list_2, color=cmap(1), ls=":", **kwargs, label=labels[1]
+        )
     if crps_list_3 is not None:
-        hist, bins, _ = ax.hist(crps_list_3,  color=cmap(2), ls='-', **kwargs, label=labels[2])
+        hist, bins, _ = ax.hist(
+            crps_list_3, color=cmap(2), ls="-", **kwargs, label=labels[2]
+        )
 
     # Add labels and a title
-    ax.set_xlabel('CRPS Scores', fontsize = 18)
-    ax.set_ylabel('Frequency', fontsize = 18)
+    ax.set_xlabel("CRPS Scores", fontsize=18)
+    ax.set_ylabel("Frequency", fontsize=18)
 
     # Add grid lines
-    ax.grid(True, linestyle='--', alpha=0.7)
+    ax.grid(True, linestyle="--", alpha=0.7)
 
     # Customize the x-axis
     ax.set_xlim(0, 0.5)
 
     # Make ticks larger
-    ax.tick_params(axis='both', which='major', labelsize=14)
+    ax.tick_params(axis="both", which="major", labelsize=14)
 
     # Calculate the mean CRPS value
     mean_crps_1 = round(np.nanmean(crps_list_1), 2)
     mean_crps_2 = round(np.nanmean(crps_list_2), 2)
     mean_crps_3 = round(np.nanmean(crps_list_3), 2)
 
-
     # Add the mean CRPS value at the top-left corner
-    ax.annotate(f"Mean CRPS {labels[0]}: {mean_crps_1}", xy=(0.57, 0.9), xycoords='axes fraction', fontsize=14, color =cmap(0))
-    ax.annotate(f"Mean CRPS {labels[1]}: {mean_crps_2}", xy=(0.57, 0.85), xycoords='axes fraction', fontsize=14, color =cmap(1))
-    ax.annotate(f"Mean CRPS {labels[2]}: {mean_crps_3}", xy=(0.57, 0.8), xycoords='axes fraction', fontsize=14, color =cmap(2))
-
-    
-    if save==True:
-        plt.savefig(f'{sample}_CRPS.pdf', bbox_inches='tight')
+    ax.annotate(
+        f"Mean CRPS {labels[0]}: {mean_crps_1}",
+        xy=(0.57, 0.9),
+        xycoords="axes fraction",
+        fontsize=14,
+        color=cmap(0),
+    )
+    ax.annotate(
+        f"Mean CRPS {labels[1]}: {mean_crps_2}",
+        xy=(0.57, 0.85),
+        xycoords="axes fraction",
+        fontsize=14,
+        color=cmap(1),
+    )
+    ax.annotate(
+        f"Mean CRPS {labels[2]}: {mean_crps_3}",
+        xy=(0.57, 0.8),
+        xycoords="axes fraction",
+        fontsize=14,
+        color=cmap(2),
+    )
+
+    if save == True:
+        plt.savefig(f"{sample}_CRPS.pdf", bbox_inches="tight")
 
     # Show the plot
     plt.show()
-
-
-
-def plot_nz(df, bins=np.arange(0,5,0.2)):
-    kwargs=dict( bins=bins,alpha=0.5)
-    plt.hist(df.zs.values, color='grey', ls='-' ,**kwargs)
-    counts, _, =np.histogram(df.z.values, bins=bins)
-    
-    plt.plot((bins[:-1]+bins[1:])*0.5,counts, color ='purple')
-    
-    #plt.legend(fontsize=14)
-    plt.xlabel(r'Redshift', fontsize=14)
-    plt.ylabel(r'Counts', fontsize=14)
-    plt.yscale('log')
-    
-    plt.show()
-    
-    return
-
-
-def plot_scatter(df, sample='specz', save=True):
-    # Calculate the point density
-    xy = np.vstack([df.zs.values,df.z.values])
-    zd = gaussian_kde(xy)(xy)
-
-    fig, ax = plt.subplots()
-    plt.scatter(df.zs.values, df.z.values,c=zd, s=1)
-    plt.xlim(0,5)
-    plt.ylim(0,5)
-
-    plt.xlabel(r'$z_{\rm s}$', fontsize = 14)
-    plt.ylabel('$z$', fontsize = 14)
-
-    plt.xticks(fontsize = 12)
-    plt.yticks(fontsize = 12)
-
-    if save==True:
-        plt.savefig(f'{sample}_scatter.pdf', dpi = 300, bbox_inches='tight')
-
-    plt.show()  
-    

temps/temps.py

@@ -6,38 +6,63 @@ from torch.optim import lr_scheduler
 from loguru import logger
 import pandas as pd
 from scipy.stats import norm
-
-from tqdm import tqdm  # Import tqdm for progress bars
+from dataclasses import dataclass, field
+from tqdm import tqdm  
+from typing import Optional, Tuple, List, Union
 
 from temps.utils import maximum_mean_discrepancy
 
 
+@dataclass
 class TempsModule:
-    """Class for managing temperature-related models and training."""
+    """Attributes:
+    model_f (nn.Module): The feature extraction model.
+    model_z (nn.Module): The model for predicting z values.
+    batch_size (int): Size of each batch for training. Default is 100.
+    rejection_param (int): Parameter for rejection sampling. Default is 1.
+    da (bool): Flag for enabling domain adaptation. Default is True.
+    verbose (bool): Flag for verbose logging. Default is False.
+    device (torch.device): Device to run the model on (CPU or GPU).
+    ngaussians (int): Number of Gaussian components in the mixture model.
+    """
+
+    model_f: nn.Module
+    model_z: nn.Module
+    batch_size: int = 100
+    rejection_param: int = 1
+    da: bool = True
+    verbose: bool = False
+    device: torch.device = field(init=False)
+    ngaussians: int = field(init=False)
+
+    def __post_init__(self) -> None:
+        """Post-initialization for setting up additional attributes."""
+        self.device: torch.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        self.ngaussians: int = (
+            self.model_z.ngaussians
+        )  # Assuming ngaussians is an integer
 
-    def __init__(
+    def _get_dataloaders(
         self,
-        model_f,
-        model_z,
-        batch_size=100,
-        rejection_param=1,
-        da=True,
-        verbose=False,
-    ):
-        self.model_z = model_z
-        self.model_f = model_f
-        self.da = da
-        self.verbose = verbose
-        self.ngaussians = model_z.ngaussians
-
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.batch_size = batch_size
-        self.rejection_parameter = rejection_param
+        input_data: np.ndarray,
+        target_data: np.ndarray,
+        input_data_da: Optional[np.ndarray] = None,
+        val_fraction: float = 0.1,
+    ) -> Tuple[DataLoader, DataLoader]:
+        """Create training and validation dataloaders.
+
+        Args:
+            input_data (np.ndarray): The input features for training.
+            target_data (np.ndarray): The target outputs for training.
+            input_data_da (Optional[np.ndarray]): Input data for domain adaptation (if any).
+            val_fraction (float): Fraction of data to use for validation. Default is 0.1.
+
+        Returns:
+            Tuple[DataLoader, DataLoader]: Training and validation data loaders.
+        """
 
-    def _get_dataloaders(
-        self, input_data, target_data, input_data_da=None, val_fraction=0.1
-    ):
-        """Create training and validation dataloaders."""
         input_data = torch.Tensor(input_data)
         target_data = torch.Tensor(target_data)
         input_data_da = (
@@ -47,10 +72,17 @@ class TempsModule:
         )
 
         dataset = TensorDataset(input_data, input_data_da, target_data)
+
+        # Calculate sizes for training and validation sets
+        total_size = len(dataset)
+        val_size = int(total_size * val_fraction)
+        train_size = total_size - val_size
+
         train_dataset, val_dataset = torch.utils.data.random_split(
             dataset,
-            [int(len(dataset) * (1 - val_fraction)), int(len(dataset) * val_fraction)+1],
+            [train_size, val_size],
         )
+
         loader_train = DataLoader(
             train_dataset, batch_size=self.batch_size, shuffle=True
         )
@@ -58,8 +90,24 @@ class TempsModule:
 
         return loader_train, loader_val
 
-    def _loss_function(self, mean, std, logmix, true):
-        """Compute the loss function."""
+    def _loss_function(
+        self,
+        mean: torch.Tensor,
+        std: torch.Tensor,
+        logmix: torch.Tensor,
+        true: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute the loss function for the model.
+
+        Args:
+            mean (torch.Tensor): Mean values predicted by the model.
+            std (torch.Tensor): Standard deviation values predicted by the model.
+            logmix (torch.Tensor): Logarithm of the mixture coefficients.
+            true (torch.Tensor): True target values.
+
+        Returns:
+            torch.Tensor: The computed loss value.
+        """
         log_prob = (
             logmix - 0.5 * (mean - true[:, None]).pow(2) / std.pow(2) - torch.log(std)
         )
@@ -67,28 +115,55 @@ class TempsModule:
         loss = -log_prob.mean()
         return loss
 
-    def _loss_function_da(self, f1, f2):
-        """Compute the KL divergence loss for domain adaptation."""
+    def _loss_function_da(self, f1: torch.Tensor, f2: torch.Tensor) -> torch.Tensor:
+        """Compute the KL divergence loss for domain adaptation.
+
+        Args:
+            f1 (torch.Tensor): Features from the primary domain.
+            f2 (torch.Tensor): Features from the domain for adaptation.
+
+        Returns:
+            torch.Tensor: The KL divergence loss value.
+        """
+
         kl_loss = nn.KLDivLoss(reduction="batchmean", log_target=True)
         loss = kl_loss(f1, f2)
         return torch.log(loss)
 
-    def _to_numpy(self, x):
-        """Convert a tensor to a NumPy array."""
+    def _to_numpy(self, x: torch.Tensor) -> np.ndarray:
+        """Convert a tensor to a NumPy array.
+
+        Args:
+            x (torch.Tensor): The input tensor to convert.
+
+        Returns:
+            np.ndarray: The converted NumPy array.
+        """
         return x.detach().cpu().numpy()
 
     def train(
         self,
-        input_data,
-        input_data_da,
-        target_data,
-        nepochs=10,
-        step_size=100,
-        val_fraction=0.1,
-        lr=1e-3,
-        weight_decay=0,
-    ):
-        """Train the models using provided data."""
+        input_data: np.ndarray,
+        input_data_da: np.ndarray,
+        target_data: np.ndarray,
+        nepochs: int = 10,
+        step_size: int = 100,
+        val_fraction: float = 0.1,
+        lr: float = 1e-3,
+        weight_decay: float = 0,
+    ) -> None:
+        """Train the models using provided data.
+
+        Args:
+            input_data (np.ndarray): The input features for training.
+            input_data_da (np.ndarray): Input data for domain adaptation.
+            target_data (np.ndarray): The target outputs for training.
+            nepochs (int): Number of training epochs. Default is 10.
+            step_size (int): Step size for learning rate scheduling. Default is 100.
+            val_fraction (float): Fraction of data to use for validation. Default is 0.1.
+            lr (float): Learning rate for the optimizer. Default is 1e-3.
+            weight_decay (float): Weight decay for regularization. Default is 0.
+        """
         self.model_z.train()
         self.model_f.train()
 
@@ -157,8 +232,11 @@ class TempsModule:
                 f"Epoch {epoch + 1}: Training Loss: {np.mean(_loss_train):.4f}, Validation Loss: {np.mean(_loss_validation):.4f}"
             )
 
-    def _validate(self, loader_val, target_data):
+    def _validate(
+        self, loader_val: DataLoader, target_data: torch.Tensor
+    ) -> List[float]:
         """Validate the model on the validation dataset."""
+
         self.model_z.eval()
         self.model_f.eval()
         _loss_validation = []
@@ -180,15 +258,49 @@ class TempsModule:
 
         return _loss_validation
 
-    def get_features(self, input_data):
-        """Get features from the model."""
+    def get_features(self, input_data: torch.Tensor) -> np.ndarray:
+        """Extract features from the model for the given input data.
+
+        Args:
+            input_data (torch.Tensor): Input tensor containing the data for which features are to be extracted.
+
+        Returns:
+            np.ndarray: Numpy array of extracted features from the model.
+        """
+
         self.model_f.eval()
         input_data = input_data.to(self.device)
         features = self.model_f(input_data)
         return self._to_numpy(features)
 
-    def get_pz(self, input_data, return_pz=True, return_flag=True, return_odds=False):
-        """Get the predicted z values and their uncertainties."""
+    def get_pz(
+        self,
+        input_data: torch.Tensor,
+        return_pz: bool = True,
+        return_flag: bool = True,
+        return_odds: bool = False,
+        ) -> Union[
+                Tuple[np.ndarray, np.ndarray],                # Return z and zerr
+                Tuple[np.ndarray, np.ndarray],                # Return z, pz
+                Tuple[np.ndarray, np.ndarray, np.ndarray]     # Return z, pz, odds
+            ]:
+        """Get the predicted redshift (z) values and their uncertainties from the model.
+
+        This function predicts the photo-z for the input galaxies, computes the mean and standard
+        deviation for the predicted redshifts, and optionally calculates the probability density function (PDF).
+
+        Args:
+            input_data (torch.Tensor): Input tensor containing galaxy data for which to predict redshifts.
+            return_pz (bool, optional): Flag indicating whether to return the probability density function. Defaults to True.
+            return_flag (bool, optional): Flag indicating whether to return additional information. Defaults to True.
+            return_odds (bool, optional): Flag indicating whether to return the odds. Defaults to False.
+
+        Returns:
+            Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+                - If return_pz is True, returns the PDF and possibly additional metrics.
+                - If return_pz is False, returns a tuple containing the predicted redshifts and their uncertainties.
+        """
+
         logger.info("Predicting photo-z for the input galaxies...")
         self.model_z.eval().to(self.device)
         self.model_f.eval().to(self.device)
@@ -206,6 +318,7 @@ class TempsModule:
             + (mix_coeff * (mu - mu.mean(dim=1, keepdim=True)) ** 2).sum(dim=1)
         )
 
+        z = self._to_numpy(z)
         mu, mix_coeff, sig = map(self._to_numpy, (mu, mix_coeff, sig))
 
         if return_pz:
@@ -214,118 +327,61 @@ class TempsModule:
         else:
             return self._to_numpy(z), self._to_numpy(zerr)
 
-    def _calculate_pdf(self, z, mu, sig, mix_coeff, return_flag):
-        """Calculate the probability density function."""
+    def _calculate_pdf(
+        self,
+        z: np.ndarray,
+        mu: np.ndarray,
+        sig: np.ndarray,
+        mix_coeff: np.ndarray,
+        return_flag: bool,
+    ) -> Union[
+        Tuple[np.ndarray, np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]
+    ]:
+        """Calculate the probability density function (PDF) for the predicted redshifts.
+
+        Args:
+            z (np.ndarray): Predicted redshift values.
+            mu (np.ndarray): Mean values for the Gaussian components.
+            sig (np.ndarray): Standard deviations for the Gaussian components.
+            mix_coeff (np.ndarray): Mixture coefficients for the Gaussian components.
+            return_flag (bool): Flag indicating whether to calculate and return odds.
+
+        Returns:
+            Union[Tuple[np.ndarray, np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]:
+                - If return_flag is True, returns a tuple containing the redshift values, PDF, and the z-grid.
+                - If return_flag is False, returns a tuple containing the redshift values and PDF.
+        """
+
         zgrid = np.linspace(0, 5, 1000)
         pz = np.zeros((len(z), len(zgrid)))
 
         for ii in range(len(z)):
             for i in range(self.ngaussians):
-                pz[ii] += mix_coeff[ii, i] * norm.pdf(
-                    zgrid, mu[ii, i], sig[ii, i]
-                )
+                pz[ii] += mix_coeff[ii, i] * norm.pdf(zgrid, mu[ii, i], sig[ii, i])
 
         if return_flag:
             logger.info("Calculating and returning ODDS")
             pz /= pz.sum(axis=1, keepdims=True)
             return self._calculate_odds(z, pz, zgrid)
-        return self._to_numpy(z), pz
-
-    def _calculate_odds(self, z, pz, zgrid):
-        """Calculate odds based on the PDF."""
-        logger.info('Calculating ODDS values')
-        diff_matrix = np.abs(self._to_numpy(z)[:, None] - zgrid[None, :])
-        idx_peak = np.argmax(pz, axis=1)
-        zpeak = zgrid[idx_peak]
-        idx_upper = np.argmin(np.abs((zpeak + 0.05)[:, None] - zgrid[None, :]), axis=1)
-        idx_lower = np.argmin(np.abs((zpeak - 0.05)[:, None] - zgrid[None, :]), axis=1)
-
-        odds = []
-        for jj in range(len(pz)):
-            odds.append(pz[jj,idx_lower[jj]:(idx_upper[jj]+1)].sum())
-    
-        odds = np.array(odds)
-        return self._to_numpy(z), pz, odds
-
-    def calculate_pit(self, input_data, target_data):
-        logger.info('Calculating PIT values')
-        
-        pit_list = []
-
-        self.model_f = self.model_f.eval()
-        self.model_f = self.model_f.to(self.device)
-        self.model_z = self.model_z.eval()
-        self.model_z = self.model_z.to(self.device)
-
-        input_data = input_data.to(self.device)
-                
-
-        features = self.model_f(input_data)
-        mu, logsig, logmix_coeff = self.model_z(features)
-        
-        logsig = torch.clamp(logsig,-6,2)
-        sig = torch.exp(logsig)
-
-        mix_coeff = torch.exp(logmix_coeff)
-        
-        mu,  mix_coeff, sig = mu.detach().cpu().numpy(),  mix_coeff.detach().cpu().numpy(), sig.detach().cpu().numpy() 
-        
-        for ii in range(len(input_data)):
-            pit = (mix_coeff[ii] * norm.cdf(target_data[ii]*np.ones(mu[ii].shape),mu[ii], sig[ii])).sum()
-            pit_list.append(pit)
-        
-        
-        return pit_list
-    
-    def calculate_crps(self, input_data, target_data):
-        logger.info('Calculating CRPS values')
-
-        def measure_crps(cdf, t):
-            zgrid = np.linspace(0,4,1000)
-            Deltaz = zgrid[None,:] - t[:,None]
-            DeltaZ_heaviside = np.where(Deltaz < 0,0,1)
-            integral = (cdf-DeltaZ_heaviside)**2
-            crps_value = integral.sum(1) / 1000
-
-            return crps_value
-
-
-        crps_list = []
-
-        self.model_f = self.model_f.eval()
-        self.model_f = self.model_f.to(self.device)
-        self.model_z = self.model_z.eval()
-        self.model_z = self.model_z.to(self.device)
-
-        input_data = input_data.to(self.device)
-
-
-        features = self.model_f(input_data)
-        mu, logsig, logmix_coeff = self.model_z(features)
-        logsig = torch.clamp(logsig,-6,2)
-        sig = torch.exp(logsig)
-
-        mix_coeff = torch.exp(logmix_coeff)
-
-
-        mu,  mix_coeff, sig = mu.detach().cpu().numpy(),  mix_coeff.detach().cpu().numpy(), sig.detach().cpu().numpy() 
-
-        z = (mix_coeff * mu).sum(1)
-
-        x = np.linspace(0, 4, 1000)
-        pz = np.zeros(shape=(len(target_data), len(x)))
-        for ii in range(len(input_data)):
-            for i in range(6):
-                pz[ii] += mix_coeff[ii,i] * norm.pdf(x, mu[ii,i], sig[ii,i])
-
-        pz = pz / pz.sum(1)[:,None]
-
-
-        cdf_z = np.cumsum(pz,1)
-
-        crps_value = measure_crps(cdf_z, target_data)
-
-
-
-        return crps_value
-
+        return z, pz
+
+    def _calculate_odds(
+        self, z: np.ndarray, pz: np.ndarray, zgrid: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Calculate the odds for the estimated redshifts based on the cumulative distribution.
+
+        Args:
+            z (np.ndarray): Predicted redshift values.
+            pz (np.ndarray): Probability density function values.
+            zgrid (np.ndarray): Grid of redshift values for evaluation.
+
+        Returns:
+            Tuple[np.ndarray, np.ndarray, np.ndarray]: A tuple containing the predicted redshift values,
+            PDF values, and calculated odds.
+        """
+
+        cumulative = np.cumsum(pz, axis=1)
+        odds = np.array(
+            [np.max(np.abs(cumulative[i] - 0.68)) for i in range(cumulative.shape[0])]
+        )
+        return z, pz, odds

temps/temps_arch.py

@@ -1,10 +1,24 @@
 import torch
-from torch import nn, optim
+from torch import nn
 import torch.nn.functional as F
 
 
 class EncoderPhotometry(nn.Module):
-    def __init__(self, input_dim=6, dropout_prob=0):
+    """Encoder for photometric data.
+
+    This neural network encodes photometric features into a lower-dimensional representation.
+
+    Attributes:
+        features (nn.Sequential): A sequential container of layers used for encoding.
+    """
+
+    def __init__(self, input_dim: int = 6, dropout_prob: float = 0) -> None:
+        """Initializes the EncoderPhotometry module.
+
+        Args:
+            input_dim (int): Number of input features (default is 6).
+            dropout_prob (float): Probability of dropout (default is 0).
+        """
         super(EncoderPhotometry, self).__init__()
 
         self.features = nn.Sequential(
@@ -23,14 +37,39 @@ class EncoderPhotometry(nn.Module):
             nn.Linear(20, 10),
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the encoder.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, input_dim).
+
+        Returns:
+            torch.Tensor: Log softmax output of shape (batch_size, 10).
+        """
         f = self.features(x)
         f = F.log_softmax(f, dim=1)
         return f
 
 
 class MeasureZ(nn.Module):
-    def __init__(self, num_gauss=10, dropout_prob=0):
+    """Model to measure redshift parameters.
+
+    This model estimates the parameters of a mixture of Gaussians used for measuring redshift.
+
+    Attributes:
+        ngaussians (int): Number of Gaussian components in the mixture.
+        measure_mu (nn.Sequential): Sequential model to measure the mean (mu).
+        measure_coeffs (nn.Sequential): Sequential model to measure the mixing coefficients.
+        measure_sigma (nn.Sequential): Sequential model to measure the standard deviation (sigma).
+    """
+
+    def __init__(self, num_gauss: int = 10, dropout_prob: float = 0) -> None:
+        """Initializes the MeasureZ module.
+
+        Args:
+            num_gauss (int): Number of Gaussian components (default is 10).
+            dropout_prob (float): Probability of dropout (default is 0).
+        """
         super(MeasureZ, self).__init__()
 
         self.ngaussians = num_gauss
@@ -55,11 +94,25 @@ class MeasureZ(nn.Module):
             nn.Linear(20, num_gauss),
         )
 
-    def forward(self, f):
+    def forward(
+        self, f: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Forward pass to measure redshift parameters.
+
+        Args:
+            f (torch.Tensor): Input tensor of shape (batch_size, 10).
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
+                - mu (torch.Tensor): Mean parameters of shape (batch_size, num_gauss).
+                - sigma (torch.Tensor): Standard deviation parameters of shape (batch_size, num_gauss).
+                - logmix_coeff (torch.Tensor): Log mixing coefficients of shape (batch_size, num_gauss).
+        """
         mu = self.measure_mu(f)
         sigma = self.measure_sigma(f)
         logmix_coeff = self.measure_coeffs(f)
 
-        logmix_coeff = logmix_coeff - torch.logsumexp(logmix_coeff, 1)[:, None]
+        # Normalize logmix_coeff to get valid mixture coefficients
+        logmix_coeff = logmix_coeff - torch.logsumexp(logmix_coeff, dim=1, keepdim=True)
 
         return mu, sigma, logmix_coeff

temps/utils.py

@@ -4,21 +4,31 @@ import matplotlib.pyplot as plt
 from scipy import stats
 import torch
 from loguru import logger
+from typing import Optional, Tuple, Union
 
 
-def caluclate_eta(df):
-    return len(df[np.abs(df.zwerr)>0.15])/len(df) *100
-    
+def calculate_eta(df: pd.DataFrame) -> float:
+    """Calculate the percentage of outliers in the DataFrame based on zwerr column."""
+    return len(df[np.abs(df.zwerr) > 0.15]) / len(df) * 100
+
 
-def nmad(data):
+def nmad(data: Union[np.ndarray, pd.Series]) -> float:
+    """Calculate the normalized median absolute deviation (NMAD) of the data."""
     return 1.4826 * np.median(np.abs(data - np.median(data)))
 
 
-def sigma68(data):
+def sigma68(data: Union[np.ndarray, pd.Series]) -> float:
+    """Calculate the sigma68 metric, a robust measure of dispersion."""
     return 0.5 * (pd.Series(data).quantile(q=0.84) - pd.Series(data).quantile(q=0.16))
 
 
-def maximum_mean_discrepancy(x, y, kernel_type="rbf", kernel_mul=2.0, kernel_num=5):
+def maximum_mean_discrepancy(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    kernel_type: str = "rbf",
+    kernel_mul: float = 2.0,
+    kernel_num: int = 5,
+) -> torch.Tensor:
     """
     Compute the Maximum Mean Discrepancy (MMD) between two sets of samples.
 
@@ -40,7 +50,13 @@ def maximum_mean_discrepancy(x, y, kernel_type="rbf", kernel_mul=2.0, kernel_num
     return mmd_loss
 
 
-def compute_kernel(x, y, kernel_type="rbf", kernel_mul=2.0, kernel_num=5):
+def compute_kernel(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    kernel_type: str = "rbf",
+    kernel_mul: float = 2.0,
+    kernel_num: int = 5,
+) -> torch.Tensor:
     """
     Compute the kernel matrix based on the chosen kernel type.
 
@@ -61,7 +77,7 @@ def compute_kernel(x, y, kernel_type="rbf", kernel_mul=2.0, kernel_num=5):
     x = x.unsqueeze(1).expand(x_size, y_size, dim)
     y = y.unsqueeze(0).expand(x_size, y_size, dim)
 
-    kernel_input = (x - y).pow(2).mean(2)  
+    kernel_input = (x - y).pow(2).mean(2)
 
     if kernel_type == "linear":
         kernel_matrix = kernel_input
@@ -80,46 +96,62 @@ def compute_kernel(x, y, kernel_type="rbf", kernel_mul=2.0, kernel_num=5):
 
 
 def select_cut(
-    df, completenss_lim=None, nmad_lim=None, outliers_lim=None, return_df=False
-):
+    df: pd.DataFrame,
+    completenss_lim: Optional[float] = None,
+    nmad_lim: Optional[float] = None,
+    outliers_lim: Optional[float] = None,
+    return_df: bool = False,
+) -> Union[Tuple[pd.DataFrame, float, pd.DataFrame], Tuple[float, pd.DataFrame]]:
+    """
+    Selects a cut based on one of the provided limits (completeness, NMAD, or outliers).
+
+    Args:
+    - df: DataFrame, containing the data
+    - completenss_lim: float, optional limit on completeness
+    - nmad_lim: float, optional limit on NMAD
+    - outliers_lim: float, optional limit on outliers (eta)
+    - return_df: bool, whether to return the filtered DataFrame
+
+    Returns:
+    - selected_cut: If return_df is False, returns the cut value and a DataFrame of cuts.
+                    If return_df is True, returns the filtered DataFrame, cut value, and cuts DataFrame.
+    """
 
-    if (completenss_lim is None) & (nmad_lim is None) & (outliers_lim is None):
-        raise (ValueError("Select at least one cut"))
+    if (completenss_lim is None) and (nmad_lim is None) and (outliers_lim is None):
+        raise ValueError("Select at least one cut")
     elif sum(c is not None for c in [completenss_lim, nmad_lim, outliers_lim]) > 1:
         raise ValueError("Select only one cut at a time")
 
-    else:
-        bin_edges = stats.mstats.mquantiles(df.odds, np.arange(0, 1.01, 0.1))
-        scatter, eta, cmptnss, nobj = [], [], [], []
-
-        for k in range(len(bin_edges) - 1):
-            edge_min = bin_edges[k]
-            edge_max = bin_edges[k + 1]
-
-            df_bin = df[(df.odds > edge_min)]
-
-            cmptnss.append(np.round(len(df_bin) / len(df), 2) * 100)
-            scatter.append(nmad(df_bin.zwerr))
-            eta.append(len(df_bin[np.abs(df_bin.zwerr) > 0.15]) / len(df_bin) * 100)
-            nobj.append(len(df_bin))
-
-        dfcuts = pd.DataFrame(
-            data=np.c_[
-                np.round(bin_edges[:-1], 5),
-                np.round(nobj, 1),
-                np.round(cmptnss, 1),
-                np.round(scatter, 3),
-                np.round(eta, 2),
-            ],
-            columns=["flagcut", "Nobj", "completeness", "nmad", "eta"],
-        )
+    bin_edges = stats.mstats.mquantiles(df.odds, np.arange(0, 1.01, 0.1))
+    scatter, eta, cmptnss, nobj = [], [], [], []
+
+    for k in range(len(bin_edges) - 1):
+        edge_min = bin_edges[k]
+        edge_max = bin_edges[k + 1]
+
+        df_bin = df[(df.odds > edge_min)]
+        cmptnss.append(np.round(len(df_bin) / len(df), 2) * 100)
+        scatter.append(nmad(df_bin.zwerr))
+        eta.append(len(df_bin[np.abs(df_bin.zwerr) > 0.15]) / len(df_bin) * 100)
+        nobj.append(len(df_bin))
+
+    dfcuts = pd.DataFrame(
+        data=np.c_[
+            np.round(bin_edges[:-1], 5),
+            np.round(nobj, 1),
+            np.round(cmptnss, 1),
+            np.round(scatter, 3),
+            np.round(eta, 2),
+        ],
+        columns=["flagcut", "Nobj", "completeness", "nmad", "eta"],
+    )
 
     if completenss_lim is not None:
         logger.info("Selecting cut based on completeness")
         selected_cut = dfcuts[dfcuts["completeness"] <= completenss_lim].iloc[0]
 
     elif nmad_lim is not None:
-        logger.info("Selecting cut based on nmad")
+        logger.info("Selecting cut based on NMAD")
         selected_cut = dfcuts[dfcuts["nmad"] <= nmad_lim].iloc[0]
 
     elif outliers_lim is not None:
@@ -127,11 +159,104 @@ def select_cut(
         selected_cut = dfcuts[dfcuts["eta"] <= outliers_lim].iloc[0]
 
     logger.info(
-        f"This cut provides completeness of {selected_cut['completeness']}, nmad={selected_cut['nmad']} and eta={selected_cut['eta']}"
+        f"This cut provides completeness of {selected_cut['completeness']}, "
+        f"nmad={selected_cut['nmad']} and eta={selected_cut['eta']}"
     )
 
     df_cut = df[(df.odds > selected_cut["flagcut"])]
-    if return_df == True:
+
+    if return_df:
         return df_cut, selected_cut["flagcut"], dfcuts
     else:
         return selected_cut["flagcut"], dfcuts
+
+def calculate_pit(model_f: nn.Module, 
+                  model_z: nn.Module,
+                  input_data: Tensor,
+                  target_data: Tensor,
+    ) -> List[float]:
+    
+    logger.info('Calculating PIT values')
+    
+    pit_list = []
+
+    model_f = model_f.eval()
+    model_f = model_f.to(self.device)
+    model_z = model_z.eval()
+    model_z = model_z.to(self.device)
+
+    input_data = input_data.to(self.device)
+            
+
+    features = model_f(input_data)
+    mu, logsig, logmix_coeff = model_z(features)
+    
+    logsig = torch.clamp(logsig,-6,2)
+    sig = torch.exp(logsig)
+
+    mix_coeff = torch.exp(logmix_coeff)
+    
+    mu,  mix_coeff, sig = mu.detach().cpu().numpy(),  mix_coeff.detach().cpu().numpy(), sig.detach().cpu().numpy() 
+    
+    for ii in range(len(input_data)):
+        pit = (mix_coeff[ii] * norm.cdf(target_data[ii]*np.ones(mu[ii].shape),mu[ii], sig[ii])).sum()
+        pit_list.append(pit)
+    
+    
+    return pit_list
+
+def calculate_crps(model_f: nn.Module, 
+                  model_z: nn.Module,
+                  input_data: Tensor,
+                  target_data: Tensor,
+    ) -> List[float]:
+    logger.info('Calculating CRPS values')
+
+    def measure_crps(cdf, t):
+        zgrid = np.linspace(0,4,1000)
+        Deltaz = zgrid[None,:] - t[:,None]
+        DeltaZ_heaviside = np.where(Deltaz < 0,0,1)
+        integral = (cdf-DeltaZ_heaviside)**2
+        crps_value = integral.sum(1) / 1000
+
+        return crps_value
+
+
+    crps_list = []
+
+    model_f = model_f.eval()
+    model_f = model_f.to(self.device)
+    model_z = model_z.eval()
+    model_z = model_z.to(self.device)
+
+    input_data = input_data.to(self.device)
+
+
+    features = model_f(input_data)
+    mu, logsig, logmix_coeff = model_z(features)
+    logsig = torch.clamp(logsig,-6,2)
+    sig = torch.exp(logsig)
+
+    mix_coeff = torch.exp(logmix_coeff)
+
+
+    mu,  mix_coeff, sig = mu.detach().cpu().numpy(),  mix_coeff.detach().cpu().numpy(), sig.detach().cpu().numpy() 
+
+    z = (mix_coeff * mu).sum(1)
+
+    x = np.linspace(0, 4, 1000)
+    pz = np.zeros(shape=(len(target_data), len(x)))
+    for ii in range(len(input_data)):
+        for i in range(6):
+            pz[ii] += mix_coeff[ii,i] * norm.pdf(x, mu[ii,i], sig[ii,i])
+
+    pz = pz / pz.sum(1)[:,None]
+
+
+    cdf_z = np.cumsum(pz,1)
+
+    crps_value = measure_crps(cdf_z, target_data)
+
+
+
+    return crps_value