From 131369fc4abfc7c36e7545af9159827cbeab63dc Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Tue, 4 Aug 2026 15:53:13 +0200 Subject: [PATCH 01/12] fix tests with hardcopy --- ctlearn/conftest.py | 9 +- ctlearn/tools/__init__.py | 6 +- ctlearn/tools/keras/__init__.py | 6 + ctlearn/tools/keras/predict_model.py | 2369 +++++++++++++++++++++ ctlearn/tools/tests/test_predict_model.py | 68 +- pyproject.toml | 4 +- 6 files changed, 2419 insertions(+), 43 deletions(-) create mode 100644 ctlearn/tools/keras/predict_model.py diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index ee30af49..41dc2115 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -403,14 +403,10 @@ def ctlearn_trained_dl1_stereo_models( # Loop over reconstruction tasks and train models for each combination ctlearn_trained_dl1_stereo_models = {} for reco_task in ["type", "energy", "skydirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - # Build command-line arguments argv = [ f"--signal={signal_dir}", "--pattern-signal=*.dl1.h5", - f"--output={output_dir}", f"--reco={reco_task}", "--TrainCTLearnModel.n_epochs=1", "--TrainCTLearnModel.batch_size=2", @@ -432,7 +428,10 @@ def ctlearn_trained_dl1_stereo_models( # Run training tools for framework, training_tool in TRAINING_TOOLS.items(): - assert run_tool(training_tool(config=config), argv=argv, cwd=tmp_path) == 0 + framework_argv = argv.copy() + output_dir = tmp_path / f"ctlearn_{framework}_{telescope_type}_{reco_task}" + framework_argv.append(f"--output={output_dir}") + assert run_tool(training_tool(config=config), argv=framework_argv, cwd=tmp_path) == 0 ctlearn_trained_dl1_stereo_models[f"{framework}_{telescope_type}_{reco_task}"] = ( output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}" ) diff --git a/ctlearn/tools/__init__.py b/ctlearn/tools/__init__.py index d05c5556..145d76ff 100644 --- a/ctlearn/tools/__init__.py +++ b/ctlearn/tools/__init__.py @@ -2,10 +2,10 @@ """ from ctlearn.tools.predict_LST1 import LST1PredictionTool -from ctlearn.tools.predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel +from ctlearn.tools.keras.predict_model import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel __all__ = [ "LST1PredictionTool", - "MonoPredictCTLearnModel", - "StereoPredictCTLearnModel", + "MonoPredictCTLearnKerasModel", + "StereoPredictCTLearnKerasModel", ] \ No newline at end of file diff --git a/ctlearn/tools/keras/__init__.py b/ctlearn/tools/keras/__init__.py index e69de29b..06d24e06 100644 --- a/ctlearn/tools/keras/__init__.py +++ b/ctlearn/tools/keras/__init__.py @@ -0,0 +1,6 @@ +from .predict_model import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel + +__all__ = [ + "MonoPredictCTLearnKerasModel", + "StereoPredictCTLearnKerasModel", +] \ No newline at end of file diff --git a/ctlearn/tools/keras/predict_model.py b/ctlearn/tools/keras/predict_model.py new file mode 100644 index 00000000..f05b8203 --- /dev/null +++ b/ctlearn/tools/keras/predict_model.py @@ -0,0 +1,2369 @@ +""" +Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``. +""" + +import atexit +import uuid +import warnings + +import numpy as np +import tables +import tensorflow as tf +import keras + +from astropy import units as u +from astropy.coordinates.earth import EarthLocation +from astropy.coordinates import AltAz, SkyCoord +from astropy.table import ( + Table, + hstack, + vstack, + join, + setdiff, + unique, +) +from astropy.time import Time + +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) +from ctapipe.coordinates import CameraFrame, NominalFrame +from ctapipe.core import Tool +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + Int, + Path, + flag, + Dict, + ComponentName, + classes_with_traits, +) +from ctapipe.monitoring.interpolation import PointingInterpolator +from ctapipe.instrument import SubarrayDescription +from ctapipe.io import read_table, write_table, HDF5Merger +from ctapipe.io.datalevels import DataLevel +from ctapipe.io.hdf5dataformat import ( + DL0_TEL_POINTING_GROUP, + DL1_SUBARRAY_GROUP, + DL1_SUBARRAY_POINTING_GROUP, + DL1_SUBARRAY_TRIGGER_TABLE, + DL1_TEL_GROUP, + DL1_TEL_CALIBRATION_GROUP, + DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, + DL1_TEL_IMAGES_GROUP, + DL1_TEL_MUON_GROUP, + DL1_TEL_MUON_THROUGHPUT_GROUP, + DL1_TEL_OPTICAL_PSF_GROUP, + DL1_TEL_PARAMETERS_GROUP, + DL1_TEL_POINTING_GROUP, + DL1_TEL_TRIGGER_TABLE, + DL2_EVENT_STATISTICS_GROUP, + FIXED_POINTING_GROUP, + R0_TEL_GROUP, + R1_TEL_GROUP, + SIMULATION_IMAGES_GROUP, + SIMULATION_IMPACT_GROUP, + SIMULATION_PARAMETERS_GROUP, + SIMULATION_RUN_TABLE, + SIMULATION_SHOWER_TABLE, + DL2_TEL_PARTICLETYPE_GROUP, + DL2_TEL_ENERGY_GROUP, + DL2_TEL_GEOMETRY_GROUP, + DL2_SUBARRAY_GROUP, + DL2_SUBARRAY_PARTICLETYPE_GROUP, + DL2_SUBARRAY_ENERGY_GROUP, + DL2_SUBARRAY_GEOMETRY_GROUP, +) +from ctapipe.reco.reconstructor import ReconstructionProperty +from ctapipe.reco.stereo_combination import StereoCombiner +from ctapipe.reco.utils import add_defaults_and_meta +from dl1_data_handler.reader import ( + DLDataReader, + ProcessType, + LST_EPOCH, +) +from ctlearn import __version__ as ctlearn_version +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.utils import validate_trait_dict + +# Convienient constants for column names and table keys +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TEL_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] +TEL_ITER_GROUPS = [ + R0_TEL_GROUP, + R1_TEL_GROUP, + FIXED_POINTING_GROUP, + DL0_TEL_POINTING_GROUP, + DL1_TEL_POINTING_GROUP, + DL1_TEL_CALIBRATION_GROUP, + DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, + DL1_TEL_MUON_THROUGHPUT_GROUP, + DL1_TEL_OPTICAL_PSF_GROUP, + DL1_TEL_PARAMETERS_GROUP, + DL1_TEL_IMAGES_GROUP, + DL1_TEL_MUON_GROUP, + SIMULATION_IMAGES_GROUP, + SIMULATION_IMPACT_GROUP, + SIMULATION_PARAMETERS_GROUP, +] +DATALEVEL_TO_GROUP = { + DataLevel.R0: R0_TEL_GROUP, + DataLevel.R1: R1_TEL_GROUP, + DataLevel.DL1_IMAGES: DL1_TEL_IMAGES_GROUP, + DataLevel.DL1_PARAMETERS: DL1_TEL_PARAMETERS_GROUP, + DataLevel.DL1_MUON: DL1_TEL_MUON_GROUP, + DataLevel.DL2: DL2_SUBARRAY_GROUP, +} + + +class CannotPredict(OSError): + """Raised when trying to predict an incompatible file""" + + +class PredictCTLearnKerasModel(Tool): + """ + Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. + + This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image + or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to + store them in the output file. The input data is loaded from the input url using the + ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.keras.sequence.KerasSequence``. + The prediction is performed using the CTLearn models. The data is stored in the output file + following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to + handle the prediction for mono and stereo mode. + + Attributes + ---------- + input_url : pathlib.Path + Input ctapipe HDF5 files including pixel-wise image or waveform data. + dl1_features : bool + Set whether to include the dl1 feature vectors in the output file. + dl2_telescope : bool + Set whether to include dl2 telescope-event-wise data in the output file. + dl2_subarray : bool + Set whether to include dl2 subarray-event-wise data in the output file. + dl1dh_reader : dl1_data_handler.reader.DLDataReader + DLDataReader object to read the data. + dl1dh_reader_type : str + Type of the DLDataReader to use for the prediction. + stack_telescope_images : bool + Set whether to stack the telescope images in the data loader. Requires ``stereo``. + sort_by_intensity : bool + Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. + prefix : str + Name of the reconstruction algorithm used to generate the dl2 data. + load_type_model_from : pathlib.Path + Path to a Keras model file (Keras3) for the classification of the primary particle type. + load_energy_model_from : pathlib.Path + Path to a Keras model file (Keras3) for the regression of the primary particle energy. + load_cameradirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) for the regression + of the primary particle arrival direction based on camera coordinate offsets. + load_skydirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) for the regression + of the primary particle arrival direction based on spherical coordinate offsets. + output_path : pathlib.Path + Output path to save the dl2 prediction results. + keras_verbose : int + Verbosity mode of Keras during the prediction. + strategy : tf.distribute.Strategy + MirroredStrategy to distribute the prediction. + data_loader : ctlearn.core.keras.sequence.KerasSequence + KerasSequence object to load the data. + indices : list of int + List of indices for the data loaders. + batch_size : int + Size of the batch to perform inference of the neural network. + last_batch_size : int + Size of the last batch in the data loaders. + + Methods + ------- + setup() + Set up the tool. + finish() + Finish the tool. + _predict_with_model(model_path) + Load and predict with a CTLearn model. + _predict_particletype(example_identifiers) + Predict the classification of the primary particle type. + _predict_energy(example_identifiers) + Predict the energy of the primary particle. + _predict_cameradirection(example_identifiers) + Predict the arrival direction of the primary particle based on camera coordinate offsets. + _predict_skydirection(example_identifiers) + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + _transform_cam_coord_offsets_to_sky(table) + Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _transform_spher_coord_offsets_to_sky(table) + Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _create_nan_table(nonexample_identifiers, columns, shapes, reco_task) + Create a table with NaNs for missing predictions. + _store_pointing(all_identifiers) + Store the telescope pointing table from to the output file. + _create_feature_vectors_table(example_identifiers, nonexample_identifiers, particletype_feature_vectors, energy_feature_vectors, direction_feature_vectors) + Create the table for the DL1 feature vectors. + """ + + input_url = Path( + help="Input ctapipe HDF5 files including pixel-wise image or waveform data", + allow_none=True, + exists=True, + directory_ok=False, + file_ok=True, + ).tag(config=True) + + dl1_features = Bool( + default_value=False, + allow_none=False, + help="Set whether to include the dl1 feature vectors in the output file.", + ).tag(config=True) + + dl2_telescope = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 telescope-event-wise data in the output file.", + ).tag(config=True) + + dl2_subarray = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 subarray-event-wise data in the output file.", + ).tag(config=True) + + dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( + config=True + ) + + stack_telescope_images = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to stack the telescope images in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + sort_by_intensity = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to sort the telescope images by intensity in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + prefixes = Dict( + default_value={ + "type": "CTLearnClassifier", + "energy": "CTLearnRegressor", + "cameradirection": "CTLearnCameraReconstructor", + "skydirection": "CTLearnSkyReconstructor", + "all": "CTLearn", + }, + allow_none=False, + help=( + "Name of the reconstruction algorithm used " + "to generate the dl2 data for each task." + ), + ).tag(config=True) + + load_type_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) for the classification " + "of the primary particle type." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_energy_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) for the regression " + "of the primary particle energy." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_cameradirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) for the reconstruction " + "of the primary particle arrival direction based on camera coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_skydirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) for the reconstruction " + "of the primary particle arrival direction based on spherical coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + batch_size = Int( + default_value=64, + allow_none=False, + help="Size of the batch to perform inference of the neural network.", + ).tag(config=True) + + output_path = Path( + default_value="./output.dl2.h5", + allow_none=False, + help="Output path to save the dl2 prediction results", + ).tag(config=True) + + keras_verbose = Int( + default_value=1, + min=0, + max=2, + allow_none=False, + help=( + "Verbosity mode of Keras during the prediction: " + "0 = silent, 1 = progress bar, 2 = one line per call." + ), + ).tag(config=True) + + aliases = { + ("i", "input_url"): "PredictCTLearnKerasModel.input_url", + ("t", "type_model"): "PredictCTLearnKerasModel.load_type_model_from", + ("e", "energy_model"): "PredictCTLearnKerasModel.load_energy_model_from", + ( + "d", + "cameradirection_model", + ): "PredictCTLearnKerasModel.load_cameradirection_model_from", + ("s", "skydirection_model"): "PredictCTLearnKerasModel.load_skydirection_model_from", + ("o", "output"): "PredictCTLearnKerasModel.output_path", + } + + flags = { + **flag( + "overwrite", + "HDF5Merger.overwrite", + "Overwrite the output file if it exists", + "Do not overwrite the output file if it exists", + ), + **flag( + "dl1-features", + "PredictCTLearnKerasModel.dl1_features", + "Include dl1 features", + "Exclude dl1 features", + ), + **flag( + "dl2-telescope", + "PredictCTLearnKerasModel.dl2_telescope", + "Include dl2 telescope-event-wise data in the output file", + "Exclude dl2 telescope-event-wise data in the output file", + ), + **flag( + "dl2-subarray", + "PredictCTLearnKerasModel.dl2_subarray", + "Include dl2 subarray-event-wise data in the output file", + "Exclude dl2 subarray-event-wise data in the output file", + ), + **flag( + "r0-waveforms", + "HDF5Merger.r0_waveforms", + "Include r0 waveforms", + "Exclude r0 waveforms", + ), + **flag( + "r1-waveforms", + "HDF5Merger.r1_waveforms", + "Include r1 waveforms", + "Exclude r1 waveforms", + ), + **flag( + "dl1-parameters", + "HDF5Merger.dl1_parameters", + "Include dl1 parameters", + "Exclude dl1 parameters", + ), + **flag( + "dl1-images", + "HDF5Merger.dl1_images", + "Include dl1 images", + "Exclude dl1 images", + ), + **flag( + "true-parameters", + "HDF5Merger.true_parameters", + "Include true parameters", + "Exclude true parameters", + ), + **flag( + "true-images", + "HDF5Merger.true_images", + "Include true images", + "Exclude true images", + ), + } + + classes = classes_with_traits(DLDataReader) + + def setup(self): + self.activity_start_time = Time.now() + self.log.info("ctlearn version %s", ctlearn_version) + # Validate the prefixes trait dictionary + validate_trait_dict( + self.prefixes, ["type", "energy", "cameradirection", "skydirection", "all"] + ) + # Copy selected tables from the input file to the output file + self.log.info("Copying to output destination.") + with HDF5Merger( + self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self + ) as merger: + merger(self.input_url) + # Create a MirroredStrategy. + self.strategy = tf.distribute.MirroredStrategy() + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) + + # Set up the data reader + self.log.info("Loading data reader:") + self.log.info("For a large dataset, this may take a while...") + self.dl1dh_reader = DLDataReader.from_name( + self.dl1dh_reader_type, + input_url_signal=[self.input_url], + parent=self, + ) + self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) + # Check if the number of events is enough to form a batch + if self.dl1dh_reader._get_n_events() < self.batch_size: + raise ToolConfigurationError( + f"{self.dl1dh_reader._get_n_events()} events are not enough " + f"to form a batch of size {self.batch_size}. Reduce the batch size." + ) + # Set the indices for the data loaders + self.indices = list(range(self.dl1dh_reader._get_n_events())) + self.last_batch_size = len(self.indices) % ( + self.batch_size * self.strategy.num_replicas_in_sync + ) + # Ensure subarray consistency in the output file + self._ensure_subarray_consistency() + + def finish(self): + # Overwrite CTAO reference metadata to the output file + self._overwrite_meta() + self.log.info("Tool is shutting down") + + def _overwrite_meta(self): + """Overwrite CTAO metadata in the output file.""" + # TODO: Upgrade to new CTAO metatdata standard when available + with warnings.catch_warnings(): + warnings.simplefilter("ignore", tables.NaturalNameWarning) + with tables.open_file(self.output_path, mode="r+") as h5_file: + # Update CTA Activity metadata + h5_file.root._v_attrs["CTA ACTIVITY ID"] = str(uuid.uuid4()) + h5_file.root._v_attrs["CTA ACTIVITY NAME"] = self.name + h5_file.root._v_attrs["CTA ACTIVITY SOFTWARE NAME"] = "ctlearn" + h5_file.root._v_attrs["CTA ACTIVITY SOFTWARE VERSION"] = ctlearn_version + h5_file.root._v_attrs["CTA ACTIVITY START TIME"] = ( + self.activity_start_time.iso + ) + h5_file.root._v_attrs["CTA ACTIVITY STOP TIME"] = Time.now().iso + # Update CTA Product metadata + h5_file.root._v_attrs["CTA PRODUCT DATA LEVELS"] = ( + self._get_data_levels(h5_file) + ) + h5_file.root._v_attrs["CTA PRODUCT CREATION TIME"] = ( + self.activity_start_time.iso + ) + h5_file.root._v_attrs["CTA PRODUCT ID"] = str(uuid.uuid4()) + h5_file.flush() + + def _get_data_levels(self, h5file): + """Get the data levels present in the HDF5 file.""" + data_levels = { + level.name + for level, group in DATALEVEL_TO_GROUP.items() + if hasattr(h5file.root, group) + } + return ",".join(sorted(data_levels)) + + def _ensure_subarray_consistency(self): + """ + Align subarray metadata and trigger tables with the selected telescopes. + + When only a subset of telescopes is processed, overwrite the output file's + SubarrayDescription and trim the DL1 trigger tables to keep events that + involve the selected telescopes. Also rebuild the subarray trigger table + with the corresponding telescope participation masks. + """ + + input_subarray = SubarrayDescription.from_hdf( + self.input_url, + focal_length_choice=self.dl1dh_reader.focal_length_choice, + ) + if input_subarray == self.dl1dh_reader.subarray: + return + + # From the merger tool a SubarrayDescription for the full array is already stored + # in the output file. We need to remove it to avoid conflicts when storing + # the new SubarrayDescription for the selected telescopes. + with tables.open_file(self.output_path, mode="a") as h5file: + h5file.remove_node("/configuration/instrument", recursive=True) + selected_subarray = input_subarray.select_subarray(set(self.dl1dh_reader.tel_ids)) + selected_subarray.to_hdf(self.output_path) + self.log.info("SubarrayDescription was stored in '%s'", self.output_path) + + tel_trigger_table = read_table( + self.output_path, + DL1_TEL_TRIGGER_TABLE, + ) + mask = np.isin(tel_trigger_table["tel_id"], self.dl1dh_reader.tel_ids) + tel_trigger_table = tel_trigger_table[mask] + tel_trigger_table.sort(TEL_EVENT_KEYS) + + write_table( + tel_trigger_table, + self.output_path, + DL1_TEL_TRIGGER_TABLE, + overwrite=True, + ) + + subarray_trigger_table = tel_trigger_table.copy() + subarray_columns = SUBARRAY_EVENT_KEYS + ["time"] + # In older data formats the event type is not included in the trigger table, so we need to + # check if it is present before keeping the column to be backwards compatible. + if "event_type" in subarray_trigger_table.colnames: + subarray_columns.append("event_type") + subarray_trigger_table.keep_columns(subarray_columns) + subarray_trigger_table = unique( + subarray_trigger_table, keys=SUBARRAY_EVENT_KEYS + ) + + tel_trigger_groups = tel_trigger_table.group_by(SUBARRAY_EVENT_KEYS) + tel_with_trigger = [] + for tel_trigger in tel_trigger_groups.groups: + tel_with_trigger_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) + tel_with_trigger_mask[ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_trigger["tel_id"]) + ] = True + tel_with_trigger.append(tel_with_trigger_mask) + + subarray_trigger_table.add_column( + tel_with_trigger, index=-2, name="tels_with_trigger" + ) + + write_table( + subarray_trigger_table, + self.output_path, + DL1_SUBARRAY_TRIGGER_TABLE, + overwrite=True, + ) + # Update the simulation shower table to keep only events present in the subarray trigger table + subarray_trigger_table.keep_columns(SUBARRAY_EVENT_KEYS) + sim_shower_table = read_table( + self.output_path, + SIMULATION_SHOWER_TABLE, + ) + sim_shower_table = join( + sim_shower_table, + subarray_trigger_table, + keys=SUBARRAY_EVENT_KEYS, + join_type="right", + ) + sim_shower_table.sort(SUBARRAY_EVENT_KEYS) + write_table( + sim_shower_table, + self.output_path, + SIMULATION_SHOWER_TABLE, + overwrite=True, + ) + # Delete telescope-specific tables for unselected telescopes + self._delete_unselected_telescope_tables() + + def _delete_unselected_telescope_tables(self): + """ + Delete telescope-specific tables for unselected telescopes from the output file. + + Iterates through all telescope-related groups in the HDF5 file and removes + tables corresponding to telescopes that are not in the selected telescope list. + This ensures the output file only contains data for the telescopes that were + processed. The camera configuration tables are also pruned based on the camera indices. + """ + # Open the HDF5 file to prune the unselected telescope tables and camera configurations + with tables.open_file(self.output_path, mode="r+") as h5_file: + + def prune_group(group, valid_ids): + for table in group._f_iter_nodes("Table"): + idx = int(table._v_name.split("_")[-1]) + if idx not in valid_ids: + table._f_remove() + + # Telescope-specific tables + tel_ids = set(self.dl1dh_reader.tel_ids) + for group_name in TEL_ITER_GROUPS: + group = getattr(h5_file.root, group_name, None) + if group is not None: + prune_group(group, tel_ids) + + + def _create_nan_table(self, nonexample_identifiers, columns, shapes, reco_task): + """ + Create a table with NaNs for missing predictions. + + This method creates a table with NaNs for missing predictions for the non-example identifiers. + In stereo mode, the table also a column for the valid telescopes is added with all False values. + + Parameters: + ----------- + nonexample_identifiers : astropy.table.Table + Table containing the non-example identifiers. + columns : list of str + List of column names to create in the table. + shapes : list of shapes + List of shapes for the columns to create in the table. + reco_task : str + Reconstruction task name. + + Returns: + -------- + nan_table : astropy.table.Table + Table containing NaNs for missing predictions. + """ + # Create a table with NaNs for missing predictions + nan_table = nonexample_identifiers.copy() + for column_name, shape in zip(columns, shapes): + nan_table.add_column(np.full(shape, np.nan), name=column_name) + # Add that no telescope is valid for the non-example identifiers in stereo mode + if self.dl1dh_reader.mode == "stereo": + nan_table.add_column( + np.zeros( + (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ), + name=f"{self.prefixes[reco_task]}_telescopes", + ) + return nan_table + + def deduplicate_first_valid( + self, + table: Table, + keys=("obs_id", "event_id"), + valid_col="CTLearn_is_valid", + ): + """ + Return a deduplicated Astropy Table. + + For each group defined by `keys`, keep the first row where + `valid_col` is True. If none are valid, keep the first row. + """ + + t = table.copy() + + t.sort(list(keys) + [valid_col], reverse=[False] * len(keys) + [True]) + + return unique(t, keys=list(keys), keep="first") + + def _predict_with_model(self, model_path): + """ + Load and predict with a CTLearn model. + + Load a model from the specified path and predict the data using the loaded model. + If a last batch loader is provided, predict the last batch and stack the results. + + Parameters + ---------- + model_path : str + Path to a Keras model file (Keras3). + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + # Create a new KerasSequence for each task + # It turned out to be more robust to initialize the KerasSequence separately. + data_loader = KerasSequence( + self.dl1dh_reader, + self.indices, + tasks=[], + batch_size=self.batch_size * self.strategy.num_replicas_in_sync, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Keras is only considering the last complete batch. + # In prediction mode we don't want to loose the last + # uncomplete batch, so we are creating an additional + # batch generator for the remaining events. + data_loader_last_batch = None + if self.last_batch_size > 0: + last_batch_indices = self.indices[-self.last_batch_size :] + data_loader_last_batch = KerasSequence( + self.dl1dh_reader, + last_batch_indices, + tasks=[], + batch_size=self.last_batch_size, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Load the model from the specified path + model = keras.saving.load_model(model_path) + prediction_colname = ( + "type" + if isinstance(model.layers[-1], keras.layers.Softmax) + else model.layers[-1].name + ) + backbone_model, feature_vectors = None, None + if self.dl1_features: + # Get the backbone model which is the second layer of the model + backbone_model = model.get_layer(index=1) + # Create a new head model with the same layers as the original model. + # The output of the backbone model is the input of the head model. + backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) + x = backbone_output_shape + for layer in model.layers[2:]: + x = layer(x) + head = keras.Model(inputs=backbone_output_shape, outputs=x) + # Apply the backbone model with the data loader to retrieve the feature vectors + try: + feature_vectors = backbone_model.predict( + data_loader, verbose=self.keras_verbose + ) + except ValueError as err: + if str(err).startswith("Input 0 of layer"): + raise ToolConfigurationError( + "Model input shape does not match the prediction data. " + "This is usually caused by selecting the wrong telescope_id. " + "Please ensure the telescope configuration matches the one used for training." + ) from err + raise + # Apply the head model with the feature vectors to retrieve the prediction + predict_data = Table( + { + prediction_colname: head.predict( + feature_vectors, verbose=self.keras_verbose + ) + } + ) + # Predict the last batch and stack the results to the prediction data + if data_loader_last_batch is not None: + feature_vectors_last_batch = backbone_model.predict( + data_loader_last_batch, verbose=self.keras_verbose + ) + feature_vectors = np.concatenate( + (feature_vectors, feature_vectors_last_batch) + ) + predict_data = vstack( + [ + predict_data, + Table( + { + prediction_colname: head.predict( + feature_vectors_last_batch, + verbose=self.keras_verbose, + ) + } + ), + ] + ) + else: + # Predict the data using the loaded model + try: + predict_data = model.predict(data_loader, verbose=self.keras_verbose) + except ValueError as err: + if str(err).startswith("Input 0 of layer"): + raise ToolConfigurationError( + "Model input shape does not match the prediction data. " + "This is usually caused by selecting the wrong telescope_id. " + "Please ensure the telescope configuration matches the one used for training." + ) from err + raise + # Create a astropy table with the prediction results + # The classification task has a softmax layer as the last layer + # which returns the probabilities for each class in an array, while + # the regression tasks have output neurons which returns the + # predicted value for the task in a dictionary. + if prediction_colname == "type": + predict_data = Table({prediction_colname: predict_data}) + else: + predict_data = Table(predict_data) + # Predict the last batch and stack the results to the prediction data + if data_loader_last_batch is not None: + predict_data_last_batch = model.predict( + data_loader_last_batch, verbose=self.keras_verbose + ) + if model.layers[-1].name == "type": + predict_data_last_batch = Table( + {prediction_colname: predict_data_last_batch} + ) + else: + predict_data_last_batch = Table(predict_data_last_batch) + predict_data = vstack([predict_data, predict_data_last_batch]) + return predict_data, feature_vectors + + def _predict_particletype(self, example_identifiers): + """ + Predict the classification of the primary particle type. + + This method uses a pre-trained type model to predict the type of the primary particle + for a given set of example identifiers. The predicted classification score ('gammaness') + is added to the example identifiers table. + + Parameters: + ----------- + particletype_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + predicted classification score ('gammaness'). + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the classification of the primary particle type..." + ) + # Predict the data using the loaded type_model + predict_data, feature_vectors = self._predict_with_model( + self.load_type_model_from + ) + # Create prediction table and add the predicted classification score ('gammaness') + particletype_table = example_identifiers.copy() + particletype_table.add_column( + predict_data["type"].T[1], name=f"{self.prefixes['type']}_tel_prediction" + ) + return particletype_table, feature_vectors + + def _predict_energy(self, example_identifiers): + """ + Predict the energy of the primary particle. + + This method uses a pre-trained energy model to predict the energy of the primary particle + for a given set of example identifiers. The predicted energy is then converted from + log10(TeV) to TeV and added to the example identifiers table. + + Parameters: + ----------- + energy_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed energy in TeV. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info("Predicting for the regression of the primary particle energy...") + # Predict the data using the loaded energy_model + predict_data, feature_vectors = self._predict_with_model( + self.load_energy_model_from + ) + # Convert the reconstructed energy from log10(TeV) to TeV + reco_energy = u.Quantity( + np.power(10, np.squeeze(predict_data["energy"])), + unit=u.TeV, + ) + # Create prediction table and add the reconstructed energy in TeV + energy_table = example_identifiers.copy() + energy_table.add_column( + reco_energy, name=f"{self.prefixes['energy']}_tel_energy" + ) + return energy_table, feature_vectors + + def _predict_cameradirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on camera coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the + primary particle for a given set of example identifiers. The predicted camera coordinate offsets + is added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + cameradirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed camera coordinate offsets in x and y. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the reconstruction of the primary particle arrival direction based on camera coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_cameradirection_model_from + ) + # For the direction task, the prediction is the camera coordinate offset in x and y + # from the telescope pointing. + cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) + cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) + # Create prediction table and add the reconstructed energy in TeV + cameradirection_table = example_identifiers.copy() + cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") + cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") + return cameradirection_table, feature_vectors + + def _predict_skydirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the primary + particle for a given set of example identifiers. The predicted spherical coordinate offsets is + added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + skydirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed spherical coordinate offsets in fov_lon and fov_lat. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the reconstruction of the primary particle arrival direction based on spherical coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_skydirection_model_from + ) + # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat + # from the telescope pointing. + fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) + fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) + # Create prediction table and add the reconstructed fov_lon and fov_lat + skydirection_table = example_identifiers.copy() + skydirection_table.add_column(fov_lon, name="fov_lon") + skydirection_table.add_column(fov_lat, name="fov_lat") + return skydirection_table, feature_vectors + + def _transform_cam_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + # Get the telescope ID from the table + tel_id = table["tel_id"][0] + # Set the telescope position + tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the telescope pointing with the trigger timestamp and the telescope position + altaz = AltAz( + location=tel_ground_frame.to_earth_location(), + obstime=trigger_time, + ) + # Set the telescope pointing + tel_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the camera frame with the focal length and rotation of the camera + camera_frame = CameraFrame( + focal_length=self.dl1dh_reader.subarray.tel[ + tel_id + ].camera.geometry.frame.focal_length, + rotation=self.dl1dh_reader.pix_rotation[tel_id], + telescope_pointing=tel_pointing, + ) + # Set the camera coordinate offset + cam_coord_offset = SkyCoord( + x=table["cam_coord_offset_x"], + y=table["cam_coord_offset_y"], + frame=camera_frame, + ) + # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Transform the true Alt/Az coordinates to camera coordinates + reco_direction = cam_coord_offset.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column( + reco_direction.az.to(u.deg), + name=f"{self.prefixes['cameradirection']}_tel_az", + ) + table.add_column( + reco_direction.alt.to(u.deg), + name=f"{self.prefixes['cameradirection']}_tel_alt", + ) + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "cam_coord_offset_x", + "cam_coord_offset_y", + ] + ) + return table + + def _transform_spher_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted spherical offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted spherical offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the AltAz frame with the reference location and time + altaz = AltAz( + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the array pointing + array_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the nominal frame with the array pointing + nom_frame = NominalFrame( + origin=array_pointing, + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the reco direction in (fov_lon, fov_lat) coordinates + reco_direction = SkyCoord( + fov_lon=table["fov_lon"], + fov_lat=table["fov_lat"], + frame=nom_frame, + ) + # Transform the reco direction from nominal frame to the AltAz frame + sky_coord = reco_direction.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column( + sky_coord.az.to(u.deg), name=f"{self.prefixes['skydirection']}_az" + ) + table.add_column( + sky_coord.alt.to(u.deg), name=f"{self.prefixes['skydirection']}_alt" + ) + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "fov_lon", + "fov_lat", + ] + ) + return table + + def _store_pointing(self, all_identifiers): + """ + Store the telescope pointing table from to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + + # Initialize the pointing interpolator from ctapipe + pointing_interpolator = PointingInterpolator( + bounds_error=False, extrapolate=True + ) + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Get the telescope pointing from the dl1dh reader + tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] + # Add the telescope pointing table to the pointing interpolator + pointing_interpolator.add_table(tel_id, tel_pointing) + tel_identifiers = all_identifiers.copy() + if self.dl1dh_reader.mode == "mono": + tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Interpolate the telescope pointing + tel_altitude, tel_azimuth = pointing_interpolator( + tel_id, tel_identifiers["time"] + ) + tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") + tel_identifiers.add_column(tel_altitude, name="pointing_altitude") + pointing_info.append(tel_identifiers) + if self.dl1dh_reader.mode == "mono": + tel_pointing_table = Table( + { + "time": tel_identifiers["time"], + "azimuth": tel_identifiers["pointing_azimuth"], + "altitude": tel_identifiers["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info = vstack(pointing_info) + if self.dl1dh_reader.mode == "stereo": + # Group the pointing information by subarray event keys + # TODO: This needs to be debugged with SST1M data + pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) + pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) + pointing_info = join( + all_identifiers, + pointing_mean, + keys=SUBARRAY_EVENT_KEYS, + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + DL1_SUBARRAY_POINTING_GROUP, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + DL1_SUBARRAY_POINTING_GROUP, + ) + return pointing_info + + def _create_feature_vectors_table( + self, + example_identifiers, + nonexample_identifiers=None, + particletype_feature_vectors=None, + energy_feature_vectors=None, + direction_feature_vectors=None, + ): + """ + Create the table for the DL1 feature vectors. + + This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for + non-example identifiers. The feature vectors are stored in the columns of the table. The table also + contains a column for the valid predictions. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + nonexample_identifiers : astropy.table.Table or None + Table containing the non-example identifiers to fill the NaNs. + particletype_feature_vectors : np.ndarray or None + Array containing the particletype feature vectors. + energy_feature_vectors : np.ndarray or None + Array containing the energy feature vectors. + direction_feature_vectors : np.ndarray or None + Array containing the direction feature vectors. + + Returns: + -------- + feature_vector_table : astropy.table.Table + Table containing the DL1 feature vectors for the example and non-example identifiers. + """ + # Create the feature vector table + feature_vector_table = example_identifiers.copy() + columns_list, shapes_list = [], [] + if particletype_feature_vectors is not None: + is_valid_col = ~np.isnan( + np.min(particletype_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + particletype_feature_vectors, + name=f"{self.prefixes['all']}_tel_particletype_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append( + f"{self.prefixes['all']}_tel_particletype_feature_vectors" + ) + shapes_list.append( + ( + len(nonexample_identifiers), + particletype_feature_vectors.shape[1], + ) + ) + if energy_feature_vectors is not None: + is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) + feature_vector_table.add_column( + energy_feature_vectors, + name=f"{self.prefixes['all']}_tel_energy_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append( + f"{self.prefixes['all']}_tel_energy_feature_vectors" + ) + shapes_list.append( + ( + len(nonexample_identifiers), + energy_feature_vectors.shape[1], + ) + ) + if direction_feature_vectors is not None: + feature_vector_table.remove_columns( + ["pointing_azimuth", "pointing_altitude", "time"] + ) + is_valid_col = ~np.isnan( + np.min(direction_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + direction_feature_vectors, + name=f"{self.prefixes['all']}_tel_geometry_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append( + f"{self.prefixes['all']}_tel_geometry_feature_vectors" + ) + shapes_list.append( + ( + len(nonexample_identifiers), + direction_feature_vectors.shape[1], + ) + ) + # Produce output table with NaNs for missing predictions + if nonexample_identifiers is not None: + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=columns_list, + shapes=shapes_list, + reco_task="all", + ) + feature_vector_table = vstack([feature_vector_table, nan_table]) + is_valid_col = np.concatenate( + (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) + ) + # Add is_valid column to the feature vector table + feature_vector_table.add_column( + is_valid_col, + name=f"{self.prefixes['all']}_tel_is_valid", + ) + return feature_vector_table + + +class MonoPredictCTLearnKerasModel(PredictCTLearnKerasModel): + """ + Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. + + This tool extends the ``PredictCTLearnKerasModel`` to specifically handle monoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_telescope_pointing(all_identifiers) + Store the telescope pointing table for the mono mode for MC simulation. + """ + + name = "ctlearn-predict-mono-keras-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnKerasModel.batch_size=64 \\ + --PredictCTLearnKerasModel.dl1dh_reader_type=DLImageReader \\ + --DLImageReader.channels=cleaned_image \\ + --DLImageReader.channels=cleaned_relative_peak_time \\ + --DLImageReader.image_mapper_type=BilinearMapper \\ + --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ + --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ + --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ + --dl1-features \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + + To predict from pixel-wise waveform data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.r1.h5 \\ + --PredictCTLearnKerasModel.dl1dh_reader_type=DLWaveformReader \\ + --DLWaveformReader.sequnce_length=20 \\ + --DLWaveformReader.image_mapper_type=BilinearMapper \\ + --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ + --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ + --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ + --no-r0-waveforms \\ + --no-r1-waveforms \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + """ + + stereo_combiner_cls = ComponentName( + StereoCombiner, + default_value="StereoMeanCombiner", + help="Which stereo combination method to use after the monoscopic reconstruction.", + ).tag(config=True) + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.example_identifiers.copy() + example_identifiers.keep_columns(TEL_EVENT_KEYS) + all_identifiers = read_table( + self.output_path, + DL1_TEL_TRIGGER_TABLE, + ) + all_identifiers.keep_columns(TEL_EVENT_KEYS + ["time"]) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=TEL_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Pointing table for the mono mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_telescope_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + particletype_feature_vectors = None + if self.load_type_model_from is not None: + # Predict the type of the primary particle + particletype_table, particletype_feature_vectors = ( + super()._predict_particletype(example_identifiers) + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefixes['type']}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + reco_task="type", + ) + particletype_table = vstack([particletype_table, nan_table]) + # Add is_valid column to the particle type table + particletype_table.add_column( + ~np.isnan( + particletype_table[f"{self.prefixes['type']}_tel_prediction"].data, + dtype=bool, + ), + name=f"{self.prefixes['type']}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + particletype_table, + ParticleClassificationContainer, + prefix=self.prefixes["type"], + add_tel_prefix=True, + ) + if self.dl2_telescope: + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = particletype_table["tel_id"] == tel_id + particletype_tel_table = particletype_table[telescope_mask] + particletype_tel_table.sort(TEL_EVENT_KEYS) + # Save the prediction to the output file for the selected telescope + write_table( + particletype_tel_table, + self.output_path, + f"{DL2_TEL_PARTICLETYPE_GROUP}/{self.prefixes['type']}/tel_{tel_id:03d}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TEL_PARTICLETYPE_GROUP}/{self.prefixes['type']}/tel_{tel_id:03d}", + ) + + if self.dl2_subarray: + self.log.info("Processing and storing the subarray type prediction...") + # If only one telescope is used, copy the particletype table + # and modify it to subarray format + if len(self.dl1dh_reader.tel_ids) == 1: + particletype_subarray_table = particletype_table.copy() + telescope_mask = ( + particletype_subarray_table["tel_id"] + == self.dl1dh_reader.tel_ids[0] + ) + particletype_subarray_table = particletype_subarray_table[ + telescope_mask + ] + particletype_subarray_table.remove_column("tel_id") + for colname in particletype_subarray_table.colnames: + if "_tel_" in colname: + particletype_subarray_table.rename_column( + colname, colname.replace("_tel", "") + ) + particletype_subarray_table.add_column( + [ + [val] + for val in particletype_subarray_table[ + f"{self.prefixes['type']}_is_valid" + ] + ], + name=f"{self.prefixes['type']}_telescopes", + ) + else: + self.type_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefixes["type"], + property=ReconstructionProperty.PARTICLE_TYPE, + parent=self, + ) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + particletype_subarray_table = ( + self.type_stereo_combiner.predict_table(particletype_table) + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + particletype_subarray_table[ + f"{self.prefixes['type']}_telescopes" + ].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + ( + len(particletype_subarray_table), + len(self.dl1dh_reader.tel_ids), + ), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + particletype_subarray_table[ + f"{self.prefixes['type']}_telescopes" + ] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices( + tel_id + ) + ] = True + # Overwrite the column with the boolean mask with fix length + particletype_subarray_table[ + f"{self.prefixes['type']}_telescopes" + ] = reco_telescopes + # Deduplicate the subarray particletype table to have only one entry per event + particletype_subarray_table = super().deduplicate_first_valid( + table=particletype_subarray_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['type']}_is_valid", + ) + # Sort the subarray particletype table + particletype_subarray_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + particletype_subarray_table, + self.output_path, + f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefixes['energy']}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + reco_task="energy", + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefixes['energy']}_tel_energy"].data, + dtype=bool, + ), + name=f"{self.prefixes['energy']}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefixes["energy"], + add_tel_prefix=True, + ) + if self.dl2_telescope: + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = energy_table["tel_id"] == tel_id + energy_tel_table = energy_table[telescope_mask] + energy_tel_table.sort(TEL_EVENT_KEYS) + # Save the prediction to the output file + write_table( + energy_tel_table, + self.output_path, + f"{DL2_TEL_ENERGY_GROUP}/{self.prefixes['energy']}/tel_{tel_id:03d}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TEL_ENERGY_GROUP}/{self.prefixes['energy']}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray energy prediction..." + ) + # If only one telescope is used, copy the particletype table + # and modify it to subarray format + if len(self.dl1dh_reader.tel_ids) == 1: + energy_subarray_table = energy_table.copy() + telescope_mask = ( + energy_subarray_table["tel_id"] == self.dl1dh_reader.tel_ids[0] + ) + energy_subarray_table = energy_subarray_table[telescope_mask] + energy_subarray_table.remove_column("tel_id") + for colname in energy_subarray_table.colnames: + if "_tel_" in colname: + energy_subarray_table.rename_column( + colname, colname.replace("_tel", "") + ) + energy_subarray_table.add_column( + [ + [val] + for val in energy_subarray_table[ + f"{self.prefixes['energy']}_is_valid" + ] + ], + name=f"{self.prefixes['energy']}_telescopes", + ) + else: + self.energy_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefixes["energy"], + property=ReconstructionProperty.ENERGY, + parent=self, + ) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + energy_subarray_table = self.energy_stereo_combiner.predict_table( + energy_table + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + energy_subarray_table[ + f"{self.prefixes['energy']}_telescopes" + ].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + ( + len(energy_subarray_table), + len(self.dl1dh_reader.tel_ids), + ), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + energy_subarray_table[ + f"{self.prefixes['energy']}_telescopes" + ] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices( + tel_id + ) + ] = True + # Overwrite the column with the boolean mask with fix length + energy_subarray_table[ + f"{self.prefixes['energy']}_telescopes" + ] = reco_telescopes + # Deduplicate the subarray energy table to have only one entry per event + energy_subarray_table = super().deduplicate_first_valid( + table=energy_subarray_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['energy']}_is_valid", + ) + # Sort the subarray energy table + energy_subarray_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + energy_subarray_table, + self.output_path, + f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", + ) + direction_feature_vectors = None + if self.load_cameradirection_model_from is not None: + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=TEL_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = ( + super()._predict_cameradirection(example_identifiers) + ) + direction_tel_tables = [] + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = direction_table["tel_id"] == tel_id + direction_tel_table = direction_table[telescope_mask] + direction_tel_table = super()._transform_cam_coord_offsets_to_sky( + direction_tel_table + ) + # Produce output table with NaNs for missing predictions + nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id + nonexample_identifiers_tel = nonexample_identifiers[nan_telescope_mask] + if len(nonexample_identifiers_tel) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers_tel, + columns=[ + f"{self.prefixes['cameradirection']}_tel_alt", + f"{self.prefixes['cameradirection']}_tel_az", + ], + shapes=[ + (len(nonexample_identifiers_tel),), + (len(nonexample_identifiers_tel),), + ], + reco_task="cameradirection", + ) + direction_tel_table = vstack([direction_tel_table, nan_table]) + direction_tel_table.sort(TEL_EVENT_KEYS) + # Add is_valid column to the direction table + direction_tel_table.add_column( + ~np.isnan( + direction_tel_table[ + f"{self.prefixes['cameradirection']}_tel_alt" + ].data, + dtype=bool, + ), + name=f"{self.prefixes['cameradirection']}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_tel_table, + ReconstructedGeometryContainer, + prefix=self.prefixes["cameradirection"], + add_tel_prefix=True, + ) + direction_tel_tables.append(direction_tel_table) + if self.dl2_telescope: + # Save the prediction to the output file + write_table( + direction_tel_table, + self.output_path, + f"{DL2_TEL_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}/tel_{tel_id:03d}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TEL_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray geometry prediction..." + ) + # Stack the telescope tables to the subarray table + direction_tel_tables = vstack(direction_tel_tables) + # Sort the table by the telescope event keys + direction_tel_tables.sort(TEL_EVENT_KEYS) + # If only one telescope is used, copy the classification table + # and modify it to subarray format + if len(self.dl1dh_reader.tel_ids) == 1: + direction_subarray_table = direction_tel_tables.copy() + telescope_mask = ( + direction_subarray_table["tel_id"] + == self.dl1dh_reader.tel_ids[0] + ) + direction_subarray_table = direction_subarray_table[telescope_mask] + direction_subarray_table.remove_column("tel_id") + for colname in direction_subarray_table.colnames: + if "_tel_" in colname: + direction_subarray_table.rename_column( + colname, colname.replace("_tel", "") + ) + direction_subarray_table.add_column( + [ + [val] + for val in direction_subarray_table[ + f"{self.prefixes['cameradirection']}_is_valid" + ] + ], + name=f"{self.prefixes['cameradirection']}_telescopes", + ) + else: + self.geometry_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefixes["cameradirection"], + property=ReconstructionProperty.GEOMETRY, + parent=self, + ) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + direction_subarray_table = ( + self.geometry_stereo_combiner.predict_table( + direction_tel_tables + ) + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + direction_subarray_table[ + f"{self.prefixes['cameradirection']}_telescopes" + ].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + ( + len(direction_subarray_table), + len(self.dl1dh_reader.tel_ids), + ), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + direction_subarray_table[ + f"{self.prefixes['cameradirection']}_telescopes" + ] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices( + tel_id + ) + ] = True + # Overwrite the column with the boolean mask with fix length + direction_subarray_table[ + f"{self.prefixes['cameradirection']}_telescopes" + ] = reco_telescopes + # Deduplicate the subarray direction table to have only one entry per event + direction_subarray_table = super().deduplicate_first_valid( + table=direction_subarray_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['cameradirection']}_is_valid", + ) + # Sort the subarray geometry table + direction_subarray_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + direction_subarray_table, + self.output_path, + f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['cameradirection']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['cameradirection']}", + ) + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + particletype_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TEL_GROUP/features/{prefix}/tel_{tel_id:03d} table. + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = feature_vector_table["tel_id"] == tel_id + feature_vectors_tel_table = feature_vector_table[telescope_mask] + feature_vectors_tel_table.sort(TEL_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vectors_tel_table, + self.output_path, + f"{DL1_TEL_GROUP}/features/{self.prefixes['all']}/tel_{tel_id:03d}", + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TEL_GROUP}/features/{self.prefixes['all']}/tel_{tel_id:03d}", + ) + + def _store_mc_telescope_pointing(self, all_identifiers): + """ + Store the telescope pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + # Create the pointing table for each telescope + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Pointing table for the mono mode + tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) + tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") + tel_pointing.rename_column( + "telescope_pointing_altitude", "pointing_altitude" + ) + # Join the prediction table with the telescope pointing table + tel_pointing = join( + left=tel_pointing, + right=all_identifiers, + keys=["obs_id", "tel_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + tel_pointing.sort(TEL_EVENT_KEYS) + # Retrieve the example identifiers for the selected telescope + tel_pointing_table = Table( + { + "time": tel_pointing["time"], + "azimuth": tel_pointing["pointing_azimuth"], + "altitude": tel_pointing["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info.append(tel_pointing) + pointing_info = vstack(pointing_info) + return pointing_info + + +class StereoPredictCTLearnKerasModel(PredictCTLearnKerasModel): + """ + Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. + + This tool extends the ``PredictCTLearnKerasModel`` to specifically handle stereoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_subarray_pointing(all_identifiers) + Store the subarray pointing table for the stereo mode for MC simulation. + """ + + name = "ctlearn-predict-stereo-keras-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in stereo mode using trained CTLearn models: + > ctlearn-predict-stereo-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnKerasModel.batch_size=16 \\ + --PredictCTLearnKerasModel.dl1dh_reader_type=DLImageReader \\ + --DLImageReader.channels=cleaned_image \\ + --DLImageReader.channels=cleaned_relative_peak_time \\ + --DLImageReader.image_mapper_type=BilinearMapper \\ + --DLImageReader.mode=stereo \\ + --DLImageReader.min_telescopes=2 \\ + --PredictCTLearnKerasModel.stack_telescope_images=True \\ + --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ + --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ + --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ + --output output.dl2.h5 \\ + """ + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() + example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) + all_identifiers = read_table( + self.output_path, + DL1_TEL_TRIGGER_TABLE, + ) + all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) + # Unique example identifiers by events + all_identifiers = unique(all_identifiers, keys=SUBARRAY_EVENT_KEYS) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Construct the survival telescopes for each event of the example_identifiers + survival_telescopes = [] + for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: + survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) + survival_tels = [ + self.dl1dh_reader.subarray.tel_indices[tel_id] + for tel_id in subarray_event["tel_id"].data + ] + survival_mask[survival_tels] = True + survival_telescopes.append(survival_mask) + # Add the survival telescopes to the example_identifiers + example_identifiers.add_column( + survival_telescopes, name=f"{self.prefixes['all']}_telescopes" + ) + # Pointing table for the stereo mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_subarray_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + particletype_feature_vectors = None + if self.load_type_model_from is not None: + # Predict the classification of the primary particle + particletype_table, particletype_feature_vectors = ( + super()._predict_particletype(example_identifiers) + ) + if self.dl2_subarray: + particletype_table.rename_column( + f"{self.prefixes['all']}_telescopes", + f"{self.prefixes['type']}_telescopes", + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefixes['type']}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + reco_task="type", + ) + particletype_table = vstack([particletype_table, nan_table]) + # Add is_valid column to the particletype table + particletype_table.add_column( + ~np.isnan( + particletype_table[ + f"{self.prefixes['type']}_tel_prediction" + ].data, + dtype=bool, + ), + name=f"{self.prefixes['type']}_is_valid", + ) + # Rename the columns for the stereo mode + particletype_table.rename_column( + f"{self.prefixes['type']}_tel_prediction", + f"{self.prefixes['type']}_prediction", + ) + # Deduplicate the subarray particletype table to have only one entry per event + particletype_table = super().deduplicate_first_valid( + table=particletype_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['type']}_is_valid", + ) + particletype_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + particletype_table, + ParticleClassificationContainer, + prefix=self.prefixes["type"], + ) + # Save the prediction to the output file + write_table( + particletype_table, + self.output_path, + f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + if self.dl2_subarray: + energy_table.rename_column( + f"{self.prefixes['all']}_telescopes", + f"{self.prefixes['energy']}_telescopes", + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefixes['energy']}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + reco_task="energy", + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefixes['energy']}_tel_energy"].data, + dtype=bool, + ), + name=f"{self.prefixes['energy']}_is_valid", + ) + # Rename the columns for the stereo mode + energy_table.rename_column( + f"{self.prefixes['energy']}_tel_energy", + f"{self.prefixes['energy']}_energy", + ) + # Deduplicate the subarray energy table to have only one entry per event + energy_table = super().deduplicate_first_valid( + table=energy_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['energy']}_is_valid", + ) + energy_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefixes["energy"], + ) + # Save the prediction to the output file + write_table( + energy_table, + self.output_path, + f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", + ) + direction_feature_vectors = None + if self.load_skydirection_model_from is not None: + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=SUBARRAY_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = super()._predict_skydirection( + example_identifiers + ) + if self.dl2_subarray: + direction_table.rename_column( + f"{self.prefixes['all']}_telescopes", + f"{self.prefixes['skydirection']}_telescopes", + ) + # Transform the spherical coordinate offsets to sky coordinates + direction_table = super()._transform_spher_coord_offsets_to_sky( + direction_table + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[ + f"{self.prefixes['skydirection']}_alt", + f"{self.prefixes['skydirection']}_az", + ], + shapes=[ + (len(nonexample_identifiers),), + (len(nonexample_identifiers),), + ], + reco_task="skydirection", + ) + direction_table = vstack([direction_table, nan_table]) + # Add is_valid column to the direction table + direction_table.add_column( + ~np.isnan( + direction_table[f"{self.prefixes['skydirection']}_alt"].data, + dtype=bool, + ), + name=f"{self.prefixes['skydirection']}_is_valid", + ) + # Deduplicate the subarray direction table to have only one entry per event + direction_table = super().deduplicate_first_valid( + table=direction_table, + keys=SUBARRAY_EVENT_KEYS, + valid_col=f"{self.prefixes['skydirection']}_is_valid", + ) + direction_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_table, + ReconstructedGeometryContainer, + prefix=self.prefixes["skydirection"], + ) + # Save the prediction to the output file + write_table( + direction_table, + self.output_path, + f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['skydirection']}", + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['skydirection']}", + ) + # Store the telescope event statistics table + write_table( + self.dl1dh_reader.quality_query.to_table(functions=True), + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['skydirection']}", + append=True, + ) + self.log.info( + "DL2 service telescope event statistics data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['skydirection']}", + ) + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + particletype_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TEL_GROUP/features/{self.prefixes['all']}/tel_{tel_id:03d} table. + # Rename the columns for the stereo mode + feature_vector_table.rename_column( + f"{self.prefixes['all']}_tel_particletype_feature_vectors", + f"{self.prefixes['all']}_particletype_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefixes['all']}_tel_energy_feature_vectors", + f"{self.prefixes['all']}_energy_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefixes['all']}_tel_geometry_feature_vectors", + f"{self.prefixes['all']}_geometry_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefixes['all']}_tel_is_valid", + f"{self.prefixes['all']}_is_valid", + ) + feature_vector_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vector_table, + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefixes['all']}", + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefixes['all']}", + ) + + def _store_mc_subarray_pointing(self, all_identifiers): + """ + Store the subarray pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the subarray pointing information. + """ + # Read the subarray pointing table + pointing_info = read_table( + self.input_url, + SIMULATION_RUN_TABLE, + ) + # Assuming min_az = max_az and min_alt = max_alt + pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) + pointing_info.rename_column("min_az", "pointing_azimuth") + pointing_info.rename_column("min_alt", "pointing_altitude") + # Join the prediction table with the telescope pointing table + pointing_info = join( + left=pointing_info, + right=all_identifiers, + keys=["obs_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + DL1_SUBARRAY_POINTING_GROUP, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + DL1_SUBARRAY_POINTING_GROUP, + ) + return pointing_info + + +def mono_tool(): + # Run the tool + mono_tool = MonoPredictCTLearnKerasModel() + mono_tool.run() + + +def stereo_tool(): + # Run the tool + stereo_tool = StereoPredictCTLearnKerasModel() + stereo_tool.run() + + +if __name__ == "mono_tool": + mono_tool() + +if __name__ == "stereo_tool": + stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/tests/test_predict_model.py b/ctlearn/tools/tests/test_predict_model.py index 81093c6a..69f4b060 100644 --- a/ctlearn/tools/tests/test_predict_model.py +++ b/ctlearn/tools/tests/test_predict_model.py @@ -2,11 +2,10 @@ import numpy as np import pytest -pytest.importorskip("tensorflow") - from ctapipe.core import run_tool from ctapipe.io import TableLoader -from ctlearn.tools import MonoPredictCTLearnModel, StereoPredictCTLearnModel +from ctlearn.conftest import MODEL_FILE_FORMATS +from ctlearn.tools.keras import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel # Columns that should be present in the output DL2 file REQUIRED_COLUMNS = [ @@ -38,8 +37,9 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2") +@pytest.mark.parametrize("framework", ["Keras"]) def test_predict_mono_model_with_r1_waveforms( - tmp_path, ctlearn_trained_r1_mono_models, r1_gamma_file + tmp_path, ctlearn_trained_r1_mono_models, r1_gamma_file, framework ): """ Test training CTLearn mono model using the R1 gamma and proton files for all reconstruction tasks @@ -57,7 +57,7 @@ def test_predict_mono_model_with_r1_waveforms( # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_r1_mono_models[key], model_dir / f"ctlearn_mono_model_{key}.keras", @@ -67,24 +67,24 @@ def test_predict_mono_model_with_r1_waveforms( # Build command-line arguments argv = [ f"--input_url={r1_gamma_file}", - "--PredictCTLearnModel.batch_size=2", - "--PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader", + "--PredictCTLearnKerasModel.batch_size=2", + "--PredictCTLearnKerasModel.dl1dh_reader_type=DLWaveformReader", "--DLWaveformReader.sequence_length=5", "--DLWaveformReader.focal_length_choice=EQUIVALENT", "--no-r1-waveforms", "--dl2-telescope", ] - output_file = dl2_dir / f"gamma_{telescope_type}_mono_from_waveforms.dl2.h5" + output_file = dl2_dir / f"gamma_{framework}_{telescope_type}_mono_from_waveforms.dl2.h5" # Run Prediction tool assert ( run_tool( - MonoPredictCTLearnModel(), + MonoPredictCTLearnKerasModel(), argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) @@ -131,9 +131,10 @@ def test_predict_mono_model_with_r1_waveforms( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras"]) @pytest.mark.parametrize("dl2_tel_flag", ["dl2-telescope", "no-dl2-telescope"]) def test_predict_mono_model_with_dl1_images( - tmp_path, ctlearn_trained_dl1_mono_models, dl1_gamma_file, dl2_tel_flag + tmp_path, ctlearn_trained_dl1_mono_models, dl1_gamma_file, framework, dl2_tel_flag ): """ Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks @@ -159,17 +160,17 @@ def test_predict_mono_model_with_dl1_images( # Hardcopy the trained models to the model directory for telescope_type in telescope_types.keys(): for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" # Build command-line arguments argv = [ f"--input_url={dl1_gamma_file}", - "--PredictCTLearnModel.batch_size=2", + "--PredictCTLearnKerasModel.batch_size=2", "--DLImageReader.focal_length_choice=EQUIVALENT", "--no-dl1-images", "--no-true-images", @@ -177,20 +178,20 @@ def test_predict_mono_model_with_dl1_images( ] for telescope_type, allowed_tels in telescope_types.items(): output_file = ( - dl2_dir / f"gamma_{dl2_tel_flag}_{telescope_type}_mono_from_images.dl2.h5" + dl2_dir / f"gamma_{dl2_tel_flag}_{telescope_type}_{framework}_mono_from_images.dl2.h5" ) # Run Prediction tool assert ( run_tool( - MonoPredictCTLearnModel(), + MonoPredictCTLearnKerasModel(), argv=argv + [ f"--output={output_file}", f"--DLImageReader.allowed_tels={allowed_tels}", f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) @@ -239,8 +240,9 @@ def test_predict_mono_model_with_dl1_images( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras"]) def test_predict_stereo_model_with_dl1_images( - tmp_path, ctlearn_trained_dl1_stereo_models, dl1_gamma_file + tmp_path, ctlearn_trained_dl1_stereo_models, dl1_gamma_file, framework ): """ Test training CTLearn stereo model using the DL1 gamma and proton files for all reconstruction tasks @@ -258,35 +260,35 @@ def test_predict_stereo_model_with_dl1_images( # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "skydirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_stereo_models[key], - model_dir / f"ctlearn_stereo_model_{key}.keras", + model_dir / f"ctlearn_stereo_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_stereo_model_{key}.keras" + model_file = model_dir / f"ctlearn_stereo_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained stereo model file not found for {key}" # Build command-line arguments argv = [ f"--input_url={dl1_gamma_file}", - "--PredictCTLearnModel.batch_size=2", - "--PredictCTLearnModel.stack_telescope_images=True", + "--PredictCTLearnKerasModel.batch_size=2", + "--PredictCTLearnKerasModel.stack_telescope_images=True", "--DLImageReader.mode=stereo", "--DLImageReader.focal_length_choice=EQUIVALENT", f"--DLImageReader.allowed_tels={allowed_tels}", "--no-dl1-images", "--no-true-images", ] - output_file = dl2_dir / f"gamma_{telescope_type}_stereo_from_images.dl2.h5" + output_file = dl2_dir / f"gamma_{framework}_{telescope_type}_stereo_from_images.dl2.h5" # Run Prediction tool assert ( run_tool( - StereoPredictCTLearnModel(), + StereoPredictCTLearnKerasModel(), argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_skydirection.keras", + f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnKerasModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_skydirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) diff --git a/pyproject.toml b/pyproject.toml index 2d21fd5c..79d2b35a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,8 +89,8 @@ documentation = "https://ctlearn.readthedocs.io/en/latest/" [project.scripts] ctlearn-train-keras-model = "ctlearn.tools.keras.train_model:main" ctlearn-train-pytorch-model = "ctlearn.tools.pytorch.train_model:main" -ctlearn-predict-mono-model = "ctlearn.tools.predict.predict_mono:main" -ctlearn-predict-stereo-model = "ctlearn.tools.predict_stereo:main" +ctlearn-predict-mono-keras-model = "ctlearn.tools.keras.predict_model:mono_tool" +ctlearn-predict-stereo-keras-model = "ctlearn.tools.predict_model:stereo_tool" ctlearn-predict-LST1= "ctlearn.tools.predict_LST1:main" [tool.setuptools_scm] From 1fcd30351598379944e322040d21a0bfa850379e Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Tue, 4 Aug 2026 16:17:54 +0200 Subject: [PATCH 02/12] polish loader tests --- ctlearn/core/tests/test_loader_pytorch.py | 60 ----------------------- ctlearn/core/tests/test_loaders.py | 55 +++++++++++++++++++++ 2 files changed, 55 insertions(+), 60 deletions(-) delete mode 100644 ctlearn/core/tests/test_loader_pytorch.py create mode 100644 ctlearn/core/tests/test_loaders.py diff --git a/ctlearn/core/tests/test_loader_pytorch.py b/ctlearn/core/tests/test_loader_pytorch.py deleted file mode 100644 index 47fbe8cf..00000000 --- a/ctlearn/core/tests/test_loader_pytorch.py +++ /dev/null @@ -1,60 +0,0 @@ -import pytest -pytest.importorskip("dl1_data_handler") -torch = pytest.importorskip("torch") -from traitlets.config.loader import Config - -from dl1_data_handler.reader import DLImageReader -from ctlearn.core.data_loader.loader import DLDataLoader - -from ctlearn.tools.train.pytorch.utils import read_configuration -from ctlearn.tools.train.pytorch.utils import get_absolute_config_path - - -def test_data_loader(dl1_gamma_file): - """check""" - # Create a configuration suitable for the test - config = Config( - { - "DLImageReader": { - "allowed_tels": [4], - "focal_length_choice": "EQUIVALENT", - }, - } - ) - # Create an image reader - dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) - # Create a data loader - config_file_dir = get_absolute_config_path() - print(config_file_dir) - parameters = read_configuration(config_file_dir) - - dl1_loader = DLDataLoader.create( - framework = "pytorch", - DLDataReader=dl1_reader, - indices=[0], - tasks=["type", "energy", "cameradirection", "skydirection"], - batch_size=1, - parameters=parameters, - use_augmentation=parameters["augmentation"]["use_augmentation"], - is_training=True - ) - # Get the features and labels fgrom the data loader for one batch - print(len(dl1_loader[0])) - print(dl1_loader[0][0]) - print(dl1_loader[0][1]) - print(dl1_loader[0][2]) - - - features, labels, _ = dl1_loader[0] - # Check that all the correct labels are present - assert ( - "type" in labels - and "energy" in labels - and "cameradirection" in labels - and "skydirection" in labels - ) - # Check the shape of the features - assert features["image"].shape == (1, 2, 110, 110) - -if __name__ == "__main__": - test_data_loader() \ No newline at end of file diff --git a/ctlearn/core/tests/test_loaders.py b/ctlearn/core/tests/test_loaders.py new file mode 100644 index 00000000..7e9c639e --- /dev/null +++ b/ctlearn/core/tests/test_loaders.py @@ -0,0 +1,55 @@ +import pytest +from torch.utils.data import DataLoader +from traitlets.config.loader import Config + +from dl1_data_handler.reader import DLImageReader +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.core.pytorch.dataset import PyTorchDataset + + +@pytest.mark.parametrize( + ("dataloader_cls", "expected_features_shape"), + [ + (KerasSequence, (1, 110, 110, 2)), + (PyTorchDataset, (1, 2, 110, 110)), + ], + ids=["Keras", "PyTorch"], +) +def test_keras_sequence(dl1_gamma_file, dataloader_cls, expected_features_shape): + """check""" + # Create a configuration suitable for the test + config = Config( + { + "DLImageReader": { + "allowed_tels": [4], + "focal_length_choice": "EQUIVALENT", + }, + } + ) + # Create an image reader + dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) + # Initialize the Dataset or Sequence + dl1_dataset = dataloader_cls( + DLDataReader=dl1_reader, + indices=[0], + tasks=["type", "energy", "cameradirection", "skydirection"], + ) + # For PyTorch, wrap in DataLoader to apply batch dimension (dim=0) + if issubclass(dataloader_cls, PyTorchDataset): + batch_loader = DataLoader(dl1_dataset, batch_size=1) + features, labels = next(iter(batch_loader)) + # Optional: convert PyTorch tensor to numpy if labels/features are Tensors + if hasattr(features, "numpy"): + features = features.numpy() + else: + # Keras Sequence handles batching directly via __getitem__ + features, labels = dl1_dataset[0] + # Check that all the correct labels are present + assert ( + "type" in labels + and "energy" in labels + and "cameradirection" in labels + and "skydirection" in labels + ) + # Check the shape of the features match the expected ones + assert features.shape == expected_features_shape From 52ab4e745f80898e327c28f147ff6d9a77f79a8f Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Tue, 4 Aug 2026 16:23:30 +0200 Subject: [PATCH 03/12] polish core testing --- ctlearn/core/tests/test_attention.py | 3 --- ctlearn/core/tests/test_loaders.py | 7 +++---- ctlearn/core/tests/test_models.py | 5 +---- 3 files changed, 4 insertions(+), 11 deletions(-) diff --git a/ctlearn/core/tests/test_attention.py b/ctlearn/core/tests/test_attention.py index 072a9279..f519943c 100644 --- a/ctlearn/core/tests/test_attention.py +++ b/ctlearn/core/tests/test_attention.py @@ -1,7 +1,4 @@ import pytest -pytest.importorskip("keras") -import numpy as np -import pytest import torch import keras import tensorflow as tf diff --git a/ctlearn/core/tests/test_loaders.py b/ctlearn/core/tests/test_loaders.py index 7e9c639e..94590829 100644 --- a/ctlearn/core/tests/test_loaders.py +++ b/ctlearn/core/tests/test_loaders.py @@ -15,8 +15,8 @@ ], ids=["Keras", "PyTorch"], ) -def test_keras_sequence(dl1_gamma_file, dataloader_cls, expected_features_shape): - """check""" +def test_data_loading(dl1_gamma_file, dataloader_cls, expected_features_shape): + """Check the data loading for the Keras and PyTorch frameworks""" # Create a configuration suitable for the test config = Config( { @@ -28,7 +28,7 @@ def test_keras_sequence(dl1_gamma_file, dataloader_cls, expected_features_shape) ) # Create an image reader dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) - # Initialize the Dataset or Sequence + # Initialize the PyTorch Dataset or the Keras Sequence dl1_dataset = dataloader_cls( DLDataReader=dl1_reader, indices=[0], @@ -42,7 +42,6 @@ def test_keras_sequence(dl1_gamma_file, dataloader_cls, expected_features_shape) if hasattr(features, "numpy"): features = features.numpy() else: - # Keras Sequence handles batching directly via __getitem__ features, labels = dl1_dataset[0] # Check that all the correct labels are present assert ( diff --git a/ctlearn/core/tests/test_models.py b/ctlearn/core/tests/test_models.py index 6eaa1962..2aa4ed1a 100644 --- a/ctlearn/core/tests/test_models.py +++ b/ctlearn/core/tests/test_models.py @@ -1,9 +1,6 @@ import pytest -pytest.importorskip("keras") -import re -import keras import numpy as np -import pytest +import keras import torch import torch.nn as nn From 33b16cb1c77c0d5d3fbcaabb08a138de1bbe2b47 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Tue, 4 Aug 2026 16:29:15 +0200 Subject: [PATCH 04/12] fix lst1 predict tool test for Keras --- ctlearn/tools/tests/test_predict_LST1.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/ctlearn/tools/tests/test_predict_LST1.py b/ctlearn/tools/tests/test_predict_LST1.py index da72dded..5bffc8dc 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -2,10 +2,9 @@ import numpy as np import pytest -pytest.importorskip("tensorflow") - from ctapipe.core import run_tool from ctapipe.io import TableLoader +from ctlearn.conftest import MODEL_FILE_FORMATS from ctlearn.tools import LST1PredictionTool # Columns that should be present in the output DL2 file @@ -33,8 +32,9 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras"]) def test_predict_mono_model_with_lst1_mock_data( - tmp_path, ctlearn_trained_dl1_mono_models, mock_lst1_dl1_file + tmp_path, ctlearn_trained_dl1_mono_models, mock_lst1_dl1_file, framework ): """ Test LST1PredictionTool using trained mono models and mock LST-1 DL1 files. @@ -50,7 +50,7 @@ def test_predict_mono_model_with_lst1_mock_data( # Hardcopy the trained models to the model directory telescope_type = "LST" for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_mono_models[key], model_dir / f"ctlearn_mono_model_{key}.keras", @@ -61,7 +61,7 @@ def test_predict_mono_model_with_lst1_mock_data( # Check that the mock LST1 DL1 file was created assert mock_lst1_dl1_file.exists(), "Mock LST1 DL1 file not found" - output_file = dl2_dir / "mock_lst1_predictions.dl2.h5" + output_file = dl2_dir / f"mock_lst1_{framework}_predictions.dl2.h5" # Build command-line arguments for LST1PredictionTool argv = [ @@ -71,9 +71,9 @@ def test_predict_mono_model_with_lst1_mock_data( "--LST1PredictionTool.channels=cleaned_image", "--LST1PredictionTool.channels=cleaned_relative_peak_time", "--LST1PredictionTool.image_mapper_type=BilinearMapper", - f"--type_model={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--energy_model={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--cameradirection_model={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--type_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--energy_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--cameradirection_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", "--dl2-telescope", "--overwrite", ] From c1fd58e761ce8a8021ad24200fc28cada85a5f46 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Tue, 4 Aug 2026 16:41:06 +0200 Subject: [PATCH 05/12] polish lst1 predict tool test --- ctlearn/tools/tests/test_predict_LST1.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ctlearn/tools/tests/test_predict_LST1.py b/ctlearn/tools/tests/test_predict_LST1.py index 5bffc8dc..ffc12b2b 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -53,9 +53,9 @@ def test_predict_mono_model_with_lst1_mock_data( key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" # Check that the mock LST1 DL1 file was created From f2dc264e163c99918031b411a74aee71fc8b7873 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Wed, 5 Aug 2026 16:39:37 +0200 Subject: [PATCH 06/12] add pytorch prediction in existing tools and add the test for pytorch --- ctlearn/core/ctlearn_enum.py | 6 +- ctlearn/core/pytorch/model.py | 4 +- ctlearn/tools/__init__.py | 6 +- ctlearn/tools/keras/__init__.py | 6 - ctlearn/tools/keras/predict_model.py | 2369 --------------------- ctlearn/tools/predict_model.py | 537 +++-- ctlearn/tools/tests/test_predict_model.py | 46 +- pyproject.toml | 4 +- 8 files changed, 450 insertions(+), 2528 deletions(-) delete mode 100644 ctlearn/tools/keras/predict_model.py diff --git a/ctlearn/core/ctlearn_enum.py b/ctlearn/core/ctlearn_enum.py index 506abd57..1854525b 100644 --- a/ctlearn/core/ctlearn_enum.py +++ b/ctlearn/core/ctlearn_enum.py @@ -37,10 +37,10 @@ class FrameworkType(Enum): >>> from ctlearn.core.ctlearn_enum import FrameworkType >>> framework = FrameworkType.PYTORCH >>> print(framework.name) # 'PYTORCH' - >>> print(framework.value) # 2 + >>> print(framework.value) # 'PyTorch' """ - KERAS = 1 - PYTORCH = 2 + KERAS = "Keras" + PYTORCH = "PyTorch" class Task(Enum): diff --git a/ctlearn/core/pytorch/model.py b/ctlearn/core/pytorch/model.py index 2d9c9bf5..07d3f3e1 100644 --- a/ctlearn/core/pytorch/model.py +++ b/ctlearn/core/pytorch/model.py @@ -36,10 +36,12 @@ class MultiHeadClassifier(nn.Module): def __init__(self, heads_dict, single_output_task=None): super().__init__() + # Save the dict with tasks info in an attribute + self.heads_dict = heads_dict # Sanitize keys because 'type' conflicts with nn.Module.type() method self._task_mapping = { task: f"head_{task}" if hasattr(nn.Module, task) else task - for task in heads_dict.keys() + for task in self.heads_dict.keys() } sanitized_heads = { self._task_mapping[task]: module for task, module in heads_dict.items() diff --git a/ctlearn/tools/__init__.py b/ctlearn/tools/__init__.py index 145d76ff..d05c5556 100644 --- a/ctlearn/tools/__init__.py +++ b/ctlearn/tools/__init__.py @@ -2,10 +2,10 @@ """ from ctlearn.tools.predict_LST1 import LST1PredictionTool -from ctlearn.tools.keras.predict_model import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel +from ctlearn.tools.predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel __all__ = [ "LST1PredictionTool", - "MonoPredictCTLearnKerasModel", - "StereoPredictCTLearnKerasModel", + "MonoPredictCTLearnModel", + "StereoPredictCTLearnModel", ] \ No newline at end of file diff --git a/ctlearn/tools/keras/__init__.py b/ctlearn/tools/keras/__init__.py index 06d24e06..e69de29b 100644 --- a/ctlearn/tools/keras/__init__.py +++ b/ctlearn/tools/keras/__init__.py @@ -1,6 +0,0 @@ -from .predict_model import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel - -__all__ = [ - "MonoPredictCTLearnKerasModel", - "StereoPredictCTLearnKerasModel", -] \ No newline at end of file diff --git a/ctlearn/tools/keras/predict_model.py b/ctlearn/tools/keras/predict_model.py deleted file mode 100644 index f05b8203..00000000 --- a/ctlearn/tools/keras/predict_model.py +++ /dev/null @@ -1,2369 +0,0 @@ -""" -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``. -""" - -import atexit -import uuid -import warnings - -import numpy as np -import tables -import tensorflow as tf -import keras - -from astropy import units as u -from astropy.coordinates.earth import EarthLocation -from astropy.coordinates import AltAz, SkyCoord -from astropy.table import ( - Table, - hstack, - vstack, - join, - setdiff, - unique, -) -from astropy.time import Time - -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) -from ctapipe.coordinates import CameraFrame, NominalFrame -from ctapipe.core import Tool -from ctapipe.core.tool import ToolConfigurationError -from ctapipe.core.traits import ( - Bool, - Int, - Path, - flag, - Dict, - ComponentName, - classes_with_traits, -) -from ctapipe.monitoring.interpolation import PointingInterpolator -from ctapipe.instrument import SubarrayDescription -from ctapipe.io import read_table, write_table, HDF5Merger -from ctapipe.io.datalevels import DataLevel -from ctapipe.io.hdf5dataformat import ( - DL0_TEL_POINTING_GROUP, - DL1_SUBARRAY_GROUP, - DL1_SUBARRAY_POINTING_GROUP, - DL1_SUBARRAY_TRIGGER_TABLE, - DL1_TEL_GROUP, - DL1_TEL_CALIBRATION_GROUP, - DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, - DL1_TEL_IMAGES_GROUP, - DL1_TEL_MUON_GROUP, - DL1_TEL_MUON_THROUGHPUT_GROUP, - DL1_TEL_OPTICAL_PSF_GROUP, - DL1_TEL_PARAMETERS_GROUP, - DL1_TEL_POINTING_GROUP, - DL1_TEL_TRIGGER_TABLE, - DL2_EVENT_STATISTICS_GROUP, - FIXED_POINTING_GROUP, - R0_TEL_GROUP, - R1_TEL_GROUP, - SIMULATION_IMAGES_GROUP, - SIMULATION_IMPACT_GROUP, - SIMULATION_PARAMETERS_GROUP, - SIMULATION_RUN_TABLE, - SIMULATION_SHOWER_TABLE, - DL2_TEL_PARTICLETYPE_GROUP, - DL2_TEL_ENERGY_GROUP, - DL2_TEL_GEOMETRY_GROUP, - DL2_SUBARRAY_GROUP, - DL2_SUBARRAY_PARTICLETYPE_GROUP, - DL2_SUBARRAY_ENERGY_GROUP, - DL2_SUBARRAY_GEOMETRY_GROUP, -) -from ctapipe.reco.reconstructor import ReconstructionProperty -from ctapipe.reco.stereo_combination import StereoCombiner -from ctapipe.reco.utils import add_defaults_and_meta -from dl1_data_handler.reader import ( - DLDataReader, - ProcessType, - LST_EPOCH, -) -from ctlearn import __version__ as ctlearn_version -from ctlearn.core.keras.sequence import KerasSequence -from ctlearn.utils import validate_trait_dict - -# Convienient constants for column names and table keys -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TEL_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] -TEL_ITER_GROUPS = [ - R0_TEL_GROUP, - R1_TEL_GROUP, - FIXED_POINTING_GROUP, - DL0_TEL_POINTING_GROUP, - DL1_TEL_POINTING_GROUP, - DL1_TEL_CALIBRATION_GROUP, - DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, - DL1_TEL_MUON_THROUGHPUT_GROUP, - DL1_TEL_OPTICAL_PSF_GROUP, - DL1_TEL_PARAMETERS_GROUP, - DL1_TEL_IMAGES_GROUP, - DL1_TEL_MUON_GROUP, - SIMULATION_IMAGES_GROUP, - SIMULATION_IMPACT_GROUP, - SIMULATION_PARAMETERS_GROUP, -] -DATALEVEL_TO_GROUP = { - DataLevel.R0: R0_TEL_GROUP, - DataLevel.R1: R1_TEL_GROUP, - DataLevel.DL1_IMAGES: DL1_TEL_IMAGES_GROUP, - DataLevel.DL1_PARAMETERS: DL1_TEL_PARAMETERS_GROUP, - DataLevel.DL1_MUON: DL1_TEL_MUON_GROUP, - DataLevel.DL2: DL2_SUBARRAY_GROUP, -} - - -class CannotPredict(OSError): - """Raised when trying to predict an incompatible file""" - - -class PredictCTLearnKerasModel(Tool): - """ - Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. - - This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image - or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to - store them in the output file. The input data is loaded from the input url using the - ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.keras.sequence.KerasSequence``. - The prediction is performed using the CTLearn models. The data is stored in the output file - following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to - handle the prediction for mono and stereo mode. - - Attributes - ---------- - input_url : pathlib.Path - Input ctapipe HDF5 files including pixel-wise image or waveform data. - dl1_features : bool - Set whether to include the dl1 feature vectors in the output file. - dl2_telescope : bool - Set whether to include dl2 telescope-event-wise data in the output file. - dl2_subarray : bool - Set whether to include dl2 subarray-event-wise data in the output file. - dl1dh_reader : dl1_data_handler.reader.DLDataReader - DLDataReader object to read the data. - dl1dh_reader_type : str - Type of the DLDataReader to use for the prediction. - stack_telescope_images : bool - Set whether to stack the telescope images in the data loader. Requires ``stereo``. - sort_by_intensity : bool - Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. - prefix : str - Name of the reconstruction algorithm used to generate the dl2 data. - load_type_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the classification of the primary particle type. - load_energy_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression of the primary particle energy. - load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression - of the primary particle arrival direction based on camera coordinate offsets. - load_skydirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression - of the primary particle arrival direction based on spherical coordinate offsets. - output_path : pathlib.Path - Output path to save the dl2 prediction results. - keras_verbose : int - Verbosity mode of Keras during the prediction. - strategy : tf.distribute.Strategy - MirroredStrategy to distribute the prediction. - data_loader : ctlearn.core.keras.sequence.KerasSequence - KerasSequence object to load the data. - indices : list of int - List of indices for the data loaders. - batch_size : int - Size of the batch to perform inference of the neural network. - last_batch_size : int - Size of the last batch in the data loaders. - - Methods - ------- - setup() - Set up the tool. - finish() - Finish the tool. - _predict_with_model(model_path) - Load and predict with a CTLearn model. - _predict_particletype(example_identifiers) - Predict the classification of the primary particle type. - _predict_energy(example_identifiers) - Predict the energy of the primary particle. - _predict_cameradirection(example_identifiers) - Predict the arrival direction of the primary particle based on camera coordinate offsets. - _predict_skydirection(example_identifiers) - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - _transform_cam_coord_offsets_to_sky(table) - Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _transform_spher_coord_offsets_to_sky(table) - Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _create_nan_table(nonexample_identifiers, columns, shapes, reco_task) - Create a table with NaNs for missing predictions. - _store_pointing(all_identifiers) - Store the telescope pointing table from to the output file. - _create_feature_vectors_table(example_identifiers, nonexample_identifiers, particletype_feature_vectors, energy_feature_vectors, direction_feature_vectors) - Create the table for the DL1 feature vectors. - """ - - input_url = Path( - help="Input ctapipe HDF5 files including pixel-wise image or waveform data", - allow_none=True, - exists=True, - directory_ok=False, - file_ok=True, - ).tag(config=True) - - dl1_features = Bool( - default_value=False, - allow_none=False, - help="Set whether to include the dl1 feature vectors in the output file.", - ).tag(config=True) - - dl2_telescope = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 telescope-event-wise data in the output file.", - ).tag(config=True) - - dl2_subarray = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 subarray-event-wise data in the output file.", - ).tag(config=True) - - dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( - config=True - ) - - stack_telescope_images = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to stack the telescope images in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - sort_by_intensity = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to sort the telescope images by intensity in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - prefixes = Dict( - default_value={ - "type": "CTLearnClassifier", - "energy": "CTLearnRegressor", - "cameradirection": "CTLearnCameraReconstructor", - "skydirection": "CTLearnSkyReconstructor", - "all": "CTLearn", - }, - allow_none=False, - help=( - "Name of the reconstruction algorithm used " - "to generate the dl2 data for each task." - ), - ).tag(config=True) - - load_type_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) for the classification " - "of the primary particle type." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_energy_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) for the regression " - "of the primary particle energy." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_cameradirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on camera coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_skydirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on spherical coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - batch_size = Int( - default_value=64, - allow_none=False, - help="Size of the batch to perform inference of the neural network.", - ).tag(config=True) - - output_path = Path( - default_value="./output.dl2.h5", - allow_none=False, - help="Output path to save the dl2 prediction results", - ).tag(config=True) - - keras_verbose = Int( - default_value=1, - min=0, - max=2, - allow_none=False, - help=( - "Verbosity mode of Keras during the prediction: " - "0 = silent, 1 = progress bar, 2 = one line per call." - ), - ).tag(config=True) - - aliases = { - ("i", "input_url"): "PredictCTLearnKerasModel.input_url", - ("t", "type_model"): "PredictCTLearnKerasModel.load_type_model_from", - ("e", "energy_model"): "PredictCTLearnKerasModel.load_energy_model_from", - ( - "d", - "cameradirection_model", - ): "PredictCTLearnKerasModel.load_cameradirection_model_from", - ("s", "skydirection_model"): "PredictCTLearnKerasModel.load_skydirection_model_from", - ("o", "output"): "PredictCTLearnKerasModel.output_path", - } - - flags = { - **flag( - "overwrite", - "HDF5Merger.overwrite", - "Overwrite the output file if it exists", - "Do not overwrite the output file if it exists", - ), - **flag( - "dl1-features", - "PredictCTLearnKerasModel.dl1_features", - "Include dl1 features", - "Exclude dl1 features", - ), - **flag( - "dl2-telescope", - "PredictCTLearnKerasModel.dl2_telescope", - "Include dl2 telescope-event-wise data in the output file", - "Exclude dl2 telescope-event-wise data in the output file", - ), - **flag( - "dl2-subarray", - "PredictCTLearnKerasModel.dl2_subarray", - "Include dl2 subarray-event-wise data in the output file", - "Exclude dl2 subarray-event-wise data in the output file", - ), - **flag( - "r0-waveforms", - "HDF5Merger.r0_waveforms", - "Include r0 waveforms", - "Exclude r0 waveforms", - ), - **flag( - "r1-waveforms", - "HDF5Merger.r1_waveforms", - "Include r1 waveforms", - "Exclude r1 waveforms", - ), - **flag( - "dl1-parameters", - "HDF5Merger.dl1_parameters", - "Include dl1 parameters", - "Exclude dl1 parameters", - ), - **flag( - "dl1-images", - "HDF5Merger.dl1_images", - "Include dl1 images", - "Exclude dl1 images", - ), - **flag( - "true-parameters", - "HDF5Merger.true_parameters", - "Include true parameters", - "Exclude true parameters", - ), - **flag( - "true-images", - "HDF5Merger.true_images", - "Include true images", - "Exclude true images", - ), - } - - classes = classes_with_traits(DLDataReader) - - def setup(self): - self.activity_start_time = Time.now() - self.log.info("ctlearn version %s", ctlearn_version) - # Validate the prefixes trait dictionary - validate_trait_dict( - self.prefixes, ["type", "energy", "cameradirection", "skydirection", "all"] - ) - # Copy selected tables from the input file to the output file - self.log.info("Copying to output destination.") - with HDF5Merger( - self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self - ) as merger: - merger(self.input_url) - # Create a MirroredStrategy. - self.strategy = tf.distribute.MirroredStrategy() - atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore - self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) - - # Set up the data reader - self.log.info("Loading data reader:") - self.log.info("For a large dataset, this may take a while...") - self.dl1dh_reader = DLDataReader.from_name( - self.dl1dh_reader_type, - input_url_signal=[self.input_url], - parent=self, - ) - self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) - # Check if the number of events is enough to form a batch - if self.dl1dh_reader._get_n_events() < self.batch_size: - raise ToolConfigurationError( - f"{self.dl1dh_reader._get_n_events()} events are not enough " - f"to form a batch of size {self.batch_size}. Reduce the batch size." - ) - # Set the indices for the data loaders - self.indices = list(range(self.dl1dh_reader._get_n_events())) - self.last_batch_size = len(self.indices) % ( - self.batch_size * self.strategy.num_replicas_in_sync - ) - # Ensure subarray consistency in the output file - self._ensure_subarray_consistency() - - def finish(self): - # Overwrite CTAO reference metadata to the output file - self._overwrite_meta() - self.log.info("Tool is shutting down") - - def _overwrite_meta(self): - """Overwrite CTAO metadata in the output file.""" - # TODO: Upgrade to new CTAO metatdata standard when available - with warnings.catch_warnings(): - warnings.simplefilter("ignore", tables.NaturalNameWarning) - with tables.open_file(self.output_path, mode="r+") as h5_file: - # Update CTA Activity metadata - h5_file.root._v_attrs["CTA ACTIVITY ID"] = str(uuid.uuid4()) - h5_file.root._v_attrs["CTA ACTIVITY NAME"] = self.name - h5_file.root._v_attrs["CTA ACTIVITY SOFTWARE NAME"] = "ctlearn" - h5_file.root._v_attrs["CTA ACTIVITY SOFTWARE VERSION"] = ctlearn_version - h5_file.root._v_attrs["CTA ACTIVITY START TIME"] = ( - self.activity_start_time.iso - ) - h5_file.root._v_attrs["CTA ACTIVITY STOP TIME"] = Time.now().iso - # Update CTA Product metadata - h5_file.root._v_attrs["CTA PRODUCT DATA LEVELS"] = ( - self._get_data_levels(h5_file) - ) - h5_file.root._v_attrs["CTA PRODUCT CREATION TIME"] = ( - self.activity_start_time.iso - ) - h5_file.root._v_attrs["CTA PRODUCT ID"] = str(uuid.uuid4()) - h5_file.flush() - - def _get_data_levels(self, h5file): - """Get the data levels present in the HDF5 file.""" - data_levels = { - level.name - for level, group in DATALEVEL_TO_GROUP.items() - if hasattr(h5file.root, group) - } - return ",".join(sorted(data_levels)) - - def _ensure_subarray_consistency(self): - """ - Align subarray metadata and trigger tables with the selected telescopes. - - When only a subset of telescopes is processed, overwrite the output file's - SubarrayDescription and trim the DL1 trigger tables to keep events that - involve the selected telescopes. Also rebuild the subarray trigger table - with the corresponding telescope participation masks. - """ - - input_subarray = SubarrayDescription.from_hdf( - self.input_url, - focal_length_choice=self.dl1dh_reader.focal_length_choice, - ) - if input_subarray == self.dl1dh_reader.subarray: - return - - # From the merger tool a SubarrayDescription for the full array is already stored - # in the output file. We need to remove it to avoid conflicts when storing - # the new SubarrayDescription for the selected telescopes. - with tables.open_file(self.output_path, mode="a") as h5file: - h5file.remove_node("/configuration/instrument", recursive=True) - selected_subarray = input_subarray.select_subarray(set(self.dl1dh_reader.tel_ids)) - selected_subarray.to_hdf(self.output_path) - self.log.info("SubarrayDescription was stored in '%s'", self.output_path) - - tel_trigger_table = read_table( - self.output_path, - DL1_TEL_TRIGGER_TABLE, - ) - mask = np.isin(tel_trigger_table["tel_id"], self.dl1dh_reader.tel_ids) - tel_trigger_table = tel_trigger_table[mask] - tel_trigger_table.sort(TEL_EVENT_KEYS) - - write_table( - tel_trigger_table, - self.output_path, - DL1_TEL_TRIGGER_TABLE, - overwrite=True, - ) - - subarray_trigger_table = tel_trigger_table.copy() - subarray_columns = SUBARRAY_EVENT_KEYS + ["time"] - # In older data formats the event type is not included in the trigger table, so we need to - # check if it is present before keeping the column to be backwards compatible. - if "event_type" in subarray_trigger_table.colnames: - subarray_columns.append("event_type") - subarray_trigger_table.keep_columns(subarray_columns) - subarray_trigger_table = unique( - subarray_trigger_table, keys=SUBARRAY_EVENT_KEYS - ) - - tel_trigger_groups = tel_trigger_table.group_by(SUBARRAY_EVENT_KEYS) - tel_with_trigger = [] - for tel_trigger in tel_trigger_groups.groups: - tel_with_trigger_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) - tel_with_trigger_mask[ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_trigger["tel_id"]) - ] = True - tel_with_trigger.append(tel_with_trigger_mask) - - subarray_trigger_table.add_column( - tel_with_trigger, index=-2, name="tels_with_trigger" - ) - - write_table( - subarray_trigger_table, - self.output_path, - DL1_SUBARRAY_TRIGGER_TABLE, - overwrite=True, - ) - # Update the simulation shower table to keep only events present in the subarray trigger table - subarray_trigger_table.keep_columns(SUBARRAY_EVENT_KEYS) - sim_shower_table = read_table( - self.output_path, - SIMULATION_SHOWER_TABLE, - ) - sim_shower_table = join( - sim_shower_table, - subarray_trigger_table, - keys=SUBARRAY_EVENT_KEYS, - join_type="right", - ) - sim_shower_table.sort(SUBARRAY_EVENT_KEYS) - write_table( - sim_shower_table, - self.output_path, - SIMULATION_SHOWER_TABLE, - overwrite=True, - ) - # Delete telescope-specific tables for unselected telescopes - self._delete_unselected_telescope_tables() - - def _delete_unselected_telescope_tables(self): - """ - Delete telescope-specific tables for unselected telescopes from the output file. - - Iterates through all telescope-related groups in the HDF5 file and removes - tables corresponding to telescopes that are not in the selected telescope list. - This ensures the output file only contains data for the telescopes that were - processed. The camera configuration tables are also pruned based on the camera indices. - """ - # Open the HDF5 file to prune the unselected telescope tables and camera configurations - with tables.open_file(self.output_path, mode="r+") as h5_file: - - def prune_group(group, valid_ids): - for table in group._f_iter_nodes("Table"): - idx = int(table._v_name.split("_")[-1]) - if idx not in valid_ids: - table._f_remove() - - # Telescope-specific tables - tel_ids = set(self.dl1dh_reader.tel_ids) - for group_name in TEL_ITER_GROUPS: - group = getattr(h5_file.root, group_name, None) - if group is not None: - prune_group(group, tel_ids) - - - def _create_nan_table(self, nonexample_identifiers, columns, shapes, reco_task): - """ - Create a table with NaNs for missing predictions. - - This method creates a table with NaNs for missing predictions for the non-example identifiers. - In stereo mode, the table also a column for the valid telescopes is added with all False values. - - Parameters: - ----------- - nonexample_identifiers : astropy.table.Table - Table containing the non-example identifiers. - columns : list of str - List of column names to create in the table. - shapes : list of shapes - List of shapes for the columns to create in the table. - reco_task : str - Reconstruction task name. - - Returns: - -------- - nan_table : astropy.table.Table - Table containing NaNs for missing predictions. - """ - # Create a table with NaNs for missing predictions - nan_table = nonexample_identifiers.copy() - for column_name, shape in zip(columns, shapes): - nan_table.add_column(np.full(shape, np.nan), name=column_name) - # Add that no telescope is valid for the non-example identifiers in stereo mode - if self.dl1dh_reader.mode == "stereo": - nan_table.add_column( - np.zeros( - (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ), - name=f"{self.prefixes[reco_task]}_telescopes", - ) - return nan_table - - def deduplicate_first_valid( - self, - table: Table, - keys=("obs_id", "event_id"), - valid_col="CTLearn_is_valid", - ): - """ - Return a deduplicated Astropy Table. - - For each group defined by `keys`, keep the first row where - `valid_col` is True. If none are valid, keep the first row. - """ - - t = table.copy() - - t.sort(list(keys) + [valid_col], reverse=[False] * len(keys) + [True]) - - return unique(t, keys=list(keys), keep="first") - - def _predict_with_model(self, model_path): - """ - Load and predict with a CTLearn model. - - Load a model from the specified path and predict the data using the loaded model. - If a last batch loader is provided, predict the last batch and stack the results. - - Parameters - ---------- - model_path : str - Path to a Keras model file (Keras3). - - Returns - ------- - predict_data : astropy.table.Table - Table containing the prediction results. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - # Create a new KerasSequence for each task - # It turned out to be more robust to initialize the KerasSequence separately. - data_loader = KerasSequence( - self.dl1dh_reader, - self.indices, - tasks=[], - batch_size=self.batch_size * self.strategy.num_replicas_in_sync, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Keras is only considering the last complete batch. - # In prediction mode we don't want to loose the last - # uncomplete batch, so we are creating an additional - # batch generator for the remaining events. - data_loader_last_batch = None - if self.last_batch_size > 0: - last_batch_indices = self.indices[-self.last_batch_size :] - data_loader_last_batch = KerasSequence( - self.dl1dh_reader, - last_batch_indices, - tasks=[], - batch_size=self.last_batch_size, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Load the model from the specified path - model = keras.saving.load_model(model_path) - prediction_colname = ( - "type" - if isinstance(model.layers[-1], keras.layers.Softmax) - else model.layers[-1].name - ) - backbone_model, feature_vectors = None, None - if self.dl1_features: - # Get the backbone model which is the second layer of the model - backbone_model = model.get_layer(index=1) - # Create a new head model with the same layers as the original model. - # The output of the backbone model is the input of the head model. - backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) - x = backbone_output_shape - for layer in model.layers[2:]: - x = layer(x) - head = keras.Model(inputs=backbone_output_shape, outputs=x) - # Apply the backbone model with the data loader to retrieve the feature vectors - try: - feature_vectors = backbone_model.predict( - data_loader, verbose=self.keras_verbose - ) - except ValueError as err: - if str(err).startswith("Input 0 of layer"): - raise ToolConfigurationError( - "Model input shape does not match the prediction data. " - "This is usually caused by selecting the wrong telescope_id. " - "Please ensure the telescope configuration matches the one used for training." - ) from err - raise - # Apply the head model with the feature vectors to retrieve the prediction - predict_data = Table( - { - prediction_colname: head.predict( - feature_vectors, verbose=self.keras_verbose - ) - } - ) - # Predict the last batch and stack the results to the prediction data - if data_loader_last_batch is not None: - feature_vectors_last_batch = backbone_model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - feature_vectors = np.concatenate( - (feature_vectors, feature_vectors_last_batch) - ) - predict_data = vstack( - [ - predict_data, - Table( - { - prediction_colname: head.predict( - feature_vectors_last_batch, - verbose=self.keras_verbose, - ) - } - ), - ] - ) - else: - # Predict the data using the loaded model - try: - predict_data = model.predict(data_loader, verbose=self.keras_verbose) - except ValueError as err: - if str(err).startswith("Input 0 of layer"): - raise ToolConfigurationError( - "Model input shape does not match the prediction data. " - "This is usually caused by selecting the wrong telescope_id. " - "Please ensure the telescope configuration matches the one used for training." - ) from err - raise - # Create a astropy table with the prediction results - # The classification task has a softmax layer as the last layer - # which returns the probabilities for each class in an array, while - # the regression tasks have output neurons which returns the - # predicted value for the task in a dictionary. - if prediction_colname == "type": - predict_data = Table({prediction_colname: predict_data}) - else: - predict_data = Table(predict_data) - # Predict the last batch and stack the results to the prediction data - if data_loader_last_batch is not None: - predict_data_last_batch = model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - if model.layers[-1].name == "type": - predict_data_last_batch = Table( - {prediction_colname: predict_data_last_batch} - ) - else: - predict_data_last_batch = Table(predict_data_last_batch) - predict_data = vstack([predict_data, predict_data_last_batch]) - return predict_data, feature_vectors - - def _predict_particletype(self, example_identifiers): - """ - Predict the classification of the primary particle type. - - This method uses a pre-trained type model to predict the type of the primary particle - for a given set of example identifiers. The predicted classification score ('gammaness') - is added to the example identifiers table. - - Parameters: - ----------- - particletype_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - predicted classification score ('gammaness'). - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the classification of the primary particle type..." - ) - # Predict the data using the loaded type_model - predict_data, feature_vectors = self._predict_with_model( - self.load_type_model_from - ) - # Create prediction table and add the predicted classification score ('gammaness') - particletype_table = example_identifiers.copy() - particletype_table.add_column( - predict_data["type"].T[1], name=f"{self.prefixes['type']}_tel_prediction" - ) - return particletype_table, feature_vectors - - def _predict_energy(self, example_identifiers): - """ - Predict the energy of the primary particle. - - This method uses a pre-trained energy model to predict the energy of the primary particle - for a given set of example identifiers. The predicted energy is then converted from - log10(TeV) to TeV and added to the example identifiers table. - - Parameters: - ----------- - energy_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed energy in TeV. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info("Predicting for the regression of the primary particle energy...") - # Predict the data using the loaded energy_model - predict_data, feature_vectors = self._predict_with_model( - self.load_energy_model_from - ) - # Convert the reconstructed energy from log10(TeV) to TeV - reco_energy = u.Quantity( - np.power(10, np.squeeze(predict_data["energy"])), - unit=u.TeV, - ) - # Create prediction table and add the reconstructed energy in TeV - energy_table = example_identifiers.copy() - energy_table.add_column( - reco_energy, name=f"{self.prefixes['energy']}_tel_energy" - ) - return energy_table, feature_vectors - - def _predict_cameradirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on camera coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the - primary particle for a given set of example identifiers. The predicted camera coordinate offsets - is added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - cameradirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed camera coordinate offsets in x and y. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the reconstruction of the primary particle arrival direction based on camera coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_cameradirection_model_from - ) - # For the direction task, the prediction is the camera coordinate offset in x and y - # from the telescope pointing. - cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) - cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) - # Create prediction table and add the reconstructed energy in TeV - cameradirection_table = example_identifiers.copy() - cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") - cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") - return cameradirection_table, feature_vectors - - def _predict_skydirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the primary - particle for a given set of example identifiers. The predicted spherical coordinate offsets is - added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - skydirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed spherical coordinate offsets in fov_lon and fov_lat. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the reconstruction of the primary particle arrival direction based on spherical coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_skydirection_model_from - ) - # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat - # from the telescope pointing. - fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) - fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) - # Create prediction table and add the reconstructed fov_lon and fov_lat - skydirection_table = example_identifiers.copy() - skydirection_table.add_column(fov_lon, name="fov_lon") - skydirection_table.add_column(fov_lat, name="fov_lat") - return skydirection_table, feature_vectors - - def _transform_cam_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - # Get the telescope ID from the table - tel_id = table["tel_id"][0] - # Set the telescope position - tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the telescope pointing with the trigger timestamp and the telescope position - altaz = AltAz( - location=tel_ground_frame.to_earth_location(), - obstime=trigger_time, - ) - # Set the telescope pointing - tel_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the camera frame with the focal length and rotation of the camera - camera_frame = CameraFrame( - focal_length=self.dl1dh_reader.subarray.tel[ - tel_id - ].camera.geometry.frame.focal_length, - rotation=self.dl1dh_reader.pix_rotation[tel_id], - telescope_pointing=tel_pointing, - ) - # Set the camera coordinate offset - cam_coord_offset = SkyCoord( - x=table["cam_coord_offset_x"], - y=table["cam_coord_offset_y"], - frame=camera_frame, - ) - # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Transform the true Alt/Az coordinates to camera coordinates - reco_direction = cam_coord_offset.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column( - reco_direction.az.to(u.deg), - name=f"{self.prefixes['cameradirection']}_tel_az", - ) - table.add_column( - reco_direction.alt.to(u.deg), - name=f"{self.prefixes['cameradirection']}_tel_alt", - ) - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "cam_coord_offset_x", - "cam_coord_offset_y", - ] - ) - return table - - def _transform_spher_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted spherical offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted spherical offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the AltAz frame with the reference location and time - altaz = AltAz( - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the array pointing - array_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the nominal frame with the array pointing - nom_frame = NominalFrame( - origin=array_pointing, - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the reco direction in (fov_lon, fov_lat) coordinates - reco_direction = SkyCoord( - fov_lon=table["fov_lon"], - fov_lat=table["fov_lat"], - frame=nom_frame, - ) - # Transform the reco direction from nominal frame to the AltAz frame - sky_coord = reco_direction.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column( - sky_coord.az.to(u.deg), name=f"{self.prefixes['skydirection']}_az" - ) - table.add_column( - sky_coord.alt.to(u.deg), name=f"{self.prefixes['skydirection']}_alt" - ) - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "fov_lon", - "fov_lat", - ] - ) - return table - - def _store_pointing(self, all_identifiers): - """ - Store the telescope pointing table from to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - - # Initialize the pointing interpolator from ctapipe - pointing_interpolator = PointingInterpolator( - bounds_error=False, extrapolate=True - ) - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Get the telescope pointing from the dl1dh reader - tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] - # Add the telescope pointing table to the pointing interpolator - pointing_interpolator.add_table(tel_id, tel_pointing) - tel_identifiers = all_identifiers.copy() - if self.dl1dh_reader.mode == "mono": - tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Interpolate the telescope pointing - tel_altitude, tel_azimuth = pointing_interpolator( - tel_id, tel_identifiers["time"] - ) - tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") - tel_identifiers.add_column(tel_altitude, name="pointing_altitude") - pointing_info.append(tel_identifiers) - if self.dl1dh_reader.mode == "mono": - tel_pointing_table = Table( - { - "time": tel_identifiers["time"], - "azimuth": tel_identifiers["pointing_azimuth"], - "altitude": tel_identifiers["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info = vstack(pointing_info) - if self.dl1dh_reader.mode == "stereo": - # Group the pointing information by subarray event keys - # TODO: This needs to be debugged with SST1M data - pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) - pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) - pointing_info = join( - all_identifiers, - pointing_mean, - keys=SUBARRAY_EVENT_KEYS, - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - DL1_SUBARRAY_POINTING_GROUP, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - DL1_SUBARRAY_POINTING_GROUP, - ) - return pointing_info - - def _create_feature_vectors_table( - self, - example_identifiers, - nonexample_identifiers=None, - particletype_feature_vectors=None, - energy_feature_vectors=None, - direction_feature_vectors=None, - ): - """ - Create the table for the DL1 feature vectors. - - This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for - non-example identifiers. The feature vectors are stored in the columns of the table. The table also - contains a column for the valid predictions. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - nonexample_identifiers : astropy.table.Table or None - Table containing the non-example identifiers to fill the NaNs. - particletype_feature_vectors : np.ndarray or None - Array containing the particletype feature vectors. - energy_feature_vectors : np.ndarray or None - Array containing the energy feature vectors. - direction_feature_vectors : np.ndarray or None - Array containing the direction feature vectors. - - Returns: - -------- - feature_vector_table : astropy.table.Table - Table containing the DL1 feature vectors for the example and non-example identifiers. - """ - # Create the feature vector table - feature_vector_table = example_identifiers.copy() - columns_list, shapes_list = [], [] - if particletype_feature_vectors is not None: - is_valid_col = ~np.isnan( - np.min(particletype_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - particletype_feature_vectors, - name=f"{self.prefixes['all']}_tel_particletype_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append( - f"{self.prefixes['all']}_tel_particletype_feature_vectors" - ) - shapes_list.append( - ( - len(nonexample_identifiers), - particletype_feature_vectors.shape[1], - ) - ) - if energy_feature_vectors is not None: - is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) - feature_vector_table.add_column( - energy_feature_vectors, - name=f"{self.prefixes['all']}_tel_energy_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append( - f"{self.prefixes['all']}_tel_energy_feature_vectors" - ) - shapes_list.append( - ( - len(nonexample_identifiers), - energy_feature_vectors.shape[1], - ) - ) - if direction_feature_vectors is not None: - feature_vector_table.remove_columns( - ["pointing_azimuth", "pointing_altitude", "time"] - ) - is_valid_col = ~np.isnan( - np.min(direction_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - direction_feature_vectors, - name=f"{self.prefixes['all']}_tel_geometry_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append( - f"{self.prefixes['all']}_tel_geometry_feature_vectors" - ) - shapes_list.append( - ( - len(nonexample_identifiers), - direction_feature_vectors.shape[1], - ) - ) - # Produce output table with NaNs for missing predictions - if nonexample_identifiers is not None: - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=columns_list, - shapes=shapes_list, - reco_task="all", - ) - feature_vector_table = vstack([feature_vector_table, nan_table]) - is_valid_col = np.concatenate( - (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) - ) - # Add is_valid column to the feature vector table - feature_vector_table.add_column( - is_valid_col, - name=f"{self.prefixes['all']}_tel_is_valid", - ) - return feature_vector_table - - -class MonoPredictCTLearnKerasModel(PredictCTLearnKerasModel): - """ - Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. - - This tool extends the ``PredictCTLearnKerasModel`` to specifically handle monoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_telescope_pointing(all_identifiers) - Store the telescope pointing table for the mono mode for MC simulation. - """ - - name = "ctlearn-predict-mono-keras-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnKerasModel.batch_size=64 \\ - --PredictCTLearnKerasModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ - --dl1-features \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - - To predict from pixel-wise waveform data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.r1.h5 \\ - --PredictCTLearnKerasModel.dl1dh_reader_type=DLWaveformReader \\ - --DLWaveformReader.sequnce_length=20 \\ - --DLWaveformReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ - --no-r0-waveforms \\ - --no-r1-waveforms \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - """ - - stereo_combiner_cls = ComponentName( - StereoCombiner, - default_value="StereoMeanCombiner", - help="Which stereo combination method to use after the monoscopic reconstruction.", - ).tag(config=True) - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.example_identifiers.copy() - example_identifiers.keep_columns(TEL_EVENT_KEYS) - all_identifiers = read_table( - self.output_path, - DL1_TEL_TRIGGER_TABLE, - ) - all_identifiers.keep_columns(TEL_EVENT_KEYS + ["time"]) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=TEL_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Pointing table for the mono mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_telescope_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - particletype_feature_vectors = None - if self.load_type_model_from is not None: - # Predict the type of the primary particle - particletype_table, particletype_feature_vectors = ( - super()._predict_particletype(example_identifiers) - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefixes['type']}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - reco_task="type", - ) - particletype_table = vstack([particletype_table, nan_table]) - # Add is_valid column to the particle type table - particletype_table.add_column( - ~np.isnan( - particletype_table[f"{self.prefixes['type']}_tel_prediction"].data, - dtype=bool, - ), - name=f"{self.prefixes['type']}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - particletype_table, - ParticleClassificationContainer, - prefix=self.prefixes["type"], - add_tel_prefix=True, - ) - if self.dl2_telescope: - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = particletype_table["tel_id"] == tel_id - particletype_tel_table = particletype_table[telescope_mask] - particletype_tel_table.sort(TEL_EVENT_KEYS) - # Save the prediction to the output file for the selected telescope - write_table( - particletype_tel_table, - self.output_path, - f"{DL2_TEL_PARTICLETYPE_GROUP}/{self.prefixes['type']}/tel_{tel_id:03d}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TEL_PARTICLETYPE_GROUP}/{self.prefixes['type']}/tel_{tel_id:03d}", - ) - - if self.dl2_subarray: - self.log.info("Processing and storing the subarray type prediction...") - # If only one telescope is used, copy the particletype table - # and modify it to subarray format - if len(self.dl1dh_reader.tel_ids) == 1: - particletype_subarray_table = particletype_table.copy() - telescope_mask = ( - particletype_subarray_table["tel_id"] - == self.dl1dh_reader.tel_ids[0] - ) - particletype_subarray_table = particletype_subarray_table[ - telescope_mask - ] - particletype_subarray_table.remove_column("tel_id") - for colname in particletype_subarray_table.colnames: - if "_tel_" in colname: - particletype_subarray_table.rename_column( - colname, colname.replace("_tel", "") - ) - particletype_subarray_table.add_column( - [ - [val] - for val in particletype_subarray_table[ - f"{self.prefixes['type']}_is_valid" - ] - ], - name=f"{self.prefixes['type']}_telescopes", - ) - else: - self.type_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefixes["type"], - property=ReconstructionProperty.PARTICLE_TYPE, - parent=self, - ) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - particletype_subarray_table = ( - self.type_stereo_combiner.predict_table(particletype_table) - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - particletype_subarray_table[ - f"{self.prefixes['type']}_telescopes" - ].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - ( - len(particletype_subarray_table), - len(self.dl1dh_reader.tel_ids), - ), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - particletype_subarray_table[ - f"{self.prefixes['type']}_telescopes" - ] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices( - tel_id - ) - ] = True - # Overwrite the column with the boolean mask with fix length - particletype_subarray_table[ - f"{self.prefixes['type']}_telescopes" - ] = reco_telescopes - # Deduplicate the subarray particletype table to have only one entry per event - particletype_subarray_table = super().deduplicate_first_valid( - table=particletype_subarray_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['type']}_is_valid", - ) - # Sort the subarray particletype table - particletype_subarray_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - particletype_subarray_table, - self.output_path, - f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefixes['energy']}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - reco_task="energy", - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefixes['energy']}_tel_energy"].data, - dtype=bool, - ), - name=f"{self.prefixes['energy']}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefixes["energy"], - add_tel_prefix=True, - ) - if self.dl2_telescope: - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = energy_table["tel_id"] == tel_id - energy_tel_table = energy_table[telescope_mask] - energy_tel_table.sort(TEL_EVENT_KEYS) - # Save the prediction to the output file - write_table( - energy_tel_table, - self.output_path, - f"{DL2_TEL_ENERGY_GROUP}/{self.prefixes['energy']}/tel_{tel_id:03d}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TEL_ENERGY_GROUP}/{self.prefixes['energy']}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray energy prediction..." - ) - # If only one telescope is used, copy the particletype table - # and modify it to subarray format - if len(self.dl1dh_reader.tel_ids) == 1: - energy_subarray_table = energy_table.copy() - telescope_mask = ( - energy_subarray_table["tel_id"] == self.dl1dh_reader.tel_ids[0] - ) - energy_subarray_table = energy_subarray_table[telescope_mask] - energy_subarray_table.remove_column("tel_id") - for colname in energy_subarray_table.colnames: - if "_tel_" in colname: - energy_subarray_table.rename_column( - colname, colname.replace("_tel", "") - ) - energy_subarray_table.add_column( - [ - [val] - for val in energy_subarray_table[ - f"{self.prefixes['energy']}_is_valid" - ] - ], - name=f"{self.prefixes['energy']}_telescopes", - ) - else: - self.energy_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefixes["energy"], - property=ReconstructionProperty.ENERGY, - parent=self, - ) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - energy_subarray_table = self.energy_stereo_combiner.predict_table( - energy_table - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - energy_subarray_table[ - f"{self.prefixes['energy']}_telescopes" - ].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - ( - len(energy_subarray_table), - len(self.dl1dh_reader.tel_ids), - ), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - energy_subarray_table[ - f"{self.prefixes['energy']}_telescopes" - ] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices( - tel_id - ) - ] = True - # Overwrite the column with the boolean mask with fix length - energy_subarray_table[ - f"{self.prefixes['energy']}_telescopes" - ] = reco_telescopes - # Deduplicate the subarray energy table to have only one entry per event - energy_subarray_table = super().deduplicate_first_valid( - table=energy_subarray_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['energy']}_is_valid", - ) - # Sort the subarray energy table - energy_subarray_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - energy_subarray_table, - self.output_path, - f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", - ) - direction_feature_vectors = None - if self.load_cameradirection_model_from is not None: - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=TEL_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = ( - super()._predict_cameradirection(example_identifiers) - ) - direction_tel_tables = [] - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = direction_table["tel_id"] == tel_id - direction_tel_table = direction_table[telescope_mask] - direction_tel_table = super()._transform_cam_coord_offsets_to_sky( - direction_tel_table - ) - # Produce output table with NaNs for missing predictions - nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id - nonexample_identifiers_tel = nonexample_identifiers[nan_telescope_mask] - if len(nonexample_identifiers_tel) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers_tel, - columns=[ - f"{self.prefixes['cameradirection']}_tel_alt", - f"{self.prefixes['cameradirection']}_tel_az", - ], - shapes=[ - (len(nonexample_identifiers_tel),), - (len(nonexample_identifiers_tel),), - ], - reco_task="cameradirection", - ) - direction_tel_table = vstack([direction_tel_table, nan_table]) - direction_tel_table.sort(TEL_EVENT_KEYS) - # Add is_valid column to the direction table - direction_tel_table.add_column( - ~np.isnan( - direction_tel_table[ - f"{self.prefixes['cameradirection']}_tel_alt" - ].data, - dtype=bool, - ), - name=f"{self.prefixes['cameradirection']}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_tel_table, - ReconstructedGeometryContainer, - prefix=self.prefixes["cameradirection"], - add_tel_prefix=True, - ) - direction_tel_tables.append(direction_tel_table) - if self.dl2_telescope: - # Save the prediction to the output file - write_table( - direction_tel_table, - self.output_path, - f"{DL2_TEL_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}/tel_{tel_id:03d}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TEL_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray geometry prediction..." - ) - # Stack the telescope tables to the subarray table - direction_tel_tables = vstack(direction_tel_tables) - # Sort the table by the telescope event keys - direction_tel_tables.sort(TEL_EVENT_KEYS) - # If only one telescope is used, copy the classification table - # and modify it to subarray format - if len(self.dl1dh_reader.tel_ids) == 1: - direction_subarray_table = direction_tel_tables.copy() - telescope_mask = ( - direction_subarray_table["tel_id"] - == self.dl1dh_reader.tel_ids[0] - ) - direction_subarray_table = direction_subarray_table[telescope_mask] - direction_subarray_table.remove_column("tel_id") - for colname in direction_subarray_table.colnames: - if "_tel_" in colname: - direction_subarray_table.rename_column( - colname, colname.replace("_tel", "") - ) - direction_subarray_table.add_column( - [ - [val] - for val in direction_subarray_table[ - f"{self.prefixes['cameradirection']}_is_valid" - ] - ], - name=f"{self.prefixes['cameradirection']}_telescopes", - ) - else: - self.geometry_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefixes["cameradirection"], - property=ReconstructionProperty.GEOMETRY, - parent=self, - ) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - direction_subarray_table = ( - self.geometry_stereo_combiner.predict_table( - direction_tel_tables - ) - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - direction_subarray_table[ - f"{self.prefixes['cameradirection']}_telescopes" - ].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - ( - len(direction_subarray_table), - len(self.dl1dh_reader.tel_ids), - ), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - direction_subarray_table[ - f"{self.prefixes['cameradirection']}_telescopes" - ] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices( - tel_id - ) - ] = True - # Overwrite the column with the boolean mask with fix length - direction_subarray_table[ - f"{self.prefixes['cameradirection']}_telescopes" - ] = reco_telescopes - # Deduplicate the subarray direction table to have only one entry per event - direction_subarray_table = super().deduplicate_first_valid( - table=direction_subarray_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['cameradirection']}_is_valid", - ) - # Sort the subarray geometry table - direction_subarray_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - direction_subarray_table, - self.output_path, - f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['cameradirection']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['cameradirection']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['cameradirection']}", - ) - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - particletype_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TEL_GROUP/features/{prefix}/tel_{tel_id:03d} table. - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = feature_vector_table["tel_id"] == tel_id - feature_vectors_tel_table = feature_vector_table[telescope_mask] - feature_vectors_tel_table.sort(TEL_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vectors_tel_table, - self.output_path, - f"{DL1_TEL_GROUP}/features/{self.prefixes['all']}/tel_{tel_id:03d}", - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TEL_GROUP}/features/{self.prefixes['all']}/tel_{tel_id:03d}", - ) - - def _store_mc_telescope_pointing(self, all_identifiers): - """ - Store the telescope pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - # Create the pointing table for each telescope - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Pointing table for the mono mode - tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) - tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") - tel_pointing.rename_column( - "telescope_pointing_altitude", "pointing_altitude" - ) - # Join the prediction table with the telescope pointing table - tel_pointing = join( - left=tel_pointing, - right=all_identifiers, - keys=["obs_id", "tel_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - tel_pointing.sort(TEL_EVENT_KEYS) - # Retrieve the example identifiers for the selected telescope - tel_pointing_table = Table( - { - "time": tel_pointing["time"], - "azimuth": tel_pointing["pointing_azimuth"], - "altitude": tel_pointing["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TEL_POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info.append(tel_pointing) - pointing_info = vstack(pointing_info) - return pointing_info - - -class StereoPredictCTLearnKerasModel(PredictCTLearnKerasModel): - """ - Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. - - This tool extends the ``PredictCTLearnKerasModel`` to specifically handle stereoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_subarray_pointing(all_identifiers) - Store the subarray pointing table for the stereo mode for MC simulation. - """ - - name = "ctlearn-predict-stereo-keras-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in stereo mode using trained CTLearn models: - > ctlearn-predict-stereo-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnKerasModel.batch_size=16 \\ - --PredictCTLearnKerasModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --DLImageReader.mode=stereo \\ - --DLImageReader.min_telescopes=2 \\ - --PredictCTLearnKerasModel.stack_telescope_images=True \\ - --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ - --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ - --output output.dl2.h5 \\ - """ - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() - example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) - all_identifiers = read_table( - self.output_path, - DL1_TEL_TRIGGER_TABLE, - ) - all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) - # Unique example identifiers by events - all_identifiers = unique(all_identifiers, keys=SUBARRAY_EVENT_KEYS) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Construct the survival telescopes for each event of the example_identifiers - survival_telescopes = [] - for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: - survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) - survival_tels = [ - self.dl1dh_reader.subarray.tel_indices[tel_id] - for tel_id in subarray_event["tel_id"].data - ] - survival_mask[survival_tels] = True - survival_telescopes.append(survival_mask) - # Add the survival telescopes to the example_identifiers - example_identifiers.add_column( - survival_telescopes, name=f"{self.prefixes['all']}_telescopes" - ) - # Pointing table for the stereo mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_subarray_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - particletype_feature_vectors = None - if self.load_type_model_from is not None: - # Predict the classification of the primary particle - particletype_table, particletype_feature_vectors = ( - super()._predict_particletype(example_identifiers) - ) - if self.dl2_subarray: - particletype_table.rename_column( - f"{self.prefixes['all']}_telescopes", - f"{self.prefixes['type']}_telescopes", - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefixes['type']}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - reco_task="type", - ) - particletype_table = vstack([particletype_table, nan_table]) - # Add is_valid column to the particletype table - particletype_table.add_column( - ~np.isnan( - particletype_table[ - f"{self.prefixes['type']}_tel_prediction" - ].data, - dtype=bool, - ), - name=f"{self.prefixes['type']}_is_valid", - ) - # Rename the columns for the stereo mode - particletype_table.rename_column( - f"{self.prefixes['type']}_tel_prediction", - f"{self.prefixes['type']}_prediction", - ) - # Deduplicate the subarray particletype table to have only one entry per event - particletype_table = super().deduplicate_first_valid( - table=particletype_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['type']}_is_valid", - ) - particletype_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - particletype_table, - ParticleClassificationContainer, - prefix=self.prefixes["type"], - ) - # Save the prediction to the output file - write_table( - particletype_table, - self.output_path, - f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_PARTICLETYPE_GROUP}/{self.prefixes['type']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['type']}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - if self.dl2_subarray: - energy_table.rename_column( - f"{self.prefixes['all']}_telescopes", - f"{self.prefixes['energy']}_telescopes", - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefixes['energy']}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - reco_task="energy", - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefixes['energy']}_tel_energy"].data, - dtype=bool, - ), - name=f"{self.prefixes['energy']}_is_valid", - ) - # Rename the columns for the stereo mode - energy_table.rename_column( - f"{self.prefixes['energy']}_tel_energy", - f"{self.prefixes['energy']}_energy", - ) - # Deduplicate the subarray energy table to have only one entry per event - energy_table = super().deduplicate_first_valid( - table=energy_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['energy']}_is_valid", - ) - energy_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefixes["energy"], - ) - # Save the prediction to the output file - write_table( - energy_table, - self.output_path, - f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_ENERGY_GROUP}/{self.prefixes['energy']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['energy']}", - ) - direction_feature_vectors = None - if self.load_skydirection_model_from is not None: - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=SUBARRAY_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = super()._predict_skydirection( - example_identifiers - ) - if self.dl2_subarray: - direction_table.rename_column( - f"{self.prefixes['all']}_telescopes", - f"{self.prefixes['skydirection']}_telescopes", - ) - # Transform the spherical coordinate offsets to sky coordinates - direction_table = super()._transform_spher_coord_offsets_to_sky( - direction_table - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[ - f"{self.prefixes['skydirection']}_alt", - f"{self.prefixes['skydirection']}_az", - ], - shapes=[ - (len(nonexample_identifiers),), - (len(nonexample_identifiers),), - ], - reco_task="skydirection", - ) - direction_table = vstack([direction_table, nan_table]) - # Add is_valid column to the direction table - direction_table.add_column( - ~np.isnan( - direction_table[f"{self.prefixes['skydirection']}_alt"].data, - dtype=bool, - ), - name=f"{self.prefixes['skydirection']}_is_valid", - ) - # Deduplicate the subarray direction table to have only one entry per event - direction_table = super().deduplicate_first_valid( - table=direction_table, - keys=SUBARRAY_EVENT_KEYS, - valid_col=f"{self.prefixes['skydirection']}_is_valid", - ) - direction_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_table, - ReconstructedGeometryContainer, - prefix=self.prefixes["skydirection"], - ) - # Save the prediction to the output file - write_table( - direction_table, - self.output_path, - f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['skydirection']}", - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GEOMETRY_GROUP}/{self.prefixes['skydirection']}", - ) - # Store the telescope event statistics table - write_table( - self.dl1dh_reader.quality_query.to_table(functions=True), - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['skydirection']}", - append=True, - ) - self.log.info( - "DL2 service telescope event statistics data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_EVENT_STATISTICS_GROUP}/{self.prefixes['skydirection']}", - ) - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - particletype_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TEL_GROUP/features/{self.prefixes['all']}/tel_{tel_id:03d} table. - # Rename the columns for the stereo mode - feature_vector_table.rename_column( - f"{self.prefixes['all']}_tel_particletype_feature_vectors", - f"{self.prefixes['all']}_particletype_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefixes['all']}_tel_energy_feature_vectors", - f"{self.prefixes['all']}_energy_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefixes['all']}_tel_geometry_feature_vectors", - f"{self.prefixes['all']}_geometry_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefixes['all']}_tel_is_valid", - f"{self.prefixes['all']}_is_valid", - ) - feature_vector_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vector_table, - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefixes['all']}", - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefixes['all']}", - ) - - def _store_mc_subarray_pointing(self, all_identifiers): - """ - Store the subarray pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the subarray pointing information. - """ - # Read the subarray pointing table - pointing_info = read_table( - self.input_url, - SIMULATION_RUN_TABLE, - ) - # Assuming min_az = max_az and min_alt = max_alt - pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) - pointing_info.rename_column("min_az", "pointing_azimuth") - pointing_info.rename_column("min_alt", "pointing_altitude") - # Join the prediction table with the telescope pointing table - pointing_info = join( - left=pointing_info, - right=all_identifiers, - keys=["obs_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - DL1_SUBARRAY_POINTING_GROUP, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - DL1_SUBARRAY_POINTING_GROUP, - ) - return pointing_info - - -def mono_tool(): - # Run the tool - mono_tool = MonoPredictCTLearnKerasModel() - mono_tool.run() - - -def stereo_tool(): - # Run the tool - stereo_tool = StereoPredictCTLearnKerasModel() - stereo_tool.run() - - -if __name__ == "mono_tool": - mono_tool() - -if __name__ == "stereo_tool": - stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index be6f1180..2328caf3 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -1,19 +1,20 @@ """ -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. +Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``. """ import atexit import uuid +import pathlib import warnings import numpy as np import tables -try: - import tensorflow as tf - import keras -except ImportError: - tf = None - keras = None +import keras +import tensorflow as tf +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from traitlets import TraitError from astropy import units as u from astropy.coordinates.earth import EarthLocation @@ -43,45 +44,44 @@ flag, Dict, ComponentName, - CaselessStrEnum, classes_with_traits, ) from ctapipe.monitoring.interpolation import PointingInterpolator from ctapipe.instrument import SubarrayDescription from ctapipe.io import read_table, write_table, HDF5Merger from ctapipe.io.datalevels import DataLevel - -DL0_TEL_POINTING_GROUP = "/dl0/event/telescope/pointing" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL1_SUBARRAY_POINTING_GROUP = "/dl1/event/subarray/pointing" -DL1_SUBARRAY_TRIGGER_TABLE = "/dl1/event/subarray/trigger" -DL1_TEL_GROUP = "/dl1/event/telescope" -DL1_TEL_CALIBRATION_GROUP = "/dl1/event/telescope/calibration" -DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP = "/dl1/event/telescope/illuminator_throughput" -DL1_TEL_IMAGES_GROUP = "/dl1/event/telescope/image" -DL1_TEL_MUON_GROUP = "/dl1/event/telescope/muon" -DL1_TEL_MUON_THROUGHPUT_GROUP = "/dl1/event/telescope/muon_throughput" -DL1_TEL_OPTICAL_PSF_GROUP = "/dl1/event/telescope/optical_psf" -DL1_TEL_PARAMETERS_GROUP = "/dl1/event/telescope/parameters" -DL1_TEL_POINTING_GROUP = "/dl1/event/telescope/pointing" -DL1_TEL_TRIGGER_TABLE = "/dl1/event/telescope/trigger" -DL2_EVENT_STATISTICS_GROUP = "/dl2/event/subarray/statistics" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -R0_TEL_GROUP = "/r0/event/telescope" -R1_TEL_GROUP = "/r1/event/telescope" -SIMULATION_IMAGES_GROUP = "/simulation/event/telescope/images" -SIMULATION_IMPACT_GROUP = "/simulation/event/telescope/impact" -SIMULATION_PARAMETERS_GROUP = "/simulation/event/telescope/parameters" -SIMULATION_RUN_TABLE = "/simulation/run_config" -SIMULATION_SHOWER_TABLE = "/simulation/event/subarray/shower" -DL2_TEL_PARTICLETYPE_GROUP = "/dl2/event/telescope/classification" -DL2_TEL_ENERGY_GROUP = "/dl2/event/telescope/energy" -DL2_TEL_GEOMETRY_GROUP = "/dl2/event/telescope/geometry" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_SUBARRAY_PARTICLETYPE_GROUP = "/dl2/event/subarray/classification" -DL2_SUBARRAY_ENERGY_GROUP = "/dl2/event/subarray/energy" -DL2_SUBARRAY_GEOMETRY_GROUP = "/dl2/event/subarray/geometry" - +from ctapipe.io.hdf5dataformat import ( + DL0_TEL_POINTING_GROUP, + DL1_SUBARRAY_GROUP, + DL1_SUBARRAY_POINTING_GROUP, + DL1_SUBARRAY_TRIGGER_TABLE, + DL1_TEL_GROUP, + DL1_TEL_CALIBRATION_GROUP, + DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, + DL1_TEL_IMAGES_GROUP, + DL1_TEL_MUON_GROUP, + DL1_TEL_MUON_THROUGHPUT_GROUP, + DL1_TEL_OPTICAL_PSF_GROUP, + DL1_TEL_PARAMETERS_GROUP, + DL1_TEL_POINTING_GROUP, + DL1_TEL_TRIGGER_TABLE, + DL2_EVENT_STATISTICS_GROUP, + FIXED_POINTING_GROUP, + R0_TEL_GROUP, + R1_TEL_GROUP, + SIMULATION_IMAGES_GROUP, + SIMULATION_IMPACT_GROUP, + SIMULATION_PARAMETERS_GROUP, + SIMULATION_RUN_TABLE, + SIMULATION_SHOWER_TABLE, + DL2_TEL_PARTICLETYPE_GROUP, + DL2_TEL_ENERGY_GROUP, + DL2_TEL_GEOMETRY_GROUP, + DL2_SUBARRAY_GROUP, + DL2_SUBARRAY_PARTICLETYPE_GROUP, + DL2_SUBARRAY_ENERGY_GROUP, + DL2_SUBARRAY_GEOMETRY_GROUP, +) from ctapipe.reco.reconstructor import ReconstructionProperty from ctapipe.reco.stereo_combination import StereoCombiner from ctapipe.reco.utils import add_defaults_and_meta @@ -91,7 +91,9 @@ LST_EPOCH, ) from ctlearn import __version__ as ctlearn_version -from ctlearn.core.data_loader.loader import DLDataLoader +from ctlearn.core.ctlearn_enum import FrameworkType +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.core.pytorch.dataset import PyTorchDataset from ctlearn.utils import validate_trait_dict # Convienient constants for column names and table keys @@ -135,7 +137,8 @@ class PredictCTLearnModel(Tool): This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to store them in the output file. The input data is loaded from the input url using the - ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.loader.DLDataLoader``. + ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.keras.sequence.KerasSequence`` or + ``~ctlearn.core.pytorch.dataset.PyTorchDataset``. The prediction is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to handle the prediction for mono and stereo mode. @@ -161,14 +164,14 @@ class PredictCTLearnModel(Tool): prefix : str Name of the reconstruction algorithm used to generate the dl2 data. load_type_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the classification of the primary particle type. + Path to a Keras or PyTorch model file for the classification of the primary particle type. load_energy_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression of the primary particle energy. + Path to a Keras or PyTorch model file for the regression of the primary particle energy. load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression + Path to a Keras or PyTorch model file for the regression of the primary particle arrival direction based on camera coordinate offsets. load_skydirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression + Path to a Keras or PyTorch model file for the regression of the primary particle arrival direction based on spherical coordinate offsets. output_path : pathlib.Path Output path to save the dl2 prediction results. @@ -176,8 +179,8 @@ class PredictCTLearnModel(Tool): Verbosity mode of Keras during the prediction. strategy : tf.distribute.Strategy MirroredStrategy to distribute the prediction. - data_loader : ctlearn.core.loader.DLDataLoader - DLDataLoader object to load the data. + data_loader : ctlearn.core.keras.sequence.KerasSequence + KerasSequence object to load the data. indices : list of int List of indices for the data loaders. batch_size : int @@ -279,8 +282,8 @@ class PredictCTLearnModel(Tool): load_type_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the classification " - "of the primary particle type." + "Path to a Keras or PyTorch model file for the classification of the primary particle type. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -291,8 +294,8 @@ class PredictCTLearnModel(Tool): load_energy_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the regression " - "of the primary particle energy." + "Path to a Keras or PyTorch model file for the regression of the primary particle energy. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -303,8 +306,9 @@ class PredictCTLearnModel(Tool): load_cameradirection_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on camera coordinate offsets." + "Path to a Keras or PyTorch model file for the reconstruction " + "of the primary particle arrival direction based on camera coordinate offsets. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -315,8 +319,9 @@ class PredictCTLearnModel(Tool): load_skydirection_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on spherical coordinate offsets." + "Path to a Keras or PyTorch model file for the reconstruction " + "of the primary particle arrival direction based on spherical coordinate offsets. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -336,23 +341,6 @@ class PredictCTLearnModel(Tool): help="Output path to save the dl2 prediction results", ).tag(config=True) - keras_verbose = Int( - default_value=1, - min=0, - max=2, - allow_none=False, - help=( - "Verbosity mode of Keras during the prediction: " - "0 = silent, 1 = progress bar, 2 = one line per call." - ), - ).tag(config=True) - - framework_type = CaselessStrEnum( - ["pytorch", "keras"], - default_value="keras", - help="Framework to use: pytorch or keras", - ).tag(config=True) - aliases = { ("i", "input_url"): "PredictCTLearnModel.input_url", ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", @@ -363,7 +351,6 @@ class PredictCTLearnModel(Tool): ): "PredictCTLearnModel.load_cameradirection_model_from", ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", ("o", "output"): "PredictCTLearnModel.output_path", - ("f", "framework"): "PredictCTLearnModel.framework_type", } flags = { @@ -429,13 +416,7 @@ class PredictCTLearnModel(Tool): ), } - @property - def classes(self): - return [ - type(self), - PredictCTLearnModel, - HDF5Merger, - ] + classes_with_traits(DLDataReader) + classes = classes_with_traits(DLDataReader) def setup(self): self.activity_start_time = Time.now() @@ -450,14 +431,9 @@ def setup(self): self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self ) as merger: merger(self.input_url) - - if tf is None: - raise ImportError("TensorFlow is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") - - # Create a MirroredStrategy. - self.strategy = tf.distribute.MirroredStrategy() - atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore - self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) + + # Reads the model paths and set up the Keras or PyTorch framework + self._setup_framework() # Set up the data reader self.log.info("Loading data reader:") @@ -477,7 +453,7 @@ def setup(self): # Set the indices for the data loaders self.indices = list(range(self.dl1dh_reader._get_n_events())) self.last_batch_size = len(self.indices) % ( - self.batch_size * self.strategy.num_replicas_in_sync + self.batch_size * self.num_devices ) # Ensure subarray consistency in the output file self._ensure_subarray_consistency() @@ -512,6 +488,87 @@ def _overwrite_meta(self): h5_file.root._v_attrs["CTA PRODUCT ID"] = str(uuid.uuid4()) h5_file.flush() + def _setup_framework(self): + """ + Detect framework from model paths, check consistency, and configure hardware devices. + Raises an error if no valid framework (Keras or PyTorch) is identified. + """ + # Collect all configured model path attributes + model_paths = [ + self.load_type_model_from, + self.load_energy_model_from, + self.load_cameradirection_model_from, + self.load_skydirection_model_from, + ] + # Detect frameworks from all non-None paths + detected_frameworks = {} + for path in model_paths: + if path is not None: + fw = self._detect_framework(path) + detected_frameworks[path] = fw + # Fail immediately if no valid model paths are provided + if not detected_frameworks: + raise ToolConfigurationError( + "No model paths were specified. At least one valid model path " + "(.keras for Keras, or .pt/.pth for PyTorch) must be provided." + ) + # Verify consistency across specified paths + unique_frameworks = set(detected_frameworks.values()) + if len(unique_frameworks) > 1: + details = ", ".join(f"{p}: {fw}" for p, fw in detected_frameworks.items()) + raise ToolConfigurationError( + f"Inconsistent model frameworks detected across paths: {details}. " + "All specified model files must belong to the same framework." + ) + # Convert detected framework string (e.g., 'Keras' / 'PyTorch') to FrameworkType enum + framework_str = next(iter(unique_frameworks)) + try: + self.framework_type = FrameworkType[framework_str.upper()] + except KeyError: + raise ToolConfigurationError( + f"Unsupported framework type '{framework_str}'. " + "Must be either FrameworkType.KERAS or FrameworkType.PYTORCH." + ) + self.log.info("Framework selected: %s", self.framework_type) + # Configure framework-specific hardware/device setup + self.device = None + if self.framework_type == FrameworkType.PYTORCH: + if torch.cuda.is_available(): + self.device = torch.device("cuda") + self.num_devices = torch.cuda.device_count() + self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) + else: + self.device = torch.device("cpu") + self.num_devices = 1 + self.log.info("Using PyTorch CPU device.") + elif self.framework_type == FrameworkType.KERAS: + self.strategy = tf.distribute.MirroredStrategy() + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.num_devices = self.strategy.num_replicas_in_sync + self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) + + @staticmethod + def _detect_framework(path_val): + """ + Determines framework based on file extension. + Returns 'Keras', 'PyTorch', or raises TraitError. + """ + if path_val is None: + return None + + path = pathlib.Path(path_val) + ext = path.suffix.lower() + + if ext in [".keras", ".h5"]: + return "Keras" + elif ext in [".pt", ".pth"]: + return "PyTorch" + else: + raise TraitError( + f"Invalid model extension '{ext}' for file '{path}'. " + "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." + ) + def _get_data_levels(self, h5file): """Get the data levels present in the HDF5 file.""" data_levels = { @@ -697,9 +754,16 @@ def deduplicate_first_valid( return unique(t, keys=list(keys), keep="first") - def _predict_with_model(self, model_path, task): + def _predict_with_model(self, model_path): + """ Select the framework to load and predict the data. """ + if self.framework_type == FrameworkType.KERAS: + return self._predict_with_keras_model(model_path) + else: + return self._predict_with_pytorch_model(model_path) + + def _predict_with_keras_model(self, model_path): """ - Load and predict with a CTLearn model. + Load and predict with a CTLearn Keras-based model. Load a model from the specified path and predict the data using the loaded model. If a last batch loader is provided, predict the last batch and stack the results. @@ -708,8 +772,6 @@ def _predict_with_model(self, model_path, task): ---------- model_path : str Path to a Keras model file (Keras3). - task : str - The task for which prediction is being made. Returns ------- @@ -718,12 +780,249 @@ def _predict_with_model(self, model_path, task): feature_vectors : np.ndarray Feature vectors extracted from the backbone model. """ - if self.framework_type == "keras": - from ctlearn.tools.predict.keras.predic_model_keras import predict_with_model - predict_data, feature_vectors = predict_with_model(self, model_path, task) + # Create a new KerasSequence for each task + # It turned out to be more robust to initialize the KerasSequence separately. + data_loader = KerasSequence( + self.dl1dh_reader, + self.indices, + tasks=[], + batch_size=self.batch_size * self.num_devices, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Keras is only considering the last complete batch. + # In prediction mode we don't want to loose the last + # uncomplete batch, so we are creating an additional + # batch generator for the remaining events. + data_loader_last_batch = None + if self.last_batch_size > 0: + last_batch_indices = self.indices[-self.last_batch_size :] + data_loader_last_batch = KerasSequence( + self.dl1dh_reader, + last_batch_indices, + tasks=[], + batch_size=self.last_batch_size, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Load the model from the specified path + model = keras.saving.load_model(model_path) + prediction_colname = ( + "type" + if isinstance(model.layers[-1], keras.layers.Softmax) + else model.layers[-1].name + ) + backbone_model, feature_vectors = None, None + if self.dl1_features: + # Get the backbone model which is the second layer of the model + backbone_model = model.get_layer(index=1) + # Create a new head model with the same layers as the original model. + # The output of the backbone model is the input of the head model. + backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) + x = backbone_output_shape + for layer in model.layers[2:]: + x = layer(x) + head = keras.Model(inputs=backbone_output_shape, outputs=x) + # Apply the backbone model with the data loader to retrieve the feature vectors + try: + feature_vectors = backbone_model.predict( + data_loader + ) + except ValueError as err: + if str(err).startswith("Input 0 of layer"): + raise ToolConfigurationError( + "Model input shape does not match the prediction data. " + "This is usually caused by selecting the wrong telescope_id. " + "Please ensure the telescope configuration matches the one used for training." + ) from err + raise + # Apply the head model with the feature vectors to retrieve the prediction + predict_data = Table( + { + prediction_colname: head.predict( + feature_vectors + ) + } + ) + # Predict the last batch and stack the results to the prediction data + if data_loader_last_batch is not None: + feature_vectors_last_batch = backbone_model.predict( + data_loader_last_batch + ) + feature_vectors = np.concatenate( + (feature_vectors, feature_vectors_last_batch) + ) + predict_data = vstack( + [ + predict_data, + Table( + { + prediction_colname: head.predict( + feature_vectors_last_batch + ) + } + ), + ] + ) + else: + # Predict the data using the loaded model + try: + predict_data = model.predict(data_loader) + except ValueError as err: + if str(err).startswith("Input 0 of layer"): + raise ToolConfigurationError( + "Model input shape does not match the prediction data. " + "This is usually caused by selecting the wrong telescope_id. " + "Please ensure the telescope configuration matches the one used for training." + ) from err + raise + # Create a astropy table with the prediction results + # The classification task has a softmax layer as the last layer + # which returns the probabilities for each class in an array, while + # the regression tasks have output neurons which returns the + # predicted value for the task in a dictionary. + if prediction_colname == "type": + predict_data = Table({prediction_colname: predict_data}) + else: + predict_data = Table(predict_data) + # Predict the last batch and stack the results to the prediction data + if data_loader_last_batch is not None: + predict_data_last_batch = model.predict( + data_loader_last_batch + ) + if model.layers[-1].name == "type": + predict_data_last_batch = Table( + {prediction_colname: predict_data_last_batch} + ) + else: + predict_data_last_batch = Table(predict_data_last_batch) + predict_data = vstack([predict_data, predict_data_last_batch]) + return predict_data, feature_vectors + + + def _predict_with_pytorch_model(self, model_path): + """ + Load and predict with a CTLearn PyTorch-based model. - return predict_data, feature_vectors - + Parameters + ---------- + model_path : str + Path to a PyTorch model checkpoint/file. + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + # Load the model and make sure it is not a state dict + model = torch.load(model_path, map_location=self.device, weights_only=False) + if not isinstance(model, nn.Module): + raise TypeError( + f"Expected a PyTorch 'nn.Module' object at '{model_path}', " + f"but got '{type(model).__name__}'. " + "Ensure the model was saved via 'torch.save(model, path)' rather than 'torch.save(model.state_dict(), path)'." + ) + model.eval() + + # PyTorch DataLoaders natively handle drop_last=False, + # so we don't need a separate generator for incomplete batches. + # Replace 'PyTorchDataset' with your actual PyTorch Dataset implementation + dataset = PyTorchDataset( + self.dl1dh_reader, + self.indices, + tasks=[], + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + data_loader = DataLoader( + dataset, + batch_size=self.batch_size * self.num_devices, + shuffle=False, + pin_memory=torch.cuda.is_available(), + drop_last=False + ) + + feature_vectors_list = [] + predictions_dict = {} + with torch.no_grad(): + for batch in data_loader: # Iterate over DataLoader without subscripting + inputs = ( + batch[0].to(self.device) + if isinstance(batch, (tuple, list)) + else batch.to(self.device) + ) + + try: + if self.dl1_features: + # Extract features from backbone + features = ( + model.backbone(inputs) + if hasattr(model, "backbone") + else model.extract_features(inputs) + ) + # Retrieve classifier / head module + head_model = ( + model.head + if hasattr(model, "head") + else model.classifier + ) + preds = head_model(features) + + feature_vectors_list.append(features.cpu().numpy()) + else: + preds = model(inputs) + + except RuntimeError as err: + if "shape" in str(err).lower() or "mat1 and mat2" in str(err).lower(): + raise ToolConfigurationError( + "Model input shape does not match the prediction data. " + "This is usually caused by selecting the wrong telescope_id. " + "Please ensure the telescope configuration matches the one used for training." + ) from err + raise + + # Standardize model outputs into predictions_dict using self.heads_dict + if isinstance(preds, torch.Tensor): + # Check for heads_dict on model or head_model, fallback to single_output_task or default + heads_dict = ( + getattr(model, "heads_dict", None) + or getattr(getattr(model, "head", None), "heads_dict", None) + or getattr(getattr(model, "classifier", None), "heads_dict", None) + ) + + if heads_dict: + task_name = list(heads_dict.keys())[0] + elif hasattr(model, "single_output_task") and model.single_output_task: + task_name = model.single_output_task + else: + task_name = getattr(self, "prediction_type", "type") + + predictions_dict.setdefault(task_name, []).append(preds.cpu().numpy()) + + elif isinstance(preds, dict): + for task_name, output_tensor in preds.items(): + predictions_dict.setdefault(task_name, []).append( + output_tensor.cpu().numpy() + ) + + # Concatenate batched prediction arrays + formatted_predictions = { + task_name: np.concatenate(batches, axis=0) + for task_name, batches in predictions_dict.items() + } + + # Format into Astropy Table + predict_data = Table(formatted_predictions) + + # Concatenate features if extracted + feature_vectors = ( + np.concatenate(feature_vectors_list, axis=0) + if self.dl1_features and feature_vectors_list + else None + ) return predict_data, feature_vectors @@ -748,7 +1047,7 @@ def _predict_particletype(self, example_identifiers): ) # Predict the data using the loaded type_model predict_data, feature_vectors = self._predict_with_model( - self.load_type_model_from, "type" + self.load_type_model_from ) # Create prediction table and add the predicted classification score ('gammaness') particletype_table = example_identifiers.copy() @@ -776,7 +1075,7 @@ def _predict_energy(self, example_identifiers): self.log.info("Predicting for the regression of the primary particle energy...") # Predict the data using the loaded energy_model predict_data, feature_vectors = self._predict_with_model( - self.load_energy_model_from, "energy" + self.load_energy_model_from ) # Convert the reconstructed energy from log10(TeV) to TeV reco_energy = u.Quantity( @@ -816,7 +1115,7 @@ def _predict_cameradirection(self, example_identifiers): ) # Predict the data using the loaded direction_model predict_data, feature_vectors = self._predict_with_model( - self.load_cameradirection_model_from, "cameradirection" + self.load_cameradirection_model_from ) # For the direction task, the prediction is the camera coordinate offset in x and y # from the telescope pointing. @@ -854,7 +1153,7 @@ def _predict_skydirection(self, example_identifiers): ) # Predict the data using the loaded direction_model predict_data, feature_vectors = self._predict_with_model( - self.load_skydirection_model_from, "skydirection" + self.load_skydirection_model_from ) # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat # from the telescope pointing. @@ -1245,9 +1544,9 @@ class MonoPredictCTLearnModel(PredictCTLearnModel): --DLImageReader.channels=cleaned_image \\ --DLImageReader.channels=cleaned_relative_peak_time \\ --DLImageReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/mono/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/mono/energy/ctlearn_model(.keras/.pth)" \\ + --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model(.keras/.pth)" \\ --dl1-features \\ --no-dl1-images \\ --no-true-images \\ @@ -1259,9 +1558,9 @@ class MonoPredictCTLearnModel(PredictCTLearnModel): --PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader \\ --DLWaveformReader.sequnce_length=20 \\ --DLWaveformReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/mono_waveform/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model(.keras/.pth)" \\ + --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model(.keras/.pth)" \\ --no-r0-waveforms \\ --no-r1-waveforms \\ --no-dl1-images \\ @@ -1905,9 +2204,9 @@ class StereoPredictCTLearnModel(PredictCTLearnModel): --DLImageReader.mode=stereo \\ --DLImageReader.min_telescopes=2 \\ --PredictCTLearnModel.stack_telescope_images=True \\ - --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ - --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/stereo/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/stereo/energy/ctlearn_model(.keras/.pth)" \\ + --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model(.keras/.pth)" \\ --output output.dl2.h5 \\ """ @@ -2223,14 +2522,10 @@ def _store_mc_subarray_pointing(self, all_identifiers): Table containing the subarray pointing information. """ # Read the subarray pointing table - try: - pointing_info = read_table( - self.input_url, - SIMULATION_RUN_TABLE, - ) - except Exception as e: - self.log.warning("Could not read simulation run table: %s", e) - return None + pointing_info = read_table( + self.input_url, + SIMULATION_RUN_TABLE, + ) # Assuming min_az = max_az and min_alt = max_alt pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) pointing_info.rename_column("min_az", "pointing_azimuth") @@ -2283,4 +2578,4 @@ def stereo_tool(): mono_tool() if __name__ == "stereo_tool": - stereo_tool() + stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/tests/test_predict_model.py b/ctlearn/tools/tests/test_predict_model.py index 69f4b060..862bdf31 100644 --- a/ctlearn/tools/tests/test_predict_model.py +++ b/ctlearn/tools/tests/test_predict_model.py @@ -5,7 +5,7 @@ from ctapipe.core import run_tool from ctapipe.io import TableLoader from ctlearn.conftest import MODEL_FILE_FORMATS -from ctlearn.tools.keras import MonoPredictCTLearnKerasModel, StereoPredictCTLearnKerasModel +from ctlearn.tools import MonoPredictCTLearnModel, StereoPredictCTLearnModel # Columns that should be present in the output DL2 file REQUIRED_COLUMNS = [ @@ -37,7 +37,7 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2") -@pytest.mark.parametrize("framework", ["Keras"]) +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_mono_model_with_r1_waveforms( tmp_path, ctlearn_trained_r1_mono_models, r1_gamma_file, framework ): @@ -60,15 +60,15 @@ def test_predict_mono_model_with_r1_waveforms( key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_r1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" # Build command-line arguments argv = [ f"--input_url={r1_gamma_file}", - "--PredictCTLearnKerasModel.batch_size=2", - "--PredictCTLearnKerasModel.dl1dh_reader_type=DLWaveformReader", + "--PredictCTLearnModel.batch_size=2", + "--PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader", "--DLWaveformReader.sequence_length=5", "--DLWaveformReader.focal_length_choice=EQUIVALENT", "--no-r1-waveforms", @@ -78,13 +78,13 @@ def test_predict_mono_model_with_r1_waveforms( # Run Prediction tool assert ( run_tool( - MonoPredictCTLearnKerasModel(), + MonoPredictCTLearnModel(), argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) @@ -131,7 +131,7 @@ def test_predict_mono_model_with_r1_waveforms( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") -@pytest.mark.parametrize("framework", ["Keras"]) +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) @pytest.mark.parametrize("dl2_tel_flag", ["dl2-telescope", "no-dl2-telescope"]) def test_predict_mono_model_with_dl1_images( tmp_path, ctlearn_trained_dl1_mono_models, dl1_gamma_file, framework, dl2_tel_flag @@ -170,7 +170,7 @@ def test_predict_mono_model_with_dl1_images( # Build command-line arguments argv = [ f"--input_url={dl1_gamma_file}", - "--PredictCTLearnKerasModel.batch_size=2", + "--PredictCTLearnModel.batch_size=2", "--DLImageReader.focal_length_choice=EQUIVALENT", "--no-dl1-images", "--no-true-images", @@ -183,15 +183,15 @@ def test_predict_mono_model_with_dl1_images( # Run Prediction tool assert ( run_tool( - MonoPredictCTLearnKerasModel(), + MonoPredictCTLearnModel(), argv=argv + [ f"--output={output_file}", f"--DLImageReader.allowed_tels={allowed_tels}", f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", - f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) @@ -240,7 +240,7 @@ def test_predict_mono_model_with_dl1_images( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") -@pytest.mark.parametrize("framework", ["Keras"]) +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_stereo_model_with_dl1_images( tmp_path, ctlearn_trained_dl1_stereo_models, dl1_gamma_file, framework ): @@ -270,8 +270,8 @@ def test_predict_stereo_model_with_dl1_images( # Build command-line arguments argv = [ f"--input_url={dl1_gamma_file}", - "--PredictCTLearnKerasModel.batch_size=2", - "--PredictCTLearnKerasModel.stack_telescope_images=True", + "--PredictCTLearnModel.batch_size=2", + "--PredictCTLearnModel.stack_telescope_images=True", "--DLImageReader.mode=stereo", "--DLImageReader.focal_length_choice=EQUIVALENT", f"--DLImageReader.allowed_tels={allowed_tels}", @@ -282,13 +282,13 @@ def test_predict_stereo_model_with_dl1_images( # Run Prediction tool assert ( run_tool( - StereoPredictCTLearnKerasModel(), + StereoPredictCTLearnModel(), argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnKerasModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", - f"--PredictCTLearnKerasModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_skydirection.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_skydirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) diff --git a/pyproject.toml b/pyproject.toml index 79d2b35a..7c662447 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,8 +89,8 @@ documentation = "https://ctlearn.readthedocs.io/en/latest/" [project.scripts] ctlearn-train-keras-model = "ctlearn.tools.keras.train_model:main" ctlearn-train-pytorch-model = "ctlearn.tools.pytorch.train_model:main" -ctlearn-predict-mono-keras-model = "ctlearn.tools.keras.predict_model:mono_tool" -ctlearn-predict-stereo-keras-model = "ctlearn.tools.predict_model:stereo_tool" +ctlearn-predict-mono-model = "ctlearn.tools.predict_model:mono_tool" +ctlearn-predict-stereo-model = "ctlearn.tools.predict_model:stereo_tool" ctlearn-predict-LST1= "ctlearn.tools.predict_LST1:main" [tool.setuptools_scm] From bcb8b6ab3beca2475c48d1f13b8ed5bb92f1cc78 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 11:30:17 +0200 Subject: [PATCH 07/12] polish prediction pytorch part --- ctlearn/core/pytorch/model.py | 8 +-- ctlearn/tools/predict_model.py | 87 ++++++---------------------- ctlearn/tools/pytorch/train_model.py | 4 +- 3 files changed, 23 insertions(+), 76 deletions(-) diff --git a/ctlearn/core/pytorch/model.py b/ctlearn/core/pytorch/model.py index 07d3f3e1..99eda31c 100644 --- a/ctlearn/core/pytorch/model.py +++ b/ctlearn/core/pytorch/model.py @@ -21,7 +21,7 @@ __all__ = [ "BasicBlock", "BottleneckBlock", - "MultiHeadClassifier", + "MultiFullyConnectedHead", "build_fully_connect_pytorch_head", "PyTorchSingleCNN", "PyTorchResNet", @@ -29,7 +29,7 @@ ] -class MultiHeadClassifier(nn.Module): +class MultiFullyConnectedHead(nn.Module): """ A PyTorch container module to hold the multi-task fully connected heads. """ @@ -99,7 +99,7 @@ def build_fully_connect_pytorch_head(in_features, layers, activation_function, t heads[task] = nn.Sequential(*task_layers) single_output_task = tasks[0] if (len(tasks) == 1 and tasks[0] == "type") else None - return MultiHeadClassifier(heads, single_output_task=single_output_task) + return MultiFullyConnectedHead(heads, single_output_task=single_output_task) class FullModelPipeline(nn.Module): @@ -113,7 +113,7 @@ def __init__(self, backbone, head): def forward(self, x): features = self.backbone(x) - return self.head(features) + return self.head(features), features class PyTorchSingleCNN(SingleCNN): diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index 2328caf3..da5ab5bd 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -925,7 +925,6 @@ def _predict_with_pytorch_model(self, model_path): "Ensure the model was saved via 'torch.save(model, path)' rather than 'torch.save(model.state_dict(), path)'." ) model.eval() - # PyTorch DataLoaders natively handle drop_last=False, # so we don't need a separate generator for incomplete batches. # Replace 'PyTorchDataset' with your actual PyTorch Dataset implementation @@ -936,7 +935,6 @@ def _predict_with_pytorch_model(self, model_path): sort_by_intensity=self.sort_by_intensity, stack_telescope_images=self.stack_telescope_images, ) - data_loader = DataLoader( dataset, batch_size=self.batch_size * self.num_devices, @@ -944,86 +942,35 @@ def _predict_with_pytorch_model(self, model_path): pin_memory=torch.cuda.is_available(), drop_last=False ) - - feature_vectors_list = [] - predictions_dict = {} + predictions_dict, feature_vectors_list = {}, [] with torch.no_grad(): for batch in data_loader: # Iterate over DataLoader without subscripting - inputs = ( - batch[0].to(self.device) - if isinstance(batch, (tuple, list)) - else batch.to(self.device) - ) - - try: - if self.dl1_features: - # Extract features from backbone - features = ( - model.backbone(inputs) - if hasattr(model, "backbone") - else model.extract_features(inputs) - ) - # Retrieve classifier / head module - head_model = ( - model.head - if hasattr(model, "head") - else model.classifier - ) - preds = head_model(features) - - feature_vectors_list.append(features.cpu().numpy()) - else: - preds = model(inputs) - - except RuntimeError as err: - if "shape" in str(err).lower() or "mat1 and mat2" in str(err).lower(): - raise ToolConfigurationError( - "Model input shape does not match the prediction data. " - "This is usually caused by selecting the wrong telescope_id. " - "Please ensure the telescope configuration matches the one used for training." - ) from err - raise - + inputs = batch[0].to(self.device) + predictions, features = model(inputs) + if self.dl1_features: + feature_vectors_list.append(features.cpu().numpy()) # Standardize model outputs into predictions_dict using self.heads_dict - if isinstance(preds, torch.Tensor): - # Check for heads_dict on model or head_model, fallback to single_output_task or default - heads_dict = ( - getattr(model, "heads_dict", None) - or getattr(getattr(model, "head", None), "heads_dict", None) - or getattr(getattr(model, "classifier", None), "heads_dict", None) - ) - - if heads_dict: - task_name = list(heads_dict.keys())[0] - elif hasattr(model, "single_output_task") and model.single_output_task: - task_name = model.single_output_task - else: - task_name = getattr(self, "prediction_type", "type") - - predictions_dict.setdefault(task_name, []).append(preds.cpu().numpy()) - - elif isinstance(preds, dict): - for task_name, output_tensor in preds.items(): + if isinstance(predictions, torch.Tensor): + task_name = list(model.head.heads_dict.keys())[0] + predictions_dict.setdefault(task_name, []).append(predictions.cpu().numpy()) + elif isinstance(predictions, dict): + for task_name, output_tensor in predictions.items(): predictions_dict.setdefault(task_name, []).append( output_tensor.cpu().numpy() ) - - # Concatenate batched prediction arrays - formatted_predictions = { - task_name: np.concatenate(batches, axis=0) - for task_name, batches in predictions_dict.items() - } - - # Format into Astropy Table - predict_data = Table(formatted_predictions) - + # Concatenate batched prediction arrays and format into Astropy Table + predict_data = Table( + { + task_name: np.concatenate(batches, axis=0) + for task_name, batches in predictions_dict.items() + } + ) # Concatenate features if extracted feature_vectors = ( np.concatenate(feature_vectors_list, axis=0) if self.dl1_features and feature_vectors_list else None ) - return predict_data, feature_vectors def _predict_particletype(self, example_identifiers): diff --git a/ctlearn/tools/pytorch/train_model.py b/ctlearn/tools/pytorch/train_model.py index 37564dae..b0460405 100644 --- a/ctlearn/tools/pytorch/train_model.py +++ b/ctlearn/tools/pytorch/train_model.py @@ -296,7 +296,7 @@ def _train_epoch(self): batch_y = self._to_device(batch_y) self.opt.zero_grad() - outputs = self.model(batch_x) + outputs, _ = self.model(batch_x) loss = self._compute_combined_loss(outputs, batch_y) loss.backward() @@ -321,7 +321,7 @@ def _validate_epoch(self): batch_x = self._to_device(batch_x) batch_y = self._to_device(batch_y) - outputs = self.model(batch_x) + outputs, _ = self.model(batch_x) loss = self._compute_combined_loss(outputs, batch_y) total_loss += loss.item() From 0b0ede5e109aa0f97b91972f8e6b67fb7834b72a Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 12:46:08 +0200 Subject: [PATCH 08/12] further polishing --- ctlearn/conftest.py | 4 +- ctlearn/core/ctlearn_enum.py | 28 -------------- ctlearn/core/model.py | 2 +- ctlearn/tools/predict_model.py | 51 ++++-------------------- ctlearn/tools/train_model.py | 2 +- ctlearn/{ => tools}/utils.py | 71 +++++++++++++++++++++++++++++++--- 6 files changed, 77 insertions(+), 81 deletions(-) rename ctlearn/{ => tools}/utils.py (59%) diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index 41dc2115..46ced9bb 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -16,13 +16,11 @@ from ctapipe.utils import get_dataset_path from ctlearn.tools.keras.train_model import TrainCTLearnKerasModel from ctlearn.tools.pytorch.train_model import TrainCTLearnPyTorchModel -from ctlearn.utils import get_lst1_subarray_description +from ctlearn.tools.utils import get_lst1_subarray_description -# TODO: ADD PyTorch here TRAINING_TOOLS = {"Keras": TrainCTLearnKerasModel, "PyTorch": TrainCTLearnPyTorchModel} MODEL_FILE_FORMATS = {"Keras": "keras", "PyTorch": "pth"} - @pytest.fixture(scope="session") def gamma_simtel_path(): return get_dataset_path("gamma_test_large.simtel.gz") diff --git a/ctlearn/core/ctlearn_enum.py b/ctlearn/core/ctlearn_enum.py index 1854525b..c102c623 100644 --- a/ctlearn/core/ctlearn_enum.py +++ b/ctlearn/core/ctlearn_enum.py @@ -14,34 +14,6 @@ from enum import Enum -class FrameworkType(Enum): - """ - Deep learning framework type enumeration. - - This enumeration specifies which deep learning framework to use for - model training and inference. CTLearn supports both Keras (TensorFlow backend) - and PyTorch frameworks. - - Attributes: - KERAS (int): Use Keras/TensorFlow framework (value: 1) - - Advantages: High-level API, easy to use, good for prototyping - - TensorFlow 2.x with Keras API - - Suitable for production deployment - - PYTORCH (int): Use PyTorch framework (value: 2) - - Advantages: Dynamic computation graphs, flexible, research-friendly - - PyTorch 1.x or 2.x - - Better for custom architectures and experimental models - - Example: - >>> from ctlearn.core.ctlearn_enum import FrameworkType - >>> framework = FrameworkType.PYTORCH - >>> print(framework.name) # 'PYTORCH' - >>> print(framework.value) # 'PyTorch' - """ - KERAS = "Keras" - PYTORCH = "PyTorch" - class Task(Enum): """ diff --git a/ctlearn/core/model.py b/ctlearn/core/model.py index 153f3f18..5f4bf31e 100644 --- a/ctlearn/core/model.py +++ b/ctlearn/core/model.py @@ -6,7 +6,7 @@ from ctapipe.core import Component from ctapipe.core.traits import Bool, Int, CaselessStrEnum, List, Dict, Unicode, Path -from ctlearn.utils import validate_trait_dict +from ctlearn.tools.utils import validate_trait_dict __all__ = [ "CTLearnModel", diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index da5ab5bd..92b06637 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -1,10 +1,9 @@ """ -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``. +Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``/``PyTorchDataset``. """ import atexit import uuid -import pathlib import warnings import numpy as np @@ -14,14 +13,11 @@ import torch import torch.nn as nn from torch.utils.data import DataLoader -from traitlets import TraitError from astropy import units as u -from astropy.coordinates.earth import EarthLocation from astropy.coordinates import AltAz, SkyCoord from astropy.table import ( Table, - hstack, vstack, join, setdiff, @@ -91,10 +87,9 @@ LST_EPOCH, ) from ctlearn import __version__ as ctlearn_version -from ctlearn.core.ctlearn_enum import FrameworkType from ctlearn.core.keras.sequence import KerasSequence from ctlearn.core.pytorch.dataset import PyTorchDataset -from ctlearn.utils import validate_trait_dict +from ctlearn.tools.utils import detect_framework, FrameworkType, validate_trait_dict # Convienient constants for column names and table keys SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] @@ -504,32 +499,24 @@ def _setup_framework(self): detected_frameworks = {} for path in model_paths: if path is not None: - fw = self._detect_framework(path) - detected_frameworks[path] = fw + detected_frameworks[path] = detect_framework(path) # Fail immediately if no valid model paths are provided if not detected_frameworks: raise ToolConfigurationError( "No model paths were specified. At least one valid model path " - "(.keras for Keras, or .pt/.pth for PyTorch) must be provided." + "(.keras/.h5 for Keras, or .pt/.pth for PyTorch) must be provided." ) # Verify consistency across specified paths unique_frameworks = set(detected_frameworks.values()) if len(unique_frameworks) > 1: - details = ", ".join(f"{p}: {fw}" for p, fw in detected_frameworks.items()) + details = ", ".join(f"{p}: {fw.value}" for p, fw in detected_frameworks.items()) raise ToolConfigurationError( f"Inconsistent model frameworks detected across paths: {details}. " "All specified model files must belong to the same framework." ) - # Convert detected framework string (e.g., 'Keras' / 'PyTorch') to FrameworkType enum - framework_str = next(iter(unique_frameworks)) - try: - self.framework_type = FrameworkType[framework_str.upper()] - except KeyError: - raise ToolConfigurationError( - f"Unsupported framework type '{framework_str}'. " - "Must be either FrameworkType.KERAS or FrameworkType.PYTORCH." - ) - self.log.info("Framework selected: %s", self.framework_type) + # Set framework directly from the detected FrameworkType enum + self.framework_type = next(iter(unique_frameworks)) + self.log.info("Framework selected: %s", self.framework_type.value) # Configure framework-specific hardware/device setup self.device = None if self.framework_type == FrameworkType.PYTORCH: @@ -547,28 +534,6 @@ def _setup_framework(self): self.num_devices = self.strategy.num_replicas_in_sync self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) - @staticmethod - def _detect_framework(path_val): - """ - Determines framework based on file extension. - Returns 'Keras', 'PyTorch', or raises TraitError. - """ - if path_val is None: - return None - - path = pathlib.Path(path_val) - ext = path.suffix.lower() - - if ext in [".keras", ".h5"]: - return "Keras" - elif ext in [".pt", ".pth"]: - return "PyTorch" - else: - raise TraitError( - f"Invalid model extension '{ext}' for file '{path}'. " - "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." - ) - def _get_data_levels(self, h5file): """Get the data levels present in the HDF5 file.""" data_levels = { diff --git a/ctlearn/tools/train_model.py b/ctlearn/tools/train_model.py index e26eca67..b6d8e9b1 100644 --- a/ctlearn/tools/train_model.py +++ b/ctlearn/tools/train_model.py @@ -22,7 +22,7 @@ ) from ctlearn import __version__ as ctlearn_version from ctlearn.core.model import CTLearnModel -from ctlearn.utils import validate_trait_dict +from ctlearn.tools.utils import validate_trait_dict from dl1_data_handler.reader import DLDataReader diff --git a/ctlearn/utils.py b/ctlearn/tools/utils.py similarity index 59% rename from ctlearn/utils.py rename to ctlearn/tools/utils.py index ee16420e..82139ed0 100644 --- a/ctlearn/utils.py +++ b/ctlearn/tools/utils.py @@ -1,15 +1,27 @@ +""" +Utility functions for the CTLearn tools. +""" + +from enum import Enum +import pathlib from importlib.resources import files, as_file +import os +import time +from tqdm import tqdm from ctapipe.core import Provenance from ctapipe.core.traits import TraitError from ctapipe.instrument.optics import FocalLengthKind from ctapipe.instrument import SubarrayDescription -import os -import time -from tqdm import tqdm -__all__ = ["validate_trait_dict","get_lst1_subarray_description","monitor_progress"] +__all__ = [ + "monitor_progress", + "validate_trait_dict", + "get_lst1_subarray_description", + "FrameworkType", + "detect_framework", +] def monitor_progress(src_path, dst_path, stop_event, logger): try: @@ -85,4 +97,53 @@ def get_lst1_subarray_description(focal_length_choice=FocalLengthKind.EFFECTIVE) """ with as_file(files("ctlearn") / "resources/LST-1_SubarrayDescription.h5") as path: Provenance().add_input_file(path, role="SubarrayDescription") - return SubarrayDescription.from_hdf(path, focal_length_choice=focal_length_choice) \ No newline at end of file + return SubarrayDescription.from_hdf(path, focal_length_choice=focal_length_choice) + +class FrameworkType(Enum): + """ + Deep learning framework type enumeration. + + This enumeration specifies which deep learning framework to use for + model training and inference. CTLearn supports both Keras (TensorFlow backend) + and PyTorch frameworks. + + Attributes: + KERAS (int): Use Keras/TensorFlow framework (value: 1) + - Advantages: High-level API, easy to use, good for prototyping + - TensorFlow 2.x with Keras API + - Suitable for production deployment + + PYTORCH (int): Use PyTorch framework (value: 2) + - Advantages: Dynamic computation graphs, flexible, research-friendly + - PyTorch 1.x or 2.x + - Better for custom architectures and experimental models + + Example: + >>> from ctlearn.core.ctlearn_enum import FrameworkType + >>> framework = FrameworkType.PYTORCH + >>> print(framework.name) # 'PYTORCH' + >>> print(framework.value) # 'PyTorch' + """ + KERAS = "Keras" + PYTORCH = "PyTorch" + +def detect_framework(path_val): + """ + Determines framework based on file extension. + Returns 'Keras', 'PyTorch', or raises TraitError. + """ + if path_val is None: + return None + + path = pathlib.Path(path_val) + ext = path.suffix.lower() + + if ext in [".keras", ".h5"]: + return FrameworkType["KERAS"] + elif ext in [".pt", ".pth"]: + return FrameworkType["PYTORCH"] + else: + raise TraitError( + f"Invalid model extension '{ext}' for file '{path}'. " + "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." + ) From 9fea9ba39432c6d2f64878b1022b950b13cb3386 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 13:15:47 +0200 Subject: [PATCH 09/12] move framework setup to utils to be reused in LST1 prediction tool --- ctlearn/tools/predict_model.py | 29 +++++++++++--- ctlearn/tools/utils.py | 70 ++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 5 deletions(-) diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index 92b06637..2dec24c9 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -89,7 +89,12 @@ from ctlearn import __version__ as ctlearn_version from ctlearn.core.keras.sequence import KerasSequence from ctlearn.core.pytorch.dataset import PyTorchDataset -from ctlearn.tools.utils import detect_framework, FrameworkType, validate_trait_dict +from ctlearn.tools.utils import ( + FrameworkType, + detect_framework, + setup_framework, + validate_trait_dict, +) # Convienient constants for column names and table keys SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] @@ -426,10 +431,24 @@ def setup(self): self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self ) as merger: merger(self.input_url) - - # Reads the model paths and set up the Keras or PyTorch framework - self._setup_framework() - + # Collect all configured model path attributes + model_paths = [ + self.load_type_model_from, + self.load_energy_model_from, + self.load_cameradirection_model_from, + self.load_skydirection_model_from, + ] + # Reads the model paths and set up the Keras or PyTorch framework + self.framework_type, self.num_devices, self.strategy, self.device = setup_framework(model_paths) + self.log.info("Framework selected: %s", self.framework_type.value) + if self.framework_type == FrameworkType.KERAS: + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) + elif self.framework_type == FrameworkType.PYTORCH: + if self.device == torch.device("cpu"): + self.log.info("Using PyTorch CPU device.") + else: + self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) # Set up the data reader self.log.info("Loading data reader:") self.log.info("For a large dataset, this may take a while...") diff --git a/ctlearn/tools/utils.py b/ctlearn/tools/utils.py index 82139ed0..0ce29029 100644 --- a/ctlearn/tools/utils.py +++ b/ctlearn/tools/utils.py @@ -9,7 +9,11 @@ import time from tqdm import tqdm +import tensorflow as tf +import torch + from ctapipe.core import Provenance +from ctapipe.core.tool import ToolConfigurationError from ctapipe.core.traits import TraitError from ctapipe.instrument.optics import FocalLengthKind from ctapipe.instrument import SubarrayDescription @@ -21,6 +25,7 @@ "get_lst1_subarray_description", "FrameworkType", "detect_framework", + "setup_framework", ] def monitor_progress(src_path, dst_path, stop_event, logger): @@ -147,3 +152,68 @@ def detect_framework(path_val): f"Invalid model extension '{ext}' for file '{path}'. " "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." ) + +def setup_framework(model_paths): + """ + Detects the deep learning framework from model paths, ensures consistency, + and configures the hardware devices for distributed or single-device execution. + + This function iterates through the provided model paths, infers the framework + (Keras or PyTorch) based on file extensions, and ensures all models belong + to the same framework. It then initializes the appropriate hardware setup + (e.g., MirroredStrategy for Keras, CUDA/CPU device for PyTorch). + + Parameters + ---------- + model_paths : list of str or pathlib.Path or None + A list of paths pointing to the saved model files. `None` values are ignored. + + Returns + ------- + tuple + A tuple containing: + - framework_type (FrameworkType): The identified framework enum (FrameworkType.KERAS or FrameworkType.PYTORCH). + - num_devices (int): The number of devices available/configured for the framework. + - strategy (tf.distribute.Strategy or None): The TensorFlow distribution strategy if Keras is detected, else None. + - device (torch.device or None): The PyTorch target device if PyTorch is detected, else None. + + Raises + ------ + ToolConfigurationError + If no valid model paths are provided, or if multiple inconsistent frameworks + are detected across the provided paths. + """ + # Detect frameworks from all non-None paths + detected_frameworks = {} + for path in model_paths: + if path is not None: + detected_frameworks[path] = detect_framework(path) + # Fail immediately if no valid model paths are provided + if not detected_frameworks: + raise ToolConfigurationError( + "No model paths were specified. At least one valid model path " + "(.keras/.h5 for Keras, or .pt/.pth for PyTorch) must be provided." + ) + # Verify consistency across specified paths + unique_frameworks = set(detected_frameworks.values()) + if len(unique_frameworks) > 1: + details = ", ".join(f"{p}: {fw.value}" for p, fw in detected_frameworks.items()) + raise ToolConfigurationError( + f"Inconsistent model frameworks detected across paths: {details}. " + "All specified model files must belong to the same framework." + ) + # Set framework directly from the detected FrameworkType enum + framework_type = next(iter(unique_frameworks)) + # Configure framework-specific hardware/device setup + strategy, device = None, None + if framework_type == FrameworkType.KERAS: + strategy = tf.distribute.MirroredStrategy() + num_devices = strategy.num_replicas_in_sync + elif framework_type == FrameworkType.PYTORCH: + if torch.cuda.is_available(): + device = torch.device("cuda") + num_devices = torch.cuda.device_count() + else: + device = torch.device("cpu") + num_devices = 1 + return framework_type, num_devices, strategy, device \ No newline at end of file From ec8c1d26290570967d96669670efa461093ce93a Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 14:04:21 +0200 Subject: [PATCH 10/12] more polishing --- ctlearn/tools/predict/predict_LST1.py | 2 +- ctlearn/tools/predict_model.py | 54 +-------------------------- ctlearn/tools/utils.py | 39 +++++++++---------- 3 files changed, 19 insertions(+), 76 deletions(-) diff --git a/ctlearn/tools/predict/predict_LST1.py b/ctlearn/tools/predict/predict_LST1.py index ab3671bc..a9ee4029 100644 --- a/ctlearn/tools/predict/predict_LST1.py +++ b/ctlearn/tools/predict/predict_LST1.py @@ -35,7 +35,7 @@ from ctapipe.io import read_table, write_table from ctapipe.reco.utils import add_defaults_and_meta -from ctlearn.utils import get_lst1_subarray_description +from ctlearn.tools.utils import get_lst1_subarray_description from dl1_data_handler.image_mapper import ImageMapper from dl1_data_handler.reader import TableQualityQuery from ctlearn.tools.predict.utils.load_model import load_model diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index 2dec24c9..4186b676 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -91,7 +91,6 @@ from ctlearn.core.pytorch.dataset import PyTorchDataset from ctlearn.tools.utils import ( FrameworkType, - detect_framework, setup_framework, validate_trait_dict, ) @@ -502,57 +501,6 @@ def _overwrite_meta(self): h5_file.root._v_attrs["CTA PRODUCT ID"] = str(uuid.uuid4()) h5_file.flush() - def _setup_framework(self): - """ - Detect framework from model paths, check consistency, and configure hardware devices. - Raises an error if no valid framework (Keras or PyTorch) is identified. - """ - # Collect all configured model path attributes - model_paths = [ - self.load_type_model_from, - self.load_energy_model_from, - self.load_cameradirection_model_from, - self.load_skydirection_model_from, - ] - # Detect frameworks from all non-None paths - detected_frameworks = {} - for path in model_paths: - if path is not None: - detected_frameworks[path] = detect_framework(path) - # Fail immediately if no valid model paths are provided - if not detected_frameworks: - raise ToolConfigurationError( - "No model paths were specified. At least one valid model path " - "(.keras/.h5 for Keras, or .pt/.pth for PyTorch) must be provided." - ) - # Verify consistency across specified paths - unique_frameworks = set(detected_frameworks.values()) - if len(unique_frameworks) > 1: - details = ", ".join(f"{p}: {fw.value}" for p, fw in detected_frameworks.items()) - raise ToolConfigurationError( - f"Inconsistent model frameworks detected across paths: {details}. " - "All specified model files must belong to the same framework." - ) - # Set framework directly from the detected FrameworkType enum - self.framework_type = next(iter(unique_frameworks)) - self.log.info("Framework selected: %s", self.framework_type.value) - # Configure framework-specific hardware/device setup - self.device = None - if self.framework_type == FrameworkType.PYTORCH: - if torch.cuda.is_available(): - self.device = torch.device("cuda") - self.num_devices = torch.cuda.device_count() - self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) - else: - self.device = torch.device("cpu") - self.num_devices = 1 - self.log.info("Using PyTorch CPU device.") - elif self.framework_type == FrameworkType.KERAS: - self.strategy = tf.distribute.MirroredStrategy() - atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore - self.num_devices = self.strategy.num_replicas_in_sync - self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) - def _get_data_levels(self, h5file): """Get the data levels present in the HDF5 file.""" data_levels = { @@ -742,7 +690,7 @@ def _predict_with_model(self, model_path): """ Select the framework to load and predict the data. """ if self.framework_type == FrameworkType.KERAS: return self._predict_with_keras_model(model_path) - else: + elif self.framework_type == FrameworkType.PYTORCH: return self._predict_with_pytorch_model(model_path) def _predict_with_keras_model(self, model_path): diff --git a/ctlearn/tools/utils.py b/ctlearn/tools/utils.py index 0ce29029..d9d88e1f 100644 --- a/ctlearn/tools/utils.py +++ b/ctlearn/tools/utils.py @@ -24,7 +24,6 @@ "validate_trait_dict", "get_lst1_subarray_description", "FrameworkType", - "detect_framework", "setup_framework", ] @@ -132,26 +131,6 @@ class FrameworkType(Enum): KERAS = "Keras" PYTORCH = "PyTorch" -def detect_framework(path_val): - """ - Determines framework based on file extension. - Returns 'Keras', 'PyTorch', or raises TraitError. - """ - if path_val is None: - return None - - path = pathlib.Path(path_val) - ext = path.suffix.lower() - - if ext in [".keras", ".h5"]: - return FrameworkType["KERAS"] - elif ext in [".pt", ".pth"]: - return FrameworkType["PYTORCH"] - else: - raise TraitError( - f"Invalid model extension '{ext}' for file '{path}'. " - "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." - ) def setup_framework(model_paths): """ @@ -183,11 +162,27 @@ def setup_framework(model_paths): If no valid model paths are provided, or if multiple inconsistent frameworks are detected across the provided paths. """ + def _detect_framework(path_val): + """ + Determines framework based on file extension. + Returns 'Keras', 'PyTorch', or raises TraitError. + """ + path = pathlib.Path(path_val) + ext = path.suffix.lower() + if ext in [".keras", ".h5"]: + return FrameworkType["KERAS"] + elif ext in [".pt", ".pth"]: + return FrameworkType["PYTORCH"] + else: + raise TraitError( + f"Invalid model extension '{ext}' for file '{path}'. " + "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." + ) # Detect frameworks from all non-None paths detected_frameworks = {} for path in model_paths: if path is not None: - detected_frameworks[path] = detect_framework(path) + detected_frameworks[path] = _detect_framework(path) # Fail immediately if no valid model paths are provided if not detected_frameworks: raise ToolConfigurationError( From e329d5affdb0976b965f9dbd4918c45e99e24570 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 15:24:57 +0200 Subject: [PATCH 11/12] added pytorch models into LST1 prediction tool --- ctlearn/conftest.py | 10 +- ctlearn/tools/predict/utils/predict_model.py | 1196 ----------- ctlearn/tools/predict_LST1.py | 318 +-- ctlearn/tools/predict_model.py | 1 - ctlearn/tools/predict_model_main.py | 1884 ----------------- .../pytorch/config/default_config_file.yml | 152 -- ctlearn/tools/tests/test_predict_LST1.py | 2 +- ctlearn/tools/train_model_current.py | 267 --- ctlearn/tools/utils.py | 2 + 9 files changed, 189 insertions(+), 3643 deletions(-) delete mode 100644 ctlearn/tools/predict/utils/predict_model.py delete mode 100644 ctlearn/tools/predict_model_main.py delete mode 100644 ctlearn/tools/pytorch/config/default_config_file.yml delete mode 100644 ctlearn/tools/train_model_current.py diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index 46ced9bb..803189a3 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -7,13 +7,14 @@ import numpy as np import pytest import shutil -from astropy import units as u from astropy.table import Column, Table from traitlets.config.loader import Config from ctapipe.core import run_tool from ctapipe.io import write_table +from ctapipe.tools.process import ProcessorTool from ctapipe.utils import get_dataset_path + from ctlearn.tools.keras.train_model import TrainCTLearnKerasModel from ctlearn.tools.pytorch.train_model import TrainCTLearnPyTorchModel from ctlearn.tools.utils import get_lst1_subarray_description @@ -136,8 +137,6 @@ def dl1_gamma_file(dl1_tmp_path, gamma_simtel_path): """ DL1 file containing both images and parameters from a gamma simulation set. """ - from ctapipe.tools.process import ProcessorTool - output = dl1_tmp_path / "gamma.dl1.h5" argv = [ f"--input={gamma_simtel_path}", @@ -172,10 +171,7 @@ def r1_gamma_file(r1_tmp_path, gamma_simtel_path): """ R1 file containing both waveforms and parameters from a gamma simulation set. """ - from ctapipe.tools.process import ProcessorTool - output = r1_tmp_path / "gamma.r1.h5" - allowed_tels = [1, 2] argv = [ f"--input={gamma_simtel_path}", @@ -193,8 +189,6 @@ def r1_proton_file(r1_tmp_path, proton_simtel_path): """ R1 file containing both waveforms and parameters from a proton simulation set. """ - from ctapipe.tools.process import ProcessorTool - # Restrict to two LSTs for R1 tests to reduce computational load allowed_tels = [1, 2] output = r1_tmp_path / "proton.r1.h5" diff --git a/ctlearn/tools/predict/utils/predict_model.py b/ctlearn/tools/predict/utils/predict_model.py deleted file mode 100644 index 2ca79a25..00000000 --- a/ctlearn/tools/predict/utils/predict_model.py +++ /dev/null @@ -1,1196 +0,0 @@ -""" -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. -""" - -import atexit -import pathlib -import numpy as np -import os -import tensorflow as tf -import keras -import threading -from ctlearn.core.ctlearn_enum import Task - -from astropy import units as u -from astropy.coordinates.earth import EarthLocation -from astropy.coordinates import AltAz, SkyCoord -from astropy.table import ( - Table, - hstack, - vstack, - join, - setdiff, -) -from ctlearn.tools.train.pytorch.utils import ( - sanity_check, - read_configuration, - expected_structure, -) - -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) -from ctapipe.coordinates import CameraFrame, NominalFrame -from ctapipe.core import Tool -from ctapipe.core.tool import ToolConfigurationError -from ctapipe.core.traits import ( - Bool, - Float, - Int, - Path, - flag, - Set, - Dict, - List, - CaselessStrEnum, - ComponentName, - Unicode, - classes_with_traits, -) -from ctapipe.monitoring.interpolation import PointingInterpolator -from ctapipe.io import read_table, write_table, HDF5Merger -from ctapipe.reco.reconstructor import ReconstructionProperty -from ctapipe.reco.stereo_combination import StereoCombiner -from ctapipe.reco.utils import add_defaults_and_meta -from dl1_data_handler.reader import ( - DLDataReader, - ProcessType, - LST_EPOCH, -) -from ctlearn.core.data_loader.loader import DLDataLoader -from ctlearn.utils import monitor_progress - -SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -POINTING_GROUP = "/dl1/monitoring/telescope/pointing" -SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" -DL1_TELESCOPE_GROUP = "/dl1/event/telescope" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_TELESCOPE_GROUP = "/dl2/event/telescope" -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] - -__all__ = ["PredictCTLearnModel"] - -class PredictCTLearnModel(Tool): - """ - Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. - - This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image - or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to - store them in the output file. The input data is loaded from the input url using the - ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.loader.DLDataLoader``. - The prediction is performed using the CTLearn models. The data is stored in the output file - following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to - handle the prediction for mono and stereo mode. - - Attributes - ---------- - input_url : pathlib.Path - Input ctapipe HDF5 files including pixel-wise image or waveform data. - use_HDF5Merger : bool - Set whether to use the HDF5Merger component to copy the selected tables from the input file to the output file. - dl1_features : bool - Set whether to include the dl1 feature vectors in the output file. - dl2_telescope : bool - Set whether to include dl2 telescope-event-wise data in the output file. - dl2_subarray : bool - Set whether to include dl2 subarray-event-wise data in the output file. - dl1dh_reader : dl1_data_handler.reader.DLDataReader - DLDataReader object to read the data. - dl1dh_reader_type : str - Type of the DLDataReader to use for the prediction. - stack_telescope_images : bool - Set whether to stack the telescope images in the data loader. Requires ``stereo``. - sort_by_intensity : bool - Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. - prefix : str - Name of the reconstruction algorithm used to generate the dl2 data. - load_type_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the classification of the primary particle type. - load_energy_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression of the primary particle energy. - load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression - of the primary particle arrival direction based on camera coordinate offsets. - load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression - of the primary particle arrival direction based on spherical coordinate offsets. - output_path : pathlib.Path - Output path to save the dl2 prediction results. - overwrite_tables : bool - Overwrite the table in the output file if it exists. - keras_verbose : int - Verbosity mode of Keras during the prediction. - strategy : tf.distribute.Strategy - MirroredStrategy to distribute the prediction. - data_loader : ctlearn.core.loader.DLDataLoader - DLDataLoader object to load the data. - indices : list of int - List of indices for the data loaders. - batch_size : int - Size of the batch to perform inference of the neural network. - last_batch_size : int - Size of the last batch in the data loaders. - - Methods - ------- - setup() - Set up the tool. - finish() - Finish the tool. - _predict_with_model(model_path) - Load and predict with a CTLearn model. - _predict_classification(example_identifiers) - Predict the classification of the primary particle type. - _predict_energy(example_identifiers) - Predict the energy of the primary particle. - _predict_cameradirection(example_identifiers) - Predict the arrival direction of the primary particle based on camera coordinate offsets. - _predict_skydirection(example_identifiers) - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - _transform_cam_coord_offsets_to_sky(table) - Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _transform_spher_coord_offsets_to_sky(table) - Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _create_nan_table(nonexample_identifiers, columns, shapes) - Create a table with NaNs for missing predictions. - _store_pointing(all_identifiers) - Store the telescope pointing table from to the output file. - _create_feature_vectors_table(example_identifiers, nonexample_identifiers, classification_feature_vectors, energy_feature_vectors, direction_feature_vectors) - Create the table for the DL1 feature vectors. - """ - - input_url = Path( - help="Input ctapipe HDF5 files including pixel-wise image or waveform data", - allow_none=True, - exists=True, - directory_ok=False, - file_ok=True, - ).tag(config=True) - - use_HDF5Merger = Bool( - default_value=True, - allow_none=False, - help=( - "Set whether to use the HDF5Merger component to copy the selected tables " - "from the input file to the output file. CAUTION: This can only be used " - "if the output file not exists." - ), - ).tag(config=True) - - dl1_features = Bool( - default_value=False, - allow_none=False, - help="Set whether to include the dl1 feature vectors in the output file.", - ).tag(config=True) - - dl2_telescope = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 telescope-event-wise data in the output file.", - ).tag(config=True) - - dl2_subarray = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 subarray-event-wise data in the output file.", - ).tag(config=True) - - dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( - config=True - ) - - stack_telescope_images = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to stack the telescope images in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - sort_by_intensity = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to sort the telescope images by intensity in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - prefix = Unicode( - default_value="CTLearn", - allow_none=False, - help="Name of the reconstruction algorithm used to generate the dl2 data.", - ).tag(config=True) - - load_type_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the classification " - "of the primary particle type." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_energy_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle energy." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_cameradirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle arrival direction based on camera coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_skydirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle arrival direction based on spherical coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - batch_size = Int( - default_value=64, - allow_none=False, - help="Size of the batch to perform inference of the neural network.", - ).tag(config=True) - - output_path = Path( - default_value="./output.dl2.h5", - allow_none=False, - help="Output path to save the dl2 prediction results", - ).tag(config=True) - - overwrite_tables = Bool( - default_value=True, - allow_none=False, - help="Overwrite the table in the output file if it exists", - ).tag(config=True) - - pytorch_config_file = Path( - default_value=None, - allow_none=True, - help="Pytorch config file", - ).tag(config=True) - - # Unified Hardware Architecture & Execution Strategy - device = CaselessStrEnum( - ["cuda", "cpu", "mps"], - default_value="cuda", - help="Device to use: cuda, cpu, or mps", - ).tag(config=True) - - devices = List( - trait=Int(), - default_value=[0], - help="List of GPU device IDs to use", - ).tag(config=True) - - strategy = Unicode( - default_value="auto", - help="Multi-GPU strategy", - ).tag(config=True) - - # Event Cut-offs - leakage_intensity_cutoff = Float( - default_value=0.2, - help="Events with leakage intensity greater than this value are removed", - ).tag(config=True) - - intensity_cutoff = Float( - default_value=50.0, - help="Events with Hillas intensity below this value are removed", - ).tag(config=True) - - # Normalizations - apply_log_scaling = List( - trait=Bool(), - default_value=[True, True], - help="List specifying whether to apply log10(X+1.0) scaling for [charge, peak_time]", - ).tag(config=True) - - use_clean = Bool( - default_value=True, - help="Use the image with the applied mask", - ).tag(config=True) - - use_clean_dvr = Bool( - default_value=False, - help="Use clean DVR mask", - ).tag(config=True) - - type_mu = Float( - default_value=0.0, - help="Mean for type channel normalization", - ).tag(config=True) - - type_sigma = Float( - default_value=1000.0, - help="Std dev for type channel normalization", - ).tag(config=True) - - dir_mu = Float( - default_value=0.0, - help="Mean for direction channel normalization", - ).tag(config=True) - - dir_sigma = Float( - default_value=1000.0, - help="Std dev for direction channel normalization", - ).tag(config=True) - - energy_mu = Float( - default_value=0.0, - help="Mean for energy channel normalization", - ).tag(config=True) - - energy_sigma = Float( - default_value=1000.0, - help="Std dev for energy channel normalization", - ).tag(config=True) - - keras_verbose = Int( - default_value=1, - min=0, - max=2, - allow_none=False, - help=( - "Verbosity mode of Keras during the prediction: " - "0 = silent, 1 = progress bar, 2 = one line per call." - ), - ).tag(config=True) - - framework_type = CaselessStrEnum( - ["pytorch", "keras"], - default_value="keras", - help="Framework to use: pytorch or keras", - ).tag(config=True) - - aliases = { - ("i", "input_url"): "PredictCTLearnModel.input_url", - ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", - ("e", "energy_model"): "PredictCTLearnModel.load_energy_model_from", - ( - "d", - "cameradirection_model", - ): "PredictCTLearnModel.load_cameradirection_model_from", - ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", - ("o", "output"): "PredictCTLearnModel.output_path", - ("f", "framework"): "PredictCTLearnModel.framework_type", - ("p", "pytorch_config_file"): "PredictCTLearnModel.pytorch_config_file", - "device": "PredictCTLearnModel.device", - "devices": "PredictCTLearnModel.devices", - "strategy": "PredictCTLearnModel.strategy", - "type_mu": "PredictCTLearnModel.type_mu", - "type_sigma": "PredictCTLearnModel.type_sigma", - "dir_mu": "PredictCTLearnModel.dir_mu", - "dir_sigma": "PredictCTLearnModel.dir_sigma", - "energy_mu": "PredictCTLearnModel.energy_mu", - "energy_sigma": "PredictCTLearnModel.energy_sigma", - } - - flags = { - **flag( - "dl1-features", - "PredictCTLearnModel.dl1_features", - "Include dl1 features", - "Exclude dl1 features", - ), - **flag( - "dl2-telescope", - "PredictCTLearnModel.dl2_telescope", - "Include dl2 telescope-event-wise data in the output file", - "Exclude dl2 telescope-event-wise data in the output file", - ), - **flag( - "dl2-subarray", - "PredictCTLearnModel.dl2_subarray", - "Include dl2 telescope-event-wise data in the output file", - "Exclude dl2 telescope-event-wise data in the output file", - ), - **flag( - "use-HDF5Merger", - "PredictCTLearnModel.use_HDF5Merger", - "Copy data using the HDF5Merger component (CAUTION: This can not be used if the output file already exists)", - "Do not copy data using the HDF5Merger component", - ), - **flag( - "r0-waveforms", - "HDF5Merger.r0_waveforms", - "Include r0 waveforms", - "Exclude r0 waveforms", - ), - **flag( - "r1-waveforms", - "HDF5Merger.r1_waveforms", - "Include r1 waveforms", - "Exclude r1 waveforms", - ), - **flag( - "dl1-parameters", - "HDF5Merger.dl1_parameters", - "Include dl1 parameters", - "Exclude dl1 parameters", - ), - **flag( - "dl1-images", - "HDF5Merger.dl1_images", - "Include dl1 images", - "Exclude dl1 images", - ), - **flag( - "true-parameters", - "HDF5Merger.true_parameters", - "Include true parameters", - "Exclude true parameters", - ), - **flag( - "true-images", - "HDF5Merger.true_images", - "Include true images", - "Exclude true images", - ), - } - - classes = classes_with_traits(DLDataReader) - - def setup(self): - if self.framework_type == "pytorch": - import torch - if self.pytorch_config_file is not None: - self.log.info(f"Using {self.pytorch_config_file} config file for pytorch framework") - legacy_params = read_configuration(self.pytorch_config_file) - sanity_check(legacy_params, expected_structure) - - def get_conf_val(key1, key2, trait_name, default_val): - in_config = False - if "PredictCTLearnModel" in self.config and trait_name in self.config["PredictCTLearnModel"]: - in_config = True - if not in_config: - return legacy_params.get(key1, {}).get(key2, default_val) - return getattr(self, trait_name) - - self.device_str = get_conf_val("arch", "device", "device", self.device.name if hasattr(self.device, "name") else str(self.device)) - self.type_mu = get_conf_val("normalization", "type_mu", "type_mu", self.type_mu) - self.type_sigma = get_conf_val("normalization", "type_sigma", "type_sigma", self.type_sigma) - self.dir_mu = get_conf_val("normalization", "dir_mu", "dir_mu", self.dir_mu) - self.dir_sigma = get_conf_val("normalization", "dir_sigma", "dir_sigma", self.dir_sigma) - self.energy_mu = get_conf_val("normalization", "energy_mu", "energy_mu", self.energy_mu) - self.energy_sigma = get_conf_val("normalization", "energy_sigma", "energy_sigma", self.energy_sigma) - self.leakage_intensity_cutoff = get_conf_val("cut-off", "leakage_intensity", "leakage_intensity_cutoff", self.leakage_intensity_cutoff) - self.intensity_cutoff = get_conf_val("cut-off", "intensity", "intensity_cutoff", self.intensity_cutoff) - - self.parameters = legacy_params - else: - self.log.info("No legacy config file provided. Using standard Traitlets configuration for PyTorch.") - self.device_str = self.device.name if hasattr(self.device, "name") else str(self.device) - - self.parameters = { - "data": { - "type_checkpoint": self.load_type_model_from, - "energy_checkpoint": self.load_energy_model_from, - "direction_checkpoint": self.load_cameradirection_model_from or self.load_skydirection_model_from, - }, - "hyp": { - "batches": self.batch_size, - "dynamic_batches": True, - }, - "cut-off": { - "leakage_intensity": self.leakage_intensity_cutoff, - "intensity": self.intensity_cutoff, - }, - "normalization": { - "apply_log_scaling": self.apply_log_scaling, - "use_clean": self.use_clean, - "use_clean_dvr": self.use_clean_dvr, - "type_mu": self.type_mu, - "type_sigma": self.type_sigma, - "dir_mu": self.dir_mu, - "dir_sigma": self.dir_sigma, - "energy_mu": self.energy_mu, - "energy_sigma": self.energy_sigma, - }, - "arch": { - "device": self.device_str, - "devices": self.devices, - "strategy": self.strategy, - }, - "model": { - "model_type": { - "model_name": "DoubleBBEfficientNet", - "parameters": { - "model_variant": "efficientnet-b3", - "task": "type", - "num_outputs": 2, - "device_str": self.device_str, - "energy_bins": None, - } - }, - "model_energy": { - "model_name": "ThinResNet", - "parameters": { - "task": "energy", - "num_inputs": 1, - "num_outputs": 1, - "num_blocks": [3, 4, 6, 3], - "dropout": 0.1, - "use_bn": False, - } - }, - "model_direction": { - "model_name": "ThinResNet_DBB", - "parameters": { - "task": "direction", - "num_inputs": 1, - "num_outputs": 3, - "num_blocks": [3, 4, 6, 3], - "dropout": 0.1, - "use_bn": False, - } - } - } - } - self.device = torch.device(self.device_str) - self.tasks = [] - if self.load_type_model_from is not None: - self.tasks.append(Task.type) - if self.load_energy_model_from is not None: - self.tasks.append(Task.energy) - if self.load_cameradirection_model_from is not None or self.load_skydirection_model_from is not None: - self.tasks.append(Task.direction) - - # Check if the ctapipe HDF5Merger component is enabled - if os.path.exists(self.output_path): - self.log.warning( - "The output file '{self.output_path}' already exists. Disabling HDF5Merger the flag '--no-use-HDF5Merger' will be not use." - ) - self.use_HDF5Merger = False - - if self.use_HDF5Merger: - if os.path.exists(self.output_path): - raise ToolConfigurationError( - f"The output file '{self.output_path}' already exists. Please use " - "'--no-use-HDF5Merger' to disable the usage of the HDF5Merger component." - ) - # Copy selected tables from the input file to the output file - self.log.info("Copying to output destination.") - stop_event = threading.Event() - monitor_thread = threading.Thread(target=monitor_progress, args=(self.input_url, self.output_path, stop_event, self.log)) - monitor_thread.start() - - try: - with HDF5Merger(self.output_path, parent=self) as merger: - merger(self.input_url) - finally: - stop_event.set() - monitor_thread.join() - - else: - self.log.info( - "No copy to output destination, since the usage of the HDF5Merger component is disabled." - ) - - # Create a MirroredStrategy. - self.strategy = tf.distribute.MirroredStrategy() - atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore - self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) - - # Set up the data reader - self.log.info("Loading data reader:") - self.log.info("For a large dataset, this may take a while...") - self.dl1dh_reader = DLDataReader.from_name( - self.dl1dh_reader_type, - input_url_signal=[self.input_url], - parent=self, - ) - self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) - # Check if the number of events is enough to form a batch - if self.dl1dh_reader._get_n_events() < self.batch_size: - raise ToolConfigurationError( - f"{self.dl1dh_reader._get_n_events()} events are not enough " - f"to form a batch of size {self.batch_size}. Reduce the batch size." - ) - # Set the indices for the data loaders - self.indices = list(range(self.dl1dh_reader._get_n_events())) - self.last_batch_size = len(self.indices) % ( - self.batch_size * self.strategy.num_replicas_in_sync - ) - - def finish(self): - self.log.info("Tool is shutting down") - - def _predict_with_model(self, model_path, task): - """ - Load and predict with a CTLearn model. - - Load a model from the specified path and predict the data using the loaded model. - If a last batch loader is provided, predict the last batch and stack the results. - - Parameters - ---------- - model_path : str - Path to a Keras model file (Keras3) or directory (Keras2). - - Returns - ------- - predict_data : astropy.table.Table - Table containing the prediction results. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - predict_data = None - feature_vectors = None - - if self.framework_type == "keras": - from ctlearn.tools.predict.keras.predic_model_keras import predict_with_model - predict_data, feature_vectors = predict_with_model(self,model_path) - return predict_data, feature_vectors - - if self.framework_type == "pytorch": - from ctlearn.tools.predict.pytorch.predic_model_pytorch import predict_with_model_pytorch - - task = None - if model_path == self.load_type_model_from: - task = Task.type - elif model_path == self.load_energy_model_from: - task = Task.energy - elif model_path in [self.load_cameradirection_model_from, self.load_skydirection_model_from]: - task = Task.direction - else: - task = Task.type - - predict_data, feature_vectors = predict_with_model_pytorch(self, task) - return predict_data, feature_vectors - return predict_data, feature_vectors - - def _predict_classification(self, example_identifiers): - """ - Predict the classification of the primary particle type. - - This method uses a pre-trained type model to predict the type of the primary particle - for a given set of example identifiers. The predicted classification score ('gammaness') - is added to the example identifiers table. - - Parameters: - ----------- - classification_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - predicted classification score ('gammaness'). - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the classification of the primary particle type..." - ) - # Predict the data using the loaded type_model - predict_data, feature_vectors = self._predict_with_model( - self.load_type_model_from, Task.type - ) - # Create prediction table and add the predicted classification score ('gammaness') - classification_table = example_identifiers.copy() - classification_table.add_column( - predict_data["type"].T[1], name=f"{self.prefix}_tel_prediction" - ) - return classification_table, feature_vectors - - def _predict_energy(self, example_identifiers): - """ - Predict the energy of the primary particle. - - This method uses a pre-trained energy model to predict the energy of the primary particle - for a given set of example identifiers. The predicted energy is then converted from - log10(TeV) to TeV and added to the example identifiers table. - - Parameters: - ----------- - energy_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed energy in TeV. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info("Predicting for the regression of the primary particle energy...") - # Predict the data using the loaded energy_model - predict_data, feature_vectors = self._predict_with_model( - self.load_energy_model_from, Task.energy - ) - # Convert the reconstructed energy from log10(TeV) to TeV - reco_energy = u.Quantity( - np.power(10, np.squeeze(predict_data["energy"])), - unit=u.TeV, - ) - print(reco_energy) - print(reco_energy.shape) - print(example_identifiers) - # Create prediction table and add the reconstructed energy in TeV - energy_table = example_identifiers.copy() - energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") - return energy_table, feature_vectors - - def _predict_cameradirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on camera coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the - primary particle for a given set of example identifiers. The predicted camera coordinate offsets - is added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - cameradirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed camera coordinate offsets in x and y. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the regression of the primary particle arrival direction based on camera coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_cameradirection_model_from, Task.direction - ) - # For the direction task, the prediction is the camera coordinate offset in x and y - # from the telescope pointing. - cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) - cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) - # Create prediction table and add the reconstructed energy in TeV - cameradirection_table = example_identifiers.copy() - cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") - cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") - return cameradirection_table, feature_vectors - - def _predict_skydirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the primary - particle for a given set of example identifiers. The predicted spherical coordinate offsets is - added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - skydirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed spherical coordinate offsets in fov_lon and fov_lat. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the regression of the primary particle arrival direction based on spherical coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_skydirection_model_from - ) - # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat - # from the telescope pointing. - fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) - fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) - # Create prediction table and add the reconstructed energy in TeV - skydirection_table = example_identifiers.copy() - skydirection_table.add_column(fov_lon, name="fov_lon") - skydirection_table.add_column(fov_lat, name="fov_lat") - return skydirection_table, feature_vectors - - def _transform_cam_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - # Get the telescope ID from the table - tel_id = table["tel_id"][0] - # Set the telescope position - tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the telescope pointing with the trigger timestamp and the telescope position - altaz = AltAz( - location=tel_ground_frame.to_earth_location(), - obstime=trigger_time, - ) - # Set the telescope pointing - tel_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the camera frame with the focal length and rotation of the camera - camera_frame = CameraFrame( - focal_length=self.dl1dh_reader.subarray.tel[ - tel_id - ].camera.geometry.frame.focal_length, - rotation=self.dl1dh_reader.pix_rotation[tel_id], - telescope_pointing=tel_pointing, - ) - # Set the camera coordinate offset - cam_coord_offset = SkyCoord( - x=table["cam_coord_offset_x"], - y=table["cam_coord_offset_y"], - frame=camera_frame, - ) - # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Transform the true Alt/Az coordinates to camera coordinates - reco_direction = cam_coord_offset.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column(reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az") - table.add_column(reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt") - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "cam_coord_offset_x", - "cam_coord_offset_y", - ] - ) - return table - - def _transform_spher_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted spherical offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted spherical offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the AltAz frame with the reference location and time - altaz = AltAz( - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the array pointing - array_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the nominal frame with the array pointing - nom_frame = NominalFrame( - origin=array_pointing, - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the reco direction in (fov_lon, fov_lat) coordinates - reco_direction = SkyCoord( - fov_lon=table["fov_lon"], - fov_lat=table["fov_lat"], - frame=nom_frame, - ) - # Transform the reco direction from nominal frame to the AltAz frame - sky_coord = reco_direction.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column(sky_coord.az.to(u.deg), name=f"{self.prefix}_az") - table.add_column(sky_coord.alt.to(u.deg), name=f"{self.prefix}_alt") - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "fov_lon", - "fov_lat", - ] - ) - return table - - def _create_nan_table(self, nonexample_identifiers, columns, shapes): - """ - Create a table with NaNs for missing predictions. - - This method creates a table with NaNs for missing predictions for the non-example identifiers. - In stereo mode, the table also a column for the valid telescopes is added with all False values. - - Parameters: - ----------- - nonexample_identifiers : astropy.table.Table - Table containing the non-example identifiers. - columns : list of str - List of column names to create in the table. - shapes : list of shapes - List of shapes for the columns to create in the table. - - Returns: - -------- - nan_table : astropy.table.Table - Table containing NaNs for missing predictions. - """ - # Create a table with NaNs for missing predictions - nan_table = nonexample_identifiers.copy() - for column_name, shape in zip(columns, shapes): - nan_table.add_column(np.full(shape, np.nan), name=column_name) - # Add that no telescope is valid for the non-example identifiers in stereo mode - if self.dl1dh_reader.mode == "stereo": - nan_table.add_column( - np.zeros( - (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ), - name=f"{self.prefix}_telescopes", - ) - return nan_table - - def _store_pointing(self, all_identifiers): - """ - Store the telescope pointing table from to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - - # Initialize the pointing interpolator from ctapipe - pointing_interpolator = PointingInterpolator( - bounds_error=False, extrapolate=True - ) - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Get the telescope pointing from the dl1dh reader - tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] - # Add the telescope pointing table to the pointing interpolator - pointing_interpolator.add_table(tel_id, tel_pointing) - tel_identifiers = all_identifiers.copy() - if self.dl1dh_reader.mode == "mono": - tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Interpolate the telescope pointing - tel_altitude, tel_azimuth = pointing_interpolator( - tel_id, tel_identifiers["time"] - ) - tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") - tel_identifiers.add_column(tel_altitude, name="pointing_altitude") - pointing_info.append(tel_identifiers) - if self.dl1dh_reader.mode == "mono": - tel_pointing_table = Table( - { - "time": tel_identifiers["time"], - "azimuth": tel_identifiers["pointing_azimuth"], - "altitude": tel_identifiers["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info = vstack(pointing_info) - if self.dl1dh_reader.mode == "stereo": - # Group the pointing information by subarray event keys - # TODO: This needs to be debugged with SST1M data - pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) - pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) - pointing_info = join( - all_identifiers, - pointing_mean, - keys=SUBARRAY_EVENT_KEYS, - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - ) - return pointing_info - - def _create_feature_vectors_table( - self, - example_identifiers, - nonexample_identifiers=None, - classification_feature_vectors=None, - energy_feature_vectors=None, - direction_feature_vectors=None, - ): - """ - Create the table for the DL1 feature vectors. - - This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for - non-example identifiers. The feature vectors are stored in the columns of the table. The table also - contains a column for the valid predictions. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - nonexample_identifiers : astropy.table.Table or None - Table containing the non-example identifiers to fill the NaNs. - classification_feature_vectors : np.ndarray or None - Array containing the classification feature vectors. - energy_feature_vectors : np.ndarray or None - Array containing the energy feature vectors. - direction_feature_vectors : np.ndarray or None - Array containing the direction feature vectors. - - Returns: - -------- - feature_vector_table : astropy.table.Table - Table containing the DL1 feature vectors for the example and non-example identifiers. - """ - # Create the feature vector table - feature_vector_table = example_identifiers.copy() - feature_vector_table.remove_columns( - ["pointing_azimuth", "pointing_altitude", "time"] - ) - columns_list, shapes_list = [], [] - if classification_feature_vectors is not None: - is_valid_col = ~np.isnan( - np.min(classification_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - classification_feature_vectors, - name=f"{self.prefix}_tel_classification_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - classification_feature_vectors.shape[1], - ) - ) - if energy_feature_vectors is not None: - is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) - feature_vector_table.add_column( - energy_feature_vectors, name=f"{self.prefix}_tel_energy_feature_vectors" - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - energy_feature_vectors.shape[1], - ) - ) - if direction_feature_vectors is not None: - is_valid_col = ~np.isnan( - np.min(direction_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - direction_feature_vectors, - name=f"{self.prefix}_tel_geometry_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_geometry_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - direction_feature_vectors.shape[1], - ) - ) - # Produce output table with NaNs for missing predictions - if nonexample_identifiers is not None: - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=columns_list, - shapes=shapes_list, - ) - feature_vector_table = vstack([feature_vector_table, nan_table]) - is_valid_col = np.concatenate( - (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) - ) - # Add is_valid column to the feature vector table - feature_vector_table.add_column( - is_valid_col, - name=f"{self.prefix}_tel_is_valid", - ) - return feature_vector_table - - -if __name__ == "__main__": - pass \ No newline at end of file diff --git a/ctlearn/tools/predict_LST1.py b/ctlearn/tools/predict_LST1.py index e814f63d..5a1e6510 100644 --- a/ctlearn/tools/predict_LST1.py +++ b/ctlearn/tools/predict_LST1.py @@ -2,12 +2,13 @@ Predict the gammaness, energy and arrival direction from lstchain DL1 data. """ +import atexit import numpy as np import tables -try: - import keras -except ImportError: - keras = None +import keras +import torch +import torch.nn as nn + from astropy import units as u from astropy.coordinates import AltAz, SkyCoord from astropy.table import Table, join, setdiff, vstack @@ -19,7 +20,6 @@ ReconstructedEnergyContainer, ) from ctapipe.coordinates import CameraFrame -from ctapipe.coordinates import CameraFrame from ctapipe.core import Tool from ctapipe.core.tool import ToolConfigurationError from ctapipe.core.traits import ( @@ -32,50 +32,31 @@ ComponentName, Dict, UseEnum, - UseEnum, classes_with_traits, ) from ctapipe.instrument.optics import FocalLengthKind -from ctapipe.instrument.optics import FocalLengthKind from ctapipe.io import read_table, write_table - -DL0_TEL_POINTING_GROUP = "/dl0/event/telescope/pointing" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL1_SUBARRAY_POINTING_GROUP = "/dl1/event/subarray/pointing" -DL1_SUBARRAY_TRIGGER_TABLE = "/dl1/event/subarray/trigger" -DL1_TEL_GROUP = "/dl1/event/telescope" -DL1_TEL_CALIBRATION_GROUP = "/dl1/event/telescope/calibration" -DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP = "/dl1/event/telescope/illuminator_throughput" -DL1_TEL_IMAGES_GROUP = "/dl1/event/telescope/image" -DL1_TEL_MUON_GROUP = "/dl1/event/telescope/muon" -DL1_TEL_MUON_THROUGHPUT_GROUP = "/dl1/event/telescope/muon_throughput" -DL1_TEL_OPTICAL_PSF_GROUP = "/dl1/event/telescope/optical_psf" -DL1_TEL_PARAMETERS_GROUP = "/dl1/event/telescope/parameters" -DL1_TEL_POINTING_GROUP = "/dl1/event/telescope/pointing" -DL1_TEL_TRIGGER_TABLE = "/dl1/event/telescope/trigger" -DL2_EVENT_STATISTICS_GROUP = "/dl2/event/subarray/statistics" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -R0_TEL_GROUP = "/r0/event/telescope" -R1_TEL_GROUP = "/r1/event/telescope" -SIMULATION_IMAGES_GROUP = "/simulation/event/telescope/images" -SIMULATION_IMPACT_GROUP = "/simulation/event/telescope/impact" -SIMULATION_PARAMETERS_GROUP = "/simulation/event/telescope/parameters" -SIMULATION_RUN_TABLE = "/simulation/run_config" -SIMULATION_SHOWER_TABLE = "/simulation/event/subarray/shower" -DL2_TEL_PARTICLETYPE_GROUP = "/dl2/event/telescope/classification" -DL2_TEL_ENERGY_GROUP = "/dl2/event/telescope/energy" -DL2_TEL_GEOMETRY_GROUP = "/dl2/event/telescope/geometry" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_SUBARRAY_PARTICLETYPE_GROUP = "/dl2/event/subarray/classification" -DL2_SUBARRAY_ENERGY_GROUP = "/dl2/event/subarray/energy" -DL2_SUBARRAY_GEOMETRY_GROUP = "/dl2/event/subarray/geometry" - +from ctapipe.io.hdf5dataformat import ( + DL1_SUBARRAY_TRIGGER_TABLE, + DL1_TEL_GROUP, + DL1_TEL_PARAMETERS_GROUP, + DL1_TEL_POINTING_GROUP, + DL1_TEL_TRIGGER_TABLE, + DL2_TEL_PARTICLETYPE_GROUP, + DL2_TEL_ENERGY_GROUP, + DL2_TEL_GEOMETRY_GROUP, + DL2_SUBARRAY_PARTICLETYPE_GROUP, + DL2_SUBARRAY_ENERGY_GROUP, + DL2_SUBARRAY_GEOMETRY_GROUP, +) from ctapipe.reco.utils import add_defaults_and_meta - -from ctlearn.core.keras.model import LoadedModel from ctlearn import __version__ as ctlearn_version -from ctlearn.utils import get_lst1_subarray_description -from ctlearn.utils import validate_trait_dict +from ctlearn.tools.utils import ( + FrameworkType, + setup_framework, + get_lst1_subarray_description, + validate_trait_dict, +) from dl1_data_handler.image_mapper import ImageMapper from dl1_data_handler.reader import ( get_unmapped_image, @@ -113,9 +94,9 @@ class LST1PredictionTool(Tool): --LST1PredictionTool.channels=cleaned_image \\ --LST1PredictionTool.channels=cleaned_relative_peak_time \\ --LST1PredictionTool.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/direction/ctlearn_model.cpk" \\ + --type_model="/path/to/your/type/ctlearn_model(.keras/.pth)"" \\ + --energy_model="/path/to/your/energy/ctlearn_model(.keras/.pth)"" \\ + --cameradirection_model="/path/to/your/direction/ctlearn_model(.keras/.pth)"" \\ --output output.dl2.h5 \\ --overwrite \\ """ @@ -210,11 +191,11 @@ class LST1PredictionTool(Tool): "cleaned_image", "peak_time", "relative_peak_time", - "cleaned_peak_time", + "cleaned_peak_time", "cleaned_relative_peak_time", ] ), - default_value=["cleaned_image", "cleaned_peak_time"], + default_value=["cleaned_image", "cleaned_relative_peak_time"], allow_none=False, help=( "Set the input channels to be loaded from the DL1 event data. " @@ -269,8 +250,6 @@ class LST1PredictionTool(Tool): ("e", "energy_model"): "LST1PredictionTool.load_energy_model_from", ("d", "cameradirection_model"): "LST1PredictionTool.load_cameradirection_model_from", ("o", "output"): "LST1PredictionTool.output_path", - ("ch", "channels"): "LST1PredictionTool.channels", - } flags = { @@ -301,13 +280,7 @@ class LST1PredictionTool(Tool): ), } - - @property - def classes(self): - return [ - type(self), - LST1PredictionTool, - ] + classes_with_traits(ImageMapper) + classes = classes_with_traits(ImageMapper) def setup(self): self.log.info("ctlearn version %s", ctlearn_version) @@ -319,39 +292,36 @@ def setup(self): self.image_table_path = "/dl1/event/telescope/image/LST_LSTCam" self.parameter_table_name = "/dl1/event/telescope/parameters/LST_LSTCam" self.tel_id = 1 - + # Collect all configured model path attributes + model_paths = [ + self.load_type_model_from, + self.load_energy_model_from, + self.load_cameradirection_model_from, + ] + # Reads the model paths and set up the Keras or PyTorch framework + self.framework_type, self.num_devices, self.strategy, self.device = setup_framework(model_paths) + self.log.info("Framework selected: %s", self.framework_type.value) + if self.framework_type == FrameworkType.KERAS: + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) + elif self.framework_type == FrameworkType.PYTORCH: + if self.device == torch.device("cpu"): + self.log.info("Using PyTorch CPU device.") + else: + self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) + # Load the models from the specified paths + input_shape_type, self.backbone_type, self.head_type = self._load_model( + self.load_type_model_from + ) + input_shape_energy, self.backbone_energy, self.head_energy = self._load_model( + self.load_energy_model_from + ) + input_shape_direction, self.backbone_direction, self.head_direction = self._load_model( + self.load_cameradirection_model_from + ) # Get the number of rows in the table with tables.open_file(self.input_url) as input_file: self.table_length = len(input_file.get_node(self.image_table_path)) - - if keras is None: - raise ImportError("TensorFlow/Keras is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") - - # Load the models from the specified paths - if self.load_type_model_from is not None: - self.log.info("Loading the type model from %s.", self.load_type_model_from) - model_type = keras.saving.load_model(self.load_type_model_from) - input_shape = model_type.input_shape[1:] - self.backbone_type, self.head_type = self._split_model(model_type) - if self.load_energy_model_from is not None: - self.log.info( - "Loading the energy model from %s.", self.load_energy_model_from - ) - model_energy = keras.saving.load_model( - self.load_energy_model_from - ) - input_shape = model_energy.input_shape[1:] - self.backbone_energy, self.head_energy = self._split_model(model_energy) - if self.load_cameradirection_model_from is not None: - self.log.info( - "Loading the cameradirection model from %s.", self.load_cameradirection_model_from - ) - model_direction = keras.saving.load_model( - self.load_cameradirection_model_from - ) - input_shape = model_direction.input_shape[1:] - self.backbone_direction, self.head_direction = self._split_model(model_direction) - # Get the SubarrayDescription of the LST-1 telescope self.subarray = get_lst1_subarray_description(focal_length_choice=self.focal_length_choice) # Write the SubarrayDescription to the output file @@ -372,13 +342,13 @@ def setup(self): parent=self, ) # Check if the input shape of the model matches the image shape of the ImageMapper - if input_shape[0] != self.image_mapper.image_shape: - raise ToolConfigurationError( - f"The input shape of the model ('{input_shape[0]}') does not match " - f"the image shape of the ImageMapper ('{self.image_mapper.image_shape}'). " - f"Use e.g. '--BilinearMapper.interpolation_image_shape={input_shape[0]}' ." - ) - + for input_shape in [input_shape_type, input_shape_energy, input_shape_direction]: + if input_shape is not None and input_shape[0] != self.image_mapper.image_shape: + raise ToolConfigurationError( + f"The input shape of the model ('{input_shape[0]}') does not match " + f"the image shape of the ImageMapper ('{self.image_mapper.image_shape}'). " + f"Use e.g. '--BilinearMapper.interpolation_image_shape={input_shape[0]}' ." + ) # Get offset and scaling of images self.transforms = {} self.transforms["image_scale"] = 0.0 @@ -388,7 +358,6 @@ def setup(self): # Get the number of rows in the table with tables.open_file(self.input_url) as input_file: img_table_v_attrs = input_file.get_node(self.image_table_path)._v_attrs - # Check the transform value used for the file compression if "CTAFIELD_3_TRANSFORM_SCALE" in img_table_v_attrs: self.transforms["image_scale"] = img_table_v_attrs[ @@ -406,7 +375,6 @@ def setup(self): ] def start(self): - all_identifiers = read_table(self.input_url, self.parameter_table_name) all_identifiers.meta = {} if self.override_obs_id is not None: @@ -578,21 +546,20 @@ def start(self): tel_altitude.extend(dl1_table["tel_alt"].data) trigger_time.extend(dl1_table["time"].mjd) if self.load_type_model_from is not None: - classification_feature_vectors = self.backbone_type.predict_on_batch(input_data) - classification_fvs.extend(classification_feature_vectors) - predict_data = self.head_type.predict_on_batch(classification_feature_vectors) - prediction.extend(predict_data[:, 1]) + fvs, preds = self._predict_batch(self.backbone_type, self.head_type, input_data) + classification_fvs.extend(fvs) + prediction.extend(preds[:, 1]) + if self.load_energy_model_from is not None: - energy_feature_vectors = self.backbone_energy.predict_on_batch(input_data) - energy_fvs.extend(energy_feature_vectors) - predict_data = self.head_energy.predict_on_batch(energy_feature_vectors) - energy.extend(predict_data.T[0]) + fvs, preds = self._predict_batch(self.backbone_energy, self.head_energy, input_data) + energy_fvs.extend(fvs) + energy.extend(preds[:, 0]) + if self.load_cameradirection_model_from is not None: - direction_feature_vectors = self.backbone_direction.predict_on_batch(input_data) - direction_fvs.extend(direction_feature_vectors) - predict_data = self.head_direction.predict_on_batch(direction_feature_vectors) - cam_coord_offset_x.extend(predict_data.T[0]) - cam_coord_offset_y.extend(predict_data.T[1]) + fvs, preds = self._predict_batch(self.backbone_direction, self.head_direction, input_data) + direction_fvs.extend(fvs) + cam_coord_offset_x.extend(preds[:, 0]) + cam_coord_offset_y.extend(preds[:, 1]) # Create the prediction tables example_identifiers = Table( @@ -923,36 +890,119 @@ def start(self): def finish(self): self.log.info("Tool is shutting down") - def _split_model(self, model): + def _predict_batch(self, backbone, head, input_data): """ - Split the model into backbone and head. + Run batch inference through a backbone and head model using Keras or PyTorch. - This method splits the model into backbone and head. The backbone is summarized - into a single layer which can be retrieved by the layer index 1. The model input - has layer index 0 and the head is the rest of the model with layer index 2 and above. + Dispatches model execution based on ``self.framework_type``. For PyTorch, + converts input data to Tensors, transposes to channels-first format (NCHW), + places inputs on the targeted device, runs inference without gradient tracking + (``torch.no_grad()``), and returns NumPy outputs on CPU. - Parameters: - ----------- - model : keras.Model - Keras model to split into backbone and head. + Parameters + ---------- + backbone : keras.Model or torch.nn.Module + The feature extractor network that processes raw input data into + feature representations. + head : keras.Model or torch.nn.Module + The task-specific head network that computes final predictions + from feature vectors. + input_data : np.ndarray + Batch of image data following the Keras convention. - Returns: - -------- - backbone : keras.Model - Backbone model of the original model. - head : keras.Model - Head model of the original model. + Returns + ------- + fvs : np.ndarray + Extracted feature vectors for the batch. + preds : np.ndarray + Model output predictions for the batch. """ - # Get the backbone model which is the second layer of the model - backbone = model.get_layer(index=1) - # Create a new head model with the same layers as the original model. - # The output of the backbone model is the input of the head model. - backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) - x = backbone_output_shape - for layer in model.layers[2:]: - x = layer(x) - head = keras.Model(inputs=backbone_output_shape, outputs=x) - return backbone, head + if self.framework_type == FrameworkType.KERAS: + fvs = backbone.predict_on_batch(input_data) + preds = head.predict_on_batch(fvs) + return fvs, preds + if self.framework_type == FrameworkType.PYTORCH: + tensor_data = torch.from_numpy(input_data).float() + # Convert Channels-Last (N, H, W, C) to Channels-First (N, C, H, W) + if tensor_data.ndim == 4: + tensor_data = tensor_data.permute(0, 3, 1, 2) + tensor_data = tensor_data.to(self.device) + + def _to_numpy(val): + if isinstance(val, dict): + # Extract the tensor value from the dictionary (e.g., {'energy': tensor(...)}) + val = next(iter(val.values())) + if isinstance(val, torch.Tensor): + return val.cpu().numpy() + return np.asarray(val) + + with torch.no_grad(): + fvs = backbone(tensor_data) + preds = head(fvs) + return fvs.cpu().numpy(), _to_numpy(preds) + + def _load_model(self, path): + """ + Load a trained model and extract its input shape, backbone, and head. + + Supports both Keras and PyTorch frameworks based on ``self.framework_type``. + For PyTorch models, the model is set to evaluation mode (``.eval()``) and its + ``backbone`` and ``head`` submodules are extracted. + + Parameters + ---------- + path : str, pathlib.Path, or None + Path to the saved model file. If ``None``, the function immediately returns ``None`` elements. + + Returns + ------- + input_shape : tuple of int or None + Input feature tensor shape (excluding batch dimension), or ``None`` if ``path`` + is ``None`` or framework is PyTorch. + backbone : keras.Model, torch.nn.Module, or None + Feature extraction backbone of the model, or ``None`` if ``path`` is ``None``. + head : keras.Model, torch.nn.Module, or None + Output prediction head of the model, or ``None`` if ``path`` is ``None``. + + Raises + ------ + TypeError + If the loaded PyTorch object is not an instance of ``torch.nn.Module`` + (e.g., when a ``state_dict`` is loaded instead of the full model). + """ + if path is None: + return None, None, None + self.log.info("Loading the model from %s.", path) + + def _split_keras_model(model): + """Split the Keras model into backbone and head.""" + # Get the backbone model which is the second layer of the model + backbone = model.get_layer(index=1) + # Create a new head model with the same layers as the original model. + # The output of the backbone model is the input of the head model. + backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) + x = backbone_output_shape + for layer in model.layers[2:]: + x = layer(x) + head = keras.Model(inputs=backbone_output_shape, outputs=x) + return backbone, head + + if self.framework_type == FrameworkType.KERAS: + model = keras.saving.load_model(path) + input_shape = model.input_shape[1:] + backbone, head = _split_keras_model(model) + return input_shape, backbone, head + if self.framework_type == FrameworkType.PYTORCH: + model = torch.load(path, map_location=self.device, weights_only=False) + if not isinstance(model, nn.Module): + raise TypeError( + f"Expected a PyTorch 'nn.Module' object at '{path}', " + f"but got '{type(model).__name__}'. " + "Ensure the model was saved via 'torch.save(model, path)' " + "rather than 'torch.save(model.state_dict(), path)'." + ) + model.eval() + return None, model.backbone, model.head def _create_nan_table(self, nonexample_identifiers, columns, shapes): """ diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index 4186b676..75138a50 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -9,7 +9,6 @@ import numpy as np import tables import keras -import tensorflow as tf import torch import torch.nn as nn from torch.utils.data import DataLoader diff --git a/ctlearn/tools/predict_model_main.py b/ctlearn/tools/predict_model_main.py deleted file mode 100644 index 4b64001b..00000000 --- a/ctlearn/tools/predict_model_main.py +++ /dev/null @@ -1,1884 +0,0 @@ -""" -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. -""" - -import atexit -import pathlib -import numpy as np -import os -try: - import tensorflow as tf - import keras -except ImportError: - tf = None - keras = None - -from astropy import units as u -from astropy.coordinates.earth import EarthLocation -from astropy.coordinates import AltAz, SkyCoord -from astropy.table import ( - Table, - hstack, - vstack, - join, - setdiff, -) - -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) -from ctapipe.coordinates import CameraFrame, NominalFrame -from ctapipe.core import Tool -from ctapipe.core.tool import ToolConfigurationError -from ctapipe.core.traits import ( - Bool, - Int, - Path, - flag, - Set, - Dict, - List, - CaselessStrEnum, - ComponentName, - Unicode, - classes_with_traits, -) -from ctapipe.monitoring.interpolation import PointingInterpolator -from ctapipe.io import read_table, write_table, HDF5Merger -from ctapipe.reco.reconstructor import ReconstructionProperty -from ctapipe.reco.stereo_combination import StereoCombiner -from ctapipe.reco.utils import add_defaults_and_meta -from dl1_data_handler.reader import ( - DLDataReader, - ProcessType, - LST_EPOCH, -) -from ctlearn.core.loader import DLDataLoader - -SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -POINTING_GROUP = "/dl1/monitoring/telescope/pointing" -SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" -DL1_TELESCOPE_GROUP = "/dl1/event/telescope" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_TELESCOPE_GROUP = "/dl2/event/telescope" -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] - -__all__ = [ - "PredictCTLearnModel", - "MonoPredictCTLearnModel", - "StereoPredictCTLearnModel", -] - - -class PredictCTLearnModel(Tool): - """ - Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. - - This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image - or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to - store them in the output file. The input data is loaded from the input url using the - ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.loader.DLDataLoader``. - The prediction is performed using the CTLearn models. The data is stored in the output file - following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to - handle the prediction for mono and stereo mode. - - Attributes - ---------- - input_url : pathlib.Path - Input ctapipe HDF5 files including pixel-wise image or waveform data. - use_HDF5Merger : bool - Set whether to use the HDF5Merger component to copy the selected tables from the input file to the output file. - dl1_features : bool - Set whether to include the dl1 feature vectors in the output file. - dl2_telescope : bool - Set whether to include dl2 telescope-event-wise data in the output file. - dl2_subarray : bool - Set whether to include dl2 subarray-event-wise data in the output file. - dl1dh_reader : dl1_data_handler.reader.DLDataReader - DLDataReader object to read the data. - dl1dh_reader_type : str - Type of the DLDataReader to use for the prediction. - stack_telescope_images : bool - Set whether to stack the telescope images in the data loader. Requires ``stereo``. - sort_by_intensity : bool - Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. - prefix : str - Name of the reconstruction algorithm used to generate the dl2 data. - load_type_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the classification of the primary particle type. - load_energy_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression of the primary particle energy. - load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression - of the primary particle arrival direction based on camera coordinate offsets. - load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) or directory (Keras2) for the regression - of the primary particle arrival direction based on spherical coordinate offsets. - output_path : pathlib.Path - Output path to save the dl2 prediction results. - overwrite_tables : bool - Overwrite the table in the output file if it exists. - keras_verbose : int - Verbosity mode of Keras during the prediction. - strategy : tf.distribute.Strategy - MirroredStrategy to distribute the prediction. - data_loader : ctlearn.core.loader.DLDataLoader - DLDataLoader object to load the data. - indices : list of int - List of indices for the data loaders. - batch_size : int - Size of the batch to perform inference of the neural network. - last_batch_size : int - Size of the last batch in the data loaders. - - Methods - ------- - setup() - Set up the tool. - finish() - Finish the tool. - _predict_with_model(model_path) - Load and predict with a CTLearn model. - _predict_classification(example_identifiers) - Predict the classification of the primary particle type. - _predict_energy(example_identifiers) - Predict the energy of the primary particle. - _predict_cameradirection(example_identifiers) - Predict the arrival direction of the primary particle based on camera coordinate offsets. - _predict_skydirection(example_identifiers) - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - _transform_cam_coord_offsets_to_sky(table) - Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _transform_spher_coord_offsets_to_sky(table) - Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - _create_nan_table(nonexample_identifiers, columns, shapes) - Create a table with NaNs for missing predictions. - _store_pointing(all_identifiers) - Store the telescope pointing table from to the output file. - _create_feature_vectors_table(example_identifiers, nonexample_identifiers, classification_feature_vectors, energy_feature_vectors, direction_feature_vectors) - Create the table for the DL1 feature vectors. - """ - - input_url = Path( - help="Input ctapipe HDF5 files including pixel-wise image or waveform data", - allow_none=True, - exists=True, - directory_ok=False, - file_ok=True, - ).tag(config=True) - - use_HDF5Merger = Bool( - default_value=True, - allow_none=False, - help=( - "Set whether to use the HDF5Merger component to copy the selected tables " - "from the input file to the output file. CAUTION: This can only be used " - "if the output file not exists." - ), - ).tag(config=True) - - dl1_features = Bool( - default_value=False, - allow_none=False, - help="Set whether to include the dl1 feature vectors in the output file.", - ).tag(config=True) - - dl2_telescope = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 telescope-event-wise data in the output file.", - ).tag(config=True) - - dl2_subarray = Bool( - default_value=True, - allow_none=False, - help="Set whether to include dl2 subarray-event-wise data in the output file.", - ).tag(config=True) - - dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( - config=True - ) - - stack_telescope_images = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to stack the telescope images in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - sort_by_intensity = Bool( - default_value=False, - allow_none=False, - help=( - "Set whether to sort the telescope images by intensity in the data loader. " - "Requires DLDataReader mode to be ``stereo``." - ), - ).tag(config=True) - - prefix = Unicode( - default_value="CTLearn", - allow_none=False, - help="Name of the reconstruction algorithm used to generate the dl2 data.", - ).tag(config=True) - - load_type_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the classification " - "of the primary particle type." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_energy_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle energy." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_cameradirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle arrival direction based on camera coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_skydirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " - "of the primary particle arrival direction based on spherical coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - batch_size = Int( - default_value=64, - allow_none=False, - help="Size of the batch to perform inference of the neural network.", - ).tag(config=True) - - output_path = Path( - default_value="./output.dl2.h5", - allow_none=False, - help="Output path to save the dl2 prediction results", - ).tag(config=True) - - overwrite_tables = Bool( - default_value=True, - allow_none=False, - help="Overwrite the table in the output file if it exists", - ).tag(config=True) - - keras_verbose = Int( - default_value=1, - min=0, - max=2, - allow_none=False, - help=( - "Verbosity mode of Keras during the prediction: " - "0 = silent, 1 = progress bar, 2 = one line per call." - ), - ).tag(config=True) - - aliases = { - ("i", "input_url"): "PredictCTLearnModel.input_url", - ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", - ("e", "energy_model"): "PredictCTLearnModel.load_energy_model_from", - ( - "d", - "cameradirection_model", - ): "PredictCTLearnModel.load_cameradirection_model_from", - ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", - ("o", "output"): "PredictCTLearnModel.output_path", - } - - flags = { - **flag( - "dl1-features", - "PredictCTLearnModel.dl1_features", - "Include dl1 features", - "Exclude dl1 features", - ), - **flag( - "dl2-telescope", - "PredictCTLearnModel.dl2_telescope", - "Include dl2 telescope-event-wise data in the output file", - "Exclude dl2 telescope-event-wise data in the output file", - ), - **flag( - "dl2-subarray", - "PredictCTLearnModel.dl2_subarray", - "Include dl2 telescope-event-wise data in the output file", - "Exclude dl2 telescope-event-wise data in the output file", - ), - **flag( - "use-HDF5Merger", - "PredictCTLearnModel.use_HDF5Merger", - "Copy data using the HDF5Merger component (CAUTION: This can not be used if the output file already exists)", - "Do not copy data using the HDF5Merger component", - ), - **flag( - "r0-waveforms", - "HDF5Merger.r0_waveforms", - "Include r0 waveforms", - "Exclude r0 waveforms", - ), - **flag( - "r1-waveforms", - "HDF5Merger.r1_waveforms", - "Include r1 waveforms", - "Exclude r1 waveforms", - ), - **flag( - "dl1-parameters", - "HDF5Merger.dl1_parameters", - "Include dl1 parameters", - "Exclude dl1 parameters", - ), - **flag( - "dl1-images", - "HDF5Merger.dl1_images", - "Include dl1 images", - "Exclude dl1 images", - ), - **flag( - "true-parameters", - "HDF5Merger.true_parameters", - "Include true parameters", - "Exclude true parameters", - ), - **flag( - "true-images", - "HDF5Merger.true_images", - "Include true images", - "Exclude true images", - ), - } - - classes = classes_with_traits(DLDataReader) - - def setup(self): - # Check if the ctapipe HDF5Merger component is enabled - if self.use_HDF5Merger: - if os.path.exists(self.output_path): - raise ToolConfigurationError( - f"The output file '{self.output_path}' already exists. Please use " - "'--no-use-HDF5Merger' to disable the usage of the HDF5Merger component." - ) - # Copy selected tables from the input file to the output file - self.log.info("Copying to output destination.") - with HDF5Merger(self.output_path, parent=self) as merger: - merger(self.input_url) - else: - self.log.info( - "No copy to output destination, since the usage of the HDF5Merger component is disabled." - ) - - if tf is None: - raise ImportError("TensorFlow is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") - - # Create a MirroredStrategy. - self.strategy = tf.distribute.MirroredStrategy() - atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore - self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) - - # Set up the data reader - self.log.info("Loading data reader:") - self.log.info("For a large dataset, this may take a while...") - self.dl1dh_reader = DLDataReader.from_name( - self.dl1dh_reader_type, - input_url_signal=[self.input_url], - parent=self, - ) - self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) - # Check if the number of events is enough to form a batch - if self.dl1dh_reader._get_n_events() < self.batch_size: - raise ToolConfigurationError( - f"{self.dl1dh_reader._get_n_events()} events are not enough " - f"to form a batch of size {self.batch_size}. Reduce the batch size." - ) - # Set the indices for the data loaders - self.indices = list(range(self.dl1dh_reader._get_n_events())) - self.last_batch_size = len(self.indices) % ( - self.batch_size * self.strategy.num_replicas_in_sync - ) - - def finish(self): - self.log.info("Tool is shutting down") - - def _predict_with_model(self, model_path): - """ - Load and predict with a CTLearn model. - - Load a model from the specified path and predict the data using the loaded model. - If a last batch loader is provided, predict the last batch and stack the results. - - Parameters - ---------- - model_path : str - Path to a Keras model file (Keras3) or directory (Keras2). - - Returns - ------- - predict_data : astropy.table.Table - Table containing the prediction results. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - # Create a new DLDataLoader for each task - # It turned out to be more robust to initialize the DLDataLoader separately. - data_loader = DLDataLoader( - self.dl1dh_reader, - self.indices, - tasks=[], - batch_size=self.batch_size * self.strategy.num_replicas_in_sync, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Keras is only considering the last complete batch. - # In prediction mode we don't want to loose the last - # uncomplete batch, so we are creating an additional - # batch generator for the remaining events. - data_loader_last_batch = None - if self.last_batch_size > 0: - last_batch_indices = self.indices[-self.last_batch_size :] - data_loader_last_batch = DLDataLoader( - self.dl1dh_reader, - last_batch_indices, - tasks=[], - batch_size=self.last_batch_size, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Load the model from the specified path - model = keras.saving.load_model(model_path) - prediction_colname = ( - model.layers[-1].name if model.layers[-1].name != "softmax" else "type" - ) - backbone_model, feature_vectors = None, None - if self.dl1_features: - # Get the backbone model which is the second layer of the model - backbone_model = model.get_layer(index=1) - # Create a new head model with the same layers as the original model. - # The output of the backbone model is the input of the head model. - backbone_output_shape = keras.Input(model.layers[2].input_shape[1:]) - x = backbone_output_shape - for layer in model.layers[2:]: - x = layer(x) - head = keras.Model(inputs=backbone_output_shape, outputs=x) - # Apply the backbone model with the data loader to retrieve the feature vectors - feature_vectors = backbone_model.predict( - data_loader, verbose=self.keras_verbose - ) - # Apply the head model with the feature vectors to retrieve the prediction - predict_data = Table( - { - prediction_colname: head.predict( - feature_vectors, verbose=self.keras_verbose - ) - } - ) - # Predict the last batch and stack the results to the prediction data - if data_loader_last_batch is not None: - feature_vectors_last_batch = backbone_model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - feature_vectors = np.concatenate( - (feature_vectors, feature_vectors_last_batch) - ) - predict_data = vstack( - [ - predict_data, - Table( - { - prediction_colname: head.predict( - feature_vectors_last_batch, - verbose=self.keras_verbose, - ) - } - ), - ] - ) - else: - # Predict the data using the loaded model - predict_data = model.predict(data_loader, verbose=self.keras_verbose) - # Create a astropy table with the prediction results - # The classification task has a softmax layer as the last layer - # which returns the probabilities for each class in an array, while - # the regression tasks have output neurons which returns the - # predicted value for the task in a dictionary. - if prediction_colname == "type": - predict_data = Table({prediction_colname: predict_data}) - else: - predict_data = Table(predict_data) - # Predict the last batch and stack the results to the prediction data - if data_loader_last_batch is not None: - predict_data_last_batch = model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - if model.layers[-1].name == "type": - predict_data_last_batch = Table( - {prediction_colname: predict_data_last_batch} - ) - else: - predict_data_last_batch = Table(predict_data_last_batch) - predict_data = vstack([predict_data, predict_data_last_batch]) - return predict_data, feature_vectors - - def _predict_classification(self, example_identifiers): - """ - Predict the classification of the primary particle type. - - This method uses a pre-trained type model to predict the type of the primary particle - for a given set of example identifiers. The predicted classification score ('gammaness') - is added to the example identifiers table. - - Parameters: - ----------- - classification_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - predicted classification score ('gammaness'). - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the classification of the primary particle type..." - ) - # Predict the data using the loaded type_model - predict_data, feature_vectors = self._predict_with_model( - self.load_type_model_from - ) - # Create prediction table and add the predicted classification score ('gammaness') - classification_table = example_identifiers.copy() - classification_table.add_column( - predict_data["type"].T[1], name=f"{self.prefix}_tel_prediction" - ) - return classification_table, feature_vectors - - def _predict_energy(self, example_identifiers): - """ - Predict the energy of the primary particle. - - This method uses a pre-trained energy model to predict the energy of the primary particle - for a given set of example identifiers. The predicted energy is then converted from - log10(TeV) to TeV and added to the example identifiers table. - - Parameters: - ----------- - energy_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed energy in TeV. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info("Predicting for the regression of the primary particle energy...") - # Predict the data using the loaded energy_model - predict_data, feature_vectors = self._predict_with_model( - self.load_energy_model_from - ) - # Convert the reconstructed energy from log10(TeV) to TeV - reco_energy = u.Quantity( - np.power(10, np.squeeze(predict_data["energy"])), - unit=u.TeV, - ) - # Create prediction table and add the reconstructed energy in TeV - energy_table = example_identifiers.copy() - energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") - return energy_table, feature_vectors - - def _predict_cameradirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on camera coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the - primary particle for a given set of example identifiers. The predicted camera coordinate offsets - is added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - cameradirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed camera coordinate offsets in x and y. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the regression of the primary particle arrival direction based on camera coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_cameradirection_model_from - ) - # For the direction task, the prediction is the camera coordinate offset in x and y - # from the telescope pointing. - cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) - cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) - # Create prediction table and add the reconstructed energy in TeV - cameradirection_table = example_identifiers.copy() - cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") - cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") - return cameradirection_table, feature_vectors - - def _predict_skydirection(self, example_identifiers): - """ - Predict the arrival direction of the primary particle based on spherical coordinate offsets. - - This method uses a pre-trained direction model to predict the arrival direction of the primary - particle for a given set of example identifiers. The predicted spherical coordinate offsets is - added to the example identifiers table. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - - Returns: - -------- - skydirection_table : astropy.table.Table - Table containing the example identifiers with an additional column for the - reconstructed spherical coordinate offsets in fov_lon and fov_lat. - feature_vectors : np.ndarray - Feature vectors extracted from the backbone model. - """ - self.log.info( - "Predicting for the regression of the primary particle arrival direction based on spherical coordinate offsets..." - ) - # Predict the data using the loaded direction_model - predict_data, feature_vectors = self._predict_with_model( - self.load_skydirection_model_from - ) - # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat - # from the telescope pointing. - fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) - fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) - # Create prediction table and add the reconstructed energy in TeV - skydirection_table = example_identifiers.copy() - skydirection_table.add_column(fov_lon, name="fov_lon") - skydirection_table.add_column(fov_lat, name="fov_lat") - return skydirection_table, feature_vectors - - def _transform_cam_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - # Get the telescope ID from the table - tel_id = table["tel_id"][0] - # Set the telescope position - tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the telescope pointing with the trigger timestamp and the telescope position - altaz = AltAz( - location=tel_ground_frame.to_earth_location(), - obstime=trigger_time, - ) - # Set the telescope pointing - tel_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the camera frame with the focal length and rotation of the camera - camera_frame = CameraFrame( - focal_length=self.dl1dh_reader.subarray.tel[ - tel_id - ].camera.geometry.frame.focal_length, - rotation=self.dl1dh_reader.pix_rotation[tel_id], - telescope_pointing=tel_pointing, - ) - # Set the camera coordinate offset - cam_coord_offset = SkyCoord( - x=table["cam_coord_offset_x"], - y=table["cam_coord_offset_y"], - frame=camera_frame, - ) - # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Transform the true Alt/Az coordinates to camera coordinates - reco_direction = cam_coord_offset.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column(reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az") - table.add_column(reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt") - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "cam_coord_offset_x", - "cam_coord_offset_y", - ] - ) - return table - - def _transform_spher_coord_offsets_to_sky(self, table) -> Table: - """ - Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. - - This method converts the predicted spherical offsets w.r.t. the telescope pointing - in the provided table to Alt/Az coordinates. It also removes the unnecessary columns - from the table that do not the ctapipe DL2 data format. - - Parameters: - ----------- - table : astropy.table.Table - A Table containing the trigger time, telescope pointing, and predicted spherical offsets. - - Returns: - -------- - table : astropy.table.Table - A Table with the Alt/Az coordinates following the ctapipe DL2 data format. - """ - - # Set the trigger timestamp based on the process type - if self.dl1dh_reader.process_type == ProcessType.Simulation: - trigger_time = LST_EPOCH - elif self.dl1dh_reader.process_type == ProcessType.Observation: - trigger_time = table["time"] - # Set the AltAz frame with the reference location and time - altaz = AltAz( - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the array pointing - array_pointing = SkyCoord( - az=table["pointing_azimuth"], - alt=table["pointing_altitude"], - frame=altaz, - ) - # Set the nominal frame with the array pointing - nom_frame = NominalFrame( - origin=array_pointing, - location=self.dl1dh_reader.subarray.reference_location, - obstime=trigger_time, - ) - # Set the reco direction in (fov_lon, fov_lat) coordinates - reco_direction = SkyCoord( - fov_lon=table["fov_lon"], - fov_lat=table["fov_lat"], - frame=nom_frame, - ) - # Transform the reco direction from nominal frame to the AltAz frame - sky_coord = reco_direction.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - table.add_column(sky_coord.az.to(u.deg), name=f"{self.prefix}_az") - table.add_column(sky_coord.alt.to(u.deg), name=f"{self.prefix}_alt") - # Remove unnecessary columns from the table that do not the ctapipe DL2 data format - table.remove_columns( - [ - "time", - "pointing_azimuth", - "pointing_altitude", - "fov_lon", - "fov_lat", - ] - ) - return table - - def _create_nan_table(self, nonexample_identifiers, columns, shapes): - """ - Create a table with NaNs for missing predictions. - - This method creates a table with NaNs for missing predictions for the non-example identifiers. - In stereo mode, the table also a column for the valid telescopes is added with all False values. - - Parameters: - ----------- - nonexample_identifiers : astropy.table.Table - Table containing the non-example identifiers. - columns : list of str - List of column names to create in the table. - shapes : list of shapes - List of shapes for the columns to create in the table. - - Returns: - -------- - nan_table : astropy.table.Table - Table containing NaNs for missing predictions. - """ - # Create a table with NaNs for missing predictions - nan_table = nonexample_identifiers.copy() - for column_name, shape in zip(columns, shapes): - nan_table.add_column(np.full(shape, np.nan), name=column_name) - # Add that no telescope is valid for the non-example identifiers in stereo mode - if self.dl1dh_reader.mode == "stereo": - nan_table.add_column( - np.zeros( - (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ), - name=f"{self.prefix}_telescopes", - ) - return nan_table - - def _store_pointing(self, all_identifiers): - """ - Store the telescope pointing table from to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - - # Initialize the pointing interpolator from ctapipe - pointing_interpolator = PointingInterpolator( - bounds_error=False, extrapolate=True - ) - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Get the telescope pointing from the dl1dh reader - tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] - # Add the telescope pointing table to the pointing interpolator - pointing_interpolator.add_table(tel_id, tel_pointing) - tel_identifiers = all_identifiers.copy() - if self.dl1dh_reader.mode == "mono": - tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] - # Interpolate the telescope pointing - tel_altitude, tel_azimuth = pointing_interpolator( - tel_id, tel_identifiers["time"] - ) - tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") - tel_identifiers.add_column(tel_altitude, name="pointing_altitude") - pointing_info.append(tel_identifiers) - if self.dl1dh_reader.mode == "mono": - tel_pointing_table = Table( - { - "time": tel_identifiers["time"], - "azimuth": tel_identifiers["pointing_azimuth"], - "altitude": tel_identifiers["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info = vstack(pointing_info) - if self.dl1dh_reader.mode == "stereo": - # Group the pointing information by subarray event keys - # TODO: This needs to be debugged with SST1M data - pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) - pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) - pointing_info = join( - all_identifiers, - pointing_mean, - keys=SUBARRAY_EVENT_KEYS, - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - ) - return pointing_info - - def _create_feature_vectors_table( - self, - example_identifiers, - nonexample_identifiers=None, - classification_feature_vectors=None, - energy_feature_vectors=None, - direction_feature_vectors=None, - ): - """ - Create the table for the DL1 feature vectors. - - This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for - non-example identifiers. The feature vectors are stored in the columns of the table. The table also - contains a column for the valid predictions. - - Parameters: - ----------- - example_identifiers : astropy.table.Table - Table containing the example identifiers. - nonexample_identifiers : astropy.table.Table or None - Table containing the non-example identifiers to fill the NaNs. - classification_feature_vectors : np.ndarray or None - Array containing the classification feature vectors. - energy_feature_vectors : np.ndarray or None - Array containing the energy feature vectors. - direction_feature_vectors : np.ndarray or None - Array containing the direction feature vectors. - - Returns: - -------- - feature_vector_table : astropy.table.Table - Table containing the DL1 feature vectors for the example and non-example identifiers. - """ - # Create the feature vector table - feature_vector_table = example_identifiers.copy() - feature_vector_table.remove_columns( - ["pointing_azimuth", "pointing_altitude", "time"] - ) - columns_list, shapes_list = [], [] - if classification_feature_vectors is not None: - is_valid_col = ~np.isnan( - np.min(classification_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - classification_feature_vectors, - name=f"{self.prefix}_tel_classification_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - classification_feature_vectors.shape[1], - ) - ) - if energy_feature_vectors is not None: - is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) - feature_vector_table.add_column( - energy_feature_vectors, name=f"{self.prefix}_tel_energy_feature_vectors" - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - energy_feature_vectors.shape[1], - ) - ) - if direction_feature_vectors is not None: - is_valid_col = ~np.isnan( - np.min(direction_feature_vectors, axis=1), dtype=bool - ) - feature_vector_table.add_column( - direction_feature_vectors, - name=f"{self.prefix}_tel_geometry_feature_vectors", - ) - if nonexample_identifiers is not None: - columns_list.append(f"{self.prefix}_tel_geometry_feature_vectors") - shapes_list.append( - ( - len(nonexample_identifiers), - direction_feature_vectors.shape[1], - ) - ) - # Produce output table with NaNs for missing predictions - if nonexample_identifiers is not None: - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=columns_list, - shapes=shapes_list, - ) - feature_vector_table = vstack([feature_vector_table, nan_table]) - is_valid_col = np.concatenate( - (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) - ) - # Add is_valid column to the feature vector table - feature_vector_table.add_column( - is_valid_col, - name=f"{self.prefix}_tel_is_valid", - ) - return feature_vector_table - - -class MonoPredictCTLearnModel(PredictCTLearnModel): - """ - Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. - - This tool extends the ``PredictCTLearnModel`` to specifically handle monoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_telescope_pointing(all_identifiers) - Store the telescope pointing table for the mono mode for MC simulation. - """ - - name = "ctlearn-predict-mono-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnModel.batch_size=64 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ - --dl1-features \\ - --use-HDF5Merger \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - - To predict from pixel-wise waveform data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.r1.h5 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader \\ - --DLWaveformReader.sequnce_length=20 \\ - --DLWaveformReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ - --use-HDF5Merger \\ - --no-r0-waveforms \\ - --no-r1-waveforms \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - """ - - stereo_combiner_cls = ComponentName( - StereoCombiner, - default_value="StereoMeanCombiner", - help="Which stereo combination method to use after the monoscopic reconstruction.", - ).tag(config=True) - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.example_identifiers.copy() - example_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) - all_identifiers = self.dl1dh_reader.tel_trigger_table.copy() - all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS + ["time"]) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Pointing table for the mono mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_telescope_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - classification_feature_vectors = None - if self.load_type_model_from is not None: - self.type_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.PARTICLE_TYPE, - parent=self, - ) - # Predict the energy of the primary particle - classification_table, classification_feature_vectors = ( - super()._predict_classification(example_identifiers) - ) - if self.dl2_telescope: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - ) - classification_table = vstack([classification_table, nan_table]) - # Add is_valid column to the energy table - classification_table.add_column( - ~np.isnan( - classification_table[f"{self.prefix}_tel_prediction"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - classification_table, - ParticleClassificationContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = classification_table["tel_id"] == tel_id - classification_tel_table = classification_table[telescope_mask] - classification_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file for the selected telescope - write_table( - classification_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info("Processing and storing the subarray type prediction...") - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_classification_table = self.type_stereo_combiner.predict_table( - classification_table - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - subarray_classification_table[f"{self.prefix}_telescopes"].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - ( - len(subarray_classification_table), - len(self.dl1dh_reader.tel_ids), - ), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_classification_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_classification_table[f"{self.prefix}_telescopes"] = ( - reco_telescopes - ) - # Save the prediction to the output file - write_table( - subarray_classification_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - self.energy_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.ENERGY, - parent=self, - ) - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - if self.dl2_telescope: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = energy_table["tel_id"] == tel_id - energy_tel_table = energy_table[telescope_mask] - energy_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file - write_table( - energy_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray energy prediction..." - ) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_energy_table = self.energy_stereo_combiner.predict_table( - energy_table - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if subarray_energy_table[f"{self.prefix}_telescopes"].dtype != np.bool_: - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - (len(subarray_energy_table), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_energy_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_energy_table[f"{self.prefix}_telescopes"] = reco_telescopes - # Save the prediction to the output file - write_table( - subarray_energy_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - ) - direction_feature_vectors = None - if self.load_cameradirection_model_from is not None: - self.geometry_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.GEOMETRY, - parent=self, - ) - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=TELESCOPE_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = ( - super()._predict_cameradirection(example_identifiers) - ) - direction_tel_tables = [] - if self.dl2_telescope: - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = direction_table["tel_id"] == tel_id - direction_tel_table = direction_table[telescope_mask] - direction_tel_table = super()._transform_cam_coord_offsets_to_sky( - direction_tel_table - ) - # Produce output table with NaNs for missing predictions - nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id - nonexample_identifiers_tel = nonexample_identifiers[ - nan_telescope_mask - ] - if len(nonexample_identifiers_tel) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers_tel, - columns=[f"{self.prefix}_tel_alt", f"{self.prefix}_tel_az"], - shapes=[ - (len(nonexample_identifiers_tel),), - (len(nonexample_identifiers_tel),), - ], - ) - direction_tel_table = vstack([direction_tel_table, nan_table]) - direction_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Add is_valid column to the direction table - direction_tel_table.add_column( - ~np.isnan( - direction_tel_table[f"{self.prefix}_tel_alt"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_tel_table, - ReconstructedGeometryContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - direction_tel_tables.append(direction_tel_table) - # Save the prediction to the output file - write_table( - direction_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray geometry prediction..." - ) - # Stack the telescope tables to the subarray table - direction_tel_tables = vstack(direction_tel_tables) - # Sort the table by the telescope event keys - direction_tel_tables.sort(TELESCOPE_EVENT_KEYS) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_direction_table = self.geometry_stereo_combiner.predict_table( - direction_tel_tables - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - subarray_direction_table[f"{self.prefix}_telescopes"].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - (len(subarray_direction_table), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_direction_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_direction_table[f"{self.prefix}_telescopes"] = ( - reco_telescopes - ) - # Save the prediction to the output file - write_table( - subarray_direction_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - ) - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - classification_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = feature_vector_table["tel_id"] == tel_id - feature_vectors_tel_table = feature_vector_table[telescope_mask] - feature_vectors_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vectors_tel_table, - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", - ) - - def _store_mc_telescope_pointing(self, all_identifiers): - """ - Store the telescope pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - # Create the pointing table for each telescope - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Pointing table for the mono mode - tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) - tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") - tel_pointing.rename_column( - "telescope_pointing_altitude", "pointing_altitude" - ) - # Join the prediction table with the telescope pointing table - tel_pointing = join( - left=tel_pointing, - right=all_identifiers, - keys=["obs_id", "tel_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - tel_pointing.sort(TELESCOPE_EVENT_KEYS) - # Retrieve the example identifiers for the selected telescope - tel_pointing_table = Table( - { - "time": tel_pointing["time"], - "azimuth": tel_pointing["pointing_azimuth"], - "altitude": tel_pointing["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info.append(tel_pointing) - pointing_info = vstack(pointing_info) - return pointing_info - - -class StereoPredictCTLearnModel(PredictCTLearnModel): - """ - Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. - - This tool extends the ``PredictCTLearnModel`` to specifically handle stereoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_subarray_pointing(all_identifiers) - Store the subarray pointing table for the stereo mode for MC simulation. - """ - - name = "ctlearn-predict-stereo-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in stereo mode using trained CTLearn models: - > ctlearn-predict-stereo-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnModel.batch_size=16 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --DLImageReader.mode=stereo \\ - --DLImageReader.min_telescopes=2 \\ - --PredictCTLearnModel.stack_telescope_images=True \\ - --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ - --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - """ - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() - example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) - all_identifiers = self.dl1dh_reader.subarray_trigger_table.copy() - all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Construct the survival telescopes for each event of the example_identifiers - survival_telescopes = [] - for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: - survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) - survival_tels = [ - self.dl1dh_reader.subarray.tel_indices[tel_id] - for tel_id in subarray_event["tel_id"].data - ] - survival_mask[survival_tels] = True - survival_telescopes.append(survival_mask) - # Add the survival telescopes to the example_identifiers - example_identifiers.add_column( - survival_telescopes, name=f"{self.prefix}_telescopes" - ) - # Pointing table for the stereo mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_subarray_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - classification_feature_vectors = None - if self.load_type_model_from is not None: - # Predict the energy of the primary particle - classification_table, classification_feature_vectors = ( - super()._predict_classification(example_identifiers) - ) - if self.dl2_subarray: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - ) - classification_table = vstack([classification_table, nan_table]) - # Add is_valid column to the energy table - classification_table.add_column( - ~np.isnan( - classification_table[f"{self.prefix}_tel_prediction"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Rename the columns for the stereo mode - classification_table.rename_column( - f"{self.prefix}_tel_prediction", f"{self.prefix}_prediction" - ) - classification_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - classification_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - classification_table, - ParticleClassificationContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - classification_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - if self.dl2_subarray: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Rename the columns for the stereo mode - energy_table.rename_column( - f"{self.prefix}_tel_energy", f"{self.prefix}_energy" - ) - energy_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - energy_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - energy_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - ) - direction_feature_vectors = None - if self.load_skydirection_model_from is not None: - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=SUBARRAY_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = super()._predict_skydirection( - example_identifiers - ) - if self.dl2_subarray: - # Transform the spherical coordinate offsets to sky coordinates - direction_table = super()._transform_spher_coord_offsets_to_sky( - direction_table - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_alt", f"{self.prefix}_az"], - shapes=[ - (len(nonexample_identifiers),), - (len(nonexample_identifiers),), - ], - ) - direction_table = vstack([direction_table, nan_table]) - # Add is_valid column to the direction table - direction_table.add_column( - ~np.isnan(direction_table[f"{self.prefix}_alt"].data, dtype=bool), - name=f"{self.prefix}_is_valid", - ) - direction_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_table, - ReconstructedGeometryContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - direction_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - ) - - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - classification_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. - # Rename the columns for the stereo mode - feature_vector_table.rename_column( - f"{self.prefix}_tel_classification_feature_vectors", - f"{self.prefix}_classification_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_energy_feature_vectors", - f"{self.prefix}_energy_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_geometry_feature_vectors", - f"{self.prefix}_geometry_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - feature_vector_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vector_table, - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", - ) - - def _store_mc_subarray_pointing(self, all_identifiers): - """ - Store the subarray pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the subarray pointing information. - """ - # Read the subarray pointing table - pointing_info = read_table( - self.input_url, - f"{SIMULATION_CONFIG_TABLE}", - ) - # Assuming min_az = max_az and min_alt = max_alt - pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) - pointing_info.rename_column("min_az", "pointing_azimuth") - pointing_info.rename_column("min_alt", "pointing_altitude") - # Join the prediction table with the telescope pointing table - pointing_info = join( - left=pointing_info, - right=all_identifiers, - keys=["obs_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - ) - return pointing_info - - -def mono_tool(): - # Run the tool - mono_tool = MonoPredictCTLearnModel() - mono_tool.run() - - -def stereo_tool(): - # Run the tool - stereo_tool = StereoPredictCTLearnModel() - stereo_tool.run() - - -if __name__ == "mono_tool": - mono_tool() - -if __name__ == "stereo_tool": - stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/pytorch/config/default_config_file.yml b/ctlearn/tools/pytorch/config/default_config_file.yml deleted file mode 100644 index abf55871..00000000 --- a/ctlearn/tools/pytorch/config/default_config_file.yml +++ /dev/null @@ -1,152 +0,0 @@ -data: - train_gamma_proton: ./data/gamma_proton_train_remix.dl1.pickle - validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle - - train_gamma: ./data/gamma_955000_train.pickle - validation_gamma: ./data/gamma_106141_validation.pickle - - test_gamma: ./data/gamma_1805522_test_gamma.pickle - test_proton: ./data/proton_130811_test_proton.pickle - test_electron: None - test_validation_gamma: ./data/gamma_180552_test_val_gamma.pickle - - test_validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle - - observation: ./run_2931.dl1.pickle - # Important: This is only for testing purpose. Set always to 0 - # when you are training, validating or estimating the dl2 files - training_reduce_factor: 0 #64 #4 - validation_reduce_factor: 0 #16 #8 - validation_test_reduce_factor: 0 #16 #8 - - # Check points - type_checkpoint: ./run/run_type_training_14/exp_14_type_train/version_0/Epoch_6_type_train_acc_80.9682309627532959.pth - energy_checkpoint: /home/cpozogonzalez/ctlearn/run/run_energy_training_14/exp_14_energy_train/version_1/Epoch_13_energy_train_loss_30.2185532478501244.pth - direction_checkpoint: /lhome/ext/ucm147/ucm1477/data/check_points/v_5/Epoch_23_cameradirection_train_loss_7296.6534562211982120.pth - -run_details: - - mode: "observation" # The option are: "train", "results", "observation" and "validate" - task: "direction" # The option are: "all", "energy" "type" and "direction" - test_type: "gamma" # The option are: "gamma" "proton" or "electron" - experiment_number: 14 # The experiment number. The experiment folder is saved into the "run" folder. - - -cut-off: - - leakage_intensity: 0.2 # bigger to this value, the event is removed - intensity: 50 # below to this value, the event is removed - -model: - - model_type: - model_name: "DoubleBBEfficientNet" - parameters: - model_variant: "efficientnet-b3" - task: 'type' - num_outputs: 2 - device_str: "cuda" - energy_bins: None - - model_energy: - model_name: "ThinResNet" - parameters: - task: 'energy' - num_inputs: 1 - num_outputs: 1 - num_blocks: [3, 4, 6, 3] #[2, 3, 3, 3] - dropout: 0.1 - use_bn: False - - model_direction: - model_name: "ThinResNet_DBB" - parameters: - task: 'direction' - num_inputs: 1 - num_outputs: 3 - num_blocks: [3, 4, 6, 3] - dropout: 0.1 - use_bn: False - - # model_direction: - # model_name: "DBBNoPropDTReg" - # parameters: - # task: 'direction' - # num_outputs: 3 - # embedding_dim: 512 - # T: 3 - # eta: 0.1 #0.1 - -# Hyper-parameters -hyp: - - epochs: 30 - batches: 128 #128 #64 - dynamic_batches: True - optimizer: Adamw - momentum: 0.957 #Yolo 0.937 # Efficient-b3 0.757 - weight_decay: 0.0005 #0.004676 #0.0001 #0.00002 Efficient-b3 0.0005 - learning_rate: 1e-4 #1e-5 #Efficient-b3 1e-5 - lrf: 0.1 - start_epoch: 0 - steps_epoch: 100 # Computed online. Must be removed - l2_lambda: 1e-7 #1e-5 #1e-5 # L2 regularization (Set to 0.0 to skip the L2 Regularization) - adam_epsilon: 1.0e-08 #7.511309034256153e-05 #1.0e-08 - gradient_clip_val: 3.0 # Avoid gradient explosion - - save_k: 200 # Save as maximum k checkpoints. - -augmentation: - # probabilities for augmentation range = [0, 1.0] - # prob = 0.0 -> Always apply the augmentation - # prob >= 1.0 -> Never apply the augmentation, i.e., Set bigger than 1.0 ( ex: 2.0) if you want disable it. - # Note: mask augmentation is always on even with flag use_augmentation = True - # To disable it, just set to 2.5 for example. - - use_augmentation: True # This apply only on training mode. - aug_prob: 0.5 # Probability of use Augmentation - rot_prob: 0.5 # Rotation probability - trans_prob: 0.5 # Translation probability - flip_hor_prob: 0.5 # Horizontal Flip probability - flip_ver_prob: 0.5 # Vertical Flip probability - mask_prob: 0.5 # Apply mask probability - mask_dvr_prob: 0.5 # Apply dvr mask probability - noise_prob: 0.5 # No implemented yet. - max_rot: 5 # Maximum rotation in augmentation - max_trans: 10 # Maximum translation in augmentation - -normalization: - - # Normalization: Im' = (Im-mu)/sigma - use_clean: True # Use the image with the applied mask (True), IOC the mask is not applied (False) - use_clean_dvr: False - type_mu: 0.0 - type_sigma: 1000.0 - - dir_mu: 0.0 - dir_sigma: 1000.0 - - energy_mu: 0.0 - energy_sigma: 1000.0 - -dataset: - num_workers: 1 # - pin_memory: True - persistent_workers: True # - -# Hardware Architecture and precision -arch: - # device: 'mps' # Apple Mx - device: 'cuda' - precision_type: "32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - precision_energy: "32-true" #"32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - precision_direction: "32-true" # "bf16-mixed" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - # (bf16 for GPU with Ampere or higher, it is better that 16 because is numerical more stability) - # devices: [0,1] # [0,1] For multiple GPUs - devices: [0,1] - # Note: Check the documentation for more information. - strategy: 'deepspeed_stage_2' # Options: auto, dpp, dpp_swap, fsdp, deepspeed, horovod, bagua, deepspeed_stage_2, deepspeed_stage_3, colossalai, hivemind, etc... - -Notes: - Note_1: Training with augmentation dvr using 1-3 dilatations - Note_2: Trainining b3 applying always the mask \ No newline at end of file diff --git a/ctlearn/tools/tests/test_predict_LST1.py b/ctlearn/tools/tests/test_predict_LST1.py index ffc12b2b..b9ee77e3 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -32,7 +32,7 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") -@pytest.mark.parametrize("framework", ["Keras"]) +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_mono_model_with_lst1_mock_data( tmp_path, ctlearn_trained_dl1_mono_models, mock_lst1_dl1_file, framework ): diff --git a/ctlearn/tools/train_model_current.py b/ctlearn/tools/train_model_current.py deleted file mode 100644 index 84b2b264..00000000 --- a/ctlearn/tools/train_model_current.py +++ /dev/null @@ -1,267 +0,0 @@ -import atexit -import pandas as pd -import numpy as np -import sys -from ctapipe.core import Tool -from ctapipe.core.traits import CaselessStrEnum, Dict -from ctlearn.core.ctlearn_enum import FrameworkType -class DLFrameWork(Tool): - """ - Tool to select and run a specific deep learning training framework (Keras or PyTorch) - for CTLearn model training. It dynamically loads the appropriate subclass based on - the user-defined --framework argument. - """ - - name = "dlframework" - - framework_type = CaselessStrEnum( - ["pytorch", "keras"], - default_value="keras", - help="Framework to use: pytorch or keras", - ).tag(config=True) - - early_stopping = Dict( - default_value=None, - allow_none=True, - help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " - ), - ).tag(config=True) - - early_stopping = Dict( - default_value=None, - allow_none=True, - help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " - ), - ).tag(config=True) - - early_stopping = Dict( - default_value=None, - allow_none=True, - help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " - ), - ).tag(config=True) - - early_stopping = Dict( - default_value=None, - allow_none=True, - help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " - ), - ).tag(config=True) - - aliases = { - "framework": "DLFrameWork.framework_type", - } - - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import TrainPyTorchModel - aliases.update(TrainPyTorchModel.aliases) - except ImportError: - pass - - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel - aliases.update(TrainKerasModel.aliases) - except ImportError: - pass - - def __init__(self, **kwargs): - """ - Initialize the DLFrameWork tool and prepare for framework injection. - """ - super().__init__(**kwargs) - self.framework_instance = None - - def setup(self): - """ - Setup method called after basic trait parsing. - This dynamically loads and prepares the correct framework subclass - (TrainKerasModel or TrainPyTorchModel). - """ - framework_enum = self.string_to_type(self.framework_type) - self.framework_instance = self.get_framework(framework_enum) - - # Inject aliases and shared config before full CLI parsing - self.framework_instance.update_config(self.config) - self.aliases.update(self.framework_instance.aliases) - DLFrameWork.aliases.update(self.framework_instance.aliases) - - def start(self): - """ - Start method called after setup. Executes the selected framework instance. - """ - print("start") - self.framework_instance.run() - - @classmethod - def string_to_type(cls, str_type: str) -> FrameworkType: - """ - Convert a string to a FrameworkType enum (case-insensitive). - - Parameters: - str_type (str): The name of the framework (e.g., 'keras', 'pytorch'). - - Returns: - FrameworkType: Corresponding enum value. - - Raises: - ValueError: If the provided string is not a valid framework type. - """ - try: - return FrameworkType[str_type.upper()] - except KeyError: - raise ValueError(f"'{str_type}' is not a valid framework type.") - - @classmethod - def get_framework(cls, framework_type: FrameworkType): - """ - Dynamically import and return the corresponding training class - based on the framework type. - - Parameters: - framework_type (FrameworkType): Enum indicating which framework to use. - - Returns: - Tool: An instance of the selected training framework (subclass of Tool). - - Raises: - ImportError: If the training module could not be imported. - ValueError: If the framework type is unknown. - """ - if framework_type == FrameworkType.KERAS: - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel - - fw = TrainKerasModel() - except ImportError as e: - raise ImportError(f"Not possible to import TrainKerasModel: {e}") from e - - elif framework_type == FrameworkType.PYTORCH: - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import ( - TrainPyTorchModel, - ) - - fw = TrainPyTorchModel() - - except ImportError as e: - raise ImportError( - f"Not possible to import TrainPyTorchModel: {e}" - ) from e - - else: - raise ValueError(f"Unknown Framework: {framework_type.name}") - - return fw - - @property - def classes(self): - from ctlearn.tools.train.base_train_model import TrainCTLearnModel - from ctapipe.core.traits import classes_with_traits - from dl1_data_handler.reader import DLDataReader - - tool_classes = [ - type(self), - TrainCTLearnModel, - ] - - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel - tool_classes.append(TrainKerasModel) - except ImportError: - pass - - try: - from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import TrainPyTorchModel - tool_classes.append(TrainPyTorchModel) - except ImportError: - pass - - return tool_classes + classes_with_traits(DLDataReader) - -def main(): - # Run the tool - tool = DLFrameWork() - - # Manually parse --framework to determine which subclass to load, as traitlets alias - # update can sometimes fail to parse it correctly before setup - framework = "keras" - for i, arg in enumerate(sys.argv[1:]): - if arg.startswith("--framework="): - framework = arg.split("=")[1].strip() - elif arg == "--framework" and i + 2 < len(sys.argv): - framework = sys.argv[i + 2].strip() - - tool.framework_type = framework - - minimal_args = [ - arg for arg in sys.argv[1:] if "--framework" in arg or arg in ["-h", "--help"] - ] - tool.initialize(argv=minimal_args) - - # Setup and inject the correct framework instance - tool.setup() - - # Parse all CLI args with the selected framework subclass - tool.framework_instance.initialize(argv=sys.argv[1:]) - - tool.run() - - -if __name__ == "__main__": - main() - - -# Example: -# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir2 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite -# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_*.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --pattern-background proton_theta_23.161_az_99.261_runs833-1250*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 - -# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 - - - - -# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 - - -# --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 - - - -# Type -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# Direction -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & - - - -# Energy -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_1.yml> nohup_energy_training.out 2>&1 & - - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs243-302.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs303-362.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs363-422.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs423-482.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs483-541.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs542-600.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_6.000_az_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_2.yml> nohup_energy_training.out 2>&1 & \ No newline at end of file diff --git a/ctlearn/tools/utils.py b/ctlearn/tools/utils.py index d9d88e1f..0ed8425a 100644 --- a/ctlearn/tools/utils.py +++ b/ctlearn/tools/utils.py @@ -162,6 +162,7 @@ def setup_framework(model_paths): If no valid model paths are provided, or if multiple inconsistent frameworks are detected across the provided paths. """ + def _detect_framework(path_val): """ Determines framework based on file extension. @@ -178,6 +179,7 @@ def _detect_framework(path_val): f"Invalid model extension '{ext}' for file '{path}'. " "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." ) + # Detect frameworks from all non-None paths detected_frameworks = {} for path in model_paths: From 773b4608c5a5c3e912599386d0a13ea9070a9b69 Mon Sep 17 00:00:00 2001 From: Tjark Miener Date: Thu, 6 Aug 2026 15:32:47 +0200 Subject: [PATCH 12/12] even more polishing --- .../tools/predict/keras/predic_LST1_keras.py | 255 ------ .../tools/predict/keras/predic_model_keras.py | 156 ---- ctlearn/tools/predict/predict_LST1.py | 834 ------------------ ctlearn/tools/predict/predict_mono.py | 541 ------------ ctlearn/tools/predict/predict_stereo.py | 375 -------- .../predict/pytorch/predic_LST1_pytorch.py | 312 ------- .../predict/pytorch/predic_model_pytorch.py | 144 --- ctlearn/tools/predict/utils/load_model.py | 55 -- .../predict/utils/optimaze_batch_size.py | 174 ---- ctlearn/tools/tests/test_predict_LST1.py | 10 - ctlearn/tools/tests/test_predict_model.py | 11 - ctlearn/tools/tests/test_train_model.py | 9 - .../pytorch/config/default_config_file.yml | 152 ---- 13 files changed, 3028 deletions(-) delete mode 100644 ctlearn/tools/predict/keras/predic_LST1_keras.py delete mode 100644 ctlearn/tools/predict/keras/predic_model_keras.py delete mode 100644 ctlearn/tools/predict/predict_LST1.py delete mode 100644 ctlearn/tools/predict/predict_mono.py delete mode 100644 ctlearn/tools/predict/predict_stereo.py delete mode 100644 ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py delete mode 100644 ctlearn/tools/predict/pytorch/predic_model_pytorch.py delete mode 100644 ctlearn/tools/predict/utils/load_model.py delete mode 100644 ctlearn/tools/predict/utils/optimaze_batch_size.py delete mode 100644 ctlearn/tools/train/pytorch/config/default_config_file.yml diff --git a/ctlearn/tools/predict/keras/predic_LST1_keras.py b/ctlearn/tools/predict/keras/predic_LST1_keras.py deleted file mode 100644 index 7118a67f..00000000 --- a/ctlearn/tools/predict/keras/predic_LST1_keras.py +++ /dev/null @@ -1,255 +0,0 @@ -""" -Keras prediction module for LST1 telescope data. -This module provides functionality to load trained Keras models and perform predictions -on DL1 level data for particle type classification, energy estimation, and direction reconstruction. -""" - -from ctapipe.io import read_table -from astropy.table import join -import keras -from dl1_data_handler.reader import get_unmapped_image -import numpy as np - - -def predictions(self): - """ - Perform predictions on input DL1 data using trained Keras models. - - This function processes the input file in batches, applies quality cuts, - and generates predictions for particle type, energy, and/or direction - depending on the configured models. The models are split into backbone - and head components to extract feature vectors. - - Returns - ------- - tuple - Contains the following arrays: - - event_id: Event identifiers - - tel_azimuth: Telescope azimuth angles - - tel_altitude: Telescope altitude angles - - trigger_time: Event trigger times in MJD - - prediction: Particle type classification scores (gammaness) - - energy: Reconstructed energy values - - cam_coord_offset_x: Camera coordinate offset in x direction - - cam_coord_offset_y: Camera coordinate offset in y direction - - classification_fvs: Classification feature vectors from backbone - - energy_fvs: Energy estimation feature vectors from backbone - - direction_fvs: Direction reconstruction feature vectors from backbone - - Notes - ----- - The function processes data in batches to manage memory efficiently and - applies quality selection criteria before making predictions. - """ - # Initialize output arrays for storing results - event_id, tel_azimuth, tel_altitude, trigger_time = [], [], [], [] - prediction, energy, cam_coord_offset_x, cam_coord_offset_y = [], [], [], [] - classification_fvs, energy_fvs, direction_fvs = [], [], [] - - # Process input file in batches - for start in range(0, self.table_length, self.batch_size): - stop = min(start + self.batch_size, self.table_length) - self.log.debug("Processing chunk from '%d' to '%d'.", start, stop - 1) - - # Read the DL1 data table for current batch - dl1_table = read_table( - self.input_url, self.image_table_path, start=start, stop=stop - ) - - # Join tables to enable quality selection - # Join with parameter table for event parameters - dl1_table = join( - left=dl1_table, - right=self.parameter_table, - keys=["event_id"], - ) - # Join with trigger table for timing information - dl1_table = join( - left=dl1_table, - right=self.trigger_table, - keys=["event_id"], - ) - - # Apply quality selection criteria - # Initialize mask to accept all events initially - passes_quality_checks = np.ones(len(dl1_table), dtype=bool) - - # Apply quality query if configured - if self.quality_query: - passes_quality_checks = self.quality_query.get_table_mask(dl1_table) - - # Filter events based on quality criteria - dl1_table = dl1_table[passes_quality_checks] - - # Skip batch if no events passed quality selection - if len(dl1_table) == 0: - self.log.debug("No events passed the quality selection.") - continue - - # Prepare input data by mapping images to model input format - data = [] - for event in dl1_table: - # Get the unmapped image with specified channels and transforms - image = get_unmapped_image(event, self.channels, self.transforms) - # Map image to model's expected input format - data.append(self.image_mapper.map_image(image)) - input_data = {"input": np.array(data)} - - # Handle compatibility between Keras 2 and Keras 3 - # Keras 3 expects direct array input, not dictionary - if int(keras.__version__.split(".")[0]) >= 3: - input_data = input_data["input"] - - # Store event metadata - event_id.extend(dl1_table["event_id"].data) - tel_azimuth.extend(dl1_table["tel_az"].data) - tel_altitude.extend(dl1_table["tel_alt"].data) - trigger_time.extend(dl1_table["time"].mjd) - - # Perform particle type classification if model is loaded - if self.load_type_model_from is not None: - # Extract feature vectors from backbone - classification_feature_vectors = self.backbone_type.predict_on_batch(input_data) - classification_fvs.extend(classification_feature_vectors) - # Generate predictions from head model - predict_data = self.head_type.predict_on_batch(classification_feature_vectors) - # Extract gammaness score (probability of being gamma) - prediction.extend(predict_data[:, 1]) - - # Perform energy estimation if model is loaded - if self.load_energy_model_from is not None: - # Extract feature vectors from backbone - energy_feature_vectors = self.backbone_energy.predict_on_batch(input_data) - energy_fvs.extend(energy_feature_vectors) - # Generate energy predictions from head model - predict_data = self.head_energy.predict_on_batch(energy_feature_vectors) - energy.extend(predict_data.T[0]) - - # Perform direction reconstruction if model is loaded - if self.load_cameradirection_model_from is not None: - # Extract feature vectors from backbone - direction_feature_vectors = self.backbone_direction.predict_on_batch(input_data) - direction_fvs.extend(direction_feature_vectors) - # Generate direction predictions from head model - predict_data = self.head_direction.predict_on_batch(direction_feature_vectors) - # Extract x and y components of camera coordinate offset - cam_coord_offset_x.extend(predict_data.T[0]) - cam_coord_offset_y.extend(predict_data.T[1]) - - return (event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, - cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs) - - -def _split_model(model): - """ - Split a Keras model into backbone and head components. - - This function separates a trained model into two parts: - - Backbone: Feature extraction layers (typically convolutional layers) - - Head: Task-specific prediction layers (typically dense layers) - - This separation allows extraction of intermediate feature representations - which can be useful for analysis or transfer learning. - - Parameters - ---------- - model : keras.Model - Complete trained Keras model to be split. The model should have: - - Layer 0: Input layer - - Layer 1: Backbone (feature extractor) - - Layers 2+: Head (prediction layers) - - Returns - ------- - backbone : keras.Model - Feature extraction model that outputs intermediate representations. - head : keras.Model - Prediction model that takes backbone outputs and produces final predictions. - - Notes - ----- - The function assumes a specific model architecture where the backbone - is the second layer (index 1) of the complete model. This is a common - pattern in CTLearn models where the backbone is wrapped as a single layer. - """ - # Extract the backbone model (second layer of the complete model) - # Layer 0 is the input, layer 1 is the backbone feature extractor - backbone = model.get_layer(index=1) - - # Create a new head model using layers after the backbone - # Define input with the same shape as backbone output - backbone_output_shape = keras.Input(model.layers[2].input_shape[1:]) - x = backbone_output_shape - - # Reconstruct head by connecting all layers after backbone - for layer in model.layers[2:]: - x = layer(x) - - # Create the head model - head = keras.Model(inputs=backbone_output_shape, outputs=x) - - return backbone, head - - -def load_keras_model(self): - """ - Load Keras models from saved files and split them into backbone and head. - - This function loads trained Keras models for different tasks (particle type - classification, energy estimation, direction reconstruction) and splits each - into backbone and head components for efficient prediction and feature extraction. - - Parameters - ---------- - self : object - Prediction handler instance containing model paths: - - load_type_model_from: Path to particle classification model - - load_energy_model_from: Path to energy estimation model - - load_cameradirection_model_from: Path to direction reconstruction model - - Returns - ------- - input_shape : tuple - Shape of the model input (height, width, channels). - Returns the shape from the last loaded model. - - Notes - ----- - The function sets the following attributes on self: - - backbone_type, head_type: Split models for particle classification - - backbone_energy, head_energy: Split models for energy estimation - - backbone_direction, head_direction: Split models for direction reconstruction - """ - input_shape = None - - # Load particle type classification model if configured - if self.load_type_model_from is not None: - self.log.info("Loading the type model from %s.", self.load_type_model_from) - model_type = keras.saving.load_model(self.load_type_model_from) - input_shape = model_type.input_shape[1:] - # Split model into backbone and head - self.backbone_type, self.head_type = _split_model(model_type) - - # Load energy estimation model if configured - if self.load_energy_model_from is not None: - self.log.info( - "Loading the energy model from %s.", self.load_energy_model_from - ) - model_energy = keras.saving.load_model(self.load_energy_model_from) - input_shape = model_energy.input_shape[1:] - # Split model into backbone and head - self.backbone_energy, self.head_energy = _split_model(model_energy) - - # Load direction reconstruction model if configured - if self.load_cameradirection_model_from is not None: - self.log.info( - "Loading the cameradirection model from %s.", self.load_cameradirection_model_from - ) - model_direction = keras.saving.load_model( - self.load_cameradirection_model_from - ) - input_shape = model_direction.input_shape[1:] - # Split model into backbone and head - self.backbone_direction, self.head_direction = _split_model(model_direction) - - return input_shape diff --git a/ctlearn/tools/predict/keras/predic_model_keras.py b/ctlearn/tools/predict/keras/predic_model_keras.py deleted file mode 100644 index 0d2b3472..00000000 --- a/ctlearn/tools/predict/keras/predic_model_keras.py +++ /dev/null @@ -1,156 +0,0 @@ -""" -Keras model prediction module for CTLearn. -This module provides functionality to load trained Keras models and perform batch predictions -on DL1 data, with optional feature vector extraction from backbone models. -""" - -from ctlearn.core.data_loader.loader import DLDataLoader -import keras -from astropy.table import Table, vstack -import numpy as np - - -def predict_with_model(self, model_path, task): - """ - Load and predict with a CTLearn Keras model. - - This function loads a trained model from the specified path and performs predictions - on the provided data. It handles both complete and incomplete batches, and optionally - extracts feature vectors from the backbone model for downstream analysis. - - Parameters - ---------- - model_path : str - Path to a Keras model file (Keras 3) or directory (Keras 2). - The model should be a complete trained CTLearn model. - task : str - The task for which prediction is being made (e.g., 'type', 'energy', 'cameradirection'). - - Returns - ------- - predict_data : astropy.table.Table - Table containing the prediction results with columns corresponding to - the model's output. - feature_vectors : np.ndarray or None - Feature vectors extracted from the backbone model if dl1_features is enabled. - Returns None if feature extraction is not requested. - """ - # Create data loader for the main batch processing - # The DLDataLoader is initialized separately for each prediction task - # Batch size is multiplied by number of replicas for distributed inference - data_loader = DLDataLoader.create( - framework="keras", - DLDataReader=self.dl1dh_reader, - indices=self.indices, - tasks=[], - batch_size=self.batch_size * self.strategy.num_replicas_in_sync, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - - # Handle incomplete last batch - # Keras only processes complete batches during prediction, so we need - # a separate data loader for remaining events that don't fill a complete batch - data_loader_last_batch = None - if self.last_batch_size > 0: - # Extract indices for the last incomplete batch - last_batch_indices = self.indices[-self.last_batch_size:] - data_loader_last_batch = DLDataLoader.create( - framework="keras", - DLDataReader=self.dl1dh_reader, - indices=last_batch_indices, - tasks=[], - batch_size=self.last_batch_size, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - - # Load the trained model from the specified path - model = keras.saving.load_model(model_path) - - # Initialize variables for optional feature extraction - backbone_model, feature_vectors = None, None - - if self.dl1_features: - # Feature extraction mode: split model into backbone and head - # This allows us to extract intermediate representations (feature vectors) - - # Extract the backbone model (second layer of the complete model) - # Layer 0: Input, Layer 1: Backbone (feature extractor), Layers 2+: Head - backbone_model = model.get_layer(index=1) - - # Reconstruct the head model from layers after the backbone - # Define input with the same shape as backbone output - backbone_output_shape = keras.Input(model.layers[2].input_shape[1:]) - x = backbone_output_shape - - # Connect all layers after backbone to create head model - for layer in model.layers[2:]: - x = layer(x) - head = keras.Model(inputs=backbone_output_shape, outputs=x) - - # Extract feature vectors from backbone - feature_vectors = backbone_model.predict( - data_loader, verbose=self.keras_verbose - ) - - # Generate predictions from head using extracted features - predict_data = Table( - { - task: head.predict( - feature_vectors, verbose=self.keras_verbose - ) - } - ) - - # Process last incomplete batch if it exists - if data_loader_last_batch is not None: - # Extract features from last batch - feature_vectors_last_batch = backbone_model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - # Concatenate feature vectors from all batches - feature_vectors = np.concatenate( - (feature_vectors, feature_vectors_last_batch) - ) - # Generate predictions for last batch and stack with main predictions - predict_data = vstack( - [ - predict_data, - Table( - { - task: head.predict( - feature_vectors_last_batch, - verbose=self.keras_verbose, - ) - } - ), - ] - ) - else: - # Standard prediction mode without feature extraction - # Use the complete model for end-to-end prediction - predict_data = model.predict(data_loader, verbose=self.keras_verbose) - - # Convert predictions to Astropy Table - if isinstance(predict_data, dict): - predict_data = Table(predict_data) - else: - predict_data = Table({task: predict_data}) - - # Process last incomplete batch if it exists - if data_loader_last_batch is not None: - # Generate predictions for last batch - predict_data_last_batch = model.predict( - data_loader_last_batch, verbose=self.keras_verbose - ) - - if isinstance(predict_data_last_batch, dict): - predict_data_last_batch = Table(predict_data_last_batch) - else: - predict_data_last_batch = Table({task: predict_data_last_batch}) - - # Stack predictions from main batches and last batch - predict_data = vstack([predict_data, predict_data_last_batch]) - - return predict_data, feature_vectors \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_LST1.py b/ctlearn/tools/predict/predict_LST1.py deleted file mode 100644 index a9ee4029..00000000 --- a/ctlearn/tools/predict/predict_LST1.py +++ /dev/null @@ -1,834 +0,0 @@ -""" -Predict the gammaness, energy and arrival direction from lstchain DL1 data. -""" - -import numpy as np -import tables -import keras -from astropy import units as u -from astropy.coordinates import AltAz,SkyCoord -from astropy.table import Table, setdiff, vstack -from astropy.coordinates import AltAz,SkyCoord -from astropy.table import Table, setdiff, vstack -from astropy.time import Time - -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) -from ctapipe.coordinates import CameraFrame -from ctapipe.core import Tool -from ctapipe.core.tool import ToolConfigurationError -from ctapipe.core.traits import ( - Bool, - Int, - Path, - List, - CaselessStrEnum, - ComponentName, - Unicode, - UseEnum, - classes_with_traits, -) -from ctapipe.instrument.optics import FocalLengthKind -from ctapipe.io import read_table, write_table -from ctapipe.reco.utils import add_defaults_and_meta - -from ctlearn.tools.utils import get_lst1_subarray_description -from dl1_data_handler.image_mapper import ImageMapper -from dl1_data_handler.reader import TableQualityQuery -from ctlearn.tools.predict.utils.load_model import load_model -from ctlearn.core.ctlearn_enum import Task, Mode -from ctlearn.tools.train.pytorch.utils import ( - sanity_check, - read_configuration, - expected_structure, -) - -POINTING_GROUP = "/dl1/monitoring/telescope/pointing" -DL1_TELESCOPE_GROUP = "/dl1/event/telescope" -DL2_TELESCOPE_GROUP = "/dl2/event/telescope" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] - - -class LST1PredictionTool(Tool): - """ - Tool to predict the gammaness, energy and arrival direction from lstchain DL1 data. - - This tool is used to predict the gammaness, energy and arrival direction - from pixel-wise image data in lstchain format. The tool loads the trained models - from the specified paths and performs inference on the input data. The - input data is expected to be in the DL1 format of lstchain and the output data is - stored in the DL2 format of ctapipe. Besides the DL2 predictions, the tool creates - the SubarrayDescription of the LST-1 telescope and stores it in the output file. - In addition, the tool also creates the trigger, pointing and DL1 parameters tables - and stores them in the output file. - - CAUTION: The tool is designed to work with the DL1 data format of lstchain only. - """ - - name = "LST1PredictionTool" - description = __doc__ - - examples = """ - To predict from DL1 lstchain data using trained CTLearn models: - > ctlearn-predict-model \\ - --input_url input.subrun.lstchain.dl1.h5 \\ - --LST1PredictionTool.batch_size=64 \\ - --LST1PredictionTool.channels=cleaned_image \\ - --LST1PredictionTool.channels=cleaned_relative_peak_time \\ - --LST1PredictionTool.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/direction/ctlearn_model.cpk" \\ - --output output.dl2.h5 \\ - --overwrite \\ - """ - - input_url = Path( - help="Input LST-1 HDF5 files including pixel-wise image data", - allow_none=True, - exists=True, - directory_ok=False, - file_ok=True, - ).tag(config=True) - - prefix = Unicode( - default_value="CTLearn", - allow_none=False, - help="Name of the reconstruction algorithm used to generate the dl2 data.", - ).tag(config=True) - - load_type_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) " - "for the classification of the primary particle type." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_energy_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) " - "for the regression of the primary particle energy." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - load_cameradirection_model_from = Path( - default_value=None, - help=( - "Path to a Keras model file (Keras3) or directory (Keras2) " - "for the regression of the primary particle arrival direction " - "based on the camera coordinate offsets." - ), - allow_none=True, - exists=True, - directory_ok=True, - file_ok=True, - ).tag(config=True) - - batch_size = Int( - default_value=128, - allow_none=False, - help="Size of the batch to perform inference of the neural network.", - ).tag(config=True) - - channels = List( - trait=CaselessStrEnum( - [ - "image", - "cleaned_image", - "peak_time", - "relative_peak_time", - "cleaned_peak_time", - "cleaned_relative_peak_time", - ] - ), - default_value=["cleaned_image", "cleaned_peak_time"], - allow_none=False, - help=( - "Set the input channels to be loaded from the DL1 event data. " - "image: integrated charges, " - "cleaned_image: integrated charges cleaned with the DL1 cleaning mask, " - "peak_time: extracted peak arrival times, " - "relative_peak_time: extracted relative peak arrival times, " - "cleaned_peak_time: extracted peak arrival times cleaned with the DL1 cleaning mask, " - "cleaned_relative_peak_time: extracted relative peak arrival times cleaned with the DL1 cleaning mask." - ), - ).tag(config=True) - - image_mapper_type = ComponentName(ImageMapper, default_value="BilinearMapper").tag( - config=True - ) - - focal_length_choice = UseEnum( - FocalLengthKind, - default_value=FocalLengthKind.EFFECTIVE, - help=( - "If both nominal and effective focal lengths are available, " - " which one to use for the `~ctapipe.coordinates.CameraFrame` attached" - " to the `~ctapipe.instrument.CameraGeometry` instances in the" - " `~ctapipe.instrument.SubarrayDescription` which will be used in" - " CameraFrame to TelescopeFrame coordinate transforms." - " The 'nominal' focal length is the one used during " - " the simulation, the 'effective' focal length is computed using specialized " - " ray-tracing from a point light source" - ), - ).tag(config=True) - - override_obs_id = Int( - default_value=None, - allow_none=True, - help=( - "Use the given obs_id instead of the default one. " - "Needed to merge subruns later with ctapipe-merge." - ), - ).tag(config=True) - - output_path = Path( - default_value="./output.dl2.h5", - allow_none=False, - help="Output path to save the dl2 prediction results", - ).tag(config=True) - - pytorch_config_file = Path( - default_value="./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml", - help="Pytorch config file", - ).tag(config=True) - - framework_type = CaselessStrEnum( - ["pytorch", "keras"], - default_value="keras", - help="Framework to use: pytorch or keras", - ).tag(config=True) - - overwrite = Bool(help="Overwrite output file if it exists").tag(config=True) - - aliases = { - ("i", "input_url"): "LST1PredictionTool.input_url", - ("t", "type_model"): "LST1PredictionTool.load_type_model_from", - ("e", "energy_model"): "LST1PredictionTool.load_energy_model_from", - ("d", "cameradirection_model"): "LST1PredictionTool.load_cameradirection_model_from", - ("o", "output"): "LST1PredictionTool.output_path", - ("f", "framework"): "LST1PredictionTool.framework_type", - ("p", "pytorch_config_file"): "LST1PredictionTool.pytorch_config_file", - - } - - flags = { - "overwrite": ( - {"LST1PredictionTool": {"overwrite": True}}, - "Overwrite existing files", - ), - } - - classes = classes_with_traits(ImageMapper) - - def _predictions(self): - if self.framework_type == "keras": - self.log.info("Using the Keras Model") - from ctlearn.tools.predict.keras.predic_LST1_keras import predictions - return predictions(self) - - elif self.framework_type == "pytorch": - self.log.info("Using the Pytorch Model") - from ctlearn.tools.predict.pytorch.predic_LST1_pytorch import predictions - return predictions(self) - - def setup(self): - # Save dl1 image and parameters tree schemas and tel id for easy access - import torch - self.image_table_path = "/dl1/event/telescope/image/LST_LSTCam" - self.parameter_table_name = "/dl1/event/telescope/parameters/LST_LSTCam" - self.tel_id = 1 - if self.framework_type == "pytorch": - self.log.info(f"Using {self.pytorch_config_file} config file for pytorch framework") - self.parameters = read_configuration(self.pytorch_config_file) - sanity_check(self.parameters, expected_structure) - self.batch_size = self.parameters["hyp"]["batches"] - self.device_str = self.parameters["arch"]["device"] - self.optim_batch_size = self.parameters["hyp"]["dynamic_batches"] - self.device = torch.device(self.device_str) - self.tasks = [] - self.type_mu = self.parameters["normalization"]["type_mu"] - self.type_sigma = self.parameters["normalization"]["type_sigma"] - self.dir_mu = self.parameters["normalization"]["dir_mu"] - self.dir_sigma = self.parameters["normalization"]["dir_sigma"] - self.energy_mu = self.parameters["normalization"]["energy_mu"] - self.energy_sigma = self.parameters["normalization"]["energy_sigma"] - - if self.load_type_model_from is not None: - self.tasks.append(Task.type) - if self.load_energy_model_from is not None: - self.tasks.append(Task.energy) - if self.load_cameradirection_model_from is not None: - self.tasks.append(Task.direction) - - # Get the number of rows in the table - with tables.open_file(self.input_url) as input_file: - self.table_length = len(input_file.get_node(self.image_table_path)) - - # Load the models from the specified paths - input_shape = load_model(self) - - # Get the SubarrayDescription of the LST-1 telescope - self.subarray = get_lst1_subarray_description(focal_length_choice=self.focal_length_choice) - # Write the SubarrayDescription to the output file - self.subarray.to_hdf(self.output_path, overwrite=self.overwrite) - self.log.info("SubarrayDescription was stored in '%s'", self.output_path) - # Initialize the Table data quality query - self.quality_query = TableQualityQuery(parent=self) - # Copy the pixel rotation of the camera geometry of the subarray in a variable - # since the ImageMapper will be derotated the pixels. The pixel rotation - # is needed to create a rotated camera frame in order to transform the - # predicted camera coordinate offsets back to the correct Alt/Az coordinates. - self.pix_rotation = self.subarray.tel[self.tel_id].camera.geometry.pix_rotation - # Create the ImageMapper - self.image_mapper = ImageMapper.from_name( - name=self.image_mapper_type, - geometry=self.subarray.tel[self.tel_id].camera.geometry, - subarray=self.subarray, - parent=self, - ) - # Check if the input shape of the model matches the image shape of the ImageMapper - if self.framework_type == "keras": - if input_shape[0] != self.image_mapper.image_shape: - raise ToolConfigurationError( - f"The input shape of the model ('{input_shape[0]}') does not match " - f"the image shape of the ImageMapper ('{self.image_mapper.image_shape}'). " - f"Use e.g. '--BilinearMapper.interpolation_image_shape={input_shape[0]}' ." - ) - - # Get offset and scaling of images - self.transforms = {} - self.transforms["image_scale"] = 0.0 - self.transforms["image_offset"] = 0 - self.transforms["peak_time_scale"] = 0.0 - self.transforms["peak_time_offset"] = 0 - - # Get the number of rows in the table - with tables.open_file(self.input_url) as input_file: - img_table_v_attrs = input_file.get_node(self.image_table_path)._v_attrs - - # Check the transform value used for the file compression - if "CTAFIELD_3_TRANSFORM_SCALE" in img_table_v_attrs: - self.transforms["image_scale"] = img_table_v_attrs[ - "CTAFIELD_3_TRANSFORM_SCALE" - ] - self.transforms["image_offset"] = img_table_v_attrs[ - "CTAFIELD_3_TRANSFORM_OFFSET" - ] - if "CTAFIELD_4_TRANSFORM_SCALE" in img_table_v_attrs: - self.transforms["peak_time_scale"] = img_table_v_attrs[ - "CTAFIELD_4_TRANSFORM_SCALE" - ] - self.transforms["peak_time_offset"] = img_table_v_attrs[ - "CTAFIELD_4_TRANSFORM_OFFSET" - ] - - def start(self): - all_identifiers = read_table(self.input_url, self.parameter_table_name) - all_identifiers.meta = {} - if self.override_obs_id is not None: - all_identifiers["obs_id"] = self.override_obs_id - self.obs_id = all_identifiers["obs_id"][0] - self.parameter_table = all_identifiers.copy() - tel_az = u.Quantity(self.parameter_table["az_tel"], unit=u.rad) - tel_alt = u.Quantity(self.parameter_table["alt_tel"], unit=u.rad) - event_type = self.parameter_table["event_type"] - time = Time(self.parameter_table["dragon_time"] * u.s, format="unix") - # Create the pointing table - # This table is used to store the telescope pointing per event - pointing_table = Table( - { - "time": time, - "azimuth": tel_az, - "altitude": tel_alt, - } - ) - write_table( - pointing_table, - self.output_path, - f"{POINTING_GROUP}/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{POINTING_GROUP}/tel_{self.tel_id:03d}", - ) - # Set the time format to MJD since in the other table we store the time in MJD - time.format = "mjd" - # Keep only the necessary columns for the creation of tables - all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) - - # Create the dl1 telescope trigger table - self.trigger_table = all_identifiers.copy() - self.trigger_table.add_column(time, name="time") - self.trigger_table.add_column(-1, name="n_trigger_pixels") - - write_table( - self.trigger_table, - self.output_path, - "/dl1/event/telescope/trigger", - overwrite=self.overwrite, - ) - self.log.info( - "DL1 telescope trigger table was stored in '%s' under '%s'", - self.output_path, - "/dl1/event/telescope/trigger", - ) - self.trigger_table.keep_columns(["obs_id", "event_id", "time"]) - self.trigger_table.add_column( - np.ones((len(self.trigger_table), 1), dtype=bool), name="tel_with_trigger" - ) - self.trigger_table.add_column(event_type, name="event_type") - - # Save the dl1 subrray trigger table to the output file - # write_table( - # self.trigger_table, - # self.output_path, - # "/dl1/event/subarray/trigger", - # overwrite=self.overwrite, - # ) - # self.log.info( - # "DL1 subarray trigger table was stored in '%s' under '%s'", - # self.output_path, - # "/dl1/event/subarray/trigger", - # ) - # Create the dl1 parameters table - self.parameter_table.rename_column("intensity", "hillas_intensity") - self.parameter_table.rename_column("x", "hillas_x") - self.parameter_table.rename_column("y", "hillas_y") - self.parameter_table.rename_column("phi", "hillas_phi") - self.parameter_table.rename_column("psi", "hillas_psi") - self.parameter_table.rename_column("length", "hillas_length") - self.parameter_table.rename_column("length_uncertainty", "hillas_length_uncertainty") - self.parameter_table.rename_column("width", "hillas_width") - self.parameter_table.rename_column("width_uncertainty", "hillas_width_uncertainty") - self.parameter_table.rename_column("skewness", "hillas_skewness") - self.parameter_table.rename_column("kurtosis", "hillas_kurtosis") - self.parameter_table.rename_column("time_gradient", "timing_deviation") - self.parameter_table.rename_column("intercept", "timing_intercept") - self.parameter_table.rename_column("n_pixels", "morphology_n_pixels") - self.parameter_table.rename_column("n_islands", "morphology_n_islands") - self.parameter_table.keep_columns( - [ - "obs_id", - "event_id", - "hillas_intensity", - "hillas_x", - "hillas_y", - "hillas_phi", - "hillas_psi", - "hillas_length", - "hillas_length_uncertainty", - "hillas_width", - "hillas_width_uncertainty", - "hillas_skewness", - "hillas_kurtosis", - "timing_deviation", - "timing_intercept", - "morphology_n_pixels", - "morphology_n_islands", - ] - ) - self.parameter_table.add_column(self.tel_id, name="tel_id", index=2) - # Save the dl1 parameters table to the output file - write_table( - self.parameter_table, - self.output_path, - f"/dl1/event/telescope/parameters/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL1 parameters table was stored in '%s' under '%s'", - self.output_path, - f"/dl1/event/telescope/parameters/tel_{self.tel_id:03d}", - ) - - # Add additional columns to the parameter table - # which are not present in the originl DL1 parameter table. - # They are needed for applying the quality selection. - self.parameter_table.add_column(event_type, name="event_type") - self.parameter_table.add_column(tel_az, name="tel_az") - self.parameter_table.add_column(tel_alt, name="tel_alt") - # Only select cosmic events for the prediction - self.parameter_table = self.parameter_table[self.parameter_table["event_type"]==32] - - self.log.info("Starting the prediction...") - # Iterate over the data in chunks based on the batch size - event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs = self._predictions() - - # Create the prediction tables - example_identifiers = Table( - { - "obs_id": np.full(len(event_id), self.obs_id, dtype=int), - "event_id": event_id, - "tel_id": np.full(len(event_id), self.tel_id, dtype=int), - } - ) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS - ) - if len(nonexample_identifiers) > 0: - nonexample_identifiers.sort(TELESCOPE_EVENT_KEYS) - # Create the feature vector table - feature_vector_table = example_identifiers.copy() - fvs_columns_list, fvs_shapes_list = [], [] - if self.load_type_model_from is not None: - classification_table = example_identifiers.copy() - classification_table.add_column( - prediction, name=f"{self.prefix}_tel_prediction" - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - ) - classification_table = vstack([classification_table, nan_table]) - classification_table.sort(TELESCOPE_EVENT_KEYS) - classification_is_valid = ~np.isnan(classification_table[f"{self.prefix}_tel_prediction"].data, dtype=bool) - classification_table.add_column( - classification_is_valid, - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - classification_table, - ParticleClassificationContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - # Save the prediction to the output file - write_table( - classification_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{self.tel_id:03d}", - ) - # Write the mono telescope prediction to the subarray prediction table - subarray_classification_table = classification_table.copy() - subarray_classification_table.remove_column("tel_id") - for colname in subarray_classification_table.colnames: - if "_tel_" in colname: - subarray_classification_table.rename_column( - colname, colname.replace("_tel", "") - ) - subarray_classification_table.add_column( - classification_is_valid, name=f"{self.prefix}_telescopes" - ) - # Save the prediction to the output file - write_table( - subarray_classification_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - ) - # Adding the feature vectors for the classification - is_valid_col = ~np.isnan( - np.min(classification_fvs, axis=1), dtype=bool - ) - feature_vector_table.add_column( - classification_fvs, - name=f"{self.prefix}_tel_classification_feature_vectors", - ) - if nonexample_identifiers is not None: - fvs_columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") - fvs_shapes_list.append( - ( - len(nonexample_identifiers), - classification_fvs[0].shape[0], - ) - ) - if self.load_energy_model_from is not None: - energy_table = example_identifiers.copy() - # Convert the reconstructed energy from log10(TeV) to TeV - reco_energy = u.Quantity(np.power(10, np.squeeze(energy)), unit=u.TeV) - # Add the reconstructed energy to the prediction table - energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - ) - energy_table = vstack([energy_table, nan_table]) - energy_table.sort(TELESCOPE_EVENT_KEYS) - energy_is_valid = ~np.isnan(energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool) - energy_table.add_column( - energy_is_valid, - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - # Save the prediction to the output file - write_table( - energy_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{self.tel_id:03d}", - ) - # Write the mono telescope prediction to the subarray prediction table - subarray_energy_table = energy_table.copy() - subarray_energy_table.remove_column("tel_id") - for colname in subarray_energy_table.colnames: - if "_tel_" in colname: - subarray_energy_table.rename_column( - colname, colname.replace("_tel", "") - ) - subarray_energy_table.add_column( - energy_is_valid, name=f"{self.prefix}_telescopes" - ) - # Save the prediction to the output file - write_table( - subarray_energy_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - ) - # Adding the feature vectors for the energy regression - is_valid_col = ~np.isnan( - np.min(energy_fvs, axis=1), dtype=bool - ) - feature_vector_table.add_column( - energy_fvs, - name=f"{self.prefix}_tel_energy_feature_vectors", - ) - if nonexample_identifiers is not None: - fvs_columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") - fvs_shapes_list.append( - ( - len(nonexample_identifiers), - energy_fvs[0].shape[0], - ) - ) - if self.load_cameradirection_model_from is not None: - direction_table = example_identifiers.copy() - # Set the telescope position - tel_ground_frame = self.subarray.tel_coords[ - self.subarray.tel_ids_to_indices(self.tel_id) - ] - # Set the telescope pointing with the trigger timestamp and the telescope position - trigger_time = Time(trigger_time, format="mjd") - altaz = AltAz( - location=tel_ground_frame.to_earth_location(), - obstime=trigger_time, - ) - # Set the telescope pointing - tel_pointing = SkyCoord( - az=u.Quantity(tel_azimuth, unit=u.rad), - alt=u.Quantity(tel_altitude, unit=u.rad), - frame=altaz, - ) - # Set a new camera frame with the pixel rotation of the camera - camera_frame = CameraFrame( - focal_length=self.subarray.tel[self.tel_id].camera.geometry.frame.focal_length, - rotation=self.pix_rotation, - telescope_pointing=tel_pointing, - ) - ## Remove (save the cam_coord_offset_x, cam_coord_offset_y) predicted by the model in a pickel file - - # with open('/lhome/ext/ucm147/ucm1477/ctlearn/test_local_cristian/cam_coord_offset_test.pkl', 'wb') as f: - # import pickle - # pickle.dump((cam_coord_offset_x, cam_coord_offset_y), f) - # print('Pickell camera_coordinates saved') - # Set the camera coordinate offset - cam_coord_offset = SkyCoord( - x=u.Quantity(cam_coord_offset_x, unit=u.m), - y=u.Quantity(cam_coord_offset_y, unit=u.m), - frame=camera_frame - ) - # Transform the true Alt/Az coordinates to camera coordinates - reco_direction = cam_coord_offset.transform_to(altaz) - # Add the reconstructed direction (az, alt) to the prediction table - direction_table.add_column( - reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az" - ) - direction_table.add_column( - reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt" - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_az", f"{self.prefix}_tel_alt"], - shapes=[(len(nonexample_identifiers),), (len(nonexample_identifiers),)], - ) - direction_table = vstack([direction_table, nan_table]) - direction_table.keep_columns( - TELESCOPE_EVENT_KEYS - + [f"{self.prefix}_tel_az", f"{self.prefix}_tel_alt"] - ) - direction_table.sort(TELESCOPE_EVENT_KEYS) - direction_is_valid = ~np.isnan(direction_table[f"{self.prefix}_tel_az"].data, dtype=bool) - direction_table.add_column( - direction_is_valid, - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_table, - ReconstructedGeometryContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - # Save the prediction to the output file - write_table( - direction_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{self.tel_id:03d}", - ) - # Write the mono telescope prediction to the subarray prediction table - subarray_direction_table = direction_table.copy() - subarray_direction_table.remove_column("tel_id") - for colname in subarray_direction_table.colnames: - if "_tel_" in colname: - subarray_direction_table.rename_column( - colname, colname.replace("_tel", "") - ) - subarray_direction_table.add_column( - direction_is_valid, name=f"{self.prefix}_telescopes" - ) - # Save the prediction to the output file - write_table( - subarray_direction_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - overwrite=self.overwrite, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - ) - # Adding the feature vectors for the arrival direction regression - is_valid_col = ~np.isnan( - np.min(direction_fvs, axis=1), dtype=bool - ) - feature_vector_table.add_column( - direction_fvs, - name=f"{self.prefix}_tel_direction_feature_vectors", - ) - if nonexample_identifiers is not None: - fvs_columns_list.append(f"{self.prefix}_tel_direction_feature_vectors") - fvs_shapes_list.append( - ( - len(nonexample_identifiers), - direction_fvs[0].shape[0], - ) - ) - # Produce output table with NaNs for missing predictions - if nonexample_identifiers is not None: - if len(nonexample_identifiers) > 0: - nan_table = self._create_nan_table( - nonexample_identifiers, - columns=fvs_columns_list, - shapes=fvs_shapes_list, - ) - feature_vector_table = vstack([feature_vector_table, nan_table]) - is_valid_col = np.concatenate( - (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) - ) - # Add is_valid column to the feature vector table - feature_vector_table.add_column( - is_valid_col, - name=f"{self.prefix}_tel_is_valid", - ) - # Save the prediction to the output file - write_table( - feature_vector_table, - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{self.tel_id:03d}", - overwrite=self.overwrite, - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{self.tel_id:03d}", - ) - - def finish(self): - self.log.info("Tool is shutting down") - - def _create_nan_table(self, nonexample_identifiers, columns, shapes): - """ - Create a table with NaNs for missing predictions. - - This method creates a table with NaNs for missing predictions for the non-example identifiers. - - Parameters: - ----------- - nonexample_identifiers : astropy.table.Table - Table containing the non-example identifiers. - columns : list of str - List of column names to create in the table. - shapes : list of shapes - List of shapes for the columns to create in the table. - - Returns: - -------- - nan_table : astropy.table.Table - Table containing NaNs for missing predictions. - """ - # Create a table with NaNs for missing predictions - nan_table = nonexample_identifiers.copy() - for column_name, shape in zip(columns, shapes): - nan_table.add_column(np.full(shape, np.nan), name=column_name) - return nan_table - - -def main(): - # Run the tool - tool = LST1PredictionTool() - tool.run() - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_mono.py b/ctlearn/tools/predict/predict_mono.py deleted file mode 100644 index 86782cbd..00000000 --- a/ctlearn/tools/predict/predict_mono.py +++ /dev/null @@ -1,541 +0,0 @@ -""" -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. -""" - - -import numpy as np -from astropy.table import ( - Table, - vstack, - join, - setdiff, -) -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) - -from ctapipe.core.traits import ComponentName -from ctapipe.io import write_table -from ctapipe.reco.reconstructor import ReconstructionProperty -from ctapipe.reco.stereo_combination import StereoCombiner -from ctapipe.reco.utils import add_defaults_and_meta -from dl1_data_handler.reader import ProcessType - - -SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -POINTING_GROUP = "/dl1/monitoring/telescope/pointing" -SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" -DL1_TELESCOPE_GROUP = "/dl1/event/telescope" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_TELESCOPE_GROUP = "/dl2/event/telescope" -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] - -__all__ = ["MonoPredictCTLearnModel"] - -from ctlearn.tools.predict.utils.predict_model import PredictCTLearnModel - -class MonoPredictCTLearnModel(PredictCTLearnModel): - """ - Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. - - This tool extends the ``PredictCTLearnModel`` to specifically handle monoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_telescope_pointing(all_identifiers) - Store the telescope pointing table for the mono mode for MC simulation. - """ - - name = "ctlearn-predict-mono-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnModel.batch_size=64 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ - --dl1-features \\ - --use-HDF5Merger \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - - To predict from pixel-wise waveform data in mono mode using trained CTLearn models: - > ctlearn-predict-mono-model \\ - --input_url input.r1.h5 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader \\ - --DLWaveformReader.sequnce_length=20 \\ - --DLWaveformReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ - --use-HDF5Merger \\ - --no-r0-waveforms \\ - --no-r1-waveforms \\ - --no-dl1-images \\ - --no-true-images \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - """ - - stereo_combiner_cls = ComponentName( - StereoCombiner, - default_value="StereoMeanCombiner", - help="Which stereo combination method to use after the monoscopic reconstruction.", - ).tag(config=True) - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.example_identifiers.copy() - example_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) - all_identifiers = self.dl1dh_reader.tel_trigger_table.copy() - all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS + ["time"]) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Pointing table for the mono mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_telescope_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - classification_feature_vectors = None - if self.load_type_model_from is not None: - self.type_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.PARTICLE_TYPE, - parent=self, - ) - # Predict the energy of the primary particle - classification_table, classification_feature_vectors = ( - super()._predict_classification(example_identifiers) - ) - if self.dl2_telescope: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - ) - classification_table = vstack([classification_table, nan_table]) - # Add is_valid column to the energy table - classification_table.add_column( - ~np.isnan( - classification_table[f"{self.prefix}_tel_prediction"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - classification_table, - ParticleClassificationContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = classification_table["tel_id"] == tel_id - classification_tel_table = classification_table[telescope_mask] - classification_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file for the selected telescope - write_table( - classification_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info("Processing and storing the subarray type prediction...") - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_classification_table = self.type_stereo_combiner.predict_table( - classification_table - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - subarray_classification_table[f"{self.prefix}_telescopes"].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - ( - len(subarray_classification_table), - len(self.dl1dh_reader.tel_ids), - ), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_classification_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_classification_table[f"{self.prefix}_telescopes"] = ( - reco_telescopes - ) - # Save the prediction to the output file - write_table( - subarray_classification_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - self.energy_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.ENERGY, - parent=self, - ) - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - if self.dl2_telescope: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = energy_table["tel_id"] == tel_id - energy_tel_table = energy_table[telescope_mask] - energy_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file - write_table( - energy_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray energy prediction..." - ) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_energy_table = self.energy_stereo_combiner.predict_table( - energy_table - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if subarray_energy_table[f"{self.prefix}_telescopes"].dtype != np.bool_: - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - (len(subarray_energy_table), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_energy_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_energy_table[f"{self.prefix}_telescopes"] = reco_telescopes - # Save the prediction to the output file - write_table( - subarray_energy_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - ) - direction_feature_vectors = None - if self.load_cameradirection_model_from is not None: - self.geometry_stereo_combiner = StereoCombiner.from_name( - self.stereo_combiner_cls, - prefix=self.prefix, - property=ReconstructionProperty.GEOMETRY, - parent=self, - ) - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=TELESCOPE_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = ( - super()._predict_cameradirection(example_identifiers) - ) - direction_tel_tables = [] - if self.dl2_telescope: - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = direction_table["tel_id"] == tel_id - direction_tel_table = direction_table[telescope_mask] - direction_tel_table = super()._transform_cam_coord_offsets_to_sky( - direction_tel_table - ) - # Produce output table with NaNs for missing predictions - nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id - nonexample_identifiers_tel = nonexample_identifiers[ - nan_telescope_mask - ] - if len(nonexample_identifiers_tel) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers_tel, - columns=[f"{self.prefix}_tel_alt", f"{self.prefix}_tel_az"], - shapes=[ - (len(nonexample_identifiers_tel),), - (len(nonexample_identifiers_tel),), - ], - ) - direction_tel_table = vstack([direction_tel_table, nan_table]) - direction_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Add is_valid column to the direction table - direction_tel_table.add_column( - ~np.isnan( - direction_tel_table[f"{self.prefix}_tel_alt"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_tel_table, - ReconstructedGeometryContainer, - prefix=self.prefix, - add_tel_prefix=True, - ) - direction_tel_tables.append(direction_tel_table) - # Save the prediction to the output file - write_table( - direction_tel_table, - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", - ) - if self.dl2_subarray: - self.log.info( - "Processing and storing the subarray geometry prediction..." - ) - # Stack the telescope tables to the subarray table - direction_tel_tables = vstack(direction_tel_tables) - # Sort the table by the telescope event keys - direction_tel_tables.sort(TELESCOPE_EVENT_KEYS) - # Combine the telescope predictions to the subarray prediction using the stereo combiner - subarray_direction_table = self.geometry_stereo_combiner.predict_table( - direction_tel_tables - ) - # TODO: Remove temporary fix once the stereo combiner returns correct table - # Check if the table has to be converted to a boolean mask - if ( - subarray_direction_table[f"{self.prefix}_telescopes"].dtype - != np.bool_ - ): - # Create boolean mask for telescopes that participate in the stereo reconstruction combination - reco_telescopes = np.zeros( - (len(subarray_direction_table), len(self.dl1dh_reader.tel_ids)), - dtype=bool, - ) - # Loop over the table and set the boolean mask for the telescopes - for index, tel_id_mask in enumerate( - subarray_direction_table[f"{self.prefix}_telescopes"] - ): - if not tel_id_mask: - continue - for tel_id in tel_id_mask: - reco_telescopes[index][ - self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) - ] = True - # Overwrite the column with the boolean mask with fix length - subarray_direction_table[f"{self.prefix}_telescopes"] = ( - reco_telescopes - ) - # Save the prediction to the output file - write_table( - subarray_direction_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - ) - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - classification_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. - for tel_id in self.dl1dh_reader.selected_telescopes[ - self.dl1dh_reader.tel_type - ]: - # Retrieve the example identifiers for the selected telescope - telescope_mask = feature_vector_table["tel_id"] == tel_id - feature_vectors_tel_table = feature_vector_table[telescope_mask] - feature_vectors_tel_table.sort(TELESCOPE_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vectors_tel_table, - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", - ) - - def _store_mc_telescope_pointing(self, all_identifiers): - """ - Store the telescope pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the telescope pointing information. - """ - # Create the pointing table for each telescope - pointing_info = [] - for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: - # Pointing table for the mono mode - tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) - tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") - tel_pointing.rename_column( - "telescope_pointing_altitude", "pointing_altitude" - ) - # Join the prediction table with the telescope pointing table - tel_pointing = join( - left=tel_pointing, - right=all_identifiers, - keys=["obs_id", "tel_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - tel_pointing.sort(TELESCOPE_EVENT_KEYS) - # Retrieve the example identifiers for the selected telescope - tel_pointing_table = Table( - { - "time": tel_pointing["time"], - "azimuth": tel_pointing["pointing_azimuth"], - "altitude": tel_pointing["pointing_altitude"], - } - ) - write_table( - tel_pointing_table, - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 telescope pointing table was stored in '%s' under '%s'", - self.output_path, - f"{POINTING_GROUP}/tel_{tel_id:03d}", - ) - pointing_info.append(tel_pointing) - pointing_info = vstack(pointing_info) - return pointing_info - -def mono_tool(): - # Run the tool - mono_tool = MonoPredictCTLearnModel() - mono_tool.run() - -if __name__ == "main": - mono_tool() - -if __name__ == "__main__": - mono_tool() \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_stereo.py b/ctlearn/tools/predict/predict_stereo.py deleted file mode 100644 index c868fb70..00000000 --- a/ctlearn/tools/predict/predict_stereo.py +++ /dev/null @@ -1,375 +0,0 @@ -""" -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. -""" - - -import numpy as np - -from astropy.table import ( - Table, - vstack, - join, - setdiff, -) - -from ctapipe.containers import ( - ParticleClassificationContainer, - ReconstructedGeometryContainer, - ReconstructedEnergyContainer, -) - -from ctapipe.io import read_table, write_table - -from ctapipe.reco.utils import add_defaults_and_meta -from dl1_data_handler.reader import ( - ProcessType, -) -from ctlearn.tools.predict.utils.predict_model import PredictCTLearnModel - - -SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" -FIXED_POINTING_GROUP = "/configuration/telescope/pointing" -POINTING_GROUP = "/dl1/monitoring/telescope/pointing" -SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" -DL1_TELESCOPE_GROUP = "/dl1/event/telescope" -DL1_SUBARRAY_GROUP = "/dl1/event/subarray" -DL2_SUBARRAY_GROUP = "/dl2/event/subarray" -DL2_TELESCOPE_GROUP = "/dl2/event/telescope" -SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] -TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] - -__all__ = ["StereoPredictCTLearnModel"] - -class StereoPredictCTLearnModel(PredictCTLearnModel): - """ - Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. - - This tool extends the ``PredictCTLearnModel`` to specifically handle stereoscopic R1/DL1 data. The prediction - is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. - It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. - - Attributes - ---------- - name : str - Name of the tool. - description : str - Description of the tool. - examples : str - Examples of how to use the tool. - - Methods - ------- - start() - Start the tool. - _store_mc_subarray_pointing(all_identifiers) - Store the subarray pointing table for the stereo mode for MC simulation. - """ - - name = "ctlearn-predict-stereo-model" - description = __doc__ - - examples = """ - To predict from pixel-wise image data in stereo mode using trained CTLearn models: - > ctlearn-predict-stereo-model \\ - --input_url input.dl1.h5 \\ - --PredictCTLearnModel.batch_size=16 \\ - --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ - --DLImageReader.channels=cleaned_image \\ - --DLImageReader.channels=cleaned_relative_peak_time \\ - --DLImageReader.image_mapper_type=BilinearMapper \\ - --DLImageReader.mode=stereo \\ - --DLImageReader.min_telescopes=2 \\ - --PredictCTLearnModel.stack_telescope_images=True \\ - --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ - --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ - --output output.dl2.h5 \\ - --PredictCTLearnModel.overwrite_tables=True \\ - """ - - def start(self): - self.log.info("Processing the telescope pointings...") - # Retrieve the IDs from the dl1dh for the prediction tables - example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() - example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) - all_identifiers = self.dl1dh_reader.subarray_trigger_table.copy() - all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) - nonexample_identifiers = setdiff( - all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS - ) - nonexample_identifiers.remove_column("time") - # Construct the survival telescopes for each event of the example_identifiers - survival_telescopes = [] - for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: - survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) - survival_tels = [ - self.dl1dh_reader.subarray.tel_indices[tel_id] - for tel_id in subarray_event["tel_id"].data - ] - survival_mask[survival_tels] = True - survival_telescopes.append(survival_mask) - # Add the survival telescopes to the example_identifiers - example_identifiers.add_column( - survival_telescopes, name=f"{self.prefix}_telescopes" - ) - # Pointing table for the stereo mode for MC simulation - if self.dl1dh_reader.process_type == ProcessType.Simulation: - pointing_info = self._store_mc_subarray_pointing(all_identifiers) - - # Pointing table for the observation mode - if self.dl1dh_reader.process_type == ProcessType.Observation: - pointing_info = super()._store_pointing(all_identifiers) - - self.log.info("Starting the prediction...") - classification_feature_vectors = None - if self.load_type_model_from is not None: - # Predict the energy of the primary particle - classification_table, classification_feature_vectors = ( - super()._predict_classification(example_identifiers) - ) - if self.dl2_subarray: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_prediction"], - shapes=[(len(nonexample_identifiers),)], - ) - classification_table = vstack([classification_table, nan_table]) - # Add is_valid column to the energy table - classification_table.add_column( - ~np.isnan( - classification_table[f"{self.prefix}_tel_prediction"].data, - dtype=bool, - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Rename the columns for the stereo mode - classification_table.rename_column( - f"{self.prefix}_tel_prediction", f"{self.prefix}_prediction" - ) - classification_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - classification_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - classification_table, - ParticleClassificationContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - classification_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", - ) - energy_feature_vectors = None - if self.load_energy_model_from is not None: - # Predict the energy of the primary particle - energy_table, energy_feature_vectors = super()._predict_energy( - example_identifiers - ) - if self.dl2_subarray: - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_tel_energy"], - shapes=[(len(nonexample_identifiers),)], - ) - energy_table = vstack([energy_table, nan_table]) - # Add is_valid column to the energy table - energy_table.add_column( - ~np.isnan( - energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool - ), - name=f"{self.prefix}_tel_is_valid", - ) - # Rename the columns for the stereo mode - energy_table.rename_column( - f"{self.prefix}_tel_energy", f"{self.prefix}_energy" - ) - energy_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - energy_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - energy_table, - ReconstructedEnergyContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - energy_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", - ) - direction_feature_vectors = None - if self.load_skydirection_model_from is not None: - # Join the prediction table with the telescope pointing table - example_identifiers = join( - left=example_identifiers, - right=pointing_info, - keys=SUBARRAY_EVENT_KEYS, - ) - # Predict the arrival direction of the primary particle - direction_table, direction_feature_vectors = super()._predict_skydirection( - example_identifiers - ) - if self.dl2_subarray: - # Transform the spherical coordinate offsets to sky coordinates - direction_table = super()._transform_spher_coord_offsets_to_sky( - direction_table - ) - # Produce output table with NaNs for missing predictions - if len(nonexample_identifiers) > 0: - nan_table = super()._create_nan_table( - nonexample_identifiers, - columns=[f"{self.prefix}_alt", f"{self.prefix}_az"], - shapes=[ - (len(nonexample_identifiers),), - (len(nonexample_identifiers),), - ], - ) - direction_table = vstack([direction_table, nan_table]) - # Add is_valid column to the direction table - direction_table.add_column( - ~np.isnan(direction_table[f"{self.prefix}_alt"].data, dtype=bool), - name=f"{self.prefix}_is_valid", - ) - direction_table.sort(SUBARRAY_EVENT_KEYS) - # Add the default values and meta data to the table - add_defaults_and_meta( - direction_table, - ReconstructedGeometryContainer, - prefix=self.prefix, - ) - # Save the prediction to the output file - write_table( - direction_table, - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL2 prediction data was stored in '%s' under '%s'", - self.output_path, - f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", - ) - - # Create the feature vector table if the DL1 features are enabled - if self.dl1_features: - self.log.info("Processing and storing dl1 feature vectors...") - feature_vector_table = super()._create_feature_vectors_table( - example_identifiers, - nonexample_identifiers, - classification_feature_vectors, - energy_feature_vectors, - direction_feature_vectors, - ) - # Loop over the selected telescopes and store the feature vectors - # for each telescope in the output file. The feature vectors are stored - # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. - # Rename the columns for the stereo mode - feature_vector_table.rename_column( - f"{self.prefix}_tel_classification_feature_vectors", - f"{self.prefix}_classification_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_energy_feature_vectors", - f"{self.prefix}_energy_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_geometry_feature_vectors", - f"{self.prefix}_geometry_feature_vectors", - ) - feature_vector_table.rename_column( - f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" - ) - feature_vector_table.sort(SUBARRAY_EVENT_KEYS) - # Save the prediction to the output file - write_table( - feature_vector_table, - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 feature vectors was stored in '%s' under '%s'", - self.output_path, - f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", - ) - - def _store_mc_subarray_pointing(self, all_identifiers): - """ - Store the subarray pointing table from MC simulation to the output file. - - Parameters: - ----------- - all_identifiers : astropy.table.Table - Table containing the subarray pointing information. - """ - # Read the subarray pointing table - pointing_info = read_table( - self.input_url, - f"{SIMULATION_CONFIG_TABLE}", - ) - # Assuming min_az = max_az and min_alt = max_alt - pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) - pointing_info.rename_column("min_az", "pointing_azimuth") - pointing_info.rename_column("min_alt", "pointing_altitude") - # Join the prediction table with the telescope pointing table - pointing_info = join( - left=pointing_info, - right=all_identifiers, - keys=["obs_id"], - ) - # TODO: use keep_order for astropy v7.0.0 - pointing_info.sort(SUBARRAY_EVENT_KEYS) - # Create the pointing table - pointing_table = Table( - { - "time": pointing_info["time"], - "array_azimuth": pointing_info["pointing_azimuth"], - "array_altitude": pointing_info["pointing_altitude"], - "array_ra": np.nan * np.ones(len(pointing_info)), - "array_dec": np.nan * np.ones(len(pointing_info)), - } - ) - # Save the pointing table to the output file - write_table( - pointing_table, - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - overwrite=self.overwrite_tables, - ) - self.log.info( - "DL1 subarray pointing table was stored in '%s' under '%s'", - self.output_path, - f"{SUBARRAY_POINTING_GROUP}", - ) - return pointing_info - -def stereo_tool(): - # Run the tool - stereo_tool = StereoPredictCTLearnModel() - stereo_tool.run() - -if __name__ == "main": - stereo_tool() -if __name__ == "__main__": - stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py b/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py deleted file mode 100644 index 19b403c1..00000000 --- a/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py +++ /dev/null @@ -1,312 +0,0 @@ -""" -PyTorch prediction module for LST1 telescope data. -This module provides functionality to load trained models and perform predictions -on DL1 level data for particle type classification, energy estimation, and direction reconstruction. -""" - -from ctlearn.core.pytorch.net_utils import create_model, ModelHelper -import torch -from ctlearn.core.ctlearn_enum import Task, Mode -from ctapipe.io import read_table -from astropy.table import join -from dl1_data_handler.reader import get_unmapped_image -import numpy as np -from tqdm import tqdm -from pytorch_lightning.callbacks import Callback - -class GPUStatsLogger(Callback): - """ - PyTorch Lightning callback to log GPU memory statistics during training. - - This callback tracks GPU memory allocation and reservation at the end of each training epoch - and logs the statistics to TensorBoard. - """ - - def on_train_epoch_end(self, trainer, pl_module): - """ - Called at the end of each training epoch to log GPU memory statistics. - - Args: - trainer: PyTorch Lightning trainer instance - pl_module: The LightningModule being trained - """ - mem_allocated = torch.cuda.memory_allocated() - mem_reserved = torch.cuda.memory_reserved() - - trainer.logger.experiment.add_scalar( - "gpu_mem_allocated", mem_allocated, global_step=trainer.current_epoch - ) - trainer.logger.experiment.add_scalar( - "gpu_mem_reserved", mem_reserved, global_step=trainer.current_epoch - ) - - -def predictions(self): - """ - Perform predictions on input DL1 data using trained models. - - This function processes the input file in batches, applies quality cuts, - and generates predictions for particle type, energy, and/or direction - depending on the configured tasks. - - Returns: - tuple: Contains the following arrays: - - event_id: Event identifiers - - tel_azimuth: Telescope azimuth angles - - tel_altitude: Telescope altitude angles - - trigger_time: Event trigger times - - prediction: Particle type classification scores - - energy: Reconstructed energy values - - cam_coord_offset_x: Camera coordinate offset in x - - cam_coord_offset_y: Camera coordinate offset in y - - classification_fvs: Classification feature vectors - - energy_fvs: Energy estimation feature vectors - - direction_fvs: Direction reconstruction feature vectors - """ - - # Update channels based on log scaling config - if "normalization" in self.parameters and "apply_log_scaling" in self.parameters["normalization"]: - new_channels = list(self.channels) - if self.parameters["normalization"]["apply_log_scaling"][0] and not new_channels[0].startswith("log_"): - new_channels[0] = "log_" + new_channels[0] - self.channels = new_channels - - # Optimize batch size if requested - if self.optim_batch_size: - batch_size_found = False - batch = 256 - step = 16 - - while not batch_size_found: - from ctlearn.tools.predict.utils.optimaze_batch_size import test_batch - - # Load a test batch to find optimal batch size - dl1_table = read_table( - self.input_url, self.image_table_path, start=0, stop=batch - ) - dl1_table = join(left=dl1_table, right=self.parameter_table, keys=["event_id"]) - dl1_table = join(left=dl1_table, right=self.trigger_table, keys=["event_id"]) - - # Prepare test data - data = [] - for event in dl1_table: - image = get_unmapped_image(dl1_table[0], self.channels, self.transforms) - data.append(self.image_mapper.map_image(image)) - input_data = {"input": np.array(data)} - - imgs = input_data['input'][:, :, :, 0] - if len(self.channels) == 2: - peak_time = input_data['input'][:, :, :, 1] - - # Test batch size for each configured task - for task in self.tasks: - if task == Task.type: - batch_size_found = not test_batch( - self.type_model, - torch.tensor(imgs).unsqueeze(1).to(self.device), - torch.tensor(peak_time).unsqueeze(1).to(self.device), - self.device - ) - if task == Task.energy: - batch_size_found = not test_batch( - self.energy_model, - torch.tensor(imgs).unsqueeze(1).to(self.device), - torch.tensor(peak_time).unsqueeze(1).to(self.device), - self.device - ) - if task in [Task.cameradirection, Task.skydirection, Task.direction]: - batch_size_found = not test_batch( - self.cameradirection_model, - torch.tensor(imgs).unsqueeze(1).to(self.device), - torch.tensor(peak_time).unsqueeze(1).to(self.device), - self.device - ) - - batch += step - if not batch_size_found: - self.log.info(f"Batch size: {batch} OK") - - self.batch_size = batch - step - self.log.info(f"Optimized batch size: {self.batch_size}") - - # Initialize output arrays - event_id, tel_azimuth, tel_altitude, trigger_time = [], [], [], [] - prediction, energy, cam_coord_offset_x, cam_coord_offset_y = [], [], [], [] - classification_fvs, energy_fvs, direction_fvs = [], [], [] - - # Process input file in batches - for start in tqdm(range(0, self.table_length, self.batch_size), desc="Processing input file"): - stop = min(start + self.batch_size, self.table_length) - self.log.debug("Processing chunk from '%d' to '%d'.", start, stop - 1) - - # Read and join tables - dl1_table = read_table(self.input_url, self.image_table_path, start=start, stop=stop) - dl1_table = join(left=dl1_table, right=self.parameter_table, keys=["event_id"]) - dl1_table = join(left=dl1_table, right=self.trigger_table, keys=["event_id"]) - - # Apply quality selection - passes_quality_checks = np.ones(len(dl1_table), dtype=bool) - if self.quality_query: - passes_quality_checks = self.quality_query.get_table_mask(dl1_table) - dl1_table = dl1_table[passes_quality_checks] - - if len(dl1_table) == 0: - self.log.debug("No events passed the quality selection.") - continue - - # Prepare input data - data = [] - for event in dl1_table: - image = get_unmapped_image(event, self.channels, self.transforms) - data.append(self.image_mapper.map_image(image)) - input_data = {"input": np.array(data)} - - # Store metadata - event_id.extend(dl1_table["event_id"].data) - tel_azimuth.extend(dl1_table["tel_az"].data) - tel_altitude.extend(dl1_table["tel_alt"].data) - trigger_time.extend(dl1_table["time"].mjd) - - # Extract and clean image data - imgs = input_data['input'][:, :, :, 0] - if len(self.channels) == 2: - peak_time = input_data['input'][:, :, :, 1] - peak_time[peak_time < 0] = 0 - peak_time[np.isnan(peak_time)] = 0 - peak_time[np.isinf(peak_time)] = 0 - - imgs[imgs < 0] = 0 - imgs[np.isnan(imgs)] = 0 - imgs[np.isinf(imgs)] = 0 - - feature_vector = True - if self.parameters["normalization"]["apply_log_scaling"][0] == True: - imgs = imgs.astype(np.float32) - imgs = np.log10(imgs + 1.0) - if self.parameters["normalization"]["apply_log_scaling"][1] == True and len(self.channels) == 2: - peak_time = peak_time.astype(np.float32) - peak_time = np.log10(peak_time + 1.0) - - # Prepare unified input tensor - if len(self.channels) == 2: - input_tensor = torch.cat([ - torch.tensor(imgs).unsqueeze(1), - torch.tensor(peak_time).unsqueeze(1) - ], dim=1).to(self.device) - else: - input_tensor = torch.tensor(imgs).unsqueeze(1).to(self.device) - - # Run predictions for each configured task - for task in self.tasks: - if task == Task.type: - # Particle type classification - classification_pred, energy_pred, direction_pred = self.type_model(input_tensor) - - prediction.extend(torch.softmax(classification_pred[0], dim=1).cpu().detach().numpy()[:, 1]) - classification_fvs.extend(classification_pred[1].cpu().detach().numpy()) - - elif task == Task.energy: - # Energy estimation - classification_pred, energy_pred, direction_pred = self.energy_model(input_tensor) - - energy.extend(energy_pred[0].cpu().detach().numpy()) - if feature_vector: - energy_fvs.extend(energy_pred[1].cpu().detach().numpy()) - else: - energy_fvs.extend(np.array([[0]] * len(energy_pred[0]))) - - elif task in [Task.cameradirection, Task.skydirection, Task.direction]: - # Direction reconstruction - classification_pred, energy_pred, direction_pred = self.dirrection_model(input_tensor) - - cam_coord_offset_x.extend(direction_pred[0][:, 0].float().cpu().detach().numpy()) - cam_coord_offset_y.extend(direction_pred[0][:, 1].float().cpu().detach().numpy()) - if feature_vector: - direction_fvs.extend(direction_pred[1].cpu().detach().numpy()) - else: - direction_fvs.extend(np.array([[0]] * len(direction_pred[0]))) - - else: - raise ValueError( - f"task:{task.name} is not supported. Task must be type, direction or energy" - ) - - return (event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, - cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs) - - -def load_pytorch_model(self): - """ - Load PyTorch models from checkpoints for the configured tasks. - - This function creates and loads models for particle type classification, - energy estimation, and/or direction reconstruction based on the tasks - specified in the configuration. - - Returns: - torch.nn.Module: The last loaded model (for compatibility) - """ - model = None - from ctlearn.core.pytorch.model_collection import CTLearnPyTorchModel - - def load_pytorch_model_net(model_info, task_name, num_inputs, num_outputs): - model_name = model_info.get("model_name", "") - try: - component_cls = CTLearnPyTorchModel.non_abstract_subclasses().get(model_name) - if component_cls is not None: - params = model_info.get("parameters", {}).copy() - params.pop("task", None) - params.pop("num_inputs", None) - params.pop("num_outputs", None) - params["parent"] = self - component = component_cls( - task=task_name, - num_inputs=num_inputs, - num_outputs=num_outputs, - **params - ) - return component.model - except Exception as e: - self.log.warning(f"Failed to load model {model_name} as Component: {e}. Falling back to create_model.") - return create_model(model_info) - - num_inputs = 1 - - for task in self.tasks: - # Create model based on task type - if task == Task.type: - model_net = load_pytorch_model_net(self.parameters["model"]["model_type"], "type", num_inputs, 2) - check_point_path = self.parameters["data"]["type_checkpoint"] - - elif task == Task.energy: - model_net = load_pytorch_model_net(self.parameters["model"]["model_energy"], "energy", num_inputs, 1) - check_point_path = self.parameters["data"]["energy_checkpoint"] - - elif task in [Task.cameradirection, Task.skydirection, Task.direction]: - model_net = load_pytorch_model_net(self.parameters["model"]["model_direction"], "direction", num_inputs, 3) - check_point_path = self.parameters["data"]["direction_checkpoint"] - - else: - raise ValueError( - f"task:{task.name} is not supported. Task must be type, direction or energy" - ) - - # Load the model from checkpoint - model = ModelHelper.loadModel( - model_net, "", check_point_path, Mode.observation, device_str=self.device_str - ) - model.eval() - - # Assign model to appropriate attribute - if task == Task.type: - self.type_model = model - elif task == Task.energy: - self.energy_model = model - elif task in [Task.cameradirection, Task.skydirection, Task.direction]: - self.dirrection_model = model - else: - raise ValueError( - f"task:{task.name} is not supported. Task must be type, direction or energy" - ) - - return model diff --git a/ctlearn/tools/predict/pytorch/predic_model_pytorch.py b/ctlearn/tools/predict/pytorch/predic_model_pytorch.py deleted file mode 100644 index 0fcd89ca..00000000 --- a/ctlearn/tools/predict/pytorch/predic_model_pytorch.py +++ /dev/null @@ -1,144 +0,0 @@ -""" -PyTorch model prediction module for CTLearn. -This module provides functionality to load trained models and perform batch predictions -on DL1 data for multiple tasks including particle classification, energy estimation, -and direction reconstruction. -""" - -from ctlearn.core.data_loader.loader import DLDataLoader -import torch -from tqdm import tqdm -import numpy as np -import inspect -from ctlearn.tools.predict.utils.load_model import load_model - - -def predict_with_model_pytorch(self, task): - """ - Load and predict with a CTLearn PyTorch model. - - This function loads a trained model from the specified path and performs predictions - on the provided data. It processes the data in batches and returns predictions for - particle type classification, energy estimation, and/or direction reconstruction - based on the configured task. - - Parameters - ---------- - task : Task - The task(s) to perform predictions for (type, energy, or direction). - - Returns - ------- - predict_data : dict - Dictionary containing prediction results with keys: - - 'type': Particle type classification probabilities (gammaness scores) - - 'energy': Reconstructed energy values - - 'cameradirection': Camera coordinate offsets for direction reconstruction - feature_vectors : None - Feature vectors (currently not extracted, placeholder for future implementation). - - Notes - ----- - The function automatically detects whether the model requires peak time information - by inspecting the model's forward method signature. Models can accept either one - input (image only) or two inputs (image and peak time). - """ - # Initialize batch size from configuration parameters - self.batch_size = self.parameters["hyp"]["batches"] - - # Create data loader for the specified task - # The DLDataLoader is initialized separately for each task to ensure robustness - channels = ["cleaned_image", "cleaned_peak_time"] - if self.parameters["normalization"]["apply_log_scaling"][0]: - channels[0] = "log_" + channels[0] - self.dl1dh_reader.channels = channels - data_loader = DLDataLoader.create( - framework="pytorch", - DLDataReader=self.dl1dh_reader, - indices=self.indices, - tasks=[task], - parameters=self.parameters, - use_augmentation=False, - batch_size=self.batch_size, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - - # Note: Handling of incomplete last batch - # In PyTorch, unlike Keras, we can process incomplete batches directly - # without needing a separate data loader. The code below is kept as reference - # for potential future use or compatibility with other frameworks. - - # data_loader_last_batch = None - # if self.last_batch_size > 0: - # last_batch_indices = self.indices[-self.last_batch_size:] - # data_loader_last_batch = DLDataLoader.create( - # framework="pytorch", - # DLDataReader=self.dl1dh_reader, - # indices=last_batch_indices, - # tasks=task, - # parameters=self.parameters, - # use_augmentation=False, - # batch_size=self.last_batch_size, - # sort_by_intensity=self.sort_by_intensity, - # stack_telescope_images=self.stack_telescope_images, - # ) - - # Load the trained model from checkpoint - model = load_model(self) - - # Inspect model signature to determine number of inputs - # This allows the code to work with models that take either: - # - Single input: image only - # - Dual input: image and peak time - sig = inspect.signature(model.forward) - num_inputs = len(sig.parameters) - - # Initialize prediction data dictionary with empty lists - predict_data = {} - predict_data['type'] = [] - predict_data['energy'] = [] - predict_data["cameradirection"] = [] - - # Set model to evaluation mode (disables dropout, batch normalization, etc.) - model.eval() - - # Perform predictions without gradient computation (faster inference) - with torch.no_grad(): - for i, x in enumerate(tqdm(data_loader, desc="Processing", total=len(data_loader))): - # Skip empty batches - if len(x[0]['image']) == 0: - continue - - # Forward pass through the model - classification_pred, energy_pred, direction_pred = model( - x[0]['image'].to(self.device) - ) - - # Collect particle type classification predictions - if classification_pred[0] is not None: - # Apply softmax to get probability distribution and extract gammaness score - gammaness = torch.softmax(classification_pred[0], dim=1).cpu().detach().numpy() - predict_data['type'].extend(gammaness) - - # Collect energy estimation predictions - if energy_pred[0] is not None: - predict_data['energy'].extend(energy_pred[0].cpu().detach().numpy()) - - # Collect direction reconstruction predictions - if direction_pred[0] is not None: - predict_data["cameradirection"].extend(direction_pred[0].cpu().detach().numpy()) - - # Log progress every 100 batches - if i % 100 == 0: - self.log.info(f"Processed {i}/{len(data_loader)} events.") - - self.log.info("Processing completed.") - - # Convert lists to numpy arrays for efficient storage and further processing - predict_data["cameradirection"] = np.array(predict_data["cameradirection"]) - predict_data["type"] = np.array(predict_data["type"]) - predict_data["energy"] = np.array(predict_data["energy"]) - - # Return predictions and placeholder for feature vectors - return predict_data, None \ No newline at end of file diff --git a/ctlearn/tools/predict/utils/load_model.py b/ctlearn/tools/predict/utils/load_model.py deleted file mode 100644 index 3963142a..00000000 --- a/ctlearn/tools/predict/utils/load_model.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -Model loading utility module for CTLearn predictions. -This module provides a framework-agnostic interface for loading trained models, -supporting both Keras and PyTorch frameworks. -""" - - -def load_model(self): - """ - Load a trained model based on the configured framework type. - - This function acts as a dispatcher that delegates model loading to the appropriate - framework-specific implementation. It supports both Keras and PyTorch frameworks - and loads the model from the checkpoint path specified in the configuration. - - Parameters - ---------- - self : PredictionHandler - The prediction handler instance containing configuration parameters including: - - framework_type: str, either "keras" or "pytorch" - - Model checkpoint paths and other configuration parameters - - Returns - ------- - model : object - The loaded model ready for inference. Type depends on the framework: - - For Keras: keras.Model - - For PyTorch: torch.nn.Module - Returns None if the framework is not recognized. - - Raises - ------ - ImportError - If the specified framework's prediction module cannot be imported. - - Notes - ----- - The function automatically detects the framework type from the configuration - and imports the appropriate loading function dynamically to avoid unnecessary - dependencies when using only one framework. - """ - if self.framework_type == "keras": - # Load Keras model using framework-specific loader - from ctlearn.tools.predict.keras.predic_LST1_keras import load_keras_model - return load_keras_model(self) - - elif self.framework_type == "pytorch": - # Load PyTorch model using framework-specific loader - from ctlearn.tools.predict.pytorch.predic_LST1_pytorch import load_pytorch_model - return load_pytorch_model(self) - - else: - # Log error if framework is not recognized - self.log.error(f"Framework '{self.framework_type}' not found! Supported frameworks: 'keras', 'pytorch'") - return None \ No newline at end of file diff --git a/ctlearn/tools/predict/utils/optimaze_batch_size.py b/ctlearn/tools/predict/utils/optimaze_batch_size.py deleted file mode 100644 index 2cfc81cb..00000000 --- a/ctlearn/tools/predict/utils/optimaze_batch_size.py +++ /dev/null @@ -1,174 +0,0 @@ -""" -Batch size optimization utilities for CTLearn predictions. -This module provides functions to test and find optimal batch sizes for model inference, -helping to maximize GPU utilization while avoiding out-of-memory errors. -""" - -import torch - - -def test_batch(model, imgs, peak_time, device): - """ - Test if a model can process a given batch without memory errors. - - This function tests whether a pre-prepared batch of images and peak times - can be successfully processed by the model without encountering out-of-memory - (OOM) errors. It's used to validate batch sizes during optimization. - - Parameters - ---------- - model : torch.nn.Module - The PyTorch model to test. - imgs : torch.Tensor or array-like - Batch of images to process. Will be converted to tensor if necessary. - peak_time : torch.Tensor or array-like - Batch of peak time information. Will be converted to tensor if necessary. - device : torch.device or str - Device to run the test on (e.g., 'cuda:0' or 'cpu'). - - Returns - ------- - bool - True if the batch can be processed successfully, False if OOM error occurs. - - Raises - ------ - RuntimeError - If a RuntimeError other than OOM occurs during processing. - - Notes - ----- - The function automatically clears the CUDA cache after each test to ensure - clean memory state for subsequent tests. - """ - # Move model to specified device and set to evaluation mode - model.to(device) - model.eval() - - # Ensure inputs are tensors and move to device - if not torch.is_tensor(imgs): - imgs = torch.as_tensor(imgs).to(device) - else: - imgs = imgs.to(device) - - if not torch.is_tensor(peak_time): - peak_time = torch.as_tensor(peak_time).to(device) - else: - peak_time = peak_time.to(device) - - try: - # Attempt forward pass without gradient computation - with torch.no_grad(): - _ = model(imgs, peak_time) - - # Clear CUDA cache to free memory - torch.cuda.empty_cache() - return True - - except RuntimeError as e: - # Check if error is due to out of memory - if "out of memory" in str(e).lower(): - torch.cuda.empty_cache() - return False - else: - # Re-raise unexpected errors after cleaning up - torch.cuda.empty_cache() - raise e - - -def find_max_batch_size(self, model, imgs, peak_time, device, start_bs=8, step=8, max_bs=512): - """ - Find the maximum batch size that can be processed without OOM errors. - - This function performs a binary-like search to find the largest batch size - that can be successfully processed by the model on the given device. It starts - with a small batch size and incrementally increases until an OOM error occurs. - - Parameters - ---------- - self : object - Reference to the parent object (for potential logging or configuration access). - model : torch.nn.Module - The PyTorch model to test. - imgs : torch.Tensor or array-like - Sample images to use for testing. Only the first image is used and replicated. - peak_time : torch.Tensor or array-like - Sample peak time data. Only the first value is used and replicated. - device : torch.device or str - Device to run tests on (e.g., 'cuda:0' or 'cpu'). - start_bs : int, optional - Initial batch size to start testing with. Default is 8. - step : int, optional - Increment step for batch size increases. Default is 8. - max_bs : int, optional - Maximum batch size to test. Default is 512. - - Returns - ------- - int - The maximum batch size that can be processed without OOM errors. - - Raises - ------ - RuntimeError - If a RuntimeError other than OOM occurs during testing. - - Notes - ----- - - The function replicates a single image/peak_time to create test batches - - CUDA cache is cleared after each test to ensure accurate memory measurements - - Progress is printed to console with emoji indicators for status - """ - batch_size = start_bs - - # Move model to device and set to evaluation mode - model.to(device) - model.eval() - - # Ensure inputs are tensors - if not torch.is_tensor(imgs): - imgs = torch.as_tensor(imgs) - if not torch.is_tensor(peak_time): - peak_time = torch.as_tensor(peak_time) - - # Disable gradient computation for efficiency - with torch.no_grad(): - while batch_size <= max_bs: - try: - # Prepare image batch - batch_imgs = imgs[:1] # Take first image as template - if batch_imgs.ndim == 3: - batch_imgs = batch_imgs.unsqueeze(1) # Add channel dimension if needed - # Replicate to create batch of desired size - batch_imgs = batch_imgs.repeat(batch_size, 1, 1, 1).to(device) - - # Prepare peak_time batch - value = peak_time[:1, 0, 0] # Extract representative value - batch_peaks = value.unsqueeze(1) # Shape [1, 1] - batch_peaks = batch_peaks.repeat(batch_size, 1) # Replicate for batch - batch_peaks = batch_peaks.unsqueeze(-1).unsqueeze(-1).to(device) # Shape [batch, 1, 1, 1] - - # Attempt forward pass - _ = model(batch_imgs, batch_peaks) - - # Clean up tensors and cache - del batch_imgs, batch_peaks, _ - torch.cuda.empty_cache() - - print(f"✅ Batch size {batch_size} OK") - batch_size += step - - except RuntimeError as e: - if "out of memory" in str(e).lower(): - # OOM encountered, return previous successful batch size - print(f"💥 OOM at batch size {batch_size}") - torch.cuda.empty_cache() - return batch_size - step - else: - # Unexpected error, clean up and re-raise - print(f"❌ Unexpected error at batch size {batch_size}: {e}") - torch.cuda.empty_cache() - raise e - - # If we reached max_bs without OOM, return it (minus step to be safe) - return batch_size - step diff --git a/ctlearn/tools/tests/test_predict_LST1.py b/ctlearn/tools/tests/test_predict_LST1.py index b9ee77e3..1fc02838 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -40,13 +40,10 @@ def test_predict_mono_model_with_lst1_mock_data( Test LST1PredictionTool using trained mono models and mock LST-1 DL1 files. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) - # Hardcopy the trained models to the model directory telescope_type = "LST" for reco_task in ["type", "energy", "cameradirection"]: @@ -57,12 +54,9 @@ def test_predict_mono_model_with_lst1_mock_data( ) model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" - # Check that the mock LST1 DL1 file was created assert mock_lst1_dl1_file.exists(), "Mock LST1 DL1 file not found" - output_file = dl2_dir / f"mock_lst1_{framework}_predictions.dl2.h5" - # Build command-line arguments for LST1PredictionTool argv = [ f"--input_url={mock_lst1_dl1_file}", @@ -77,13 +71,10 @@ def test_predict_mono_model_with_lst1_mock_data( "--dl2-telescope", "--overwrite", ] - # Run LST1PredictionTool assert run_tool(LST1PredictionTool(), argv=argv, cwd=tmp_path) == 0 - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" - # Check that the created DL2 file can be read with the TableLoader allowed_tels = [1] with TableLoader( @@ -109,7 +100,6 @@ def test_predict_mono_model_with_lst1_mock_data( assert ( tel_events[tel_id][col][0] is not np.nan ), f"{col} has NaN values in DL2 file {output_file.name}" - # Check subarray-wise data subarray_events = loader.read_subarray_events(start=0, stop=2, dl2=True) assert len(subarray_events) > 0 diff --git a/ctlearn/tools/tests/test_predict_model.py b/ctlearn/tools/tests/test_predict_model.py index 862bdf31..1434192b 100644 --- a/ctlearn/tools/tests/test_predict_model.py +++ b/ctlearn/tools/tests/test_predict_model.py @@ -45,16 +45,13 @@ def test_predict_mono_model_with_r1_waveforms( Test training CTLearn mono model using the R1 gamma and proton files for all reconstruction tasks and predicting DL2 from R1 waveforms. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their available telescopes telescope_type = "LST" available_tels = [1, 2] - # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "cameradirection"]: key = f"{framework}_{telescope_type}_{reco_task}" @@ -90,7 +87,6 @@ def test_predict_mono_model_with_r1_waveforms( ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader @@ -140,10 +136,8 @@ def test_predict_mono_model_with_dl1_images( Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks and predicting DL2 from DL1 images. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their allowed telescopes @@ -197,7 +191,6 @@ def test_predict_mono_model_with_dl1_images( ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader @@ -248,16 +241,13 @@ def test_predict_stereo_model_with_dl1_images( Test training CTLearn stereo model using the DL1 gamma and proton files for all reconstruction tasks and predicting DL2 from DL1 images. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their available telescopes telescope_type = "MST" allowed_tels = [7, 13, 15] - # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "skydirection"]: key = f"{framework}_{telescope_type}_{reco_task}" @@ -294,7 +284,6 @@ def test_predict_stereo_model_with_dl1_images( ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader diff --git a/ctlearn/tools/tests/test_train_model.py b/ctlearn/tools/tests/test_train_model.py index c7bfdecd..3466f2cf 100644 --- a/ctlearn/tools/tests/test_train_model.py +++ b/ctlearn/tools/tests/test_train_model.py @@ -13,23 +13,18 @@ def test_train_ctlearn_model(framework, model, reco_task, dl1_gamma_file, dl1_pr Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks. Each test run gets its own isolated temp directories. """ - # Restrict to MST array telescope_type = "MST" allowed_tels = [7, 13, 15, 16, 17, 19] - # Temporary directories for signal and background signal_dir = tmp_path / "gamma_dl1" signal_dir.mkdir(parents=True, exist_ok=True) - background_dir = tmp_path / "proton_dl1" background_dir.mkdir(parents=True, exist_ok=True) - # Hardcopy DL1 gamma file to the signal directory shutil.copy(dl1_gamma_file, signal_dir) # Hardcopy DL1 proton file to the background directory shutil.copy(dl1_proton_file, background_dir) - # Hardcopy the trained models to the model directory if model == "LoadedModel": model_dir = tmp_path / "pretrained_model" @@ -41,10 +36,8 @@ def test_train_ctlearn_model(framework, model, reco_task, dl1_gamma_file, dl1_pr ) model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained {framework} mono model file not found for {key}" - # Output directory for trained model output_dir = tmp_path / f"ctlearn_{framework}_{model}_{reco_task}" - # Build command-line arguments argv = [ f"--signal={signal_dir}", @@ -56,7 +49,6 @@ def test_train_ctlearn_model(framework, model, reco_task, dl1_gamma_file, dl1_pr "--DLImageReader.focal_length_choice=EQUIVALENT", f"--DLImageReader.allowed_tels={allowed_tels}", ] - # Include background only for classification task if reco_task == "type": argv.extend( @@ -66,7 +58,6 @@ def test_train_ctlearn_model(framework, model, reco_task, dl1_gamma_file, dl1_pr "--DLImageReader.enforce_subarray_equality=False", ] ) - argv.append(f"--TrainCTLearnModel.model_type={model}") if model == "LoadedModel": argv.append(f"--LoadedModel.load_model_from={model_file}") diff --git a/ctlearn/tools/train/pytorch/config/default_config_file.yml b/ctlearn/tools/train/pytorch/config/default_config_file.yml deleted file mode 100644 index abf55871..00000000 --- a/ctlearn/tools/train/pytorch/config/default_config_file.yml +++ /dev/null @@ -1,152 +0,0 @@ -data: - train_gamma_proton: ./data/gamma_proton_train_remix.dl1.pickle - validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle - - train_gamma: ./data/gamma_955000_train.pickle - validation_gamma: ./data/gamma_106141_validation.pickle - - test_gamma: ./data/gamma_1805522_test_gamma.pickle - test_proton: ./data/proton_130811_test_proton.pickle - test_electron: None - test_validation_gamma: ./data/gamma_180552_test_val_gamma.pickle - - test_validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle - - observation: ./run_2931.dl1.pickle - # Important: This is only for testing purpose. Set always to 0 - # when you are training, validating or estimating the dl2 files - training_reduce_factor: 0 #64 #4 - validation_reduce_factor: 0 #16 #8 - validation_test_reduce_factor: 0 #16 #8 - - # Check points - type_checkpoint: ./run/run_type_training_14/exp_14_type_train/version_0/Epoch_6_type_train_acc_80.9682309627532959.pth - energy_checkpoint: /home/cpozogonzalez/ctlearn/run/run_energy_training_14/exp_14_energy_train/version_1/Epoch_13_energy_train_loss_30.2185532478501244.pth - direction_checkpoint: /lhome/ext/ucm147/ucm1477/data/check_points/v_5/Epoch_23_cameradirection_train_loss_7296.6534562211982120.pth - -run_details: - - mode: "observation" # The option are: "train", "results", "observation" and "validate" - task: "direction" # The option are: "all", "energy" "type" and "direction" - test_type: "gamma" # The option are: "gamma" "proton" or "electron" - experiment_number: 14 # The experiment number. The experiment folder is saved into the "run" folder. - - -cut-off: - - leakage_intensity: 0.2 # bigger to this value, the event is removed - intensity: 50 # below to this value, the event is removed - -model: - - model_type: - model_name: "DoubleBBEfficientNet" - parameters: - model_variant: "efficientnet-b3" - task: 'type' - num_outputs: 2 - device_str: "cuda" - energy_bins: None - - model_energy: - model_name: "ThinResNet" - parameters: - task: 'energy' - num_inputs: 1 - num_outputs: 1 - num_blocks: [3, 4, 6, 3] #[2, 3, 3, 3] - dropout: 0.1 - use_bn: False - - model_direction: - model_name: "ThinResNet_DBB" - parameters: - task: 'direction' - num_inputs: 1 - num_outputs: 3 - num_blocks: [3, 4, 6, 3] - dropout: 0.1 - use_bn: False - - # model_direction: - # model_name: "DBBNoPropDTReg" - # parameters: - # task: 'direction' - # num_outputs: 3 - # embedding_dim: 512 - # T: 3 - # eta: 0.1 #0.1 - -# Hyper-parameters -hyp: - - epochs: 30 - batches: 128 #128 #64 - dynamic_batches: True - optimizer: Adamw - momentum: 0.957 #Yolo 0.937 # Efficient-b3 0.757 - weight_decay: 0.0005 #0.004676 #0.0001 #0.00002 Efficient-b3 0.0005 - learning_rate: 1e-4 #1e-5 #Efficient-b3 1e-5 - lrf: 0.1 - start_epoch: 0 - steps_epoch: 100 # Computed online. Must be removed - l2_lambda: 1e-7 #1e-5 #1e-5 # L2 regularization (Set to 0.0 to skip the L2 Regularization) - adam_epsilon: 1.0e-08 #7.511309034256153e-05 #1.0e-08 - gradient_clip_val: 3.0 # Avoid gradient explosion - - save_k: 200 # Save as maximum k checkpoints. - -augmentation: - # probabilities for augmentation range = [0, 1.0] - # prob = 0.0 -> Always apply the augmentation - # prob >= 1.0 -> Never apply the augmentation, i.e., Set bigger than 1.0 ( ex: 2.0) if you want disable it. - # Note: mask augmentation is always on even with flag use_augmentation = True - # To disable it, just set to 2.5 for example. - - use_augmentation: True # This apply only on training mode. - aug_prob: 0.5 # Probability of use Augmentation - rot_prob: 0.5 # Rotation probability - trans_prob: 0.5 # Translation probability - flip_hor_prob: 0.5 # Horizontal Flip probability - flip_ver_prob: 0.5 # Vertical Flip probability - mask_prob: 0.5 # Apply mask probability - mask_dvr_prob: 0.5 # Apply dvr mask probability - noise_prob: 0.5 # No implemented yet. - max_rot: 5 # Maximum rotation in augmentation - max_trans: 10 # Maximum translation in augmentation - -normalization: - - # Normalization: Im' = (Im-mu)/sigma - use_clean: True # Use the image with the applied mask (True), IOC the mask is not applied (False) - use_clean_dvr: False - type_mu: 0.0 - type_sigma: 1000.0 - - dir_mu: 0.0 - dir_sigma: 1000.0 - - energy_mu: 0.0 - energy_sigma: 1000.0 - -dataset: - num_workers: 1 # - pin_memory: True - persistent_workers: True # - -# Hardware Architecture and precision -arch: - # device: 'mps' # Apple Mx - device: 'cuda' - precision_type: "32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - precision_energy: "32-true" #"32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - precision_direction: "32-true" # "bf16-mixed" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" - # (bf16 for GPU with Ampere or higher, it is better that 16 because is numerical more stability) - # devices: [0,1] # [0,1] For multiple GPUs - devices: [0,1] - # Note: Check the documentation for more information. - strategy: 'deepspeed_stage_2' # Options: auto, dpp, dpp_swap, fsdp, deepspeed, horovod, bagua, deepspeed_stage_2, deepspeed_stage_3, colossalai, hivemind, etc... - -Notes: - Note_1: Training with augmentation dvr using 1-3 dilatations - Note_2: Trainining b3 applying always the mask \ No newline at end of file