warpfield.ndimage

  1import cupyx.scipy.signal.windows
  2import numpy as np
  3import cupy as cp
  4import cupyx
  5import scipy.ndimage
  6import cupyx.scipy.ndimage
  7import cupyx.scipy.signal
  8
  9
 10def dogfilter(vol, sigma_low=1, sigma_high=4, mode="reflect"):
 11    """Diffference of Gaussians filter
 12
 13    Args:
 14        vol (array_like): data to be filtered
 15        sigma_low (scalar or sequence of scalar): standard deviations
 16        sigma_high (scalar or sequence of scalar): standard deviations
 17        mode (str): The array borders are handled according to the given mode
 18
 19    Returns:
 20        (array_like): filtered data
 21
 22    See also:
 23        cupyx.scipy.ndimage.gaussian_filter
 24        skimage.filters.difference_of_gaussians
 25    """
 26    in_module = vol.__class__.__module__
 27    vol = cp.array(vol, "float32", copy=False)
 28    out = cupyx.scipy.ndimage.gaussian_filter(vol, sigma_low, mode=mode)
 29    out -= cupyx.scipy.ndimage.gaussian_filter(vol, sigma_high, mode=mode)
 30    if in_module == "numpy":
 31        out = out.get()
 32    return out
 33
 34
 35def periodic_smooth_decomposition_nd_rfft(img):
 36    """
 37    Decompose ND arrays of 2D images into periodic plus smooth components. This can help with edge artifacts in
 38    Fourier transforms.
 39
 40    Args:
 41        img (cupy.ndarray): input image or volume. The last two axes are treated as the image dimensions.
 42
 43    Returns:
 44        cupy.ndarray: periodic component
 45    """
 46    # compute border-difference
 47    B = cp.zeros_like(img)
 48    B[..., 0, :] = img[..., -1, :] - img[..., 0, :]
 49    B[..., -1, :] = -B[..., 0, :]
 50    B[..., :, 0] += img[..., :, -1] - img[..., :, 0]
 51    B[..., :, -1] -= img[..., :, -1] - img[..., :, 0]
 52
 53    # real FFT of border difference
 54    B_rfft = cp.fft.rfftn(B, axes=(-2, -1))
 55    del B
 56
 57    # build denom for full grid then slice to half-spectrum
 58    M, N = img.shape[-2:]
 59    q = cp.arange(M, dtype="float32").reshape(M, 1)
 60    r = cp.arange(N, dtype="float32").reshape(1, N)
 61    denom_full = 2 * cp.cos(2 * np.pi * q / M) + 2 * cp.cos(2 * np.pi * r / N) - 4
 62    # take only first N//2+1 columns
 63    denom_half = denom_full[:, : (N // 2 + 1)]
 64    denom_half[0, 0] = 1  # avoid divide by zero
 65
 66    # compute smooth in freq domain (half-spectrum)
 67    B_rfft /= denom_half
 68    B_rfft[..., 0, 0] = 0
 69
 70    # invert real FFT back to spatial
 71    # smooth = cp.fft.irfftn(B_rfft, s=(M, N), axes=(-2, -1))
 72    # periodic = img - smooth
 73    tmp = cp.fft.irfftn(B_rfft, s=(M, N), axes=(-2, -1))
 74    tmp *= -1
 75    tmp += img
 76    return tmp
 77
 78
 79def gausswin(shape, sigma):
 80    """Create Gaussian window of a given shape and sigma
 81
 82    Args:
 83        shape (list or tuple): shape along each dimension
 84        sigma (list or tuple): sigma along each dimension
 85
 86    Returns:
 87        (array_like): Gauss window
 88    """
 89    grid = np.indices(shape).astype("float32")
 90    for dim in range(len(grid)):
 91        grid[dim] -= shape[dim] // 2
 92        grid[dim] /= sigma[dim]
 93    out = np.exp(-(grid**2).sum(0) / 2)
 94    out /= out.sum()
 95    # out = np.fft.fftshift(out)
 96    return out
 97
 98
 99def gausskernel_sheared(sigma, shear=0, truncate=3):
100    """Create Gaussian window of a given shape and sigma. The window is sheared along the first two axes.
101
102    Args:
103        sigma (float or tuple of float): Standard deviation for Gaussian kernel.
104        shear (float): Shear factor in d_axis0 / d_axis1
105        truncate (float): Truncate the filter at this many standard deviations.
106
107    Returns:
108        window (array_like): n-dimensional window
109    """
110    # TODO: consider moving to .unshear
111
112    shape = (np.r_[sigma] * truncate * 2).astype("int")
113    shape[0] = np.maximum(shape[0], int(np.ceil(shape[1] * np.abs(shear))))
114    shape = (shape // 2) * 2 + 1
115    grid = np.indices(shape).astype("float32")
116    for dim in range(len(grid)):
117        grid[dim] -= shape[dim] // 2
118        grid[dim] /= sigma[dim]
119    grid[0] = grid[0] + shear * grid[1] * sigma[1] / sigma[0]
120    out = np.exp(-(grid**2).sum(0) / 2)
121    out /= out.sum()
122    return out
123
124
125def ndwindow(shape, window_func):
126    """Create a n-dimensional window function
127
128    Args:
129        shape (tuple): shape of the window
130        window_func (function): window function to be applied to each dimension
131
132    Returns:
133        window (array_like): n-dimensional window
134    """
135    out = 1
136    for i in range(len(shape)):
137        newshape = np.ones(len(shape), dtype="int")
138        newshape[i] = shape[i]
139        w = window_func(shape[i])
140        out = out * np.reshape(w, newshape)
141    return out
142
143
144def accumarray(coords, shape, weights=None, clip=False):
145    """Accumulate values into an array using given coordinates and weights
146
147    Args:
148        coords (array_like): 3-by-n array of coordinates
149        shape (tuple): shape of the output array
150        weights (array_like): weights to be accumulated. If None, all weights are set to 1
151        clip (bool): if True, clip coordinates to the shape of the output array, else ignore out-of-bounds coordinates. Default is False.
152
153    Returns:
154        accum (array_like): accumulated array of the given shape
155    """
156    assert coords.shape[0] == 3
157    coords = np.round(coords.reshape(3, -1)).astype("int")
158    if clip:
159        for d in len(shape):
160            coords[d] = np.minimum(np.maximum(coords[d], 0), shape[d] - 1)
161    else:
162        valid_ix = np.all((coords >= 0) & (coords < np.array(shape)[:, None]), axis=0)
163        coords = coords[:, valid_ix]
164        if weights is not None:
165            weights = weights.ravel()[valid_ix]
166    coords_as_ix = np.ravel_multi_index((*coords,), shape).ravel()
167    accum = np.bincount(coords_as_ix, minlength=np.prod(shape), weights=weights)
168    accum = accum.reshape(shape)
169    return accum
170
171
172def infill_nans(arr, sigma=0.5, truncate=50, ignore_warning=True):
173    """Infill NaNs in an array using Gaussian basis interpolation
174
175    Args:
176        arr (array_like): input array
177        sigma (float): standard deviation of the Gaussian basis function
178        truncate (float): truncate the filter at this many standard deviations. Note: values outside the truncation may still contain NaNs.
179        ignore_warning (bool): if True, ignore warnings about invalid values during division
180    """
181    nans = np.isnan(arr)
182    arr_zeros = arr.copy()
183    arr_zeros[nans] = 0
184    a = scipy.ndimage.gaussian_filter(np.array(arr_zeros, dtype="float64"), sigma=sigma, truncate=truncate)
185    b = scipy.ndimage.gaussian_filter(np.array(~nans, dtype="float64"), sigma=sigma, truncate=truncate)
186    if ignore_warning:
187        with np.errstate(invalid="ignore"):
188            out = (a / b).astype(arr.dtype)
189    else:
190        out = (a / b).astype(arr.dtype)
191    return out
192
193
194def sliding_block(data, block_size=100, block_stride=1):
195    """Create a sliding window/block view into the array with the given block shape and stride. The block slides across all dimensions of the array and extracts subsets of the array at all positions.
196
197    Args:
198        data (array_like): Array to create the sliding window view from
199        block_size (int or tuple of int): Size of window over each axis that takes part in the sliding block
200        block_stride (int or tuple of int): Stride of teh window along each axis
201
202    Returns:
203        view (ndarray): Sliding block view of the array.
204
205    See Also:
206        numpy.lib.stride_tricks.sliding_window_view
207        numpy.lib.stride_tricks.as_strided
208
209    """
210    block_stride *= np.ones(data.ndim, dtype="int")
211    block_size *= np.ones(data.ndim, dtype="int")
212    shape = np.r_[1 + (data.shape - block_size) // block_stride, block_size]
213    strides = np.r_[block_stride * data.strides, data.strides]
214    xp = cp.get_array_module(data)
215    out = xp.lib.stride_tricks.as_strided(data, shape, strides)
216    return out
217
218
219def upsampled_dft_rfftn(
220    data: cp.ndarray, upsampled_region_size, upsample_factor: int = 1, axis_offsets=None
221) -> cp.ndarray:
222    """
223    Performs an upsampled inverse DFT on a small region around given offsets,
224    taking as input the output of cupy.fft.rfftn (real-to-complex FFT).
225
226    This implements the Guizar‑Sicairos local DFT upsampling: no full zero‑padding,
227    just a small m×n patch at subpixel resolution.
228
229    Args:
230        data: A real-to-complex FFT array of shape (..., M, Nf),
231            where Nf = N//2 + 1 corresponds to an original real image width N.
232        upsampled_region_size: Size of the output patch (m, n). If an int is
233            provided, the same size is used for both dimensions.
234        upsample_factor: The integer upsampling factor in each axis.
235        axis_offsets: The center of the patch in original-pixel coordinates
236            (off_y, off_x). If None, defaults to (0, 0).
237
238    Returns:
239        A complex-valued array of shape (..., m, n) containing the
240        upsampled inverse DFT patch.
241    """
242    if data.ndim < 2:
243        raise ValueError("Input must have at least 2 dimensions")
244    *batch_shape, M, Nf = data.shape
245    # determine patch size
246    if isinstance(upsampled_region_size, int):
247        m, n = upsampled_region_size, upsampled_region_size
248    else:
249        m, n = upsampled_region_size
250    # full width of original image
251    N = (Nf - 1) * 2
252
253    # default offset: origin
254    off_y, off_x = (0.0, 0.0) if axis_offsets is None else axis_offsets
255
256    # reconstruct full complex FFT via Hermitian symmetry
257    full = cp.empty(batch_shape + [M, N], dtype=cp.complex64)
258    full[..., :Nf] = data
259    if Nf > 1:
260        tail = data[..., :, 1:-1]
261        full[..., Nf:] = tail[..., ::-1, ::-1].conj()
262
263    # frequency coordinates
264    fy = cp.fft.fftfreq(M)[None, :]  # shape (1, M)
265    fx = cp.fft.fftfreq(N)[None, :]  # shape (1, N)
266
267    # sample coordinates around offsets
268    y_idx = cp.arange(m) - (m // 2)
269    x_idx = cp.arange(n) - (n // 2)
270    y_coords = off_y[:, None] + y_idx[None, :] / upsample_factor  # (B, m)
271    x_coords = off_x[:, None] + x_idx[None, :] / upsample_factor  # (B, n)
272
273    # Build small inverse‐DFT kernels
274    ky = cp.exp(2j * cp.pi * y_coords[:, :, None] * fy[None, :, :]).astype("complex64")
275    kx = cp.exp(2j * cp.pi * x_coords[:, :, None] * fx[None, :, :]).astype("complex64")
276
277    # First apply along y: (B,m,M) × (B,M,N) -> (B,m,N)
278    out1 = cp.einsum("b m M, b M N -> b m N", ky, full)
279    # Then along x: (B,m,N) × (B,n,N)ᵀ -> (B,m,n)
280    patch = cp.einsum("b m N, b n N -> b m n", out1, kx)
281
282    return patch.real.reshape(*batch_shape, m, n)
283
284
285def zoom_chop_pad(
286    arr, target_shape=None, scale=(1, 1, 1), soft_edge=(0, 0, 0), shift=(0, 0, 0), flip=(False, False, False), cval=0
287):
288    """Zooms, softens, flips, shifts, and pads/crops a 3D array to match the target shape.
289
290    The conceptual order is as follows: zoom, soften, flip, shift, crop/pad.
291
292    Args:
293        arr (np.ndarray or cp.ndarray): The input array to be transformed
294        target_shape (tuple of int): The desired target shape to pad/crop to. Defaults to the shape of the input array.
295        scale (tuple): Zoom factors for each axis. Default: (1, 1, 1).
296        soft_edge (tuple of int): The size of the soft edge (Tukey envelope) to be applied to the input array, in voxels. Default: (0, 0, 0).
297        shift (tuple): Shifts for each axis, in voxels. Default: (0, 0, 0).
298        flip (tuple of bool): Whether to flip each axis. Default: (False, False, False).
299        cval (int, float): The value to use for padding. Default: 0.
300
301    Returns:
302        np.ndarray or cp.ndarray: The transformed array. Dtype is float32.
303    """
304
305    was_numpy = isinstance(arr, np.ndarray)
306
307    if target_shape is None:
308        target_shape = arr.shape
309
310    if any(s > 0 for s in soft_edge):
311        arr = cp.array(arr, dtype="float32", copy=True, order="C")
312        scaled_edge = np.array(soft_edge) / np.array(scale)
313        arr = soften_edges(arr, soft_edge=scaled_edge, copy=False)
314    else:
315        arr = cp.array(arr, dtype="float32", copy=False, order="C")
316
317    coords = cp.indices(target_shape, dtype=cp.float32)
318    for i in range(len(coords)):
319        coords[i] -= target_shape[i] / 2
320        coords[i] /= scale[i]
321        coords[i] += arr.shape[i] / 2
322        if flip[i]:
323            coords[i] *= -1
324            coords[i] += arr.shape[i] - 1
325        coords[i] -= shift[i]
326    result = cupyx.scipy.ndimage.map_coordinates(arr, coords, order=1, mode="constant", cval=cval)
327
328    if was_numpy:
329        result = result.get()
330    return result
331
332
333def soften_edges(arr, soft_edge=(0, 0, 0), copy=True):
334    """Apply a soft Tukey edge to the input array.
335
336    Args:
337        arr (np.ndarray or cp.ndarray): The input array
338        soft_edge (tuple of int): The size of the soft edge (Tukey envelope) to be applied to the input array, in voxels. Default: (0, 0, 0).
339        copy (bool): If True, a copy of the array is made. Default: True.
340
341    Returns:
342        np.ndarray or cp.ndarray: The transformed array. Dtype is float32.
343    """
344    was_numpy = isinstance(arr, np.ndarray)
345    input_dtype = arr.dtype
346    arr = cp.array(arr, dtype="float32", copy=copy)
347    if isinstance(soft_edge, int) or isinstance(soft_edge, float):
348        soft_edge = (soft_edge,) * arr.ndim
349    soft_edge = np.clip(soft_edge, 0, np.array(arr.shape) / 2)
350
351    for i in range(arr.ndim):
352        if soft_edge[i] > 0:
353            alpha = 2 * soft_edge[i] / arr.shape[i]
354            alpha = np.clip(alpha, 0, 1)
355            win = cupyx.scipy.signal.windows.tukey(arr.shape[i], alpha)
356            arr *= cp.moveaxis(win[:, None, None], 0, i)
357
358    arr = arr.astype(input_dtype, copy=False)
359    if was_numpy:
360        arr = arr.get()
361    return arr
362
363
364def zoom(arr, zoom_factors, order=1, mode="constant"):
365    """Zooms an array by given factors along each axis.
366
367    Args:
368        arr (np.ndarray or cp.ndarray): The input array to be zoomed.
369        zoom_factors (tuple of float): Zoom factors for each axis. Values greater than 1 result in a larger output array,
370            while values less than 1 result in a smaller array. Divide the physical voxel size of the input array by these values to get the physical voxel size of the output array.
371        order (int): The order of the spline interpolation. Default is 1 (linear).
372
373    Returns:
374        np.ndarray or cp.ndarray: The zoomed array.
375    """
376    was_numpy = isinstance(arr, np.ndarray)
377    arr = cp.array(arr, dtype="float32", copy=False, order="C")
378    out = cupyx.scipy.ndimage.zoom(arr, zoom_factors, order=order)
379    if was_numpy:
380        out = out.get()
381    return out
382
383
384def match_volumes(fixed, fixed_res, moving, moving_res, order=1, soft_edge=(0, 0, 0), cval=0, res_method="fixed"):
385    """
386    Rescale and pad the fixed and moving volumes so both have the same physical size and resolution.
387
388    Args:
389        fixed (ndarray): Fixed/reference volume.
390        fixed_res (tuple): Pixel size (resolution) for fixed, in physical units.
391        moving (ndarray): Moving volume.
392        moving_res (tuple): Pixel size (resolution) for moving, in physical units.
393        order (int): Interpolation order for zooming.
394        soft_edge (tuple): Soft edge parameter for padding.
395        cval (float): Constant value for padding.
396        res_method (str): How to choose the target resolution.
397            "fixed" (default): use fixed_res,
398            "min": use the finest (smallest) resolution,
399            "max": use the coarsest (largest) resolution,
400            "mean": use the mean of fixed_res and moving_res.
401
402    Returns:
403        fixed_out (ndarray): The fixed volume, rescaled and padded to the target resolution and physical size.
404        moving_out (ndarray): The moving volume, rescaled and padded to the target resolution and physical size.
405        target_res (tuple): The target resolution used for both volumes.
406    """
407    fixed_res = np.array(fixed_res)
408    moving_res = np.array(moving_res)
409    fixed_shape = np.array(fixed.shape)
410    moving_shape = np.array(moving.shape)
411
412    # Compute physical sizes
413    fixed_phys = fixed_shape * fixed_res
414    moving_phys = moving_shape * moving_res
415
416    # Target: match the larger physical size along each axis
417    target_phys = np.maximum(fixed_phys, moving_phys)
418
419    # Determine target resolution
420    if res_method == "fixed":
421        target_res = fixed_res
422    elif res_method == "min":
423        target_res = np.minimum(fixed_res, moving_res)
424    elif res_method == "max":
425        target_res = np.maximum(fixed_res, moving_res)
426    elif res_method == "mean":
427        target_res = (fixed_res + moving_res) / 2
428    else:
429        raise ValueError(f"Unknown target_res_type: {res_method}")
430
431    # Compute the target shape
432    target_shape = np.ceil(target_phys / target_res).astype(int)
433
434    # Rescale fixed
435    scale_fixed = fixed_res / target_res
436    fixed_out = zoom_chop_pad(
437        fixed, target_shape=target_shape, scale=scale_fixed, soft_edge=soft_edge, cval=cval, order=order
438    )
439
440    # Rescale moving
441    scale_moving = moving_res / target_res
442    moving_out = zoom_chop_pad(
443        moving, target_shape=target_shape, scale=scale_moving, soft_edge=soft_edge, cval=cval, order=order
444    )
445
446    return fixed_out, moving_out, tuple(target_res)
447
448
449def richardson_lucy_generic(img, convolve_psf, correlate_psf=None, num_iter=5, epsilon=1e-3, beta=0.0, initial_guess=None):
450    """Richardson-Lucy deconvolution using arbitrary convolution operations, with optional Biggs acceleration.
451
452    Args:
453        img (ArrayLike): input image or volume
454        convolve_psf (Callable): function to convolve with PSF. Should take an image and return a convolved image. Ensure that the PSF is non-negative and normalized.
455        correlate_psf (Callable): function to correlate with PSF. Defaults to convolve_psf (if PSF is symmetric)
456        num_iter (int): number of iterations. Default is 5.
457        epsilon (float): small constant to prevent divide-by-zero. Default is 1e-3.
458        beta (float): acceleration parameter. Default is 0.0 (no acceleration). Typically, beta is in the range [0, 0.5].
459
460    Returns:
461        ndarray: deconvolved image
462    """
463    epsilon = cp.float32(epsilon)
464    img = cp.array(img, dtype="float32", copy=False)
465    cp.clip(img, 0, None, out=img)
466    if num_iter < 1:
467        return img
468    if correlate_psf is None:
469        correlate_psf = convolve_psf
470
471    if initial_guess is not None:
472        assert initial_guess.shape == img.shape, "Initial guess must have the same shape as the input image."
473        img_decon = cp.array(initial_guess, dtype="float32", copy=False)
474        cp.clip(img_decon, 0, None, out=img_decon)
475    else:   
476        img_decon = img.copy()
477    img_decon += epsilon
478
479    for i in range(num_iter):
480        img_decon *= correlate_psf(img / (convolve_psf(img_decon) + epsilon))
481
482        if beta > 0:
483            if i == 0: 
484                img_decon_prev = img_decon.copy()
485            else:
486                img_decon_new = img_decon.copy()
487                img_decon += beta * (img_decon - img_decon_prev)
488                cp.clip(img_decon, epsilon, None, out=img_decon)
489                img_decon_prev = img_decon_new
490    
491    return img_decon
492
493
494def richardson_lucy_fft(img, psf, num_iter=5, epsilon=1e-3, beta=0.0, initial_guess=None):
495    """Richardson-Lucy deconvolution using FFT-based convolution and optional Biggs acceleration.
496
497    Args:
498        img (ndarray): input image or volume
499        psf (ndarray): point spread function (before fftshift)
500        num_iter (int): number of iterations
501        epsilon (float): small constant to avoid divide-by-zero
502        beta (float): Biggs acceleration parameter (0 = no acceleration)
503
504    Returns:
505        ndarray: deconvolved image
506    """
507    psf = cp.array(psf, dtype="float32")
508    cp.clip(psf, 0, None, out=psf)
509    psf /= psf.sum()
510
511    shape = img.shape
512    psf_ft = cp.fft.rfftn(cp.fft.ifftshift(psf), s=shape)
513    psf_ft_conj = cp.conj(psf_ft)
514
515    def convolve(x):
516        return cp.fft.irfftn(cp.fft.rfftn(x, s=shape) * psf_ft, s=shape)
517
518    def correlate(x):
519        return cp.fft.irfftn(cp.fft.rfftn(x, s=shape) * psf_ft_conj, s=shape)
520
521    out = richardson_lucy_generic(
522        img, convolve, correlate, num_iter=num_iter, epsilon=epsilon, beta=beta, initial_guess=initial_guess
523    )
524    return out
525
526
527def richardson_lucy_gaussian(img, sigmas, num_iter=5, epsilon=1e-3, beta=0.0, initial_guess=None):
528    """Richardson-Lucy deconvolution using Gaussian convolution operations
529
530    Args:
531        img (ndarray): input image or volume
532        sigmas (list or ndarray): list of Gaussian sigmas along each dimension
533        num_iter (int): number of iterations
534        epsilon (float): small constant to prevent divide-by-zero
535        beta (float): acceleration parameter (0 = no acceleration)
536
537    Returns:
538        ndarray: deconvolved image
539    """
540    conv_with_gauss = lambda x: cupyx.scipy.ndimage.gaussian_filter(x, sigmas)
541    out = richardson_lucy_generic(
542        img, conv_with_gauss, num_iter=num_iter, epsilon=epsilon, beta=beta, initial_guess=initial_guess
543    )
544    return out
545
546
547def richardson_lucy_gaussian_shear(img, sigmas, shear, num_iter=5, epsilon=1e-3, beta=0.0, initial_guess=None):
548    """Richardson-Lucy deconvolution using a sheared Gaussian psf
549
550    Args:
551        img (ndarray): input image or volume
552        sigmas (list or ndarray): list of Gaussian sigmas along each dimension
553        shear (scalar): shear ratio
554        num_iter (int): number of iterations
555        epsilon (float): small constant to prevent divide-by-zero
556        beta (float): acceleration parameter (0 = no acceleration)
557
558    Returns:
559        ndarray: deconvolved image
560    """
561    if shear == 0:
562        return richardson_lucy_gaussian(img, sigmas, num_iter)
563
564    sigmas = np.array(sigmas)
565    gw = cp.array(gausskernel_sheared(sigmas, shear=shear, truncate=4), "float32")
566    gw01 = gw.sum(2)[:, :, None]
567    gw01 /= gw01.sum()
568    gw2 = gw.sum(axis=(0, 1))[None, None, :]
569    gw2 /= gw2.sum()
570    conv_shear = lambda vol: cupyx.scipy.ndimage.convolve(cupyx.scipy.ndimage.convolve(vol, gw01), gw2)
571    out = richardson_lucy_generic(
572        img, conv_shear, num_iter=num_iter, epsilon=epsilon, beta=beta, initial_guess=initial_guess
573    )
574    return out