Add predictions in the altaz frame

aict_tools/configuration.py

@@ -339,18 +339,25 @@ def parse_cta(self, model_config, aict_config):
             aict_config.delta_column,
         }
 
-        cols.update(model_config['features'])
+        cols.update(model_config["features"])
         if self.feature_generation:
             cols.update(self.feature_generation.needed_columns)
         # Add id's because we generate new tables instead of adding columns
         # and want these to be included
         # focal_length is necessary for coordinate transformations
-        cols.update(["tel_id", "event_id", "obs_id", "equivalent_focal_length"])
-        self.columns_to_read_apply = list(cols)
         cols.update(
             {
+                "tel_id",
+                "event_id",
+                "obs_id",
+                "equivalent_focal_length",
                 aict_config.pointing_az_column,
                 aict_config.pointing_alt_column,
+            }
+        )
+        self.columns_to_read_apply = list(cols)
+        cols.update(
+            {
                 aict_config.source_az_column,
                 aict_config.source_alt_column,
             }
@@ -372,23 +379,28 @@ def parse_simple(self, model_config, aict_config):
             aict_config.delta_column,
         }
 
-        cols.update(model_config['features'])
+        cols.update(model_config["features"])
         if self.feature_generation:
             cols.update(self.feature_generation.needed_columns)
-
-        self.columns_to_read_apply = list(cols)
         cols.update(
             {
                 aict_config.pointing_az_column,
                 aict_config.pointing_zd_column,
                 aict_config.pointing_alt_column,
+            }
+        )
+        cols.discard(None)
+        self.columns_to_read_apply = list(cols)
+
+        cols.update(
+            {
                 aict_config.source_az_column,
                 aict_config.source_zd_column,
                 aict_config.source_alt_column,
             }
         )
         cols.discard(None)
-        if aict_config.coordinate_transformation == 'CTA':
+        if aict_config.coordinate_transformation == "CTA":
             cols.add(aict_config.focal_length_column)
         cols.update(get_optional_training_columns(aict_config))
         self.columns_to_read_train = list(cols)
@@ -448,14 +460,23 @@ def parse_cta(self, model_config, aict_config):
         # Add id's because we generate new tables instead of adding columns
         # and want these to be included
         # focal_length is necessary for coordinate transformations
-        cols.update(['tel_id', 'event_id', 'obs_id', 'equivalent_focal_length'])
+        cols.update(
+            {
+                "tel_id",
+                "event_id",
+                "obs_id",
+                "equivalent_focal_length",
+                aict_config.pointing_az_column,
+                aict_config.pointing_alt_column,
+            }
+        )
         self.columns_to_read_apply = list(cols)
-        cols.update({
-            aict_config.pointing_az_column,
-            aict_config.pointing_alt_column,
-            aict_config.source_az_column,
-            aict_config.source_alt_column,
-        })
+        cols.update(
+            {
+                aict_config.source_az_column,
+                aict_config.source_alt_column,
+            }
+        )
 
         cols.update(get_optional_training_columns(aict_config))
         self.columns_to_read_train = list(cols)
@@ -476,17 +497,24 @@ def parse_simple(self, model_config, aict_config):
         cols.update(model_config["features"])
         if self.feature_generation:
             cols.update(self.feature_generation.needed_columns)
+        cols.update(
+            {
+                aict_config.pointing_az_column,
+                aict_config.pointing_zd_column,
+                aict_config.pointing_alt_column,
+            }
+        )
+        cols.discard(None)
         self.columns_to_read_apply = list(cols)
-        cols.update({
-            aict_config.pointing_az_column,
-            aict_config.pointing_zd_column,
-            aict_config.pointing_alt_column,
-            aict_config.source_az_column,
-            aict_config.source_zd_column,
-            aict_config.source_alt_column,
-        })
+        cols.update(
+            {
+                aict_config.source_az_column,
+                aict_config.source_zd_column,
+                aict_config.source_alt_column,
+            }
+        )
         cols.discard(None)
-        if aict_config.coordinate_transformation == 'CTA':
+        if aict_config.coordinate_transformation == "CTA":
             cols.add(aict_config.focal_length_column)
 
         cols.update(get_optional_training_columns(aict_config))

aict_tools/cta_helpers.py

@@ -8,7 +8,7 @@
     raise ImportError("This functionality requires ctapipe to be installed")
 
 
-def horizontal_to_camera_cta_simtel(alt, az, alt_pointing, az_pointing, focal_length):
+def horizontal_to_camera(alt, az, alt_pointing, az_pointing, focal_length):
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", MissingFrameAttributeWarning)
 
@@ -29,5 +29,26 @@ def horizontal_to_camera_cta_simtel(alt, az, alt_pointing, az_pointing, focal_le
             telescope_pointing=tel_pointing,
         )
 
-        cam_coords = source_altaz.transform_to(camera_frame)
-        return cam_coords.x.to_value(u.m), cam_coords.y.to_value(u.m)
+        cam_coordinates = source_altaz.transform_to(camera_frame)
+        return cam_coordinates.x.to_value(u.m), cam_coordinates.y.to_value(u.m)
+
+
+def camera_to_horizontal(x, y, alt_pointing, az_pointing, focal_length):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", MissingFrameAttributeWarning)
+
+        altaz = AltAz()
+        tel_pointing = SkyCoord(
+            alt=u.Quantity(alt_pointing, u.deg, copy=False),
+            az=u.Quantity(az_pointing, u.deg, copy=False),
+            frame=altaz,
+        )
+        camera_coordinates = CameraFrame(
+            x=u.Quantity(x, u.m, copy=False),
+            y=u.Quantity(y, u.m, copy=False),
+            focal_length=u.Quantity(focal_length, u.m, copy=False),
+            telescope_pointing=tel_pointing,
+        )
+
+        horizontal_coordinates = camera_coordinates.transform_to(altaz)
+        return horizontal_coordinates.alt.to_value(u.deg), horizontal_coordinates.az.to_value(u.deg)

aict_tools/preprocessing.py

@@ -3,7 +3,8 @@
 import astropy.units as u
 from astropy.table import Table
 from fact.coordinates.utils import horizontal_to_camera as horizontal_to_camera_fact
-
+from fact.coordinates.utils import camera_to_horizontal as camera_to_horizontal_fact
+from fact.coordinates.utils import camera_to_equatorial as camera_to_equatorial_fact
 from .configuration import AICTConfig
 
 
@@ -127,12 +128,7 @@ def get_alt(df, config: AICTConfig):
         source_alt = 90 - df[config.source_zd_column]
     else:
         source_alt = df[config.source_alt_column]
-
-    if config.pointing_zd_column:
-        pointing_alt = 90 - df[config.pointing_zd_column]
-    else:
-        pointing_alt = df[config.pointing_alt_column]
-
+    pointing_alt = get_alt_pointing(df, config)
     return source_alt, pointing_alt
 
 
@@ -146,20 +142,32 @@ def get_zd(df, config):
     else:
         source_zd = df[config.source_zd_column]
 
+    pointing_zd = get_zd_pointing(df, config)
+    return source_zd, pointing_zd
+
+
+def get_alt_pointing(df, config):
+    if config.pointing_zd_column:
+        pointing_alt = 90 - df[config.pointing_zd_column]
+    else:
+        pointing_alt = df[config.pointing_alt_column]
+    return pointing_alt
+
+
+def get_zd_pointing(df, config):
     if config.pointing_alt_column:
         pointing_zd = 90 - df[config.pointing_alt_column]
     else:
         pointing_zd = df[config.pointing_zd_column]
-
-    return source_zd, pointing_zd
+    return pointing_zd
 
 
 def horizontal_to_camera(df, config):
     if config.coordinate_transformation == "CTA":
-        from .cta_helpers import horizontal_to_camera_cta_simtel
+        from .cta_helpers import horizontal_to_camera as horizontal_to_camera_cta
 
         alt_source, alt_pointing = get_alt(df, config)
-        source_x, source_y = horizontal_to_camera_cta_simtel(
+        source_x, source_y = horizontal_to_camera_cta(
             az=df[config.source_az_column],
             alt=alt_source,
             az_pointing=df[config.pointing_az_column],
@@ -180,6 +188,32 @@ def horizontal_to_camera(df, config):
     return source_x, source_y
 
 
+def camera_to_horizontal(df, config, source_x, source_y):
+    if config.coordinate_transformation == "CTA":
+        from .cta_helpers import camera_to_horizontal as camera_to_horizontal_cta
+
+        alt_pointing = get_alt_pointing(df, config)
+        source_alt, source_az = camera_to_horizontal_cta(
+            x=source_x,
+            y=source_y,
+            az_pointing=df[config.pointing_az_column],
+            alt_pointing=alt_pointing,
+            focal_length=df[config.focal_length_column],
+        )
+        return source_alt, source_az
+    elif config.coordinate_transformation == "FACT":
+        zd_pointing = get_zd_pointing(df, config)
+        source_zd, source_az = camera_to_horizontal_fact(
+            x=source_x,
+            y=source_y,
+            az_pointing=df[config.pointing_az_column],
+            zd_pointing=zd_pointing,
+        )
+        return source_zd, source_az
+    else:
+        raise ValueError("Unsupported value for coordinate_transformation")
+
+
 def delta_error(data_df, model_config):
     df = data_df.copy()
     source_x, source_y = horizontal_to_camera(df, model_config)

aict_tools/scripts/apply_disp_regressor.py

@@ -13,7 +13,7 @@
 from ..apply import predict_disp
 from ..configuration import AICTConfig
 from ..logging import setup_logging
-from ..preprocessing import convert_units
+from ..preprocessing import camera_to_horizontal
 
 
 @click.command()
@@ -124,11 +124,14 @@ def main(
         )
 
         if config.data_format == "CTA":
+            source_alt, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
             df_data.reset_index(inplace=True)
             for tel_id, group in df_data.groupby("tel_id"):
                 d = group[["obs_id", "event_id"]].copy()
                 d["source_y_prediction"] = source_y[group.index]
                 d["source_x_prediction"] = source_x[group.index]
+                d["source_alt_prediction"] = source_alt[group.index]
+                d["source_az_prediction"] = source_az[group.index]
                 d["disp_prediction"] = disp[group.index]
                 append_predictions_cta(
                     data_path,
@@ -141,6 +144,13 @@ def main(
             append_column_to_hdf5(data_path, source_x, key, "source_x_prediction")
             append_column_to_hdf5(data_path, source_y, key, "source_y_prediction")
             append_column_to_hdf5(data_path, disp, key, "disp_prediction")
+            if config.coordinate_transformation == "CTA":
+                source_alt, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
+                append_column_to_hdf5(data_path, source_alt, key, "source_alt_prediction")
+            elif config.coordinate_transformation == "FACT":
+                source_zd, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
+                append_column_to_hdf5(data_path, source_zd, key, "source_zd_prediction")
+            append_column_to_hdf5(data_path, source_az, key, "source_az_prediction")
 
 
 if __name__ == "__main__":

aict_tools/scripts/apply_dxdy_regressor.py

@@ -14,7 +14,7 @@
 from ..apply import predict_dxdy
 from ..configuration import AICTConfig
 from ..logging import setup_logging
-from ..preprocessing import calc_true_disp
+from ..preprocessing import calc_true_disp, camera_to_horizontal
 
 
 @click.command()
@@ -105,26 +105,37 @@ def main(
 
         source_x = df_data[config.cog_x_column] + dxdy[:, 0]
         source_y = df_data[config.cog_y_column] + dxdy[:, 1]
+
         if config.data_format == "CTA":
+            source_alt, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
             df_data.reset_index(inplace=True)
             for tel_id, group in df_data.groupby("tel_id"):
                 d = group[["obs_id", "event_id"]].copy()
                 d["source_y_prediction"] = source_y[group.index]
                 d["source_x_prediction"] = source_x[group.index]
                 d["dx_prediction"] = dxdy[:, 0][group.index]
                 d["dy_prediction"] = dxdy[:, 1][group.index]
+                d["source_alt_prediction"] = source_alt[group.index]
+                d["source_az_prediction"] = source_az[group.index]
                 append_predictions_cta(
                     data_path,
                     d,
                     f"/dl2/event/telescope/{model_config.output_name}/tel_{tel_id:03d}",
                 )
-
         elif config.data_format == "simple":
             key = config.events_key
             append_column_to_hdf5(data_path, source_x, key, "source_x_prediction")
             append_column_to_hdf5(data_path, source_y, key, "source_y_prediction")
             append_column_to_hdf5(data_path, dxdy[:, 0], key, "dx_prediction")
             append_column_to_hdf5(data_path, dxdy[:, 1], key, "dy_prediction")
+            if config.coordinate_transformation == "CTA":
+                source_alt, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
+                append_column_to_hdf5(data_path, source_alt, key, "source_alt_prediction")
+            elif config.coordinate_transformation == "FACT":
+                source_zd, source_az = camera_to_horizontal(df_data, config, source_x, source_y)
+                append_column_to_hdf5(data_path, source_zd, key, "source_zd_prediction")
+            append_column_to_hdf5(data_path, source_az, key, "source_az_prediction")
+
 
 if __name__ == "__main__":
     # pylint: disable=no-value-for-parameter

tests/test_cta_coords.py

@@ -1,14 +1,14 @@
 from numpy.testing import assert_allclose
 
 
-def test_horizontal_to_camera():
+def test_horizontal_to_camera_cta():
     from aict_tools.io import read_data
-    from aict_tools.cta_helpers import horizontal_to_camera_cta_simtel
+    from aict_tools.cta_helpers import horizontal_to_camera
 
     df = read_data("tests/cta_coord_test.hdf", "telescope_events")
     expected_x = df.x
     expected_y = df.y
-    transformed_x, transformed_y = horizontal_to_camera_cta_simtel(
+    transformed_x, transformed_y = horizontal_to_camera(
         alt=90 - df.zd,
         az=df.az,
         alt_pointing=90 - df.zd_pointing,
@@ -17,3 +17,21 @@ def test_horizontal_to_camera():
     )
     assert_allclose(expected_x, transformed_x)
     assert_allclose(expected_y, transformed_y)
+
+
+def test_camera_to_horizontal():
+    from aict_tools.io import read_data
+    from aict_tools.cta_helpers import camera_to_horizontal
+
+    df = read_data("tests/cta_coord_test.hdf", "telescope_events")
+    expected_alt = df.alt
+    expected_az = df.az
+    transformed_alt, transformed_az = camera_to_horizontal(
+        x=df.x,
+        y=df.y,
+        alt_pointing=90 - df.zd_pointing,
+        az_pointing=df.az_pointing,
+        focal_length=df.focal_length,
+    )
+    assert_allclose(expected_alt, transformed_alt)
+    assert_allclose(expected_az, transformed_az)