prtools 1.2.1__tar.gz → 1.3.0__tar.gz

{prtools-1.2.1/prtools.egg-info → prtools-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: prtools
-Version: 1.2.1
+Version: 1.3.0
 Summary: Utility functions for image-based phase retrieval
 Author: Andy Kee
 License: MIT License
@@ -36,15 +36,16 @@ Classifier: Topic :: Scientific/Engineering :: Astronomy
 Classifier: Topic :: Scientific/Engineering :: Image Processing
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Scientific/Engineering :: Physics
-Classifier: Topic :: Software Development
-Classifier: Topic :: Software Development :: Libraries
-Classifier: Topic :: Software Development :: Libraries :: Python Modules 
 Requires-Python: >=3.8
 Description-Content-Type: text/x-rst
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: numpy
 Requires-Dist: scipy
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: optax; extra == "jax"
+Dynamic: license-file
 
 prtools
 =======

prtools-1.3.0/prtools/__init__.py

@@ -0,0 +1,66 @@
+__version__ = '1.3.0'
+
+from prtools._backend import __backend__, use
+
+from prtools.array import (
+    centroid,
+    pad,
+    subarray,
+    boundary,
+    rebin,
+    rescale,
+    normpow,
+    shift,
+    register,
+    medfix
+)
+
+from prtools.cost import sserror
+
+from prtools.fourier import dft2, idft2
+
+from prtools.misc import (
+    calcpsf,
+    translation_defocus,
+    fft_shape,
+    min_sampling,
+    pixelscale_nyquist,
+    find_wrapped
+)
+
+from prtools.segmented import hex_segments
+
+from prtools.shape import (
+    mesh,
+    circle,
+    hexagon,
+    rectangle,
+    spider,
+    gauss,
+    sin,
+    waffle
+)
+
+from prtools.sparse import (
+    index,
+    dense,
+    sparse,
+    index_from_mask,
+    mask_from_index
+)
+
+from prtools.stats import (
+    ee,
+    pv,
+    rms,
+    radial_avg
+)
+
+from prtools.zernike import (
+    zernike,
+    zernike_compose,
+    zernike_basis,
+    zernike_fit,
+    zernike_remove,
+    zernike_coordinates
+)

prtools-1.3.0/prtools/_backend.py

@@ -0,0 +1,73 @@
+import sys
+import types
+
+from prtools.backend import registry
+from prtools.backend._base import _BackendLibrary
+
+
+class _BackendName:
+    """Helper class for storing the backend name - needed since strings are
+    immutable.
+    """
+
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return self.name
+
+    def __str__(self):
+        return self.name
+
+    def __eq__(self, other):
+        return self.name == other
+
+
+class _Backend(types.ModuleType):
+    __backend__ = _BackendName(None)
+    numpy = _BackendLibrary(None)
+    scipy = _BackendLibrary(None)
+
+    @classmethod
+    def use(cls, backend):
+        """Select the backend used for N-dimensional array operations.
+
+        Parameters
+        ----------
+        backend : str
+            The backend to use. This can be any of the following backends,
+            which are case-insensitive:
+            
+            * NumPy
+            * JAX
+        """
+        #if name in registry:
+        backend = backend.lower()
+        cls.__backend__.name = backend
+        cls.numpy.module = registry.load_numpy(backend)
+        cls.scipy.module = registry.load_scipy(backend)
+
+
+    @classmethod
+    def set_backend(cls, name, numpy_module, scipy_module):
+        """Change the backend
+
+        Parameters
+        ----------
+        name : str
+            Name of the backend.
+        numpy_module : module
+            Library providing numpy-like functionality
+        scipy_module : module
+            Library providing scipy-like functionality
+        """
+        cls.__backend__.name = name
+        cls.numpy.module = numpy_module
+        cls.scipy.module = scipy_module
+
+
+# use NumPy by default
+_Backend.use('numpy')
+
+# https://stackoverflow.com/a/72911884
+sys.modules[__name__].__class__ = _Backend

{prtools-1.2.1 → prtools-1.3.0}/prtools/array.py

@@ -1,8 +1,8 @@
 import warnings
 
-import numpy as np
-from numpy.lib.stride_tricks import as_strided
-from scipy.ndimage import map_coordinates
+from prtools import __backend__
+from prtools._backend import numpy as np
+from prtools._backend import scipy
 
 
 def centroid(a, where=None, kind='absolute', indexing='ij'):
@@ -16,9 +16,9 @@ def centroid(a, where=None, kind='absolute', indexing='ij'):
         Elements to include in the centroid calculation. If None (default),
         all finite and non-NaN values are used.
     kind : {'absolute', 'center'}, optional
-        Specifies the kind of centroid as a string. If 'absolute' (default), the 
-        absolute centroid within the input is returned. If 'center', the centroid
-        relative to the center of the input is returned.
+        Specifies the kind of centroid as a string. If 'absolute' (default),
+        the absolute centroid within the input is returned. If 'center', the
+        centroid relative to the center of the input is returned.
     indexing : {'ij', 'xy'}, optional
         Matrix ('ij', default) or cartesian ('xy') indexing of mesh.
 
@@ -36,7 +36,7 @@ def centroid(a, where=None, kind='absolute', indexing='ij'):
 
         >>> circ = prtools.circle(shape=(256, 256), radius=25, shift=(-80, 50))
         >>> plt.imshow(circ, cmap='gray')
-    
+
     .. code:: pycon
 
         >>> prtools.centroid(circ)
@@ -49,7 +49,7 @@ def centroid(a, where=None, kind='absolute', indexing='ij'):
     """
     if kind not in ('absolute', 'center'):
         raise ValueError(f'Unknown kind {kind}')
-    
+
     if indexing not in ('ij', 'xy'):
         raise ValueError("Valid values for indexing are 'xy' and 'ij'.")
 
@@ -63,7 +63,7 @@ def centroid(a, where=None, kind='absolute', indexing='ij'):
     if np.isnan(a[where]).any():
         warnings.warn('Unmasked NaN in input', RuntimeWarning,
                       stacklevel=2)
-    
+
     anorm = a[where]/np.sum(a[where])
 
     nr, nc = a.shape
@@ -96,7 +96,7 @@ def pad(a, shape, fill=0):
         Shape of output array in ``(nrows, ncols)``.
     fill : scalar
         Fill vlue used when pad operation increases array size.
-    
+
     Returns
     -------
     padded_array : ndarray
@@ -117,7 +117,7 @@ def pad(a, shape, fill=0):
         >>> ax[0].set_title('Original array')
         >>> ax[1].imshow(circ_pad, cmap='gray')
         >>> ax[1].set_title('Padded array')
-    
+
     .. plot::
         :include-source:
         :context: reset
@@ -133,14 +133,14 @@ def pad(a, shape, fill=0):
 
     """
     a = np.atleast_2d(a)
-    shape = np.broadcast_to(np.asarray(shape, dtype=int), (2,))    
+    shape = np.broadcast_to(np.asarray(shape, dtype=int), (2,))
 
     if a.ndim == 2:
         shape = np.append(1, shape)
-        a = a[np.newaxis,:]         
+        a = a[np.newaxis, :]
     else:  # a.ndim == 3
         shape = np.append(a.shape[0], shape)
-        
+
     out = np.ones(shape, dtype=a.dtype) * fill
 
     # The row and col indices here are 1 and 2 respectively since we've
@@ -153,14 +153,14 @@ def pad(a, shape, fill=0):
     a_slc_r = slice(a.shape[1]//2-rmin, a.shape[1]//2+rmax)
     a_slc_c = slice(a.shape[2]//2-cmin, a.shape[2]//2+cmax)
 
-    out_slc_r = slice(out.shape[1]//2-rmin, out.shape[1]//2+rmax) 
+    out_slc_r = slice(out.shape[1]//2-rmin, out.shape[1]//2+rmax)
     out_slc_c = slice(out.shape[2]//2-cmin, out.shape[2]//2+cmax)
-    
+
     out[:, out_slc_r, out_slc_c] = a[:, a_slc_r, a_slc_c]
     return np.squeeze(out)
 
 
-def subarray(a, shape, shift=(0,0)):
+def subarray(a, shape, shift=(0, 0)):
     """Extract a contiguous subarray from a larger array.
 
     The subarray is extracted about the center of the source array unless
@@ -174,7 +174,7 @@ def subarray(a, shape, shift=(0,0)):
         Shape of subarray array in ``(nrows, ncols)``.
     shift : array_like of ints
         Relative shift of the center of the subarray in ``(row, col)``.
-    
+
     Returns
     -------
     out : ndarray
@@ -194,7 +194,7 @@ def subarray(a, shape, shift=(0,0)):
         >>> ax[0].set_title('Original array')
         >>> ax[1].imshow(circ_subarray, cmap='gray')
         >>> ax[1].set_title('Subarray')
-    
+
     .. plot::
         :include-source:
         :context: reset
@@ -213,7 +213,6 @@ def subarray(a, shape, shift=(0,0)):
     a = np.asarray(a)
     shape = np.asarray(shape)
 
-
     rmin = a.shape[0]//2 - shape[0]//2 + shift[0]
     cmin = a.shape[1]//2 - shape[1]//2 + shift[1]
     rmax = rmin + shape[0]
@@ -258,7 +257,7 @@ def boundary(x, threshold=0):
 
         >>> prtools.boundary(circ)
         (50, 150, 50, 150)
-    
+
     """
     x = np.asarray(x)
     x = (x > threshold)
@@ -278,8 +277,8 @@ def rebin(img, factor):
     Parameters
     ----------
     img : array_like
-        Array or cube of arrays to rebin. If a cube is provided, the first dimension
-        should index the image slices.
+        Array or cube of arrays to rebin. If a cube is provided, the first
+        dimension should index the image slices.
 
     factor : int
         Rebinning factor
@@ -321,8 +320,10 @@ def rescale(img, scale, shape=None, mask=None, order=3, mode='nearest',
     img : array_like
         Image to rescale
 
-    scale : float
-        Scaling factor. Scale factors less than 1 will shrink the image. Scale
+    scale : float or tuple of floats
+        Scaling factor. If scale is a tuple, img is rescaled according to
+        ``scale_row, scale_col)``. A single value rescales img equally in both
+        dimensions. Scale factors less than 1 will shrink the image. Scale
         factors greater than 1 will grow the image.
 
     shape : array_like or int, optional
@@ -340,7 +341,7 @@ def rescale(img, scale, shape=None, mask=None, order=3, mode='nearest',
 
     mode : {'constant', 'nearest', 'reflect', 'wrap'}, optional
         Points outside the boundaries of the input are filled according to the
-        given mode. Default is 'constant'.
+        given mode. Default is 'nearest'.
 
     unitary : bool, optional
         Normalization flag. If True (default), a normalization is performed on
@@ -364,35 +365,33 @@ def rescale(img, scale, shape=None, mask=None, order=3, mode='nearest',
     """
 
     img = np.asarray(img)
+    scale = np.broadcast_to(scale, (2,))
+
+    if shape is None:
+        shape = img.shape
+    shape = np.broadcast_to(shape, (2,))
+    shape = np.ceil(shape * scale).astype(int)
 
     if mask is None:
-        # take the real portion to ensure that even if img is complex, mask will
-        # be real
+        # take the real portion to ensure that even if img is complex, mask
+        # will be real
         mask = np.zeros_like(img).real
         mask[img != 0] = 1
 
-    if shape is None:
-        shape = np.ceil((img.shape[0]*scale, img.shape[1]*scale)).astype(int)
-    else:
-        if np.isscalar(shape):
-            shape = np.ceil((shape*scale, shape*scale)).astype(int)
-        else:
-            shape = np.ceil((shape[0]*scale, shape[1]*scale)).astype(int)
-
-    x = (np.arange(shape[1], dtype=np.float64) - shape[1]/2.)/scale + img.shape[1]/2.
-    y = (np.arange(shape[0], dtype=np.float64) - shape[0]/2.)/scale + img.shape[0]/2.
+    r = (np.arange(shape[0], dtype=np.float64) - shape[0]/2.)/scale[0] + img.shape[0]/2.
+    c = (np.arange(shape[1], dtype=np.float64) - shape[1]/2.)/scale[1] + img.shape[1]/2.
 
-    xx, yy = np.meshgrid(x, y)
+    rr, cc = np.meshgrid(r, c, indexing='ij')
 
-    mask = map_coordinates(mask, [yy, xx], order=1, mode='nearest')
+    mask = scipy.ndimage.map_coordinates(mask, [rr, cc], order=1, mode='nearest')
     mask[mask < np.finfo(mask.dtype).eps] = 0
 
     if np.iscomplexobj(img):
         out = np.zeros(shape, dtype=np.complex128)
-        out.real = map_coordinates(img.real, [yy, xx], order=order, mode=mode)
-        out.imag = map_coordinates(img.imag, [yy, xx], order=order, mode=mode)
+        out.real = scipy.ndimage.map_coordinates(img.real, [rr, cc], order=order, mode=mode)
+        out.imag = scipy.ndimage.map_coordinates(img.imag, [rr, cc], order=order, mode=mode)
     else:
-        out = map_coordinates(img, [yy, xx], order=order, mode=mode)
+        out = scipy.ndimage.map_coordinates(img, [rr, cc], order=order, mode=mode)
 
     if unitary:
         out *= np.sum(img)/np.sum(out)
@@ -401,6 +400,7 @@ def rescale(img, scale, shape=None, mask=None, order=3, mode='nearest',
 
     return out
 
+
 def normpow(array, power=1):
     r"""Normalizie the power in an array.
 
@@ -437,12 +437,11 @@ def normpow(array, power=1):
     return array * np.sqrt(power/np.sum(np.abs(array)**2))
 
 
-
-def shift(a, shift):
+def shift(a, shift, mode='wrap', fill=0.0):
     """Shift an array via FFT.
 
-    Shift an array by (row, column). The shifts may be non-integer as the 
-    shift operation is implemented by introducing a Fourier-domain tilt. If 
+    Shift an array by (row, column). The shifts may be non-integer as the
+    shift operation is implemented by introducing a Fourier-domain tilt. If
     ``a`` is complex, the result will also be complex.
 
     Parameters
@@ -451,6 +450,26 @@ def shift(a, shift):
         The input array.
     shift : (2,) sequence
         The shift specified as (row, column).
+    mode : {'wrap', 'reflect', 'mirror', 'constant'}, optional
+        Determines how the input array is extended beyond its boundaries.
+        Default is 'wrap'.
+
+        * 'wrap' (a b c d | a b c d | a b c d)
+            The input is extended by wrapping around to the opposite edge.
+        * 'reflect' (d c b a | a b c d | d c b a)
+            The input is extended by reflecting about the edge of the last
+            pixel. This mode is also sometimes referred to as half-sample
+            symmetric.
+        * 'mirror' (d c b | a b c d | c b a)
+            The input is extended by reflecting about the center of the last
+            pixel. This mode is also sometimes referred to as whole-sample
+            symmetric.
+        * 'constant'
+            The input is extended by filling all values beyond the edge with
+            the same constant value, defined by the fill parameter.
+
+    fill : scalar, optional
+        Value to fill past edges if `mode` is 'constant'. Default is 0.0
     Returns
     -------
     shifted : ndarray
@@ -473,8 +492,11 @@ def shift(a, shift):
                [-3.12823642e-17,  2.22044605e-16, -4.18468327e-17]])
     """
     a = np.asarray(a)
-
+    r, c = a.shape
     dr, dc = shift
+
+    a = _extend(a, shift, mode, fill)
+
     R = dr * np.fft.fftfreq(a.shape[0])
     C = dc * np.fft.fftfreq(a.shape[1])
 
@@ -482,19 +504,58 @@ def shift(a, shift):
     K = np.exp(-1j*2*np.pi*(RR+CC))
     shifted = np.fft.ifft2(np.fft.fft2(a)*K)
 
+    shifted = pad(shifted, (r, c))  # crop back to original size
+
     if np.any(np.iscomplex(a)):
         return shifted
     else:
         return shifted.real
-    
+
+
+def _extend(a, shift, mode, fill):
+    r, c = a.shape
+    dr, dc = np.ceil(np.abs(shift)).astype(int)
+
+    if mode == 'wrap':  # (a b c d | a b c d | a b c d)
+        # no need to do anything since the Fourier transform
+        # already does this
+        out = a
+    else:
+        # pad the array to accomodate for
+        out = pad(a, shape=(r+2*dr, c+2*dc), fill=fill)
+
+        if mode == 'constant':  # (k k k k | a b c d | k k k k)
+            pass
+        elif mode == 'reflect':  # (d c b a | a b c d | d c b a)
+            out[0:dr, 0:dc] = np.flip(a[0:dr, 0:dc])  # upper left
+            out[0:dr, dc:c+dc] = np.flip(a[0:dr, :], axis=0)  # top
+            out[0:dr, c+dc:] = np.flip(a[0:dr, c-dc:])  # upper right
+            out[dr:r+dr, 0:dc] = np.flip(a[:, 0:dc], axis=1)  # left
+            out[dr:r+dr, c+dc:] = np.flip(a[:, c-dc:], axis=1)  # right
+            out[r+dr:, 0:dc] = np.flip(a[r-dr:r, 0:dc])  # lower left
+            out[r+dr:, dc:c+dc] = np.flip(a[r-dr:r, :], axis=0)  # bottom
+            out[r+dr:, c+dc:] = np.flip(a[r-dr:r, c-dc:c])  # lower right
+        elif mode == 'mirror':  # (d c b | a b c d | c b a)
+            out[0:dr, 0:dc] = np.flip(a[1:dr+1, 1:dc+1])  # upper left
+            out[0:dr, dc:c+dc] = np.flip(a[1:dr+1, :], axis=0)  # top
+            out[0:dr, c+dc:] = np.flip(a[1:dr+1, c-dc-1:c-1])  # upper right
+            out[dr:r+dr, 0:dc] = np.flip(a[:, 1:dc+1], axis=1)  # left
+            out[dr:r+dr, c+dc:] = np.flip(a[:, c-dc-1:c-1], axis=1)  # right
+            out[r+dr:, 0:dc] = np.flip(a[r-dr-1:r-1, 1:dc+1])  # lower left
+            out[r+dr:, dc:c+dc] = np.flip(a[r-dr-1:r-1, :], axis=0)  # bottom
+            out[r+dr:, c+dc:] = np.flip(a[r-dr-1:r-1, c-dc-1:c-1])  # lower right
+        else:
+            raise ValueError(f'Unknown mode {mode}')
+    return out
+
 
 def register(arr, ref, oversample, return_error=False):
     """Compute the subpixel image translation to register the input array to a
     reference array.
 
-    The registration shift is computed in two steps: first a coarse estimate 
+    The registration shift is computed in two steps: first a coarse estimate
     is computed from the FFT-based cross-correlation of the two input arrays.
-    This estimate is then refined to subpixel accuracy by computing the 
+    This estimate is then refined to subpixel accuracy by computing the
     upsampled DFT-based cross-correlation in a small neigborhood around the
     initial estimate.
 
@@ -519,7 +580,7 @@ def register(arr, ref, oversample, return_error=False):
 
     References
     ----------
-    Guizar-Sicairos, Thurman, and Fienup, "Efficient subpixel image 
+    Guizar-Sicairos, Thurman, and Fienup, "Efficient subpixel image
     registration algorithms". Optics Letters 33, 156-158 (2008)
 
     See also
@@ -542,7 +603,7 @@ def register(arr, ref, oversample, return_error=False):
     F = np.fft.fft2(arr)
     G = np.fft.fft2(ref)
     xcorr = np.fft.fftshift(np.fft.ifft2(G*np.conj(F)))
-    
+
     # find peak
     maxima = np.unravel_index(np.argmax(np.abs(xcorr)), xcorr.shape)
     peak = xcorr[maxima]
@@ -550,8 +611,8 @@ def register(arr, ref, oversample, return_error=False):
     # compute shifts
     center = np.array([np.fix(x/2) for x in arr.shape])
     shift = maxima - center
-    if oversample !=1:
-        # now we can set up and perform the oversampled dft on an oversampled 
+    if oversample != 1:
+        # now we can set up and perform the oversampled dft on an oversampled
         # 1.5 x 1.5 pixel region about the peak
         npix_dft = np.ceil(oversample*1.5)
         dft_shift = np.fix(npix_dft/2)
@@ -563,15 +624,15 @@ def register(arr, ref, oversample, return_error=False):
         Y = np.arange(arr.shape[0]) - np.floor(arr.shape[0]/2)
         U = np.arange(npix_dft) - cs
         V = np.arange(npix_dft) - rs
-        E1 = np.exp(-2*np.pi*1j/(arr.shape[0]*oversample)*np.outer(V,Y))
-        E2 = np.exp(-2*np.pi*1j/(arr.shape[1]*oversample)*np.outer(X,U))
-        xcorr = np.dot(np.dot(E1,np.conj(np.fft.ifftshift(G*np.conj(F)))),E2)
-        
+        E1 = np.exp(-2*np.pi*1j/(arr.shape[0]*oversample)*np.outer(V, Y))
+        E2 = np.exp(-2*np.pi*1j/(arr.shape[1]*oversample)*np.outer(X, U))
+        xcorr = np.dot(np.dot(E1, np.conj(np.fft.ifftshift(G*np.conj(F)))), E2)
+
         maxima_subpx = np.unravel_index(np.argmax(np.abs(xcorr)), xcorr.shape)
         peak = xcorr[maxima_subpx]
 
         # Combine subpixel peak coordinates with integer pixel peak coords
-        maxima_subpx -= dft_shift 
+        maxima_subpx -= dft_shift
         shift[0] += maxima_subpx[0]/oversample
         shift[1] += maxima_subpx[1]/oversample
 
@@ -588,7 +649,7 @@ def register(arr, ref, oversample, return_error=False):
     return shift
 
 
-def medfix(input, mask, kernel=(3,3), nanwarn=False):
+def medfix(input, mask, kernel=(3, 3), nanwarn=False):
     """Fix masked entries in a 2-dimensional array via median filtering.
 
     Parameters
@@ -599,8 +660,8 @@ def medfix(input, mask, kernel=(3,3), nanwarn=False):
         A 2-dimensional mask with the same shape as input. Entries which
         evaluate to True are considered masked and will be repaired.
     kernel : array_like, optional
-        A scalar or list of length 2 specifying the filter window in 
-        each dimension. Elements of *kernel* should be odd. If *kernel*
+        A scalar or list of length 2 specifying the filter window in
+        each dimension. Elements of *kernel* must be odd. If *kernel*
         is a scalar, it is used for each dimension. Default is (3,3).
     nanwarn : bool, optional
         If True, a RuntimeWarning will be raised if NaNs are present
@@ -609,7 +670,7 @@ def medfix(input, mask, kernel=(3,3), nanwarn=False):
     Returns
     -------
     ndarray
-   
+
     Notes
     -----
     Masked areas larger than the kernel size will introduce NaNs into
@@ -618,36 +679,55 @@ def medfix(input, mask, kernel=(3,3), nanwarn=False):
     """
     # force a copy by calling array instead of asarray
     input = np.array(input, dtype=float)
-
     mask = np.asarray(mask, dtype=bool)
 
+    if not mask.any():
+        # nothing to do
+        return input
+
     kernel = np.asarray(kernel)
     if kernel.shape == ():
         kernel = np.repeat(kernel, 2)
     if np.any(kernel % 2 == 0):
-        raise ValueError("Each element of kernel must be odd")
+        raise ValueError("Kernel must be odd sized")
+
+    if __backend__ == 'jax':
+        input = input.at[mask].set(np.nan)
+    else:
+        input[mask] = np.nan
 
-    input[mask] = np.nan
+    pw = (kernel - 1)//2
+    pad_width = ((pw[0], pw[0]), (pw[1], pw[1]))
 
-    pad_size = np.array((kernel-1)/2, dtype=int)
-    input_pad = np.full((input.shape[0] + pad_size[0] * 2, 
-                         input.shape[1] + pad_size[1] * 2), np.nan)
-    slc = (slice(pad_size[0], pad_size[0] + input.shape[0]), 
-           slice(pad_size[1], pad_size[1] + input.shape[1]))
-    input_pad[slc] = input
+    input_pad = np.pad(input, pad_width=pad_width, mode='constant',
+                       constant_values=np.nan)
 
-    input_pad = as_strided(input_pad, 
-                           shape=(input_pad.shape[0], input_pad.shape[1], kernel[0], kernel[1]),
-                           strides=input_pad.strides + input_pad.strides)
-    input_pad = input_pad.reshape((input_pad.shape[0], input_pad.shape[1], np.prod(kernel)))
+    i, j = np.nonzero(mask)  # indices of bad pixels
+    i_pad, j_pad = i + pw[0], j + pw[1]  # indices offset by padding width
+
+    # define neighborhood offsets
+    di = np.arange(kernel[0]) - pw[0]
+    dj = np.arange(kernel[1]) - pw[1]
+    window = np.stack(np.meshgrid(di, dj, indexing='ij'), axis=-1).reshape(-1, 2)  # shape (prod(kernel), 2)
+
+    # compute all neighborhood coordinates
+    rows = i_pad[:, None] + window[:, 0]  # shape (num_bad_px, prod(kernel))
+    cols = j_pad[:, None] + window[:, 1]  # shape (num_bad_px, prod(kernel))
+
+    # extract neighborhoods using advanced indexing and compute replacement
+    # values
+    bad_px_kernel = input_pad[rows, cols]
 
-    mask_idx = np.where(mask)
-    
     with warnings.catch_warnings():
         warnings.simplefilter('ignore', category=RuntimeWarning)
-        input[mask_idx] = np.nanmedian(input_pad[mask_idx[0], mask_idx[1], :], axis=1)
+        bad_px_vals = np.nanmedian(bad_px_kernel, axis=1)
+
+    if __backend__ == 'jax':
+        input = input.at[i, j].set(bad_px_vals)
+    else:
+        input[i, j] = bad_px_vals
 
-    if np.isnan(input).any():
+    if nanwarn and np.isnan(input).any():
         warnings.warn('Result contains NaNs', RuntimeWarning,
                       stacklevel=2)