Merge pull request #414 from PyAutoLabs/feature/point-source-chi-squared-variants

feat: thread xp through Grid2DIrregular.grid_2d_via_deflection_grid_from

Author: Jammy2211

autoarray/structures/grids/irregular_2d.py

@@ -168,7 +168,7 @@ def extent_with_buffer_from(self, buffer: float = 1.0e-8) -> List[float]:
         ]
 
     def grid_2d_via_deflection_grid_from(
-        self, deflection_grid: np.ndarray
+        self, deflection_grid: np.ndarray, xp=None
     ) -> "Grid2DIrregular":
         """
         Returns a new Grid2DIrregular from this grid coordinates, where the (y,x) coordinates of this grid have a
@@ -181,8 +181,20 @@ def grid_2d_via_deflection_grid_from(
         ----------
         deflection_grid
             The grid of (y,x) coordinates which is subtracted from this grid.
-        """
-        return Grid2DIrregular(values=self - deflection_grid, xp=self._xp)
+        xp
+            The array module (``numpy`` or ``jax.numpy``) used to construct the returned grid, mirroring
+            ``subtracted_from`` / ``subtracted_and_rotated_from``. Defaults to ``self._xp`` (the pre-existing
+            behaviour); JIT-traced call sites (where ``self`` may not carry a reliable ``use_jax`` flag) should
+            pass it explicitly so they do not silently fall back to NumPy.
+        """
+        if xp is None:
+            xp = self._xp
+        deflections = (
+            deflection_grid.array
+            if hasattr(deflection_grid, "array")
+            else deflection_grid
+        )
+        return Grid2DIrregular(values=self.array - deflections, xp=xp)
 
     def squared_distances_to_coordinate_from(
         self, coordinate: Tuple[float, float] = (0.0, 0.0)