Model selector

pyproject.toml

@@ -15,14 +15,14 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy",
     "ctapipe==0.26.0",
-    "gammapy",
     "pyirf==0.13.0",
     "ctaplot==0.6.3",
     "h5py",
     "ctadata",
     "ctlearn",
     "astropy",
     "matplotlib==3.10.0",
+    "tables<3.9",
 ]
 
 # needed for setuptools_scm, we don"t define a static version

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):
@@ -953,48 +1023,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 +1178,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)

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__ = [
@@ -50,26 +51,31 @@ 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")
         project_directories = CTLMDirectories(self.project_directory, tri_model_nickname)
 
 
         direction_reco = tri_model_parameters.get("direction_reco", "cameradirection")
         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 = ['type', 'energy', direction_reco]
         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.")
-
+
+
         for reco, training_sample in zip(recos, training_samples):
             match reco:
                 case "type":
@@ -108,6 +114,32 @@ 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,

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

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

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):
         """