nicer rotations

ctapipe/coordinates/ground_frames.py

@@ -26,7 +26,7 @@
     frame_transform_graph,
 )
 from astropy.units.quantity import Quantity
-from numpy import cos, sin
+from scipy.spatial.transform import Rotation as R
 
 __all__ = [
     "GroundFrame",
@@ -108,20 +108,11 @@ def get_shower_trans_matrix(azimuth, altitude):
     -------
     trans: 3x3 ndarray transformation matrix
     """
+    zen = 90 * u.deg - altitude
 
-    cos_z = sin(altitude)
-    sin_z = cos(altitude)
-    cos_az = cos(azimuth)
-    sin_az = sin(azimuth)
-
-    trans = np.array(
-        [
-            [cos_z * cos_az, -cos_z * sin_az, -sin_z],
-            [sin_az, cos_az, np.zeros_like(sin_z)],
-            [sin_z * cos_az, -sin_z * sin_az, cos_z],
-        ],
-        dtype=np.float64,
-    )
+    trans = R.from_euler(
+        "zy", [azimuth.to_value(u.rad), -zen.to_value(u.rad)]
+    ).as_matrix()
 
     return trans
 
@@ -142,18 +133,16 @@ def ground_to_tilted(ground_coord, tilted_frame):
     -------
     SkyCoordinate transformed to `tilted_frame` coordinates
     """
-    x_grd, y_grd, z_grd = ground_coord.cartesian.xyz
+    xyz = ground_coord.cartesian.xyz
 
-    altitude = tilted_frame.pointing_direction.alt.to_value(u.rad)
-    azimuth = tilted_frame.pointing_direction.az.to_value(u.rad)
+    altitude = tilted_frame.pointing_direction.alt
+    azimuth = tilted_frame.pointing_direction.az
 
     trans = get_shower_trans_matrix(azimuth, altitude)
 
-    x_tilt = trans[0, 0] * x_grd + trans[0, 1] * y_grd + trans[0, 2] * z_grd
-    y_tilt = trans[1, 0] * x_grd + trans[1, 1] * y_grd + trans[1, 2] * z_grd
-    z_tilt = trans[2, 0] * x_grd + trans[2, 1] * y_grd + trans[2, 2] * z_grd
+    vec = np.einsum("...ij,j...->i...", trans, xyz)
 
-    representation = CartesianRepresentation(x_tilt, y_tilt, z_tilt)
+    representation = CartesianRepresentation(*vec)
 
     return tilted_frame.realize_frame(representation)
 
@@ -174,19 +163,19 @@ def tilted_to_ground(tilted_coord, ground_frame):
     -------
     GroundFrame coordinates
     """
-    x_tilt, y_tilt, z_tilt = tilted_coord.cartesian.xyz
+    xyz = tilted_coord.cartesian.xyz
 
-    altitude = tilted_coord.pointing_direction.alt.to(u.rad)
-    azimuth = tilted_coord.pointing_direction.az.to(u.rad)
+    altitude = tilted_coord.pointing_direction.alt
+    azimuth = tilted_coord.pointing_direction.az
 
     trans = get_shower_trans_matrix(azimuth, altitude)
+    trans_inverse = np.swapaxes(
+        trans, axis1=range(trans.ndim)[-2], axis2=range(trans.ndim)[-1]
+    )
 
-    x_grd = trans[0][0] * x_tilt + trans[1][0] * y_tilt + trans[2][0] * z_tilt
-    y_grd = trans[0][1] * x_tilt + trans[1][1] * y_tilt + trans[2][1] * z_tilt
-    z_grd = trans[0][2] * x_tilt + trans[1][2] * y_tilt + trans[2][2] * z_tilt
-
-    representation = CartesianRepresentation(x_grd, y_grd, z_grd)
+    vec = np.einsum("...ij,j...->i...", trans_inverse, xyz)
 
+    representation = CartesianRepresentation(*vec)
     grd = ground_frame.realize_frame(representation)
     return grd