finish

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -331,24 +331,27 @@ def __getattr__(self, item):
 
     def __getitem__(self, item):
 
-        import jax.numpy as jnp
-
         result = self._array[item]
 
         if isinstance(item, slice):
             result = self.with_new_array(result)
-        if isinstance(result, jnp.ndarray):
-            result = self.with_new_array(result)
+
+        try:
+            import jax.numpy as jnp
+            if isinstance(result, jnp.ndarray):
+                result = self.with_new_array(result)
+        except ImportError:
+            pass
+
         return result
 
     def __setitem__(self, key, value):
-        from jax import Array
-        import jax.numpy as jnp
 
-        if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
-            self._array = jnp.where(key, value, self._array)
-        else:
+        if isinstance(self._array, np.ndarray):
             self._array[key] = value
+        else:
+            import jax.numpy as jnp
+            self._array = jnp.where(key, value, self._array)
 
     def __repr__(self):
         return repr(self._array).replace(

autoarray/fit/fit_dataset.py

@@ -163,7 +163,7 @@ def subtracted_from(grid, offset):
             if grid is None:
                 return None
 
-            return grid.subtracted_from(offset=offset)
+            return grid.subtracted_from(offset=offset, xp=self._xp)
 
         lp = subtracted_from(
             grid=self.dataset.grids.lp, offset=self.dataset_model.grid_offset

autoarray/inversion/inversion/interferometer/abstract.py

@@ -65,7 +65,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
         """
         return [
             self.transformer.transform_mapping_matrix(
-                mapping_matrix=linear_obj.mapping_matrix
+                mapping_matrix=linear_obj.mapping_matrix, xp=self._xp
             )
             for linear_obj in self.linear_obj_list
         ]

autoarray/operators/over_sampling/over_sampler.py

@@ -227,6 +227,7 @@ def binned_array_2d_from(self, array: Array2D, xp=np) -> "Array2D":
         Sub-pixels that are part of the same mask array pixel are indexed next to one another, such that the second
         sub-pixel in the first pixel has index 1, its next sub-pixel has index 2, and so forth.
         """
+
         if conf.instance["general"]["structures"]["native_binned_only"]:
             return self
 
@@ -243,16 +244,28 @@ def binned_array_2d_from(self, array: Array2D, xp=np) -> "Array2D":
 
         else:
 
-            import jax
+            if xp.__name__.startswith("jax"):
 
-            # Compute the group means
+                import jax
+
+                sums = jax.ops.segment_sum(
+                    array, self.segment_ids, self.mask.pixels_in_mask
+                )
+                counts = jax.ops.segment_sum(
+                    xp.ones_like(array), self.segment_ids, self.mask.pixels_in_mask
+                )
+
+            else:
+
+                # Sum values per segment
+                sums = np.bincount(self.segment_ids, weights=array, minlength=self.mask.pixels_in_mask)
+
+                # Count number of items per segment
+                counts = np.bincount(self.segment_ids, minlength=self.mask.pixels_in_mask)
+
+                # Avoid division by zero
+                counts[counts == 0] = 1
 
-            sums = jax.ops.segment_sum(
-                array, self.segment_ids, self.mask.pixels_in_mask
-            )
-            counts = jax.ops.segment_sum(
-                xp.ones_like(array), self.segment_ids, self.mask.pixels_in_mask
-            )
             binned_array_2d = sums / counts
 
         return Array2D(

autoarray/operators/transformer.py

@@ -39,7 +39,6 @@ def __init__(
         uv_wavelengths: np.ndarray,
         real_space_mask: Mask2D,
         preload_transform: bool = True,
-        xp=np,
     ):
         """
         A direct Fourier transform (DFT) operator for radio interferometric imaging.
@@ -112,9 +111,7 @@ def __init__(
             2.0 * self.grid.shape_native[1]
         )
 
-        self._xp = xp
-
-    def visibilities_from(self, image: Array2D) -> Visibilities:
+    def visibilities_from(self, image: Array2D, xp=np) -> Visibilities:
         """
         Computes the visibilities from a real-space image using the direct Fourier transform (DFT).
 
@@ -138,19 +135,20 @@ def visibilities_from(self, image: Array2D) -> Visibilities:
                 image_1d=image.array,
                 preloaded_reals=self.preload_real_transforms,
                 preloaded_imags=self.preload_imag_transforms,
-                xp=self._xp,
+                xp=xp,
             )
         else:
             visibilities = transformer_util.visibilities_from(
                 image_1d=image.slim.array,
                 grid_radians=self.grid.array,
                 uv_wavelengths=self.uv_wavelengths,
+                xp=xp
             )
 
-        return Visibilities(visibilities=self._xp.array(visibilities))
+        return Visibilities(visibilities=xp.array(visibilities))
 
     def image_from(
-        self, visibilities: Visibilities, use_adjoint_scaling: bool = False
+        self, visibilities: Visibilities, use_adjoint_scaling: bool = False, xp=np
     ) -> Array2D:
         """
         Computes the real-space image from a set of visibilities using the adjoint of the DFT.
@@ -178,12 +176,12 @@ def image_from(
         )
 
         image_native = array_2d_util.array_2d_native_from(
-            array_2d_slim=image_slim, mask_2d=self.real_space_mask, xp=self._xp
+            array_2d_slim=image_slim, mask_2d=self.real_space_mask, xp=xp
         )
 
         return Array2D(values=image_native, mask=self.real_space_mask)
 
-    def transform_mapping_matrix(self, mapping_matrix: np.ndarray) -> np.ndarray:
+    def transform_mapping_matrix(self, mapping_matrix: np.ndarray, xp=np) -> np.ndarray:
         """
         Applies the DFT to a mapping matrix that maps source pixels to image pixels.
 
@@ -310,8 +308,6 @@ def __init__(
             2.0 * self.grid.shape_native[1]
         )
 
-        self._xp = xp
-
     def initialize_plan(self, ratio: int = 2, interp_kernel: Tuple[int, int] = (6, 6)):
         """
         Initializes the PyNUFFT plan for performing the NUFFT operation.
@@ -394,7 +390,7 @@ def visibilities_from(self, image: Array2D) -> Visibilities:
         )
 
     def image_from(
-        self, visibilities: Visibilities, use_adjoint_scaling: bool = False
+        self, visibilities: Visibilities, use_adjoint_scaling: bool = False, xp=np
     ) -> Array2D:
         """
         Reconstructs a real-space image from visibilities using the NUFFT adjoint transform.
@@ -425,24 +421,24 @@ def image_from(
 
         return Array2D(values=image, mask=self.real_space_mask)
 
-    def transform_mapping_matrix(self, mapping_matrix: np.ndarray) -> np.ndarray:
+    def transform_mapping_matrix(self, mapping_matrix: np.ndarray, xp=np) -> np.ndarray:
         """
-            Applies the NUFFT forward transform to each column of a mapping matrix, producing transformed visibilities.
+        Applies the NUFFT forward transform to each column of a mapping matrix, producing transformed visibilities.
 
-            Parameters
-            ----------
-            mapping_matrix
-                A 2D array where each column corresponds to a source-plane pixel intensity distribution flattened into image space.
+        Parameters
+        ----------
+        mapping_matrix
+            A 2D array where each column corresponds to a source-plane pixel intensity distribution flattened into image space.
 
-            Returns
+        Returns
         -------
-            A complex-valued 2D array where each column contains the visibilities corresponding to the respective column
-            in the input mapping matrix.
+        A complex-valued 2D array where each column contains the visibilities corresponding to the respective column
+        in the input mapping matrix.
 
-            Notes
-            -----
-            - Each column of the input mapping matrix is reshaped into the native 2D image grid before transformation.
-            - This method repeatedly calls `visibilities_from` for each column, which may be computationally intensive.
+        Notes
+        -----
+        - Each column of the input mapping matrix is reshaped into the native 2D image grid before transformation.
+        - This method repeatedly calls `visibilities_from` for each column, which may be computationally intensive.
         """
         transformed_mapping_matrix = 0 + 0j * np.zeros(
             (self.uv_wavelengths.shape[0], mapping_matrix.shape[1])
@@ -452,7 +448,7 @@ def transform_mapping_matrix(self, mapping_matrix: np.ndarray) -> np.ndarray:
             image_2d = array_2d_util.array_2d_native_from(
                 array_2d_slim=mapping_matrix[:, source_pixel_1d_index],
                 mask_2d=self.grid.mask,
-                xp=self._xp,
+                xp=xp,
             )
 
             image = Array2D(values=image_2d, mask=self.grid.mask)

autoarray/operators/transformer_util.py

@@ -120,7 +120,7 @@ def visibilities_via_preload_from(
 
 
 def visibilities_from(
-    image_1d: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray
+    image_1d: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray, xp=np
 ) -> np.ndarray:
     """
     Compute complex visibilities from an input sky image using the Fourier transform,
@@ -150,19 +150,19 @@ def visibilities_from(
     # Compute the dot product for each pixel-uv pair
     phase = (
         -2.0
-        * np.pi
+        * xp.pi
         * (
-            np.outer(grid_radians[:, 1], uv_wavelengths[:, 0])
-            + np.outer(grid_radians[:, 0], uv_wavelengths[:, 1])
+            xp.outer(grid_radians[:, 1], uv_wavelengths[:, 0])
+            + xp.outer(grid_radians[:, 0], uv_wavelengths[:, 1])
         )
     )  # shape (n_pixels, n_vis)
 
     # Multiply image values with phase terms
-    vis_real = image_1d[:, None] * np.cos(phase)
-    vis_imag = image_1d[:, None] * np.sin(phase)
+    vis_real = image_1d[:, None] * xp.cos(phase)
+    vis_imag = image_1d[:, None] * xp.sin(phase)
 
     # Sum over all pixels for each visibility
-    visibilities = np.sum(vis_real + 1j * vis_imag, axis=0)
+    visibilities = xp.sum(vis_real + 1j * vis_imag, axis=0)
 
     return visibilities
 
@@ -247,7 +247,7 @@ def transformed_mapping_matrix_via_preload_from(
 
 
 def transformed_mapping_matrix_from(
-    mapping_matrix: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray
+    mapping_matrix: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray, xp=np
 ) -> np.ndarray:
     """
     Computes the Fourier-transformed mapping matrix used in radio interferometric imaging.
@@ -273,16 +273,16 @@ def transformed_mapping_matrix_from(
     # Compute phase term: (n_image_pixels, n_visibilities)
     phase = (
         -2.0
-        * np.pi
+        * xp.pi
         * (
-            np.outer(grid_radians[:, 1], uv_wavelengths[:, 0])  # y * u
-            + np.outer(grid_radians[:, 0], uv_wavelengths[:, 1])  # x * v
+            xp.outer(grid_radians[:, 1], uv_wavelengths[:, 0])  # y * u
+            + xp.outer(grid_radians[:, 0], uv_wavelengths[:, 1])  # x * v
         )
     )
 
     # Compute real and imaginary Fourier matrices
-    fourier_real = np.cos(phase)
-    fourier_imag = np.sin(phase)
+    fourier_real = xp.cos(phase)
+    fourier_imag = xp.sin(phase)
 
     # Only compute contributions from non-zero mapping entries
     # This matrix multiplication is: (n_visibilities x n_image_pixels) dot (n_image_pixels x n_source_pixels)

autoarray/preloads.py

@@ -9,9 +9,7 @@
 
 def mapper_indices_from(total_linear_light_profiles, total_mapper_pixels):
 
-    import jax.numpy as jnp
-
-    return jnp.arange(
+    return np.arange(
         total_linear_light_profiles,
         total_linear_light_profiles + total_mapper_pixels,
         dtype=int,
@@ -54,31 +52,28 @@ def __init__(
             is fixed to the maximum likelihood solution, allowing the blurred mapping matrix to be preloaded, but
             the intensity values will still be solved for during the inversion.
         """
-        import jax.numpy as jnp
-
         self.mapper_indices = None
         self.source_pixel_zeroed_indices = None
         self.source_pixel_zeroed_indices_to_keep = None
         self.linear_light_profile_blurred_mapping_matrix = None
 
         if mapper_indices is not None:
 
-            self.mapper_indices = jnp.array(mapper_indices)
+            self.mapper_indices = np.array(mapper_indices)
 
         if source_pixel_zeroed_indices is not None:
 
-            self.source_pixel_zeroed_indices = jnp.array(source_pixel_zeroed_indices)
+            self.source_pixel_zeroed_indices = np.array(source_pixel_zeroed_indices)
 
-            ids_zeros = jnp.array(source_pixel_zeroed_indices, dtype=int)
+            ids_zeros = np.array(source_pixel_zeroed_indices, dtype=int)
 
-            values_to_solve = jnp.ones(np.max(mapper_indices), dtype=bool)
-            values_to_solve = values_to_solve.at[ids_zeros].set(False)
+            values_to_solve = np.ones(np.max(mapper_indices)+1, dtype=bool)
+            values_to_solve[ids_zeros] = False
 
-            # Get the indices where values_to_solve is True
-            self.source_pixel_zeroed_indices_to_keep = jnp.where(values_to_solve)[0]
+            self.source_pixel_zeroed_indices_to_keep = np.where(values_to_solve)[0]
 
         if linear_light_profile_blurred_mapping_matrix is not None:
 
-            self.linear_light_profile_blurred_mapping_matrix = jnp.array(
+            self.linear_light_profile_blurred_mapping_matrix = np.array(
                 linear_light_profile_blurred_mapping_matrix
             )

autoarray/structures/triangles/abstract.py

@@ -12,21 +12,27 @@ def __len__(self):
         return len(self.triangles)
 
     @property
-    @abstractmethod
     def area(self) -> float:
         """
         The total area covered by the triangles.
         """
+        triangles = self.triangles
+        return (
+            0.5
+            * np.abs(
+                (triangles[:, 0, 0] * (triangles[:, 1, 1] - triangles[:, 2, 1]))
+                + (triangles[:, 1, 0] * (triangles[:, 2, 1] - triangles[:, 0, 1]))
+                + (triangles[:, 2, 0] * (triangles[:, 0, 1] - triangles[:, 1, 1]))
+            ).sum()
+        )
 
     @property
-    @abstractmethod
     def indices(self):
-        pass
+        return self._indices
 
     @property
-    @abstractmethod
     def vertices(self):
-        pass
+        return self._vertices
 
     def __str__(self):
         return f"{self.__class__.__name__} with {len(self.indices)} triangles"