Merge pull request #360 from Jammy2211/feature/jax_interferometer_linear

Feature/jax interferometer linear

Author: Jammy2211

autolens/analysis/analysis/lens.py

@@ -127,8 +127,10 @@ def log_likelihood_penalty_from(
 
             try:
                 for positions_likelihood in self.positions_likelihood_list:
-                    log_likelihood_penalty = positions_likelihood.log_likelihood_penalty_from(
-                        instance=instance, analysis=self
+                    log_likelihood_penalty = (
+                        positions_likelihood.log_likelihood_penalty_from(
+                            instance=instance, analysis=self
+                        )
                     )
 
                     log_likelihood_penalty += log_likelihood_penalty
@@ -137,4 +139,4 @@ def log_likelihood_penalty_from(
             except (ValueError, np.linalg.LinAlgError) as e:
                 raise exc.FitException from e
 
-        return log_likelihood_penalty
+        return log_likelihood_penalty

autolens/analysis/positions.py

@@ -136,7 +136,9 @@ def output_positions_info(
             )
             f.write("")
 
-    def log_likelihood_penalty_from(self, instance: af.ModelInstance, analysis: AnalysisDataset) -> jnp.array:
+    def log_likelihood_penalty_from(
+        self, instance: af.ModelInstance, analysis: AnalysisDataset
+    ) -> jnp.array:
         """
         Returns a log-likelihood penalty used to constrain lens models where multiple image-plane
         positions do not trace to within a threshold distance of one another in the source-plane.
@@ -174,7 +176,7 @@ def log_likelihood_penalty_from(self, instance: af.ModelInstance, analysis: Anal
         tracer = analysis.tracer_via_instance_from(instance=instance)
 
         if not tracer.has(cls=ag.mp.MassProfile) or len(tracer.planes) == 1:
-            return jnp.array(0.0),
+            return (jnp.array(0.0),)
 
         positions_fit = SourceMaxSeparation(
             data=self.positions,
@@ -183,11 +185,11 @@ def log_likelihood_penalty_from(self, instance: af.ModelInstance, analysis: Anal
             plane_redshift=self.plane_redshift,
         )
 
-        max_separation = jnp.max(positions_fit.furthest_separations_of_plane_positions.array)
+        max_separation = jnp.max(
+            positions_fit.furthest_separations_of_plane_positions.array
+        )
 
-        penalty = self.log_likelihood_penalty_factor * (
-                max_separation - self.threshold
-            )
+        penalty = self.log_likelihood_penalty_factor * (max_separation - self.threshold)
 
         return jax.lax.cond(
             max_separation > self.threshold,

autolens/analysis/result.py

@@ -310,9 +310,7 @@ def positions_likelihood_from(
 
         mask = np.isfinite(positions.array).all(axis=1)
 
-        positions = aa.Grid2DIrregular(
-            positions[mask]
-        )
+        positions = aa.Grid2DIrregular(positions[mask])
 
         threshold = self.positions_threshold_from(
             factor=factor,

autolens/interferometer/model/analysis.py

@@ -152,14 +152,15 @@ def log_likelihood_function(self, instance):
         """
 
         try:
-            log_likelihood_penalty = self.log_likelihood_penalty_from(
-                instance=instance
-            )
+            log_likelihood_penalty = self.log_likelihood_penalty_from(instance=instance)
         except Exception as e:
             raise e
 
         try:
-            return self.fit_from(instance=instance).figure_of_merit + log_likelihood_penalty
+            return (
+                self.fit_from(instance=instance).figure_of_merit
+                + log_likelihood_penalty
+            )
         except (
             PixelizationException,
             exc.PixelizationException,
@@ -171,6 +172,8 @@ def log_likelihood_function(self, instance):
             np.linalg.LinAlgError,
             OverflowError,
         ) as e:
+            print(e)
+            fggdfg
             raise exc.FitException from e
 
     def fit_from(

autolens/mock.py

@@ -1,4 +1,6 @@
-from autofit.jax_wrapper import register_pytree_node_class, numpy as np
+import jax.numpy as jnp
+
+from autofit.jax_wrapper import register_pytree_node_class
 from autofit.mock import *  # noqa
 from autoarray.mock import *  # noqa
 from autogalaxy.mock import *  # noqa
@@ -17,15 +19,15 @@ def __init__(self):
         super().__init__([])
 
     def deflections_yx_2d_from(self, grid):
-        return np.zeros_like(grid.array)
+        return jnp.zeros_like(grid.array)
 
     def deflections_between_planes_from(self, grid, plane_i=0, plane_j=-1):
-        return np.zeros_like(grid.array)
+        return jnp.zeros_like(grid.array)
 
     def magnification_2d_via_hessian_from(
         self, grid, buffer: float = 0.01, deflections_func=None
     ) -> aa.ArrayIrregular:
-        return aa.ArrayIrregular(values=np.ones(grid.shape[0]))
+        return aa.ArrayIrregular(values=jnp.ones(grid.shape[0]))
 
     def tree_flatten(self):
         """

autolens/point/fit/fluxes.py

@@ -102,10 +102,12 @@ def model_data(self):
         are used.
         """
         return aa.ArrayIrregular(
-            values=jnp.array([
-                magnification * self.profile.flux
-                for magnification in self.magnifications_at_positions
-            ])
+            values=jnp.array(
+                [
+                    magnification * self.profile.flux
+                    for magnification in self.magnifications_at_positions
+                ]
+            )
         )
 
     @property

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

@@ -107,16 +107,18 @@ def all_permutations_log_likelihoods(self) -> np.ndarray:
             [
                 jnp.log(
                     jnp.sum(
-                        jnp.array([
-                            jnp.exp(
-                                self.log_p(
-                                    data_position,
-                                    model_position,
-                                    sigma,
+                        jnp.array(
+                            [
+                                jnp.exp(
+                                    self.log_p(
+                                        data_position,
+                                        model_position,
+                                        sigma,
+                                    )
                                 )
-                            )
-                            for model_position in model_data
-                        ])
+                                for model_position in model_data
+                            ]
+                        )
                     )
                 )
                 for data_position, sigma in zip(self.data, self.noise_map)

autolens/point/max_separation.py

@@ -42,7 +42,9 @@ def __init__(
         except TypeError:
             plane_index = -1
 
-        self.plane_positions = aa.Grid2DIrregular(values=tracer.traced_grid_2d_list_from(grid=data)[plane_index])
+        self.plane_positions = aa.Grid2DIrregular(
+            values=tracer.traced_grid_2d_list_from(grid=data)[plane_index]
+        )
 
     @property
     def furthest_separations_of_plane_positions(self) -> aa.ArrayIrregular:

autolens/point/solver/point_solver.py

@@ -62,4 +62,4 @@ def solve(
         sentinel = jnp.full_like(solution[0], fill_value=jnp.inf)
         solution = jnp.where(is_nan[:, None], sentinel, solution)
 
-        return aa.Grid2DIrregular(solution)
+        return aa.Grid2DIrregular(solution)

autolens/point/solver/shape_solver.py

@@ -2,14 +2,14 @@
 import logging
 import math
 
-from typing import Tuple, List, Iterator, Type, Optional
+from typing import Tuple, List, Iterator, Optional
 
 import autoarray as aa
 
 from autoarray.structures.triangles.shape import Shape
 from autofit.jax_wrapper import register_pytree_node_class
 
-from autoarray.structures.triangles.coordinate_array.jax_coordinate_array import (
+from autoarray.structures.triangles.coordinate_array import (
     CoordinateArrayTriangles,
 )
 from autoarray.structures.triangles.abstract import AbstractTriangles
@@ -59,7 +59,6 @@ def for_grid(
         grid: aa.Grid2D,
         pixel_scale_precision: float,
         magnification_threshold=0.1,
-        array_triangles_cls: Type[AbstractTriangles] = CoordinateArrayTriangles,
         neighbor_degree: int = 1,
     ):
         """
@@ -75,9 +74,6 @@ def for_grid(
             The precision to which the triangles should be subdivided.
         magnification_threshold
             The threshold for the magnification under which multiple images are filtered.
-        array_triangles_cls
-            The class to use for the triangles. JAX is used implicitly if USE_JAX=1 and
-            jax is installed.
         max_containing_size
             Only applies to JAX. This is the maximum number of multiple images expected.
             We need to know this in advance to allocate memory for the JAX array.
@@ -106,7 +102,6 @@ def for_grid(
             scale=scale,
             pixel_scale_precision=pixel_scale_precision,
             magnification_threshold=magnification_threshold,
-            array_triangles_cls=array_triangles_cls,
             neighbor_degree=neighbor_degree,
         )
 
@@ -120,7 +115,6 @@ def for_limits_and_scale(
         scale=0.1,
         pixel_scale_precision: float = 0.001,
         magnification_threshold=0.1,
-        array_triangles_cls: Type[AbstractTriangles] = CoordinateArrayTriangles,
         neighbor_degree: int = 1,
     ):
         """
@@ -141,17 +135,14 @@ def for_limits_and_scale(
             The precision to which the triangles should be subdivided.
         magnification_threshold
             The threshold for the magnification under which multiple images are filtered.
-        array_triangles_cls
-            The class to use for the triangles. JAX is used implicitly if USE_JAX=1 and
-            jax is installed.
         neighbor_degree
             The number of times recursively add neighbors for the triangles that contain
 
         Returns
         -------
         The solver.
         """
-        initial_triangles = array_triangles_cls.for_limits_and_scale(
+        initial_triangles = CoordinateArrayTriangles.for_limits_and_scale(
             y_min=y_min,
             y_max=y_max,
             x_min=x_min,
@@ -310,6 +301,7 @@ def steps(
         An iterator over the steps of the triangle solver algorithm.
         """
         initial_triangles = self.initial_triangles
+
         for number in range(self.n_steps):
             plane_triangles = self._plane_triangles(
                 tracer=tracer,