black

Author: Jammy2211

autoarray/__init__.py

@@ -1,3 +1,4 @@
+from autoconf import jax_wrapper
 from autoconf.dictable import register_parser
 from autoconf import conf
 

autoarray/abstract_ndarray.py

@@ -332,6 +332,7 @@ def __getattr__(self, item):
     def __getitem__(self, item):
 
         import jax.numpy as jnp
+
         result = self._array[item]
 
         if isinstance(item, slice):

autoarray/config/general.yaml

@@ -1,5 +1,3 @@
-jax:
-  use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 fits:
   flip_for_ds9: false                 # If True, the image is flipped before output to a .fits file, which is useful for viewing in DS9.
 psf:

autoarray/structures/triangles/array.py

@@ -7,6 +7,7 @@
 
 MAX_CONTAINING_SIZE = 15
 
+
 class ArrayTriangles(AbstractTriangles):
     def __init__(
         self,

test_autoarray/structures/triangles/conftest.py

@@ -1,21 +1,19 @@
-from autoconf.jax_wrapper import np
-from autoarray.structures.triangles.array import ArrayTriangles
-from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
-
+import jax.numpy as jnp
 from matplotlib import pyplot as plt
-
-
 import pytest
 
+from autoarray.structures.triangles.array import ArrayTriangles
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
+
 
 @pytest.fixture
 def plot():
     plt.figure(figsize=(8, 8))
 
     def plot(triangles, color="black"):
         for triangle in triangles:
-            triangle = np.array(triangle)
-            triangle = np.append(triangle, np.array([triangle[0]]), axis=0)
+            triangle = jnp.array(triangle)
+            triangle = jnp.append(triangle, jnp.array([triangle[0]]), axis=0)
             plt.plot(triangle[:, 0], triangle[:, 1], "o-", color=color)
 
     yield plot
@@ -26,27 +24,27 @@ def plot(triangles, color="black"):
 @pytest.fixture
 def compare_with_nans():
     def compare_with_nans_(arr1, arr2):
-        nan_mask1 = np.isnan(arr1)
-        nan_mask2 = np.isnan(arr2)
+        nan_mask1 = jnp.isnan(arr1)
+        nan_mask2 = jnp.isnan(arr2)
 
         arr1 = arr1[~nan_mask1]
         arr2 = arr2[~nan_mask2]
 
-        return np.all(arr1 == arr2)
+        return jnp.all(arr1 == arr2)
 
     return compare_with_nans_
 
 
 @pytest.fixture
 def triangles():
     return ArrayTriangles(
-        indices=np.array(
+        indices=jnp.array(
             [
                 [0, 1, 2],
                 [1, 2, 3],
             ]
         ),
-        vertices=np.array(
+        vertices=jnp.array(
             [
                 [0.0, 0.0],
                 [1.0, 0.0],
@@ -60,14 +58,14 @@ def triangles():
 @pytest.fixture
 def one_triangle():
     return CoordinateArrayTriangles(
-        coordinates=np.array([[0, 0]]),
+        coordinates=jnp.array([[0, 0]]),
         side_length=1.0,
     )
 
 
 @pytest.fixture
 def two_triangles():
     return CoordinateArrayTriangles(
-        coordinates=np.array([[0, 0], [1, 0]]),
+        coordinates=jnp.array([[0, 0], [1, 0]]),
         side_length=1.0,
     )

test_autoarray/structures/triangles/test_coordinate.py

@@ -1,14 +1,16 @@
 import numpy as np
+from jax.tree_util import register_pytree_node_class
 
 import pytest
 
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
-from autoarray.structures.triangles.shape import Point
 
 from autoarray.structures.triangles.coordinate_array import (
     CoordinateArrayTriangles,
 )
 
+CoordinateArrayTriangles = register_pytree_node_class(CoordinateArrayTriangles)
+
 
 def test__two(two_triangles):
 

test_autoarray/structures/triangles/test_jax.py

@@ -10,6 +10,7 @@
 ArrayTriangles = register_pytree_node_class(ArrayTriangles)
 Point = register_pytree_node_class(Point)
 
+
 @pytest.fixture
 def triangles():
     return ArrayTriangles(