Vectorize operations in HillasReconstructor

ctapipe/coordinates/__init__.py

@@ -1,27 +1,26 @@
 """
 Coordinates.
 """
+import warnings
 
-import numpy as np
-
+import astropy.units as u
 from astropy.coordinates import (
+    CIRS,
     AltAz,
     FunctionTransformWithFiniteDifference,
-    CIRS,
     frame_transform_graph,
-    spherical_to_cartesian,
 )
-import warnings
-from .telescope_frame import TelescopeFrame
-from .nominal_frame import NominalFrame
+from erfa.ufunc import s2p as spherical_to_cartesian
+
+from .camera_frame import CameraFrame, EngineeringCameraFrame
 from .ground_frames import (
+    EastingNorthingFrame,
     GroundFrame,
     TiltedGroundFrame,
     project_to_ground,
-    EastingNorthingFrame,
 )
-from .camera_frame import CameraFrame, EngineeringCameraFrame
-
+from .nominal_frame import NominalFrame
+from .telescope_frame import TelescopeFrame
 
 __all__ = [
     "TelescopeFrame",
@@ -67,7 +66,7 @@ def altaz_to_altaz(from_coo, to_frame):
         location = to_frame.location
 
     if obstime is None or location is None:
-        return to_frame.realize_frame(from_coo.spherical)
+        return to_frame.realize_frame(from_coo.data)
 
     return from_coo.transform_to(CIRS(obstime=obstime)).transform_to(to_frame)
 
@@ -85,4 +84,8 @@ def altaz_to_righthanded_cartesian(alt, az):
     az: u.Quantity
         azimuth
     """
-    return np.array(spherical_to_cartesian(r=1, lat=alt, lon=-az))
+    if hasattr(az, "unit"):
+        az = az.to_value(u.rad)
+        alt = alt.to_value(u.rad)
+
+    return spherical_to_cartesian(-az, alt, 1.0)

ctapipe/coordinates/ground_frames.py

@@ -14,18 +14,18 @@
 - Tests Tests Tests!
 """
 import astropy.units as u
-from astropy.units.quantity import Quantity
 import numpy as np
 from astropy.coordinates import (
+    AffineTransform,
     AltAz,
     BaseCoordinateFrame,
     CartesianRepresentation,
     CoordinateAttribute,
     FunctionTransform,
     RepresentationMapping,
     frame_transform_graph,
-    AffineTransform,
 )
+from astropy.units.quantity import Quantity
 from numpy import cos, sin
 
 __all__ = [
@@ -216,13 +216,18 @@ def project_to_ground(tilt_system):
     z_initial = ground_system.z.value
 
     trans = get_shower_trans_matrix(
-        tilt_system.pointing_direction.az, tilt_system.pointing_direction.alt
+        tilt_system.pointing_direction.az,
+        tilt_system.pointing_direction.alt,
     )
 
     x_projected = x_initial - trans[2][0] * z_initial / trans[2][2]
     y_projected = y_initial - trans[2][1] * z_initial / trans[2][2]
 
-    return GroundFrame(x=x_projected * unit, y=y_projected * unit, z=0 * unit)
+    return GroundFrame(
+        x=u.Quantity(x_projected, unit),
+        y=u.Quantity(y_projected, unit),
+        z=u.Quantity(0, unit),
+    )
 
 
 @frame_transform_graph.transform(FunctionTransform, GroundFrame, GroundFrame)

ctapipe/coordinates/nominal_frame.py

@@ -6,23 +6,23 @@
 corresponding class.
 """
 import astropy.units as u
-from astropy.coordinates.matrix_utilities import (
-    rotation_matrix,
-    matrix_product,
-    matrix_transpose,
-)
 from astropy.coordinates import (
-    frame_transform_graph,
-    FunctionTransform,
-    DynamicMatrixTransform,
-    UnitSphericalRepresentation,
+    AltAz,
+    Angle,
     BaseCoordinateFrame,
     CoordinateAttribute,
-    TimeAttribute,
+    DynamicMatrixTransform,
     EarthLocationAttribute,
+    FunctionTransform,
     RepresentationMapping,
-    Angle,
-    AltAz,
+    TimeAttribute,
+    UnitSphericalRepresentation,
+    frame_transform_graph,
+)
+from astropy.coordinates.matrix_utilities import (
+    matrix_product,
+    matrix_transpose,
+    rotation_matrix,
 )
 
 __all__ = ["NominalFrame"]
@@ -73,31 +73,31 @@ def __init__(self, *args, **kwargs):
 
 
 @frame_transform_graph.transform(FunctionTransform, NominalFrame, NominalFrame)
-def skyoffset_to_skyoffset(from_telescope_coord, to_telescope_frame):
+def nominal_to_nominal(from_nominal_coord, to_nominal_frame):
     """Transform between two skyoffset frames."""
 
-    intermediate_from = from_telescope_coord.transform_to(from_telescope_coord.origin)
-    intermediate_to = intermediate_from.transform_to(to_telescope_frame.origin)
-    return intermediate_to.transform_to(to_telescope_frame)
+    intermediate_from = from_nominal_coord.transform_to(from_nominal_coord.origin)
+    intermediate_to = intermediate_from.transform_to(to_nominal_frame.origin)
+    return intermediate_to.transform_to(to_nominal_frame)
 
 
 @frame_transform_graph.transform(DynamicMatrixTransform, AltAz, NominalFrame)
-def reference_to_skyoffset(reference_frame, telescope_frame):
+def altaz_to_nominal(altaz_coord, nominal_frame):
     """Convert a reference coordinate to an sky offset frame."""
 
     # Define rotation matrices along the position angle vector, and
     # relative to the origin.
-    origin = telescope_frame.origin.spherical
+    origin = nominal_frame.origin.represent_as(UnitSphericalRepresentation)
     mat1 = rotation_matrix(-origin.lat, "y")
     mat2 = rotation_matrix(origin.lon, "z")
     return matrix_product(mat1, mat2)
 
 
 @frame_transform_graph.transform(DynamicMatrixTransform, NominalFrame, AltAz)
-def skyoffset_to_reference(skyoffset_coord, reference_frame):
+def nominal_to_altaz(nominal_coord, altaz_frame):
     """Convert an sky offset frame coordinate to the reference frame"""
 
     # use the forward transform, but just invert it
-    mat = reference_to_skyoffset(reference_frame, skyoffset_coord)
+    mat = altaz_to_nominal(altaz_frame, nominal_coord)
     # transpose is the inverse because mat is a rotation matrix
     return matrix_transpose(mat)

ctapipe/coordinates/telescope_frame.py

@@ -6,28 +6,31 @@
 corresponding class.
 """
 import astropy.units as u
-from astropy.coordinates.matrix_utilities import (
-    rotation_matrix,
-    matrix_product,
-    matrix_transpose,
-)
 from astropy.coordinates import (
-    frame_transform_graph,
-    FunctionTransform,
-    DynamicMatrixTransform,
+    AltAz,
+    Angle,
     BaseCoordinateFrame,
     CoordinateAttribute,
-    TimeAttribute,
+    DynamicMatrixTransform,
     EarthLocationAttribute,
+    FunctionTransform,
     RepresentationMapping,
+    TimeAttribute,
     UnitSphericalRepresentation,
-    Angle,
-    AltAz,
+    frame_transform_graph,
+)
+from astropy.coordinates.matrix_utilities import (
+    matrix_product,
+    matrix_transpose,
+    rotation_matrix,
 )
 
 __all__ = ["TelescopeFrame"]
 
 
+_wrap_angle = Angle(180, unit=u.deg)
+
+
 class TelescopeFrame(BaseCoordinateFrame):
     """
     Telescope coordinate frame.
@@ -73,11 +76,11 @@ def __init__(self, *args, **kwargs):
 
         # make sure telescope coordinate is in range [-180°, 180°]
         if isinstance(self._data, UnitSphericalRepresentation):
-            self._data.lon.wrap_angle = Angle(180, unit=u.deg)
+            self._data.lon.wrap_angle = _wrap_angle
 
 
 @frame_transform_graph.transform(FunctionTransform, TelescopeFrame, TelescopeFrame)
-def skyoffset_to_skyoffset(from_telescope_coord, to_telescope_frame):
+def telescope_to_telescope(from_telescope_coord, to_telescope_frame):
     """Transform between two skyoffset frames."""
 
     intermediate_from = from_telescope_coord.transform_to(
@@ -90,22 +93,24 @@ def skyoffset_to_skyoffset(from_telescope_coord, to_telescope_frame):
 
 
 @frame_transform_graph.transform(DynamicMatrixTransform, AltAz, TelescopeFrame)
-def reference_to_skyoffset(reference_frame, telescope_frame):
+def altaz_to_telescope(altaz_coord, telescope_frame):
     """Convert a reference coordinate to an sky offset frame."""
 
     # Define rotation matrices along the position angle vector, and
     # relative to the telescope_pointing.
-    telescope_pointing = telescope_frame.telescope_pointing.spherical
+    telescope_pointing = telescope_frame.telescope_pointing.represent_as(
+        UnitSphericalRepresentation
+    )
     mat1 = rotation_matrix(-telescope_pointing.lat, "y")
     mat2 = rotation_matrix(telescope_pointing.lon, "z")
     return matrix_product(mat1, mat2)
 
 
 @frame_transform_graph.transform(DynamicMatrixTransform, TelescopeFrame, AltAz)
-def skyoffset_to_reference(skyoffset_coord, reference_frame):
+def telescope_to_altaz(telescope_coord, altaz_frame):
     """Convert an sky offset frame coordinate to the reference frame"""
 
     # use the forward transform, but just invert it
-    mat = reference_to_skyoffset(reference_frame, skyoffset_coord)
+    mat = altaz_to_telescope(altaz_frame, telescope_coord)
     # transpose is the inverse because mat is a rotation matrix
     return matrix_transpose(mat)