diff --git a/ctlearn/core/loader.py b/ctlearn/core/loader.py index 27570d9..9fa77a3 100644 --- a/ctlearn/core/loader.py +++ b/ctlearn/core/loader.py @@ -170,12 +170,21 @@ def _get_mono_item(self, batch): ) if "energy" in self.tasks: labels["energy"] = batch["log_true_energy"].data - if "direction" in self.tasks: - labels["direction"] = np.stack( + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( ( - batch["spherical_offset_az"].data, - batch["spherical_offset_alt"].data, - batch["angular_separation"].data, + batch["sky_offset_lon"].data, + batch["sky_offset_lat"].data, + batch["sky_angular_separation"].data, + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( + ( + batch["cam_coord_offset_x"].data, + batch["cam_coord_offset_y"].data, + batch["cam_coord_distance"].data, ), axis=1, ) @@ -215,7 +224,8 @@ def _get_stereo_item(self, batch): features, mono_feature_vectors, stereo_feature_vectors = [], [], [] true_shower_primary_class = [] log_true_energy = [] - spherical_offset_az, spherical_offset_alt, angular_separation = [], [], [] + sky_offset_lon, sky_offset_lat, sky_angular_separation = [], [], [] + cam_coord_offset_x, cam_coord_offset_y, cam_coord_distance = [], [], [] for group_element in batch_grouped.groups: if "features" in batch.colnames: if self.sort_by_intensity: @@ -252,12 +262,18 @@ def _get_stereo_item(self, batch): ) if "energy" in self.tasks: log_true_energy.append(group_element["log_true_energy"].data[0]) - if "direction" in self.tasks: - spherical_offset_az.append(group_element["spherical_offset_az"].data[0]) - spherical_offset_alt.append( - group_element["spherical_offset_alt"].data[0] + if "skydirection" in self.tasks: + sky_offset_lon.append(group_element["sky_offset_lon"].data[0]) + sky_offset_lat.append( + group_element["sky_offset_lat"].data[0] + ) + sky_angular_separation.append(group_element["sky_angular_separation"].data[0]) + if "cameradirection" in self.tasks: + cam_coord_offset_x.append(group_element["cam_coord_offset_x"].data) + cam_coord_offset_y.append( + group_element["cam_coord_offset_y"].data ) - angular_separation.append(group_element["angular_separation"].data[0]) + cam_coord_distance.append(group_element["cam_coord_distance"].data) # Store the labels in the labels dictionary if "type" in self.tasks: labels["type"] = to_categorical( @@ -273,12 +289,21 @@ def _get_stereo_item(self, batch): ) if "energy" in self.tasks: labels["energy"] = np.array(log_true_energy) - if "direction" in self.tasks: - labels["direction"] = np.stack( + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( + ( + np.array(sky_offset_lon), + np.array(sky_offset_lat), + np.array(sky_angular_separation), + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( ( - np.array(spherical_offset_az), - np.array(spherical_offset_alt), - np.array(angular_separation), + np.array(cam_coord_offset_x), + np.array(cam_coord_offset_y), + np.array(cam_coord_distance), ), axis=1, ) diff --git a/ctlearn/core/model.py b/ctlearn/core/model.py index 80ff066..015f27e 100644 --- a/ctlearn/core/model.py +++ b/ctlearn/core/model.py @@ -77,11 +77,11 @@ class CTLearnModel(Component): ).tag(config=True) head = Dict( - default_value={'type': [512, 256, 2], 'energy': [512, 256, 1], 'direction': [512, 256, 3]}, + default_value={'type': [512, 256, 2], 'energy': [512, 256, 1], 'cameradirection': [512, 256, 3], 'skydirection': [512, 256, 3]}, allow_none=False, help=( "Dictionary containing the number of neurons in the fully connected head for each " - "task ('type', 'energy', 'direction'). Note: The number of neurons in the last layer " + "task ('type', 'energy', 'cameradirection', 'skydirection'). Note: The number of neurons in the last layer " "must match the number of classes or the number of reconstructed values." ), ).tag(config=True) diff --git a/ctlearn/tools/train_model.py b/ctlearn/tools/train_model.py index c4e17b5..ab79a01 100644 --- a/ctlearn/tools/train_model.py +++ b/ctlearn/tools/train_model.py @@ -40,7 +40,8 @@ class TrainCTLearnModel(Tool): ``~dl1_data_handler.loader.DLDataLoader``. The tool supports the following reconstruction tasks: - Classification of the primary particle type (gamma/proton) - Regression of the primary particle energy - - Regression of the primary particle arrival direction + - Regression of the primary particle arrival direction based on the offsets in camera coordinates + - Regression of the primary particle arrival direction based on the offsets in sky coordinates """ name = "ctlearn-train-model" @@ -65,14 +66,24 @@ class TrainCTLearnModel(Tool): --pattern-signal "gamma_*_run10.dl1.h5" \\ --output /path/to/your/energy/ \\ --reco energy \\ - - To train a CTLearn model for the regression of the primary particle arrival direction: + + To train a CTLearn model for the regression of the primary particle + arrival direction based on the offsets in camera coordinates: > ctlearn-train-model \\ --signal /path/to/your/gammas_dl1_dir/ \\ --pattern-signal "gamma_*_run1.dl1.h5" \\ --pattern-signal "gamma_*_run10.dl1.h5" \\ --output /path/to/your/direction/ \\ - --reco direction \\ + --reco cameradirection \\ + + To train a CTLearn model for the regression of the primary particle + arrival direction based on the offsets in sky coordinates: + > ctlearn-train-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco skydirection \\ """ input_dir_signal = Path( @@ -140,13 +151,14 @@ class TrainCTLearnModel(Tool): ).tag(config=True) reco_tasks = List( - trait=CaselessStrEnum(["type", "energy", "direction"]), + trait=CaselessStrEnum(["type", "energy", "cameradirection", "skydirection"]), allow_none=False, help=( "List of reconstruction tasks to perform. " "'type': classification of the primary particle type " "'energy': regression of the primary particle energy " - "'direction': regression of the primary particle arrival direction " + "'cameradirection': regression of the primary particle arrival direction in camera coordinates " + "'skydirection': regression of the primary particle arrival direction in sky coordinates" ) ).tag(config=True) @@ -501,11 +513,16 @@ def _get_losses_and_mertics(self, tasks): reduction="sum_over_batch_size" ) metrics["energy"] = keras.metrics.MeanAbsoluteError(name="mae_energy") - if "direction" in self.reco_tasks: - losses["direction"] = keras.losses.MeanAbsoluteError( + if "cameradirection" in self.reco_tasks: + losses["cameradirection"] = keras.losses.MeanAbsoluteError( + reduction="sum_over_batch_size" + ) + metrics["cameradirection"] = keras.metrics.MeanAbsoluteError(name="mae_cameradirection") + if "skydirection" in self.reco_tasks: + losses["skydirection"] = keras.losses.MeanAbsoluteError( reduction="sum_over_batch_size" ) - metrics["direction"] = keras.metrics.MeanAbsoluteError(name="mae_direction") + metrics["skydirection"] = keras.metrics.MeanAbsoluteError(name="mae_skydirection") return losses, metrics