going to take detour removing grid project

Jammy2211 · Jammy2211 · commit e69ff0726593 · 2025-11-08T15:42:19.000Z
diff --git a/autoarray/abstract_ndarray.py b/autoarray/abstract_ndarray.py
@@ -6,7 +6,7 @@
 from abc import abstractmethod
 import jax.numpy as jnp
 from jax._src.tree_util import register_pytree_node
-from jax import Array
+
 import numpy as np
 
 from autoconf.fitsable import output_to_fits
@@ -341,6 +341,7 @@ def __getitem__(self, item):
         return result
 
     def __setitem__(self, key, value):
+        from jax import Array
         if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
             self._array = jnp.where(key, value, self._array)
         else:
diff --git a/autoarray/dataset/abstract/dataset.py b/autoarray/dataset/abstract/dataset.py
@@ -4,10 +4,6 @@
 import warnings
 from typing import Optional, Union
 
-from autoconf import cached_property
-
-from autoarray.dataset.grids import GridsDataset
-
 from autoarray import exc
 from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py
@@ -3,9 +3,6 @@
 from pathlib import Path
 from typing import Optional, Union
 
-from autoconf import cached_property
-from autoconf import instance
-
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
 from autoarray.dataset.imaging.w_tilde import WTildeImaging
diff --git a/autoarray/dataset/imaging/w_tilde.py b/autoarray/dataset/imaging/w_tilde.py
@@ -1,8 +1,6 @@
 import logging
 import numpy as np
 
-from autoconf import cached_property
-
 from autoarray.dataset.abstract.w_tilde import AbstractWTilde
 
 from autoarray.inversion.inversion.imaging import inversion_imaging_util
diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py
@@ -2,7 +2,6 @@
 import numpy as np
 from pathlib import Path
 
-from autoconf import cached_property
 from autoconf.fitsable import ndarray_via_fits_from, output_to_fits
 
 from autoarray.dataset.abstract.dataset import AbstractDataset
diff --git a/autoarray/fit/fit_dataset.py b/autoarray/fit/fit_dataset.py
@@ -5,8 +5,6 @@
 
 import numpy as np
 
-from autoconf import cached_property
-
 from autoarray.dataset.grids import GridsInterface
 from autoarray.dataset.dataset_model import DatasetModel
 from autoarray.fit import fit_util
diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py
@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Dict, List, Union, Type
 
-from autoconf import cached_property
-
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList
diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py
@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Dict, List, Optional, Union
 
-from autoconf import cached_property
-
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.inversion.imaging.abstract import AbstractInversionImaging
diff --git a/autoarray/inversion/inversion/interferometer/w_tilde.py b/autoarray/inversion/inversion/interferometer/w_tilde.py
@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Dict, List, Optional, Union
 
-from autoconf import cached_property
-
 from autoarray.dataset.interferometer.dataset import Interferometer
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.inversion.interferometer.abstract import (
diff --git a/autoarray/inversion/linear_obj/func_list.py b/autoarray/inversion/linear_obj/func_list.py
@@ -1,7 +1,5 @@
 import numpy as np
-from typing import Optional, Dict
-
-from autoconf import cached_property
+from typing import Optional
 
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.linear_obj.neighbors import Neighbors
@@ -15,6 +13,7 @@ def __init__(
         self,
         grid: Grid1D2DLike,
         regularization: Optional[AbstractRegularization],
+        xp=np
     ):
         """
         A linear object which reconstructs a dataset based on mapping between the data points of that dataset and
@@ -41,7 +40,7 @@ def __init__(
             The regularization scheme which may be applied to this linear object in order to smooth its solution.
         """
 
-        super().__init__(regularization=regularization)
+        super().__init__(regularization=regularization, xp=xp)
 
         self.grid = grid
 
diff --git a/autoarray/inversion/pixelization/mappers/abstract.py b/autoarray/inversion/pixelization/mappers/abstract.py
@@ -1,9 +1,8 @@
 import itertools
 import numpy as np
-from typing import Dict, List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 from autoconf import conf
-from autoconf import cached_property
 
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.linear_obj.func_list import UniqueMappings
diff --git a/autoarray/inversion/pixelization/mappers/delaunay.py b/autoarray/inversion/pixelization/mappers/delaunay.py
@@ -1,11 +1,8 @@
 import numpy as np
 
-from autoconf import cached_property
-
 from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
 from autoarray.inversion.pixelization.mappers.abstract import PixSubWeights
 
-from autoarray.inversion.pixelization.mappers import mapper_util
 from autoarray.inversion.pixelization.mappers import mapper_numba_util
 
 
diff --git a/autoarray/inversion/pixelization/mappers/voronoi.py b/autoarray/inversion/pixelization/mappers/voronoi.py
@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Optional, Tuple
 
-from autoconf import cached_property
-
 from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
 from autoarray.inversion.pixelization.mappers.abstract import PixSubWeights
 from autoarray.structures.arrays.uniform_2d import Array2D
diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py
@@ -56,6 +56,7 @@ def contour_list(self):
         # make sure to use base numpy to convert JAX array back to a normal array
 
         from skimage import measure
+        import jax.numpy as jnp
 
         if isinstance(self.contour_array, jnp.ndarray):
             contour_array = np.array(self.contour_array)
diff --git a/autoarray/structures/decorators/project_grid.py b/autoarray/structures/decorators/project_grid.py
@@ -1,5 +1,5 @@
 from functools import wraps
-
+import numpy as np
 from typing import Union
 
 from autoarray import exc
@@ -29,6 +29,7 @@ def project_grid(func):
     def wrapper(
         obj: object,
         grid: Union[Grid1D, Grid2D, Grid2DIrregular],
+        xp=np,
         *args,
         **kwargs,
     ) -> Union[Array1D, ArrayIrregular, Grid2DIrregular]:
diff --git a/autoarray/structures/mesh/delaunay_2d.py b/autoarray/structures/mesh/delaunay_2d.py
@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Optional, Tuple
 
-from autoconf import cached_property
-
 from autoarray.inversion.linear_obj.neighbors import Neighbors
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.mesh.triangulation_2d import Abstract2DMeshTriangulation
diff --git a/autoarray/structures/mesh/triangulation_2d.py b/autoarray/structures/mesh/triangulation_2d.py
@@ -2,8 +2,6 @@
 
 from typing import List, Union, Tuple
 
-from autoconf import cached_property
-
 from autoarray.geometry.geometry_2d_irregular import Geometry2DIrregular
 from autoarray.structures.mesh.abstract_2d import Abstract2DMesh
 
diff --git a/autoarray/structures/mesh/voronoi_2d.py b/autoarray/structures/mesh/voronoi_2d.py
@@ -2,8 +2,6 @@
 
 from typing import Optional, Tuple
 
-from autoconf import cached_property
-
 from autoarray.inversion.linear_obj.neighbors import Neighbors
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.mesh.triangulation_2d import Abstract2DMeshTriangulation
diff --git a/autoarray/structures/visibilities.py b/autoarray/structures/visibilities.py
@@ -5,7 +5,6 @@
 from pathlib import Path
 from typing import List, Tuple, Union
 
-from autoconf import cached_property
 from autoconf.fitsable import ndarray_via_fits_from, output_to_fits
 
 from autoarray.structures.abstract_structure import Structure
diff --git a/test_autoarray/structures/triangles/test_coordinate.py b/test_autoarray/structures/triangles/test_coordinate.py
@@ -1,8 +1,6 @@
-from jax import numpy as np
-import jax
+
 import numpy as np
 
-jax.config.update("jax_log_compiles", True)
 import pytest
 
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
@@ -318,33 +316,12 @@ def one_triangle():
     )
 
 
-@jax.jit
-def full_routine(triangles):
-    neighborhood = triangles.neighborhood()
-    up_sampled = neighborhood.up_sample()
-    with_vertices = up_sampled.with_vertices(up_sampled.vertices)
-    indexes = with_vertices.containing_indices(Point(0.1, 0.1))
-    return up_sampled.for_indexes(indexes)
-
-
-# def test_full_routine(one_triangle, compare_with_nans):
-#     result = full_routine(one_triangle)
-#
-#     assert compare_with_nans(
-#         result.triangles,
-#         np.array(
-#             [
-#                 [
-#                     [0.0, 0.4330126941204071],
-#                     [0.25, 0.0],
-#                     [-0.25, 0.0],
-#                 ]
-#             ]
-#         ),
-#     )
+
 
 
 def test_neighborhood(one_triangle):
+    import jax
+
     assert np.allclose(
         np.array(jax.jit(one_triangle.neighborhood)().triangles),
         np.array(
@@ -375,6 +352,8 @@ def test_neighborhood(one_triangle):
 
 
 def test_up_sample(one_triangle):
+    import jax
+
     up_sampled = jax.jit(one_triangle.up_sample)()
     assert np.allclose(
         np.array(up_sampled.triangles),
