add camera and sky offset reco

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,
             )

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)

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