Merge pull request #519 from PyAutoLabs/feature/numpy-deflections-p3

perf: rotation-matrix grid transform + one Grid2D per VectorYX2D (numpy deflections phase 3)

Author: James Nightingale

autoarray/geometry/geometry_util.py

@@ -519,15 +519,20 @@ def transform_grid_2d_to_reference_frame(
 
     shifted_grid_2d = grid_2d - xp.array(centre)
 
-    radius = xp.sqrt(xp.sum(xp.square(shifted_grid_2d), axis=1))
-    theta_coordinate_to_profile = xp.arctan2(
-        shifted_grid_2d[:, 0], shifted_grid_2d[:, 1]
-    ) - xp.radians(angle)
+    # A clockwise rotation by `angle` is the rotation matrix applied to the shifted (y, x)
+    # coordinates, built from one scalar cos / sin of the angle. This is algebraically
+    # identical to the polar form (r, theta - angle) -> (r sin, r cos) but costs two
+    # multiply-adds per coordinate instead of a per-pixel sqrt, arctan2, sin and cos.
+    cos_angle = xp.cos(xp.radians(angle))
+    sin_angle = xp.sin(xp.radians(angle))
+
+    y = shifted_grid_2d[:, 0]
+    x = shifted_grid_2d[:, 1]
 
     return xp.vstack(
         [
-            radius * xp.sin(theta_coordinate_to_profile),
-            radius * xp.cos(theta_coordinate_to_profile),
+            y * cos_angle - x * sin_angle,
+            x * cos_angle + y * sin_angle,
         ]
     ).T
 

autoarray/structures/vectors/uniform.py

@@ -140,11 +140,19 @@ def __init__(
             grid_2d=values, mask_2d=mask, store_native=store_native
         )
 
-        grid = grid_2d_util.convert_grid_2d(
-            grid_2d=grid, mask_2d=mask, store_native=store_native
-        )
+        # A `Grid2D` already paired with this mask is the grid the vectors are located on: the
+        # `@to_vector_yx` decorator hands the input grid straight through, so re-converting its
+        # values and rebuilding a second `Grid2D` on every profile evaluation would only pay the
+        # constructor twice for the same object.
+        if isinstance(grid, Grid2D) and grid.mask is mask and not store_native:
+            self.grid = grid
+        else:
+            grid = grid_2d_util.convert_grid_2d(
+                grid_2d=grid, mask_2d=mask, store_native=store_native
+            )
+
+            self.grid = Grid2D(values=grid, mask=mask)
 
-        self.grid = Grid2D(values=grid, mask=mask)
         self.mask = mask
 
         super().__init__(values)