From 5c8adce747a6f92027506b8332f7cae83c7a2c5e Mon Sep 17 00:00:00 2001
From: Bastien Lacave <bast.lacave@gmail.com>
Date: Fri, 28 Nov 2025 13:41:40 +0100
Subject: [PATCH 1/3] fixes select models

---
 src/ctlearn_manager/model_manager.py          | 91 ++++++++++---------
 src/ctlearn_manager/project.py                | 31 +++++++
 .../resources/CTLearnStyleDark.mplstyle       |  2 +
 .../resources/CTLearnStyleLight.mplstyle      |  3 +-
 src/ctlearn_manager/tri_model.py              | 15 ++-
 src/ctlearn_manager/tri_model_collection.py   | 79 +++++++++++-----
 src/ctlearn_manager/utils/DL2_processing.py   | 78 +++++++++++++---
 src/ctlearn_manager/utils/utils.py            |  4 +-
 src/ctlearn_manager/utils/who_is_better.py    |  3 +-
 9 files changed, 221 insertions(+), 85 deletions(-)

diff --git a/src/ctlearn_manager/model_manager.py b/src/ctlearn_manager/model_manager.py
index e36b42e..1a613d2 100644
--- a/src/ctlearn_manager/model_manager.py
+++ b/src/ctlearn_manager/model_manager.py
@@ -953,48 +953,53 @@ def plot_zenith_azimuth_ranges(self, ax=None, plot_testing_nodes=True):
                         label="Training",
                     )
 
-        try:
-            testing_dl1_table = read_table_hdf5(
-                self.project_directories.model_index_file, path=IndexTables(self, ParticleType.GAMMA_POINT).TESTING.table_path
-            )
-            zeniths = testing_dl1_table["testing_gamma_point_zenith_distances"]
-            azimuths = testing_dl1_table["testing_gamma_point_azimuths"].to(u.rad)
-            for zenith, azimuth in zip(zeniths, azimuths):
-                if (zenith == np.nan) or (azimuth == np.nan) or not plot_testing_nodes:
-                    continue
-                else:
-                    ax.scatter(
-                        azimuth,
-                        zenith,
-                        s=50,
-                        facecolors="none",
-                        edgecolors=get_color("ctlearn_accent_1"),
-                        label="Testing DL1",
-                        zorder=3,
-                    )
-        except:
-            a = 1
-
-        try:
-            mc_dl2_table = read_table_hdf5(
-                self.project_directories.model_index_file, path=IndexTables(self, ParticleType.GAMMA_POINT).DL2_MC.table_path
-            )
-            zeniths = mc_dl2_table["testing_DL2_gamma_point_zenith_distances"]
-            azimuths = mc_dl2_table["testing_DL2_gamma_point_azimuths"].to(u.rad)
-            for zenith, azimuth in zip(zeniths, azimuths):
-                if (zenith == np.nan) or (azimuth == np.nan) or not plot_testing_nodes:
-                    continue
-                else:
-                    ax.scatter(
-                        azimuth,
-                        zenith,
-                        s=50,
-                        color=get_color("ctlearn_accent_2"),
-                        label="Testing DL2",
-                        zorder=2,
-                    )
-        except:
-            a = 1
+        for particle_type, s, color in zip(
+            (ParticleType.GAMMA_POINT, ParticleType.PROTON),
+            (90, 50),
+            (get_color("ctlearn_accent_1"), get_color("ctlearn_accent_2")),
+        ):
+            try:
+                testing_dl1_table = read_table_hdf5(
+                    self.project_directories.model_index_file, path=IndexTables(self, particle_type).TESTING.table_path
+                )
+                zeniths = testing_dl1_table[f"testing_{particle_type.value}_zenith_distances"]
+                azimuths = testing_dl1_table[f"testing_{particle_type.value}_azimuths"].to(u.rad)
+                for zenith, azimuth in zip(zeniths, azimuths):
+                    if (zenith == np.nan) or (azimuth == np.nan) or not plot_testing_nodes:
+                        continue
+                    else:
+                        ax.scatter(
+                            azimuth,
+                            zenith,
+                            s=s,
+                            facecolors="none",
+                            edgecolors=color,
+                            label=f"Testing {particle_type.value} DL1",
+                            zorder=3,
+                        )
+            except:
+                a = 1
+
+            try:
+                mc_dl2_table = read_table_hdf5(
+                    self.project_directories.model_index_file, path=IndexTables(self, particle_type).DL2_MC.table_path
+                )
+                zeniths = mc_dl2_table[f"testing_DL2_{particle_type.value}_zenith_distances"]
+                azimuths = mc_dl2_table[f"testing_DL2_{particle_type.value}_azimuths"].to(u.rad)
+                for zenith, azimuth in zip(zeniths, azimuths):
+                    if (zenith == np.nan) or (azimuth == np.nan) or not plot_testing_nodes:
+                        continue
+                    else:
+                        ax.scatter(
+                            azimuth,
+                            zenith,
+                            s=s,
+                            color=color,
+                            label=f"Testing {particle_type.value} DL2",
+                            zorder=2,
+                        )
+            except:
+                a = 1
         handles, labels = ax.get_legend_handles_labels()
         by_label = dict(zip(labels, handles))
         ax.legend(by_label.values(), by_label.keys())
@@ -1103,7 +1108,7 @@ def plot_training_nodes(self):
         ax.set_theta_zero_location("E")
         ax.set_theta_direction(-1)
         ax.set_rlabel_position(-30)
-        ax.set_ylim(0, 60)
+        # ax.set_ylim(0, 60)
         ax.set_yticks(np.arange(10, 61, 10))
         ax.set_yticklabels(["", "", "30°", "", "", "60°"], fontsize=10)
         ax.set_xlabel("Azimuth [deg]", fontsize=10)
diff --git a/src/ctlearn_manager/project.py b/src/ctlearn_manager/project.py
index b318bea..60f0a68 100644
--- a/src/ctlearn_manager/project.py
+++ b/src/ctlearn_manager/project.py
@@ -4,6 +4,7 @@
 from . import CTLearnModelManager, CTLearnTriModelManager, DataSample
 from .utils import CTLMDirectories, get_user_confirmation, ClusterConfiguration, set_mpl_style
 from ctlearn_manager.utils.utils import set_global_theme, ColorTheme
+from .utils import remove_model_from_index
 # from.io import load_model_from_index
 
 __all__ = [
@@ -51,6 +52,9 @@ def create_tri_model(
             else:
                 get_user_confirmation(prompt=f"TriModel {tri_model_nickname} already exists. Do you want to overwrite it?\n This will delete the existing model and all its data.")
                 os.system(f"rm -rf {tri_models_directory}")
+                remove_model_from_index(f"{tri_model_nickname}_type", f"{self.project_directory}model_index.h5")
+                remove_model_from_index(f"{tri_model_nickname}_energy", f"{self.project_directory}model_index.h5")
+                remove_model_from_index(f"{tri_model_nickname}_direction", f"{self.project_directory}model_index.h5")
         project_directories = CTLMDirectories(self.project_directory, tri_model_nickname)
         
 
@@ -108,6 +112,33 @@ def create_tri_model(
                 project_directories
             )
 
+    def delete_tri_model(
+        self,
+        tri_model_nickname: str,
+    ):
+        """
+        Delete an existing CTLearn model manager instance with the given nickname.
+        
+        Parameters:
+            tri_model_nickname (str): Nickname for the model.
+        """
+        tri_models_directory = f"{self.project_directory}/models/{tri_model_nickname}"
+        if not Path(tri_models_directory).exists():
+            raise FileNotFoundError(
+                f"TriModel {tri_model_nickname} does not exist in {self.project_directory}."
+            )
+        get_user_confirmation(prompt=f"Are you sure you want to delete TriModel {tri_model_nickname} and all its data? This action cannot be undone.")
+        os.system(f"rm -rf {tri_models_directory}")
+        import h5py
+
+        project_directories = CTLMDirectories(self.project_directory, tri_model_nickname)
+
+        with h5py.File(project_directories.model_index_file, "r+") as f:
+            del f[f"{tri_model_nickname}_type"]
+            del f[f"{tri_model_nickname}_energy"]
+            del f[f"{tri_model_nickname}_direction"]
+
+
     def open_tri_model(
         self,
         tri_model_nickname: str,
diff --git a/src/ctlearn_manager/resources/CTLearnStyleDark.mplstyle b/src/ctlearn_manager/resources/CTLearnStyleDark.mplstyle
index d2f8018..7d97de8 100644
--- a/src/ctlearn_manager/resources/CTLearnStyleDark.mplstyle
+++ b/src/ctlearn_manager/resources/CTLearnStyleDark.mplstyle
@@ -39,6 +39,8 @@ lines.dash_capstyle : round
 image.cmap : magma
 # SCATTER PLOT
 scatter.marker : o
+lines.markeredgecolor : "black"
+lines.markeredgewidth : 1.5
 
 # TICKS
 xtick.labelsize : 13
diff --git a/src/ctlearn_manager/resources/CTLearnStyleLight.mplstyle b/src/ctlearn_manager/resources/CTLearnStyleLight.mplstyle
index 6ef8211..0c34118 100644
--- a/src/ctlearn_manager/resources/CTLearnStyleLight.mplstyle
+++ b/src/ctlearn_manager/resources/CTLearnStyleLight.mplstyle
@@ -37,7 +37,8 @@ lines.markeredgewidth : 0.0
 lines.solid_capstyle : round
 lines.dash_capstyle : round
 image.cmap : viridis
-
+lines.markeredgecolor : "white"
+lines.markeredgewidth : 1.5
 
 # SCATTER PLOT
 scatter.marker : o
diff --git a/src/ctlearn_manager/tri_model.py b/src/ctlearn_manager/tri_model.py
index 8fbdaad..16b4bcc 100644
--- a/src/ctlearn_manager/tri_model.py
+++ b/src/ctlearn_manager/tri_model.py
@@ -323,9 +323,9 @@ def delete_table_from_index(self, path: str):
         """
         import h5py
 
-        with h5py.File(self.direction_model.model_index_file, "r+") as f:
+        with h5py.File(self.project_directories.model_index_file, "r+") as f:
             del f[path]
-            print(f"Table {path} erased from {self.direction_model.model_index_file}")
+            print(f"Table {path} erased from {self.project_directories.model_index_file}")
 
     def get_available_testing_directions(self):
         """
@@ -663,6 +663,8 @@ def launch_testing(
                     f"Output file {output_file} already exists, skipping, set overwrite=True to overwrite"
                 )
                 continue
+            if os.path.exists(output_file) and overwrite:
+                os.system(f"rm {output_file}")
             if self.stereo:
                 cmd = f"ctlearn-predict-stereo-model --input_url {input_file} \
 --PredictCTLearnModel.batch_size={batch_size} \
@@ -2284,13 +2286,18 @@ def plot_irfs(self, zenith, azimuth, cuts: Cuts = DefaultCuts.EFF_70.value):
         )
 
         irf_file = self.get_IRF_data(zenith, azimuth, cuts)["irf_file"]
+        print(f"Reading IRFs from file: {irf_file}")
         # rad_max = RadMax2D.read(irf_file, hdu="RAD MAX")
         aeff = EffectiveAreaTable2D.read(irf_file, hdu="EFFECTIVE AREA")
-        bkg = Background2D.read(irf_file, hdu="BACKGROUND")
+        
         edisp = EnergyDispersion2D.read(irf_file, hdu="ENERGY DISPERSION")
         edisp.peek()
         aeff.peek()
-        bkg.peek()
+        try:
+            bkg = Background2D.read(irf_file, hdu="BACKGROUND")
+            bkg.peek()
+        except:
+            print(f"Could not read BACKGROUND IRF from file {irf_file}")
 
     def plot_loss(self):
         """
diff --git a/src/ctlearn_manager/tri_model_collection.py b/src/ctlearn_manager/tri_model_collection.py
index 231adec..e4822a5 100644
--- a/src/ctlearn_manager/tri_model_collection.py
+++ b/src/ctlearn_manager/tri_model_collection.py
@@ -400,8 +400,10 @@ def find_closest_model_to(
             If the pointing data cannot be found in the provided pointing table.
         """
         import astropy.units as u
+        from astropy.io.misc.hdf5 import read_table_hdf5
 
         from ctlearn_manager.utils.utils import get_avg_pointing
+        from .model_manager import IndexTables
 
         try:
             avg_data_ze, avg_data_az = get_avg_pointing(
@@ -415,28 +417,50 @@ def find_closest_model_to(
             print(f"⚠️ Pointing not found at {pointing_table}, skipping : {input_file}")
             return
 
-        avg_model_azs = []
-        avg_model_zes = []
+        min_distances = []
+        
         for tri_model in self.tri_models:
-            avg_model_azs.append(
-                np.mean(tri_model.direction_model.validity.azimuth_range)
-                .to(u.deg)
-                .value
+            training_gamma_table = read_table_hdf5(
+                tri_model.direction_model.project_directories.model_index_file,
+                path=IndexTables(tri_model.direction_model, ParticleType.GAMMA_DIFFUSE).TRAINING.table_path,
             )
-            avg_model_zes.append(
-                np.mean(tri_model.direction_model.validity.zenith_range).to(u.deg).value
+            zeniths_nodes = training_gamma_table[
+                "training_gamma_diffuse_zenith_distances"
+            ]
+            azimuths_nodes = training_gamma_table["training_gamma_diffuse_azimuths"].to(
+                u.rad
             )
-        avg_model_azs = np.array(avg_model_azs) * u.deg
-        avg_model_zes = np.array(avg_model_zes) * u.deg
-        # angular_distance_matrix = angular_distance(avg_data_ze, avg_data_az, avg_model_zes, avg_model_azs)
-        closest_model_index = np.argmin(
-            angular_distance(avg_data_ze, avg_data_az, avg_model_zes, avg_model_azs)
-        )
+            angular_distances = angular_distance(
+                avg_data_ze, avg_data_az, 
+                zeniths_nodes, azimuths_nodes
+            )
+            min_distances.append(np.min(angular_distances).value)
+
+        
+
+        # avg_model_azs = []
+        # avg_model_zes = []
+        # for tri_model in self.tri_models:
+        #     avg_model_azs.append(
+        #         np.mean(tri_model.direction_model.validity.azimuth_range)
+        #         .to(u.deg)
+        #         .value
+        #     )
+        #     avg_model_zes.append(
+        #         np.mean(tri_model.direction_model.validity.zenith_range).to(u.deg).value
+        #     )
+        # avg_model_azs = np.array(avg_model_azs) * u.deg
+        # avg_model_zes = np.array(avg_model_zes) * u.deg
+        # # angular_distance_matrix = angular_distance(avg_data_ze, avg_data_az, avg_model_zes, avg_model_azs)
+        # closest_model_index = np.argmin(
+        #     angular_distance(avg_data_ze, avg_data_az, avg_model_zes, avg_model_azs)
+        # )
+        closest_model_index = np.argmin(min_distances)
         closest_model = self.tri_models[closest_model_index]
 
         if verbose:
             print(
-                f"📁 File : {input_file.split('/')[-1]}      📡 Pointing : ({avg_data_ze.value:.3f}, {avg_data_az.value:.3f})      🧠 Closest Model : ({np.mean(closest_model.direction_model.validity.zenith_range).value:.3f}, {np.mean(closest_model.direction_model.validity.azimuth_range).value:.3f})"
+                f"📁 File : {input_file.split('/')[-1]}      📡 Pointing : ({avg_data_ze.value:.3f}, {avg_data_az.value:.3f})      🧠 Closest Model {closest_model.direction_model.model_nickname} : ({np.mean(closest_model.direction_model.validity.zenith_range).value:.3f}, {np.mean(closest_model.direction_model.validity.azimuth_range).value:.3f})"
             )
         # print(f"｜📡 Average pointing of {input_file.split('/')[-1]} : ({avg_data_ze:3f}, {avg_data_az:3f})")
         # print(f"｜🔍 Closest model avg node : ({np.mean(closest_model.direction_model.validity.zenith_range).value}, {np.mean(closest_model.direction_model.validity.azimuth_range).value})")
@@ -532,14 +556,18 @@ def plot_energy_resolution_DL2(
         ]
         if compare_with is not None:
             fig, (ax, ax_rel) = plt.subplots(
-                2, 1, gridspec_kw={"height_ratios": [3, 1]}
+                2, 1, gridspec_kw={"height_ratios": [3, 1]},
+                sharex=True
             )
+            plt.subplots_adjust(hspace=0)
+            ax.set_xticks([])
+            ax.set_xlabel("")
             ax_rel.set_xlabel("True Energy (TeV)")
             ax_rel.set_ylabel("Rel. Impr. (%)")
             ax_rel.grid(True, linestyle="--", alpha=0.5)
             ax_rel.set_xscale("log")
             ax_rel.set_ymargin(0.05)
-            ax_rel.set_yticks([0, 10, 20, 30, 40, 50])
+            # ax_rel.set_yticks([0, 10, 20, 30, 40, 50])
         else:
             fig, ax = plt.subplots()
         cuts.plot_cuts_info_plt(ax)
@@ -677,7 +705,7 @@ def plot_energy_resolution_DL2(
 
         if compare_with is not None:
             ax_rel.set_xlim(ax.get_xlim())
-            ax_rel.set_ylim(bottom=0)
+            # ax_rel.set_ylim(bottom=0)
         ax.legend()
         plt.tight_layout()
         plt.subplots_adjust(hspace=0.0)
@@ -705,14 +733,18 @@ def plot_angular_resolution_DL2(
         ]
         if compare_with is not None:
             fig, (ax, ax_rel) = plt.subplots(
-                2, 1, gridspec_kw={"height_ratios": [3, 1]}
+                2, 1, gridspec_kw={"height_ratios": [3, 1]},
+                sharex=True
             )
+            plt.subplots_adjust(hspace=0)
+            ax.set_xticks([])
+            ax.set_xlabel("")
             ax_rel.set_xlabel("True Energy (TeV)")
             ax_rel.set_ylabel("Rel. Impr. (%)")
             ax_rel.grid(True, linestyle="--", alpha=0.5)
             ax_rel.set_xscale("log")
             ax_rel.set_ymargin(0.05)
-            ax_rel.set_yticks([0, 10, 20, 30, 40, 50])
+            # ax_rel.set_yticks([0, 10, 20, 30, 40, 50])
         else:
             fig, ax = plt.subplots()
         stored_efficiency_theta = cuts.efficiency_theta
@@ -769,8 +801,9 @@ def plot_angular_resolution_DL2(
             )
             if compare_with is not None:
                 # ax.set_xscale(ax_rel.get_xscale())
-                ax.set_xticks([])
-                ax.set_xlabel("")
+
+                # ax.set_xticks([])
+                # ax.set_xlabel("")
                 fig.subplots_adjust(hspace=0)
                 if len(compare_with_index) > 0:
                     ref_e_bins, ref_ang_res_err = self.tri_models[
@@ -846,7 +879,7 @@ def plot_angular_resolution_DL2(
 
         if compare_with is not None:
             ax_rel.set_xlim(ax.get_xlim())
-            ax_rel.set_ylim(bottom=0)
+            # ax_rel.set_ylim(bottom=0)
         ax.legend()
         plt.tight_layout()
         plt.subplots_adjust(hspace=0.0)
diff --git a/src/ctlearn_manager/utils/DL2_processing.py b/src/ctlearn_manager/utils/DL2_processing.py
index 7c2724b..f2f5ae7 100644
--- a/src/ctlearn_manager/utils/DL2_processing.py
+++ b/src/ctlearn_manager/utils/DL2_processing.py
@@ -304,7 +304,7 @@ class DL2DataProcessor:
         Initializes the DL2DataProcessor with the given parameters and processes the DL2 data.
     process_DL2_data(self):
         Processes the DL2 data files, applying cuts and computing sky positions.
-    plot_theta2_distribution(self, bins, n_off=5):
+    plot_theta2_distribution(self, bins, n_off=3):
         Plots the theta^2 distribution for the processed DL2 data.
     compute_off_regions(self, pointing, n_off):
         Computes the off-source regions for background estimation.
@@ -332,6 +332,7 @@ def __init__(
         # max_theta2: float=0.2,
         workers=None,
         reco_field_suffix = None,
+        fit_src_position: bool = True,
     ):
         self.workers = workers
         mp.set_start_method("fork", force=True)
@@ -401,7 +402,6 @@ def __init__(
         self.process_DL2_data()
         self.load_processed_data()
 
-
         self.cut_file_theta_cuts = []
         self.cut_file_gammaness_cuts = []
 
@@ -426,6 +426,11 @@ def __init__(
                 self.cut_file_theta_cuts.append(file_theta_cuts)
                 self.cut_file_gammaness_cuts.append(file_gammaness_cuts)
 
+        if fit_src_position:
+            center_ra, center_dec = self.find_gaussian_center(0)
+            self.source_position = SkyCoord(ra=center_ra * u.deg, dec=center_dec * u.deg, frame=self.source_position)
+
+
 
         # self.cut_file_theta_cuts = []
         # self.cut_file_gammaness_cuts = []
@@ -505,6 +510,54 @@ def __init__(
         
         # set_mpl_style()
 
+    def find_gaussian_center(self, cuts_index):
+        import concurrent.futures
+        from scipy.optimize import curve_fit
+        # Prepare arguments for parallel processing
+        file_args = list(
+            zip(
+                self.reco_directions,
+                self.cuts_masks_gammaness_only,
+                self.dl2s,
+                self.pointings,
+            )
+        )
+
+        def extract_coords(args):
+            reco, cuts_mask, dl2, pointing = args
+            cuts_mask = cuts_mask[cuts_index]
+            ra = reco[cuts_mask].ra.deg
+            dec = reco[cuts_mask].dec.deg
+            pointing_ra = pointing[cuts_mask].ra.deg
+            pointing_dec = pointing[cuts_mask].dec.deg
+            return ra, dec, pointing_ra, pointing_dec
+
+        # Parallel extraction of coordinates
+        ra_values = []
+        dec_values = []
+        pointings_ra = []
+        pointings_dec = []
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            for ra, dec, pra, pdec in executor.map(extract_coords, file_args):
+                ra_values.append(ra)
+                dec_values.append(dec)
+                pointings_ra.append(pra)
+                pointings_dec.append(pdec)
+
+        # Flatten arrays for plotting
+        ra_values = np.concatenate(ra_values)
+        dec_values = np.concatenate(dec_values)
+        pointings_ra = np.concatenate(pointings_ra)
+        pointings_dec = np.concatenate(pointings_dec)
+
+        histogram, xedges, yedges = np.histogram2d(ra_values, dec_values, bins=300)
+        # Find the position of the maximum pixel
+        max_idx = np.unravel_index(np.argmax(histogram), histogram.shape)
+        center_ra = (xedges[max_idx[0]] + xedges[max_idx[0] + 1]) / 2
+        center_dec = (yedges[max_idx[1]] + yedges[max_idx[1] + 1]) / 2
+        print(f"Maximum pixel center: RA={center_ra:.4f}°, DEC={center_dec:.4f}°")
+        return center_ra, center_dec
+
     def set_keys(self):
         self.gammaness_key = (
             f"{self.reco_field_suffix}_prediction"  # if self.CTLearn else "gammaness"
@@ -706,7 +759,6 @@ def load_processed_data(self):
             )
             for i, DL2_file in enumerate(self.DL2_files)
         ]
-
         with ProcessPoolExecutor(self.workers) as executor:
             for result in tqdm(executor.map(load_one_worker, args_list), total=n_files, desc="Loading processed data"):
                 results.append(result)
@@ -736,7 +788,8 @@ def load_processed_data(self):
                 self.I_g_off_counts[i] = I_g_off_counts
                 self.corresponding_models[i] = corresponding_model
                 self.effective_times[i] = eff_time
-
+        if n_tot == 0:
+            raise ValueError("No events. Maybe wait for processing.")
         print(f"Cut {n_saved_tot} ({(n_saved_tot/n_tot * 100):.2f}%) events in total from intensity cut {self.intensity_cut}p.e. and global gammaness cut {self.global_gammaness_cut}.")
         print(f"Remaining {n_tot - n_saved_tot} ({((n_tot - n_saved_tot)/n_tot * 100):.2f}%) events in total after cuts.")
         # Filter arrays
@@ -757,7 +810,7 @@ def load_processed_data(self):
 
 
 
-    def plot_theta2_distribution(self, bins=25, n_off=5, output_file=None, cuts_index=0, t2_max=0.4):
+    def plot_theta2_distribution(self, bins=25, n_off=3, output_file=None, cuts_index=0, t2_max=0.4):
         import concurrent.futures
 
         import matplotlib.pyplot as plt
@@ -1122,7 +1175,7 @@ def compute_on_off_counts(
 
         return on_count, off_count, on_separation, all_off_separation, significance_lima
 
-    def plot_skymap(self, output_file=None, cuts_index=0, n_off=5):
+    def plot_skymap(self, output_file=None, cuts_index=0, n_off=3):
         import concurrent.futures
 
         import matplotlib.pyplot as plt
@@ -1230,7 +1283,7 @@ def extract_coords(args):
         else:
             plt.show()
     
-    def optimize_cuts_on_crab(self, n_off=5, output_suffix="", max_gammaness_cut=1.0, max_theta2_cut=0.2, gcut_step=0.01, theta2_cut_step=0.001, E_bins=None):
+    def optimize_cuts_on_crab(self, n_off=3, output_suffix="", max_gammaness_cut=1.0, max_theta2_cut=0.2, gcut_step=0.01, theta2_cut_step=0.001, E_bins=None):
         """
         Compute and store optimal gammaness/theta2 cuts for even and odd events for each energy bin.
         """
@@ -1713,7 +1766,7 @@ def plot_cuts_optimized_on_crab(self, cuts_h5_file):
         plt.tight_layout()
         plt.show()
 
-    def plot_sensitivity(self, n_off=5, ax=None, label="CTLearn", output_file=None, export_to_h5: str=None,
+    def plot_sensitivity(self, n_off=3, ax=None, label="CTLearn", output_file=None, export_to_h5: str=None,
         import_from_h5: str = None,
         import_label: str = None,
         optimized_on_crab: bool = False, output_suffix: str=''):
@@ -2098,6 +2151,7 @@ def plot_perf_paper_sensitivity_src_indep_percentage_without_5percentbg_errors_s
                 import importlib.resources as pkg_resources
 
                 from .. import resources
+
                 with pkg_resources.path(resources, "sensitivity_src_indep_without_5percentbg.txt") as text_path:
                     data = np.loadtxt(text_path)
 
@@ -2156,7 +2210,7 @@ def plot_perf_paper_sensitivity_src_indep_percentage_without_5percentbg_errors_s
             if ax is None:
                 plt.show()
 
-    def plot_PSF(self, n_off=5, ax=None, label="CTLearn", output_file=None, plot_MC: list[str]=[], export_to_h5: str=None,
+    def plot_PSF(self, n_off=3, ax=None, label="CTLearn", output_file=None, plot_MC: list[str]=[], export_to_h5: str=None,
         import_from_h5: str = None,
         import_label: str = None, ylim=(0, 0.6)):
         import concurrent.futures
@@ -2374,7 +2428,7 @@ def get_efficiency_for_gamaness_cuts(
         return efficiencies
 
     def plot_bkg_discrimination_capability(
-        self, n_off=5, axs=None, label="CTLearn", output_file=None
+        self, n_off=3, axs=None, label="CTLearn", output_file=None
     ):
         gammaness_cuts = np.arange(0, 1.05, 0.05)
         import matplotlib.pyplot as plt
@@ -2466,7 +2520,7 @@ def plot_bkg_discrimination_capability(
             if axs is None:
                 plt.show()
 
-    def plot_excess_vs_background_rates(self, n_off=5, output_file=None):
+    def plot_excess_vs_background_rates(self, n_off=3, output_file=None):
         gammaness_cuts = np.arange(0, 1.05, 0.05)
         import matplotlib.pyplot as plt
 
@@ -2533,7 +2587,7 @@ def plot_excess_vs_background_rates(self, n_off=5, output_file=None):
             plt.show()
 
     def plot_excess_and_background_rates_vs_energy(
-        self, n_off=5, output_file=None, cuts_index=0
+        self, n_off=3, output_file=None, cuts_index=0
     ):
         import matplotlib.pyplot as plt
 
diff --git a/src/ctlearn_manager/utils/utils.py b/src/ctlearn_manager/utils/utils.py
index 60deea3..513534a 100644
--- a/src/ctlearn_manager/utils/utils.py
+++ b/src/ctlearn_manager/utils/utils.py
@@ -865,6 +865,7 @@ def __init__(
         import astropy.units as u
         from ctapipe.io import read_table
         from tqdm import tqdm
+        import gc
 
         self.directory = directory
         self.pattern = pattern
@@ -937,6 +938,7 @@ def __init__(
         print(
             f"\t -> {self.particle_type.value} @ ({self.zenith_distance}, {self.azimuth})"
         )
+        gc.collect()
 
 
 class CutType(Enum):
@@ -1566,7 +1568,7 @@ def get_dl2_mc_files(self, zenith:float, azimuth:float, particle_types: list[Par
     def get_available_MC_directions(self, particle_type: ParticleType):
         import glob
 
-        paths = glob.glob(f"{self.dl2_mc_directory}/{particle_type.value}/*/")
+        paths = glob.glob(f"{self.dl2_mc_directory}/{particle_type.value}/*/*.h5")
         zeniths = []
         azimuths = []
         for path in paths:
diff --git a/src/ctlearn_manager/utils/who_is_better.py b/src/ctlearn_manager/utils/who_is_better.py
index 48550fc..8d9fdc0 100644
--- a/src/ctlearn_manager/utils/who_is_better.py
+++ b/src/ctlearn_manager/utils/who_is_better.py
@@ -25,6 +25,7 @@ def plot_skymap(self, output_file=None):
     def plot_sensitivity(self, output_file=None,
         import_from_h5: str = None,
         import_label: str = None,
+        export_to_h5_out: str = None
         ):
         fig, ax = plt.subplots()
         if len(self.cuts) == 1:
@@ -33,7 +34,7 @@ def plot_sensitivity(self, output_file=None,
             if i < len(self.labels) - 1:
                 dl2_processor.plot_sensitivity(ax=ax, label=label)
             else:
-                dl2_processor.plot_sensitivity(ax=ax, label=label, import_from_h5=import_from_h5, import_label=import_label)
+                dl2_processor.plot_sensitivity(ax=ax, label=label, import_from_h5=import_from_h5, import_label=import_label, export_to_h5=export_to_h5_out)
                 
         if output_file:
             plt.savefig(output_file)

From 35af04bf09f19919fd87673ad50691a20b3eef67 Mon Sep 17 00:00:00 2001
From: Bastien Lacave <bast.lacave@gmail.com>
Date: Fri, 28 Nov 2025 13:44:36 +0100
Subject: [PATCH 2/3] read subarray pointing and allow incomplete trimodels

---
 pyproject.toml                     |  2 +-
 src/ctlearn_manager/project.py     | 17 +++++++++--------
 src/ctlearn_manager/utils/utils.py |  9 ++++++---
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7e01ad5..dd4894b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,6 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy",
     "ctapipe==0.26.0",
-    "gammapy",
     "pyirf==0.13.0",
     "ctaplot==0.6.3",
     "h5py",
@@ -23,6 +22,7 @@ dependencies = [
     "ctlearn",
     "astropy",
     "matplotlib==3.10.0",
+    "tables<3.9",
 ]
 
 # needed for setuptools_scm, we don"t define a static version
diff --git a/src/ctlearn_manager/project.py b/src/ctlearn_manager/project.py
index 60f0a68..9170d8a 100644
--- a/src/ctlearn_manager/project.py
+++ b/src/ctlearn_manager/project.py
@@ -62,18 +62,19 @@ def create_tri_model(
         assert direction_reco in ["cameradirection", "skydirection"], (
             f"direction_reco must be one of ['cameradirection', 'skydirection']: {direction_reco}"
         )
-        recos = ['type', 'energy', direction_reco]
+        recos = tri_model_parameters.get("training_samples").keys()
         if isinstance(tri_model_parameters.get("training_samples"), dict):
-            training_samples = [
-                tri_model_parameters.get("training_samples")['type'],
-                tri_model_parameters.get("training_samples")['energy'],
-                tri_model_parameters.get("training_samples")[direction_reco],
-            ]
+            training_samples = []
+            for reco in recos:
+                training_samples.append(tri_model_parameters.get("training_samples").get(reco))
+                
         else: # isinstance(tri_model_parameters.get("training_samples"), list[DataSample]):
-            training_samples = [tri_model_parameters.get("training_samples")]*3
+            raise ValueError("training_samples must be a dict with keys 'type', 'energy', and 'cameradirection' or 'skydirection'.")
+            # training_samples = [tri_model_parameters.get("training_samples")]*3
         # else:
         #     raise ValueError("training_samples must be a dict or a list of DataSample instances.")
-
+        # recos = ['type', 'energy', direction_reco]
+        
         for reco, training_sample in zip(recos, training_samples):
             match reco:
                 case "type":
diff --git a/src/ctlearn_manager/utils/utils.py b/src/ctlearn_manager/utils/utils.py
index 513534a..0a6824b 100644
--- a/src/ctlearn_manager/utils/utils.py
+++ b/src/ctlearn_manager/utils/utils.py
@@ -885,13 +885,16 @@ def __init__(
                     f"File {file} is not an absolute path. Please provide absolute paths for the files."
                 )
             shower_parameters = read_table(file, "simulation/event/subarray/shower")
-            pointing = read_table(file, "configuration/telescope/pointing/tel_001")
+            
+            pointing = read_table(file, "configuration/subarray/pointing")
+            # pointing = read_table(file, "dl1/monitoring/subarray/pointing")
+            # print(pointing.colnames)
             particle_id = np.unique(shower_parameters["true_shower_primary_id"])
 
             zenith_distance = np.unique(
-                90 * u.deg - pointing["telescope_pointing_altitude"].to(u.deg)
+                90 * u.deg - pointing["array_altitude"].to(u.deg)
             )
-            azimuth = np.unique(pointing["telescope_pointing_azimuth"].to(u.deg))
+            azimuth = np.unique(pointing["array_azimuth"].to(u.deg))
 
             assert len(zenith_distance) == 1, (
                 f"More than one zenith distance found in {file}"

From b40618f2c0b2cf9f1d15128d52153b13cab2861b Mon Sep 17 00:00:00 2001
From: Bastien Lacave <bast.lacave@gmail.com>
Date: Thu, 11 Dec 2025 10:50:37 +0100
Subject: [PATCH 3/3] properly overwrite models

---
 src/ctlearn_manager/model_manager.py | 70 ++++++++++++++++++++++++++++
 src/ctlearn_manager/project.py       |  8 ++--
 src/ctlearn_manager/utils/utils.py   |  9 ++--
 3 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/src/ctlearn_manager/model_manager.py b/src/ctlearn_manager/model_manager.py
index 1a613d2..7bf1a3e 100644
--- a/src/ctlearn_manager/model_manager.py
+++ b/src/ctlearn_manager/model_manager.py
@@ -7,6 +7,7 @@
 import ast
 from pathlib import Path
 
+
 import astropy.units as u
 import numpy as np
 from astropy.table import QTable
@@ -776,6 +777,75 @@ def update_model_manager_testing_data(self, testing_data_sample: DataSample):
         # print(
         #     f"Testing {particle_type.value} at ({testing_zenith_distance}, {testing_azimuth}) : {testing_dir}/{testing_pattern} updated"
         # )
+        
+    def update_model_manager_training_data(self, training_data_sample: DataSample):
+        """
+        Update the training data for the model manager with a new data sample.
+
+        Parameters
+        ----------
+        training_data_sample : DataSample
+            The data sample containing training information to be added or updated.
+            It includes the directory, zenith distance, azimuth, pattern, and particle type.
+
+        Raises
+        ------
+        Exception
+            If there is an issue reading or writing the HDF5 file.
+
+        Notes
+        -----
+        - If the training data for the given zenith distance and azimuth already exists,
+          it updates the directory and pattern for that entry.
+        - If no matching entry exists, it adds a new row to the training data table.
+        - The updated table is saved back to the HDF5 file.
+        """
+        from astropy.io.misc.hdf5 import read_table_hdf5, write_table_hdf5
+
+        training_dir = training_data_sample.directory
+        training_zenith_distance = training_data_sample.zenith_distance
+        training_azimuth = training_data_sample.azimuth
+        training_pattern = training_data_sample.pattern
+        particle_type = training_data_sample.particle_type
+
+        training_index_table = IndexTables(self, particle_type).TRAINING
+        try:
+            training_table = read_table_hdf5(
+                self.project_directories.model_index_file,
+                path=training_index_table.table_path,
+            )
+        except:
+            training_table = training_index_table.default_table
+
+        match = np.where(
+            (
+                training_table[f"training_{particle_type.value}_zenith_distances"]
+                == training_zenith_distance
+            )
+            & (
+                training_table[f"training_{particle_type.value}_azimuths"]
+                == training_azimuth
+            )
+        )[0]
+        if len(match) > 0:
+            training_table[f"training_{particle_type.value}_dirs"][match[0]] = (
+                training_dir
+            )
+            training_table[f"training_{particle_type.value}_patterns"][
+                match[0]
+            ] = training_pattern
+        else:
+            training_table.add_row(
+                [training_dir, training_pattern, training_zenith_distance, training_azimuth]
+            )
+        write_table_hdf5(
+            training_table,
+            self.project_directories.model_index_file,
+            path=training_index_table.table_path,
+            append=True,
+            overwrite=True,
+            serialize_meta=True,
+        )
 
 
     def plot_zenith_azimuth_ranges(self, ax=None, plot_testing_nodes=True):
diff --git a/src/ctlearn_manager/project.py b/src/ctlearn_manager/project.py
index 9170d8a..256b251 100644
--- a/src/ctlearn_manager/project.py
+++ b/src/ctlearn_manager/project.py
@@ -51,7 +51,7 @@ def create_tri_model(
                 )
             else:
                 get_user_confirmation(prompt=f"TriModel {tri_model_nickname} already exists. Do you want to overwrite it?\n This will delete the existing model and all its data.")
-                os.system(f"rm -rf {tri_models_directory}")
+                # os.system(f"rm -rf {tri_models_directory}")
                 remove_model_from_index(f"{tri_model_nickname}_type", f"{self.project_directory}model_index.h5")
                 remove_model_from_index(f"{tri_model_nickname}_energy", f"{self.project_directory}model_index.h5")
                 remove_model_from_index(f"{tri_model_nickname}_direction", f"{self.project_directory}model_index.h5")
@@ -62,7 +62,8 @@ def create_tri_model(
         assert direction_reco in ["cameradirection", "skydirection"], (
             f"direction_reco must be one of ['cameradirection', 'skydirection']: {direction_reco}"
         )
-        recos = tri_model_parameters.get("training_samples").keys()
+        # recos = tri_model_parameters.get("training_samples").keys()
+        recos = ['type', 'energy', direction_reco]
         if isinstance(tri_model_parameters.get("training_samples"), dict):
             training_samples = []
             for reco in recos:
@@ -73,7 +74,7 @@ def create_tri_model(
             # training_samples = [tri_model_parameters.get("training_samples")]*3
         # else:
         #     raise ValueError("training_samples must be a dict or a list of DataSample instances.")
-        # recos = ['type', 'energy', direction_reco]
+        
         
         for reco, training_sample in zip(recos, training_samples):
             match reco:
@@ -139,7 +140,6 @@ def delete_tri_model(
             del f[f"{tri_model_nickname}_energy"]
             del f[f"{tri_model_nickname}_direction"]
 
-
     def open_tri_model(
         self,
         tri_model_nickname: str,
diff --git a/src/ctlearn_manager/utils/utils.py b/src/ctlearn_manager/utils/utils.py
index 0a6824b..ce570a6 100644
--- a/src/ctlearn_manager/utils/utils.py
+++ b/src/ctlearn_manager/utils/utils.py
@@ -886,16 +886,15 @@ def __init__(
                 )
             shower_parameters = read_table(file, "simulation/event/subarray/shower")
             
-            pointing = read_table(file, "configuration/subarray/pointing")
+            pointing = read_table(file, "/configuration/telescope/pointing/tel_001")
             # pointing = read_table(file, "dl1/monitoring/subarray/pointing")
             # print(pointing.colnames)
             particle_id = np.unique(shower_parameters["true_shower_primary_id"])
-
+            
             zenith_distance = np.unique(
-                90 * u.deg - pointing["array_altitude"].to(u.deg)
+                90 * u.deg - pointing["telescope_pointing_altitude"].to(u.deg)
             )
-            azimuth = np.unique(pointing["array_azimuth"].to(u.deg))
-
+            azimuth = np.unique(pointing["telescope_pointing_azimuth"].to(u.deg))
             assert len(zenith_distance) == 1, (
                 f"More than one zenith distance found in {file}"
             )