black

Author: Jammy2211

autolens/aggregator/__init__.py

@@ -27,6 +27,8 @@
 from autolens.aggregator.subplot import SubplotFitX1Plane as subplot_fit_x1_plane
 from autolens.aggregator.subplot import SubplotFit as subplot_fit
 from autolens.aggregator.subplot import SubplotFitLog10 as subplot_fit_log10
-from autolens.aggregator.subplot import FITSModelGalaxyImages as fits_model_galaxy_images
+from autolens.aggregator.subplot import (
+    FITSModelGalaxyImages as fits_model_galaxy_images,
+)
 from autolens.aggregator.subplot import FITSTracer as fits_tracer
 from autolens.aggregator.subplot import FITSFit as fits_fit

autolens/analysis/positions.py

@@ -199,4 +199,3 @@ def log_likelihood_penalty_from(
             )
 
         return penalty if max_separation > self.threshold else np.array(0.0)
-

autolens/mock.py

@@ -23,6 +23,10 @@ def deflections_between_planes_from(self, grid, xp=np, plane_i=0, plane_j=-1):
         return xp.zeros_like(grid.array)
 
     def magnification_2d_via_hessian_from(
-        self, grid, buffer: float = 0.01, deflections_func=None, xp=np,
+        self,
+        grid,
+        buffer: float = 0.01,
+        deflections_func=None,
+        xp=np,
     ) -> aa.ArrayIrregular:
         return aa.ArrayIrregular(values=xp.ones(grid.shape[0]))

autolens/point/fit/abstract.py

@@ -133,9 +133,7 @@ def magnifications_at_positions(self) -> aa.ArrayIrregular:
         """
         return abs(
             self.tracer.magnification_2d_via_hessian_from(
-                grid=self.positions,
-                deflections_func=self.deflections_func,
-                xp=self._xp
+                grid=self.positions, deflections_func=self.deflections_func, xp=self._xp
             )
         )
 

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

@@ -104,5 +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
+            remove_infinities=False,
         )

autolens/point/mock/mock_solver.py

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

autolens/point/solver/point_solver.py

@@ -60,7 +60,9 @@ def solve(
             tracer=tracer, points=kept_triangles.means
         )
 
-        solution = aa.Grid2DIrregular([pair for pair in filtered_means], xp=self._xp).array
+        solution = aa.Grid2DIrregular(
+            [pair for pair in filtered_means], xp=self._xp
+        ).array
 
         is_nan = self._xp.isnan(solution).any(axis=1)
         sentinel = self._xp.full_like(solution[0], fill_value=self._xp.inf)

autolens/point/solver/shape_solver.py

@@ -23,7 +23,7 @@ def __init__(
         pixel_scale_precision: float,
         magnification_threshold=0.1,
         neighbor_degree: int = 1,
-        xp = np
+        xp=np,
     ):
         """
         Determine the image plane coordinates that are traced to be a source plane coordinate.
@@ -144,9 +144,13 @@ def for_limits_and_scale(
         """
 
         if xp.__name__.startswith("jax"):
-            from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles as triangle_cls
+            from autoarray.structures.triangles.coordinate_array import (
+                CoordinateArrayTriangles as triangle_cls,
+            )
         else:
-            from autoarray.structures.triangles.coordinate_array_np import CoordinateArrayTrianglesNp as triangle_cls
+            from autoarray.structures.triangles.coordinate_array_np import (
+                CoordinateArrayTrianglesNp as triangle_cls,
+            )
 
         initial_triangles = triangle_cls.for_limits_and_scale(
             y_min=y_min,
@@ -196,10 +200,7 @@ def _plane_grid(
             plane_index = tracer.plane_index_via_redshift_from(redshift=plane_redshift)
 
         deflections = tracer.deflections_between_planes_from(
-            grid=grid,
-            plane_i=0,
-            plane_j=plane_index,
-            xp=self._xp
+            grid=grid, plane_i=0, plane_j=plane_index, xp=self._xp
         )
         # noinspection PyTypeChecker
         return grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections)
@@ -265,9 +266,7 @@ def _filter_low_magnification(
         """
         points = self._xp.array(points)
         magnifications = tracer.magnification_2d_via_hessian_from(
-            grid=aa.Grid2DIrregular(points).array,
-            buffer=self.scale,
-            xp=self._xp
+            grid=aa.Grid2DIrregular(points).array, buffer=self.scale, xp=self._xp
         )
         mask = self._xp.abs(magnifications.array) > self.magnification_threshold
         return self._xp.where(mask[:, None], points, self._xp.nan)

test_autolens/point/triangles/test_solver.py

@@ -71,25 +71,16 @@ def test_real_example_jax(grid, tracer):
 
     import jax.numpy as jnp
 
-    jax_solver = PointSolver.for_grid(
-        grid=grid,
-        pixel_scale_precision=0.001,
-        xp=jnp
-    )
+    jax_solver = PointSolver.for_grid(grid=grid, pixel_scale_precision=0.001, xp=jnp)
 
     result = jax_solver.solve(
-        tracer=tracer,
-        source_plane_coordinate=(0.07, 0.07),
-        remove_infinities=True
+        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
+        tracer=tracer, source_plane_coordinate=(0.07, 0.07), remove_infinities=False
     )
 
     assert len(result) == 15