diff --git a/ctapipe/coordinates/ground_frames.py b/ctapipe/coordinates/ground_frames.py index 849fc7fad55..10b7c1896f4 100644 --- a/ctapipe/coordinates/ground_frames.py +++ b/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