fix flux fit by pasing .,array

Author: Jammy2211

autolens/point/fit/fluxes.py

@@ -128,5 +128,5 @@ def chi_squared(self) -> float:
         RMS noise-map values squared.
         """
         return ag.util.fit.chi_squared_from(
-            chi_squared_map=self.chi_squared_map,
+            chi_squared_map=self.chi_squared_map.array,
         )

autolens/point/solver/point_solver.py

@@ -4,7 +4,7 @@
 import autoarray as aa
 from autoarray.structures.triangles.shape import Point
 
-from autofit.jax_wrapper import jit, register_pytree_node_class
+from autofit.jax_wrapper import register_pytree_node_class
 from autogalaxy import OperateDeflections
 from .shape_solver import AbstractSolver
 
@@ -14,7 +14,7 @@
 
 @register_pytree_node_class
 class PointSolver(AbstractSolver):
-    @jit
+
     def solve(
         self,
         tracer: OperateDeflections,
@@ -55,21 +55,3 @@ def solve(
         )
 
         return aa.Grid2DIrregular([pair for pair in filtered_means])
-
-        # filtered_means = [
-        #     pair for pair in filtered_means if not np.any(np.isnan(pair)).all()
-        # ]
-        #
-        # difference = len(kept_triangles.means) - len(filtered_means)
-        # if difference > 0:
-        #     logger.debug(
-        #         f"Filtered one multiple-image with magnification below threshold."
-        #     )
-        # elif difference > 1:
-        #     logger.warning(
-        #         f"Filtered {difference} multiple-images with magnification below threshold."
-        #     )
-        #
-        # return aa.Grid2DIrregular(
-        #     [pair for pair in filtered_means if not np.isnan(pair).all()]
-        # )

autolens/point/solver/shape_solver.py

@@ -1,5 +1,4 @@
 import jax.numpy as jnp
-from jax import jit
 import logging
 import math
 
@@ -208,7 +207,6 @@ def _source_plane_grid(
         # noinspection PyTypeChecker
         return grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections)
 
-    @jit
     def solve_triangles(
         self,
         tracer: OperateDeflections,
@@ -270,7 +268,7 @@ def _filter_low_magnification(
         """
         points = jnp.array(points)
         magnifications = tracer.magnification_2d_via_hessian_from(
-            grid=aa.Grid2DIrregular(points),
+            grid=aa.Grid2DIrregular(points).array,
             buffer=self.scale,
         )
         mask = jnp.abs(magnifications.array) > self.magnification_threshold