Merge pull request #287 from PyAutoLabs/feature/grid-irregular-xp-propagation

Propagate xp through Grid2DIrregular.grid_2d_via_deflection_grid_from

Author: Jammy2211

autoarray/structures/grids/irregular_2d.py

@@ -182,7 +182,7 @@ 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)
+        return Grid2DIrregular(values=self - deflection_grid, xp=self._xp)
 
     def squared_distances_to_coordinate_from(
         self, coordinate: Tuple[float, float] = (0.0, 0.0)

test_autoarray/structures/grids/test_irregular_2d.py

@@ -75,6 +75,25 @@ def test__grid_2d_via_deflection_grid_from():
     assert grid.in_list == [(0.0, 1.0), (1.0, 1.0)]
 
 
+def test__grid_2d_via_deflection_grid_from__propagates_xp():
+    # numpy-backed receiver -> numpy-backed result
+    grid_np = aa.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 2.0)])
+    result_np = grid_np.grid_2d_via_deflection_grid_from(
+        deflection_grid=np.array([[1.0, 0.0], [1.0, 1.0]])
+    )
+    assert result_np._xp is np
+
+    # jax-backed receiver -> jax-backed result (so downstream .square calls use jnp)
+    jnp = pytest.importorskip("jax.numpy")
+    grid_jax = aa.Grid2DIrregular(
+        values=jnp.array([[1.0, 1.0], [2.0, 2.0]]), xp=jnp
+    )
+    result_jax = grid_jax.grid_2d_via_deflection_grid_from(
+        deflection_grid=jnp.array([[1.0, 0.0], [1.0, 1.0]])
+    )
+    assert result_jax._xp is jnp
+
+
 def test__furthest_distances_to_other_coordinates():
     grid = aa.Grid2DIrregular(values=[(0.0, 0.0), (0.0, 1.0)])