From c4962c93f109c5f778267404f1d2ff33c4ca8f48 Mon Sep 17 00:00:00 2001 From: Kolen Cheung Date: Thu, 19 Dec 2024 16:53:44 +0000 Subject: [PATCH 1/2] fix jit on mask_2d_circular_from --- autoarray/mask/mask_2d_util.py | 26 +++++++------------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py index 4103e89d..f6399d82 100644 --- a/autoarray/mask/mask_2d_util.py +++ b/autoarray/mask/mask_2d_util.py @@ -81,7 +81,7 @@ def total_pixels_2d_from(mask_2d: np.ndarray) -> int: return total_regular_pixels -@numba_util.jit() +@numba_util.jit(static_argnums=0) def mask_2d_circular_from( shape_native: Tuple[int, int], pixel_scales: ty.PixelScales, @@ -114,24 +114,12 @@ def mask_2d_circular_from( mask = mask_circular_from( shape=(10, 10), pixel_scales=0.1, radius=0.5, centre=(0.0, 0.0)) """ - - mask_2d = np.full(shape_native, True) - - centres_scaled = mask_2d_centres_from( - shape_native=mask_2d.shape, pixel_scales=pixel_scales, centre=centre - ) - - for y in range(mask_2d.shape[0]): - for x in range(mask_2d.shape[1]): - y_scaled = (y - centres_scaled[0]) * pixel_scales[0] - x_scaled = (x - centres_scaled[1]) * pixel_scales[1] - - r_scaled = np.sqrt(x_scaled**2 + y_scaled**2) - - if r_scaled <= radius: - mask_2d[y, x] = False - - return mask_2d + centres_scaled = mask_2d_centres_from(shape_native, pixel_scales, centre) + ys, xs = np.mgrid[:shape_native[0], :shape_native[1]] + return (radius * radius) < ( + np.square((ys - centres_scaled[0]) * pixel_scales[0]) + + np.square((xs - centres_scaled[1]) * pixel_scales[1]) + ) @numba_util.jit() From ba31c1c66a35f6308bababdd5b944b93453641f2 Mon Sep 17 00:00:00 2001 From: Kolen Cheung Date: Thu, 19 Dec 2024 17:17:46 +0000 Subject: [PATCH 2/2] use indices instead of mgrid --- autoarray/mask/mask_2d_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py index f6399d82..1981c317 100644 --- a/autoarray/mask/mask_2d_util.py +++ b/autoarray/mask/mask_2d_util.py @@ -115,7 +115,7 @@ def mask_2d_circular_from( shape=(10, 10), pixel_scales=0.1, radius=0.5, centre=(0.0, 0.0)) """ centres_scaled = mask_2d_centres_from(shape_native, pixel_scales, centre) - ys, xs = np.mgrid[:shape_native[0], :shape_native[1]] + ys, xs = np.indices(shape_native) return (radius * radius) < ( np.square((ys - centres_scaled[0]) * pixel_scales[0]) + np.square((xs - centres_scaled[1]) * pixel_scales[1])