adaptive_pixel_signals_from JAX-d

Author: Jammy2211

autoarray/inversion/inversion/abstract.py

@@ -68,17 +68,6 @@ def __init__(
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
         """
 
-        try:
-            import numba
-        except ModuleNotFoundError:
-            raise exc.InversionException(
-                "Inversion functionality (linear light profiles, pixelized reconstructions) is "
-                "disabled if numba is not installed.\n\n"
-                "This is because the run-times without numba are too slow.\n\n"
-                "Please install numba, which is described at the following web page:\n\n"
-                "https://pyautolens.readthedocs.io/en/latest/installation/overview.html"
-            )
-
         self.dataset = dataset
 
         self.linear_obj_list = linear_obj_list
@@ -160,17 +149,10 @@ def param_range_list_from(self, cls: Type) -> List[List[int]]:
         -------
         A list of the index range of the parameters of each linear object in the inversion of the input cls type.
         """
-        index_list = []
-
-        pixel_count = 0
-
-        for linear_obj in self.linear_obj_list:
-            if isinstance(linear_obj, cls):
-                index_list.append([pixel_count, pixel_count + linear_obj.params])
-
-            pixel_count += linear_obj.params
-
-        return index_list
+        return inversion_util.param_range_list_from(
+            cls=cls,
+            linear_obj_list=self.linear_obj_list
+        )
 
     def cls_list_from(self, cls: Type, cls_filtered: Optional[Type] = None) -> List:
         """

autoarray/inversion/inversion/imaging/abstract.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import Dict, List, Optional, Union, Type
+from typing import Dict, List, Union, Type
 
 from autoconf import cached_property
 

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -49,6 +49,17 @@ def __init__(
             the simultaneous linear equations are combined and solved simultaneously.
         """
 
+        try:
+            import numba
+        except ModuleNotFoundError:
+            raise exc.InversionException(
+                "Inversion functionality (linear light profiles, pixelized reconstructions) is "
+                "disabled if numba is not installed.\n\n"
+                "This is because the run-times without numba are too slow.\n\n"
+                "Please install numba, which is described at the following web page:\n\n"
+                "https://pyautolens.readthedocs.io/en/latest/installation/overview.html"
+            )
+
         super().__init__(
             dataset=dataset,
             linear_obj_list=linear_obj_list,

autoarray/inversion/inversion/inversion_util.py

@@ -2,7 +2,7 @@
 import jax.lax as lax
 import numpy as np
 
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from autoconf import conf
 
@@ -346,3 +346,48 @@ def preconditioner_matrix_via_mapping_matrix_from(
     return (
         preconditioner_noise_normalization * curvature_matrix
     ) + regularization_matrix
+
+
+def param_range_list_from(cls: Type, linear_obj_list) -> List[List[int]]:
+    """
+    Each linear object in the `Inversion` has N parameters, and these parameters correspond to a certain range
+    of indexing values in the matrices used to perform the inversion.
+
+    This function returns the `param_range_list` of an input type of linear object, which gives the indexing range
+    of each linear object of the input type.
+
+    For example, if an `Inversion` has:
+
+    - A `LinearFuncList` linear object with 3 `params`.
+    - A `Mapper` with 100 `params`.
+    - A `Mapper` with 200 `params`.
+
+    The corresponding matrices of this inversion (e.g. the `curvature_matrix`) have `shape=(303, 303)` where:
+
+    - The `LinearFuncList` values are in the entries `[0:3]`.
+    - The first `Mapper` values are in the entries `[3:103]`.
+    - The second `Mapper` values are in the entries `[103:303]
+
+    For this example, `param_range_list_from(cls=AbstractMapper)` therefore returns the
+    list `[[3, 103], [103, 303]]`.
+
+    Parameters
+    ----------
+    cls
+        The type of class that the list of their parameter range index values are returned for.
+
+    Returns
+    -------
+    A list of the index range of the parameters of each linear object in the inversion of the input cls type.
+    """
+    index_list = []
+
+    pixel_count = 0
+
+    for linear_obj in linear_obj_list:
+        if isinstance(linear_obj, cls):
+            index_list.append([pixel_count, pixel_count + linear_obj.params])
+
+        pixel_count += linear_obj.params
+
+    return index_list

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -498,7 +498,6 @@ def remove_bad_entries_voronoi_nn(
     return pix_weights_for_sub_slim_index, pix_indexes_for_sub_slim_index
 
 
-@numba_util.jit()
 def adaptive_pixel_signals_from(
     pixels: int,
     pixel_weights: np.ndarray,
@@ -536,30 +535,43 @@ def adaptive_pixel_signals_from(
         The image of the galaxy which is used to compute the weigghted pixel signals.
     """
 
-    pixel_signals = np.zeros((pixels,))
-    pixel_sizes = np.zeros((pixels,))
+    M_sub, B = pix_indexes_for_sub_slim_index.shape
 
-    for sub_slim_index in range(len(pix_indexes_for_sub_slim_index)):
-        vertices_indexes = pix_indexes_for_sub_slim_index[sub_slim_index]
+    # 1) Flatten the per‐mapping tables:
+    flat_pixidx  = pix_indexes_for_sub_slim_index.reshape(-1)    # (M_sub*B,)
+    flat_weights = pixel_weights.reshape(-1)   # (M_sub*B,)
 
-        mask_1d_index = slim_index_for_sub_slim_index[sub_slim_index]
+    # 2) Build a matching “parent‐slim” index for each flattened entry:
+    I_sub = jnp.repeat(jnp.arange(M_sub), B)                     # (M_sub*B,)
 
-        pix_size_tem = pix_size_for_sub_slim_index[sub_slim_index]
+    # 3) Mask out any k >= pix_size_for_sub_slim_index[i]
+    valid = (I_sub < 0)  # dummy to get shape
+    # better:
+    valid = (jnp.arange(B)[None, :] < pix_size_for_sub_slim_index[:, None]).reshape(-1)
 
-        if pix_size_tem > 1:
-            pixel_signals[vertices_indexes[:pix_size_tem]] += (
-                adapt_data[mask_1d_index] * pixel_weights[sub_slim_index]
-            )
-            pixel_sizes[vertices_indexes] += 1
-        else:
-            pixel_signals[vertices_indexes[0]] += adapt_data[mask_1d_index]
-            pixel_sizes[vertices_indexes[0]] += 1
+    flat_weights = jnp.where(valid, flat_weights, 0.0)
+    flat_pixidx  = jnp.where(valid, flat_pixidx, pixels)  # send invalid indices to an out-of-bounds slot
+
+    # 4) Look up data & multiply by mapping weights:
+    flat_data_vals = adapt_data[slim_index_for_sub_slim_index][I_sub]  # (M_sub*B,)
+    flat_contrib   = flat_data_vals * flat_weights                     # (M_sub*B,)
+
+    # 5) Scatter‐add into signal sums and counts:
+    pixel_signals = jnp.zeros((pixels+1,)).at[flat_pixidx].add(flat_contrib)
+    pixel_counts  = jnp.zeros((pixels+1,)).at[flat_pixidx].add(valid.astype(float))
+
+    # 6) Drop the extra “out-of-bounds” slot:
+    pixel_signals = pixel_signals[:pixels]
+    pixel_counts  = pixel_counts[:pixels]
 
-    pixel_sizes[pixel_sizes == 0] = 1
-    pixel_signals /= pixel_sizes
-    pixel_signals /= np.max(pixel_signals)
+    # 7) Normalize
+    pixel_counts = jnp.where(pixel_counts > 0, pixel_counts, 1.0)
+    pixel_signals = pixel_signals / pixel_counts
+    max_sig = jnp.max(pixel_signals)
+    pixel_signals = jnp.where(max_sig > 0, pixel_signals / max_sig, pixel_signals)
 
-    return pixel_signals**signal_scale
+    # 8) Exponentiate
+    return pixel_signals ** signal_scale
 
 
 def mapping_matrix_from(

test_autoarray/inversion/pixelization/mappers/test_rectangular.py

@@ -68,7 +68,7 @@ def test__pixel_signals_from__matches_util(grid_2d_sub_1_7x7, image_7x7):
         pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,
         pixel_weights=mapper.pix_weights_for_sub_slim_index,
         slim_index_for_sub_slim_index=grid_2d_sub_1_7x7.over_sampler.slim_for_sub_slim,
-        adapt_data=np.array(image_7x7),
+        adapt_data=image_7x7,
     )
 
     assert (pixel_signals == pixel_signals_util).all()