speed up JAX compiltation a lot via cleveer trick

Author: Jammy2211

autoarray/geometry/geometry_util.py

@@ -374,17 +374,20 @@ def transform_grid_2d_to_reference_frame(
     grid
         The 2d grid of (y, x) coordinates which are transformed to a new reference frame.
     """
-    shifted_grid_2d = grid_2d - jnp.array(centre)
+    grid_2d = jnp.asarray(grid_2d)
+    centre = jnp.asarray(centre)
 
-    radius = jnp.sqrt(jnp.sum(jnp.square(shifted_grid_2d), axis=1))
-    theta_coordinate_to_profile = jnp.arctan2(
-        shifted_grid_2d[:, 0], shifted_grid_2d[:, 1]
-    ) - jnp.radians(angle)
+    # Inject a tiny dynamic dependency on `angle` to prevent heavy constant folding
+    # (adds zero at runtime; negligible cost)
+    dynamic_zero = jnp.zeros_like(grid_2d) * angle
+    shifted = (grid_2d + dynamic_zero) - centre
 
+    radius = jnp.sqrt(jnp.sum(shifted * shifted, axis=1))
+    theta  = jnp.arctan2(shifted[:, 0], shifted[:, 1]) - jnp.deg2rad(angle)
     return jnp.vstack(
         [
-            radius * jnp.sin(theta_coordinate_to_profile),
-            radius * jnp.cos(theta_coordinate_to_profile),
+            radius * jnp.sin(theta),
+            radius * jnp.cos(theta),
         ]
     ).T