remove infinities from point solver solution

Author: Jammy2211

autolens/point/fit/positions/image/abstract.py

@@ -104,4 +104,5 @@ def model_data(self) -> aa.Grid2DIrregular:
             tracer=self.tracer,
             source_plane_coordinate=self.source_plane_coordinate,
             plane_redshift=self.plane_redshift,
+            remove_infinities=False
         )

autolens/point/mock/mock_solver.py

@@ -10,5 +10,6 @@ def solve(
         tracer,
         source_plane_coordinate,
         plane_redshift: Optional[float] = None,
+        remove_infinities : bool = True
     ):
         return self.model_positions

autolens/point/solver/point_solver.py

@@ -18,6 +18,7 @@ def solve(
         tracer: OperateDeflections,
         source_plane_coordinate: Tuple[float, float],
         plane_redshift: Optional[float] = None,
+        remove_infinities: bool = True,
     ) -> aa.Grid2DIrregular:
         """
         Solve for the image plane coordinates that are traced to the source plane coordinate.
@@ -26,9 +27,15 @@ def solve(
         within the triangle. The triangles are sub-sampled to increase the resolution with only the triangles that
         contain the source plane coordinate and their neighbours being kept.
 
-        The means of the triangles  are then filtered to keep only those with an absolute magnification above the
+        The means of the triangles are then filtered to keep only those with an absolute magnification above the
         threshold.
 
+        The positions are stored on an array of fixed shape defined by `MAX_CONTAINING_SIZE`. This ensures the
+        array is static, which is important for JAX compatibility. This array typically has many entries
+        which use the sentinel value of `inf`, subsequent JAX calculations incorporated. By default, these
+        sentinel values are removed from the output, for example general use outside of JAX when simulating
+        strong lenses.
+
         Parameters
         ----------
         source_plane_coordinate
@@ -59,4 +66,8 @@ def solve(
         sentinel = self._xp.full_like(solution[0], fill_value=self._xp.inf)
         solution = self._xp.where(is_nan[:, None], sentinel, solution)
 
+        if remove_infinities:
+
+            solution = solution[~self._xp.isinf(solution).any(axis=1)]
+
         return aa.Grid2DIrregular(solution)

test_autolens/point/triangles/test_solver.py

@@ -80,6 +80,16 @@ def test_real_example_jax(grid, tracer):
     result = jax_solver.solve(
         tracer=tracer,
         source_plane_coordinate=(0.07, 0.07),
+        remove_infinities=True
+    )
+
+    assert len(result) == 5
+
+
+    result = jax_solver.solve(
+        tracer=tracer,
+        source_plane_coordinate=(0.07, 0.07),
+        remove_infinities=False
     )
 
     assert len(result) == 15