prtools 1.3.2__tar.gz → 1.4.1__tar.gz

{prtools-1.3.2/prtools.egg-info → prtools-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prtools
-Version: 1.3.2
+Version: 1.4.1
 Summary: Utility functions for image-based phase retrieval
 Author: Andy Kee
 License: MIT License

{prtools-1.3.2 → prtools-1.4.1}/prtools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '1.3.2'
+__version__ = '1.4.1'
 
 from prtools.backend import __backend__, use
 
@@ -15,6 +15,16 @@ from prtools.array import (
     medfix
 )
 
+from prtools.convolve import (
+    fftconv,
+    gauss,
+    gauss_blur,
+    gauss_kernel,
+    pixel_kernel,
+    pixelate,
+    sinc,
+)
+
 from prtools.cost import sserror
 
 from prtools.fourier import dft2, idft2
@@ -38,7 +48,6 @@ from prtools.shape import (
     hexagon,
     rectangle,
     spider,
-    gauss,
     sin,
     waffle
 )

{prtools-1.3.2 → prtools-1.4.1}/prtools/array.py

@@ -3,6 +3,7 @@ import warnings
 from prtools import __backend__
 from prtools.backend import numpy as np
 from prtools.backend import scipy
+import prtools.jax
 
 
 def centroid(a, where=None, kind='absolute', indexing='ij'):
@@ -300,11 +301,14 @@ def rebin(img, factor):
         raise ValueError('rebin is not defined for complex data')
 
     if img.ndim == 3:
-        rebinned_shape = (img.shape[0], img.shape[1]//factor, img.shape[2]//factor)
-        img_rebinned = np.zeros(rebinned_shape, dtype=img.dtype)
-        for i in range(img.shape[0]):
-            img_rebinned[i] = img[i].reshape(rebinned_shape[1], factor,
-                                             rebinned_shape[2], factor).sum(-1).sum(1)
+        if __backend__ == 'jax':
+            return prtools.jax._ndrebin(img, factor)
+        else:
+            rebinned_shape = (img.shape[0], img.shape[1]//factor, img.shape[2]//factor)
+            img_rebinned = np.zeros(rebinned_shape, dtype=img.dtype)
+            for i in range(img.shape[0]):
+                img_rebinned[i] = img[i].reshape(rebinned_shape[1], factor,
+                                                 rebinned_shape[2], factor).sum(-1).sum(1)
     else:
         img_rebinned = img.reshape(img.shape[0]//factor, factor, img.shape[1]//factor,
                                    factor).sum(-1).sum(1)

{prtools-1.3.2 → prtools-1.4.1}/prtools/backend/_jax.py

@@ -6,6 +6,7 @@ from ._base import BackendLibrary
 class Numpy(BackendLibrary):
     def __init__(self):
         super().__init__(importlib.import_module('jax.numpy'))
+        self.jax = importlib.import_module('jax')
 
     def broadcast_to(self, array, shape):
         # jax numpy.broadcast_to expects an array input
@@ -43,6 +44,12 @@ class Numpy(BackendLibrary):
         a = self.module.asarray(a)
         return self.module.sum(a, *args, **kwargs)
 
+    def take(self, a, indices, *args, **kwargs):
+        # jax numpy.take expects an array input for a and indices
+        a = self.module.asarray(a)
+        indices = self.module.asarray(indices)
+        return self.module.take(a, indices, *args, **kwargs)
+
 
 class Scipy(BackendLibrary):
     def __init__(self):

prtools-1.4.1/prtools/convolve.py

@@ -0,0 +1,339 @@
+from prtools.backend import numpy as np
+
+
+def fftconv(array, kernel, normalize_kernel=True, fft_array=True,
+            fft_kernel=False, fftshift_kernel=False):
+    r"""Convolve an array with a kernel using the FFT.
+
+    The colvolution is computed as
+
+    .. math::
+
+        u*v = \mathcal{F}^{-1}\left\{\mathcal{F}\{u\}\cdot\mathcal{F}\{v\}\right\}
+
+    Parameters
+    ----------
+    array : array_like
+        Array to be convolved with ``kernel``.
+    kernel : array_like
+        Convolution kernel. Should have the same shape as ``array``.
+    normalize_kernel : bool, optional
+        If True (default), kernel is normalized so that  ``sum(kernel) == 1``.
+    fft_array : bool, optional
+        If True (default), the array is assumed to be provided in the spatial
+        domain and its FFT will be computed by this function. If False, the
+        array is assumed to be provided in the frequency domain.
+    fft_kernel : bool, optional
+        If True, the kernel is assumed to be provided in the spatial
+        domain and its FFT will be computed by this function. If False
+        (default), the kernel is assumed to be provided in the frequency
+        domain.
+    fftshift_kernel : bool, optional
+        If True, the supplied kernel will be shifted so that its
+        zero-frequency (DC) term is placed at the upper left ``(0,0)`` corner
+        of the array. Default is False.
+
+    Returns
+    -------
+    ndarray
+
+    See also
+    --------
+    :func:`~prtools.gauss_blur`
+    :func:`~prtools.pixelate`
+    """
+
+    a = np.fft.fft2(array) if fft_array else array
+    k = np.fft.fftshift(kernel) if fftshift_kernel else kernel
+    k = np.fft.fft2(k) if fft_kernel else k
+
+    if normalize_kernel:
+        k /= np.sum(k)
+
+    return np.fft.ifft2(a*k).real
+
+
+def gauss_blur(img, sigma, oversample=1):
+    """Blur an image using a Gaussian filter using the FFT. 
+
+    Parameters
+    ----------
+    img : array_like
+        Array to be blurred
+    sigma : float or (2,) array_like
+        Standard deviation for the Gaussian kernel. Providing two values
+        allows for non-symmetric Gaussian interpreted as `(sigma_row,
+        sigma_col)`
+    oversample : float, optional
+        Oversampling factor of `img`. Default is 1.
+    
+    Returns
+    -------
+    ndarray
+
+    Examples
+    --------
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> img = np.zeros((100,100))
+        >>> img[0:50,0:50] = 1
+        >>> img[50:100,50:100] = 1
+        >>> img = np.tile(img, (3,3))
+        >>> img_blurred = prtools.gauss_blur(img, sigma=3)
+        >>> fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(5,2))
+        >>> ax[0].imshow(img, cmap='gray')
+        >>> ax[0].set_title('Input image')
+        >>> ax[1].imshow(img_blurred, cmap='gray')
+        >>> ax[1].set_title('Blurred image')
+    """
+    img = np.asarray(img)
+    kernel = gauss_kernel(img.shape, sigma, oversample, fftshift=False)
+    # gauss_kernel always returns a normalized kernel
+    return fftconv(img, kernel, normalize_kernel=False, fft_array=True,
+                   fft_kernel=False, fftshift_kernel=False)
+
+
+def pixelate(img, oversample=1):
+    """Apply the aperture effects of an idealized square pixel using the FFT.
+
+    Parameters
+    ----------
+    img : array_like
+        Input image
+    oversample : float, optional
+        Oversampling factor of ``img``. Default is 1.
+
+    Returns
+    -------
+    out : ndarray
+        Image with pixel sampling effects applied.
+
+    Notes
+    -----
+    To avoid the introduction of numerical artifacts, this function should be
+    performed on data that is at least 2x oversampled.
+    """
+    img = np.asarray(img)
+    kernel = pixel_kernel(img.shape, oversample=oversample, fftshift=False)
+    return fftconv(img, kernel, normalize_kernel=False, fft_array=True,
+                   fft_kernel=False, fftshift_kernel=False)
+
+
+def gauss(x1, x2, sigma, indexing='ij', normalize=False):
+    """2D Gaussian function
+
+    Parameters
+    ----------
+    x1, x2 : array_like
+        1-D arrays representing the grid coordinates
+    sigma : float or (2,) array_like
+        Standard deviation of Gaussian. Providing two values allows for
+        non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
+    indexing : {'ij', 'xy'}, optional
+        Matrix ('ij', default) or Cartesian ('xy') indexing of output.
+    normalize : bool, optional
+        If True, the output is normalized such that its sum is equal to 1.
+        If False (default), the output has max equal to one.
+
+    Returns
+    -------
+    ndarray
+
+    Examples
+    --------
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> r = np.arange(-128,128)
+        >>> c = np.arange(-128,128)
+        >>> g = prtools.gauss(r, c, sigma=20)
+        >>> plt.imshow(g, cmap='gray')
+    """
+    xx1, xx2 = np.meshgrid(x1, x2, indexing=indexing)
+    sigma = np.broadcast_to(sigma, (2,))
+    g = np.exp(-((xx1**2/(2*sigma[0]**2)) + (xx2**2/(2*sigma[1]**2))))
+    if normalize:
+        g = g / (2*np.pi * np.prod(sigma))
+    return g
+
+
+def sinc(x1, x2, indexing='ij'):
+    r"""2D sinc function
+
+    Parameters
+    ----------
+    x1, x2 : array_like
+        1-D arrays representing the grid coordinates
+    indexing : {'ij', 'xy'}, optional
+        Matrix ('ij', default) or Cartesian ('xy') indexing of output.
+    normalize : bool, optional
+        If True, the output is normalized such that its sum is equal to 1.
+        If False (default), the output has max equal to one.
+
+    Returns
+    -------
+    ndarray
+
+    See also
+    --------
+    :func:`~prtools.pixel_kernel`
+
+    Notes
+    -----
+    This function uses ``numpy.sinc``, which computes the sinc function as
+
+    .. math::
+        \mbox{sinc}(x) = \frac{\sin \pi x}{\pi x}
+
+    Examples
+    --------
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> r = np.linspace(-3,3,512)
+        >>> c = np.linspace(-3,3,512)
+        >>> s = prtools.sinc(r, c)
+        >>> plt.imshow(s, cmap='gray', extent=[r.min(), r.max(), c.min(), c.max()])
+    """
+    xx1, xx2 = np.meshgrid(x1, x2, indexing=indexing)
+    return np.sinc(xx1) * np.sinc(xx2)
+
+
+def gauss_kernel(shape, sigma, oversample=1, pixelscale=1.0, fftshift=False):
+    """2D Gaussian filter kernel
+
+    This function returns the Fourier transform of a normalized 2D Gaussian
+    function.
+
+    Parameters
+    ----------
+    shape : int or tuple of int
+        Shape of the kernel
+    sigma : float
+        Standard deviation of Gaussian. Providing two values allows for a
+        non-symmetric Gaussian interpreted as ``(sigma_row, sigma_col)``.
+    oversample : float, optional
+        Oversampling factor to represent in the kernel. Default is 1.
+    pixelscale : scalar or tuple of scalars, optional
+        Sample spacing. Providing two values defines non-symmetric sampling
+        interpreted as ``(pixelscale_row, pixelscale_col)``. Default is 1.
+    fftshift : bool, optional
+        If True, the kernel is FFT-shifted so that the zero-frequency (DC)
+        term is placed at the center of the returned array. Default is False.
+
+    Returns
+    -------
+    ndarray
+
+    See also
+    --------
+    :func:`~prtools.fftconv`
+    :func:`~prtools.gauss_blur`
+    :func:`~prtools.gauss`
+
+    Examples
+    --------
+    Create a Gaussian kernel with 
+
+    .. plot::
+        :include-source:
+        :scale: 50
+
+        >>> G = prtools.gauss_kernel((256,256), sigma=1, fftshift=True)
+        >>> plt.imshow(G, cmap='gray')
+
+    Oversampling is in effect the same as zero-padding:
+
+    .. plot::
+        :include-source:
+        :scale: 50
+
+        >>> G = prtools.gauss_kernel((512,512), sigma=1, oversample=2,
+        ...                          fftshift=True)
+        >>> plt.imshow(G, cmap='gray')
+
+    For a kernel with custom sample spacing:
+
+    .. plot::
+        :include-source:
+        :scale: 50
+
+        >>> G = prtools.gauss_kernel((256,256), sigma=1, pixelscale=0.25,
+        ...                          fftshift=True)
+        >>> plt.imshow(G, cmap='gray')
+    """
+    
+    shape = np.broadcast_to(shape, (2,))
+    pixelscale = np.broadcast_to(pixelscale, (2,))
+    sigma = np.broadcast_to(sigma, (2,))
+
+    x1 = np.fft.fftfreq(n=shape[0], d=pixelscale[0])
+    x2 = np.fft.fftfreq(n=shape[1], d=pixelscale[1])
+
+    k = gauss(x1, x2, sigma=1/(2*np.pi*sigma*oversample), indexing='ij',
+              normalize=False)
+
+    if fftshift:
+        return np.fft.fftshift(k)
+    else:
+        return k
+
+
+def pixel_kernel(shape, oversample=1, pixelscale=1.0, fftshift=False):
+    r"""2D pixel MTF filter kernel
+    
+    This function returns a normalized 2D sinc function sized to represent
+    the transfer function of an idealized square pixel.
+
+    Parameters
+    ----------
+    shape : int or tuple of int
+        Shape of the kernel
+    oversample : float, optional
+        Oversampling factor to represent in the kernel. Default is 1.
+    pixelscale : scalar or tuple of scalars, optional
+        Sample spacing. Providing two values defines non-symmetric sampling
+        interpreted as ``(pixelscale_row, pixelscale_col)``. Default is 1.
+    fftshift : bool, optional
+        If True, the kernel is FFT-shifted so that the zero-frequency (DC)
+        term is placed at the center of the returned array. Default is False.
+
+    Returns
+    -------
+    ndarray
+
+    See also
+    --------
+    :func:`~prtools.fftconv`
+    :func:`~prtools.pixelate`
+    :func:`~prtools.sinc`
+
+    Examples
+    --------
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> pixel_mtf = prtools.pixel_kernel((256,256), 1, fftshift=True)
+        >>> plt.imshow(pixel_mtf, cmap='gray', vmin=0)
+    """
+    shape = np.broadcast_to(shape, (2,))
+    pixelscale = np.broadcast_to(pixelscale, (2,))
+
+    x1 = np.fft.fftfreq(n=shape[0], d=pixelscale[0]) * oversample
+    x2 = np.fft.fftfreq(n=shape[1], d=pixelscale[1]) * oversample
+
+    k = np.abs(sinc(x1, x2, indexing='ij'))
+
+    if fftshift:
+        return np.fft.fftshift(k)
+    else:
+        return k

{prtools-1.3.2 → prtools-1.4.1}/prtools/fourier.py

@@ -1,9 +1,10 @@
 from prtools import __backend__
 from prtools.backend import numpy as np
+import prtools.jax
 
 
-def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
-         out=None):
+def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), axes=(-2, -1),
+         unitary=True, out=None):
     r"""Compute the 2-dimensional discrete Fourier Transform.
 
     The DFT is defined in one dimension as
@@ -31,10 +32,13 @@ def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
         ``F.shape = (shape[0], shape[1])``. Default is ``f.shape``.
     shift : array_like, optional
         Number of pixels in (r,c) to shift the DC pixel in the output plane
-        with the origin centrally located in the plane. Default is ``(0,0)``.
+        with the origin centrally located in the plane. Default is ``(0, 0)``.
     offset : array_like, optional
         Number of pixels in (r,c) that the input plane is shifted relative to
-        the origin. Default is ``(0,0)``.
+        the origin. Default is ``(0, 0)``.
+    axes : (2,) array_like of ints, optional
+        Axes over which to compute the DFT. If not given, the last two axes are
+        used.
     unitary : bool, optional
         Normalization flag. If ``True``, a normalization is performed on the
         output such that the DFT operation is unitary and energy is conserved
@@ -42,8 +46,8 @@ def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
         way, the energy in in a limited-area DFT is a fraction of the total
         energy corresponding to the limited area. Default is ``True``.
     out : ndarray or None
-        A location into which the result is stored. If provided, out.shape ==
-        shape and out.dtype == np.complex. If not provided or None, a
+        A location into which the result is stored. If provided, ``out.shape ==
+        shape`` and ``out.dtype == np.complex``. If not provided or None, a
         freshly-allocated array is returned.
 
     Returns
@@ -74,38 +78,41 @@ def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
     References
     ----------
     [1] Soummer, et. al. Fast computation of Lyot-style coronagraph
-        propagation (2007)
+    propagation (2007)
     """
-    return _dftcore(f, alpha, shape, shift, offset, unitary, out, forward=True)
 
+    return _dftcore(f, alpha, shape, shift, offset, axes, unitary,
+                    forward=True, out=out)
 
-def _dftcore(f, alpha, shape, shift, offset, unitary, out, forward):
 
-    #__backend__ = prtools.__backend__
+def _dftcore(f, alpha, shape, shift, offset, axes, unitary, forward, out):
 
-    if out is not None:
-        if __backend__ == 'numpy':
+    if __backend__ == 'numpy':
+        _mtp = prtools.fourier._multi_dot_three
+        if out is not None:
             if not np.can_cast(complex, out.dtype):
                 raise TypeError(f"Cannot cast complex output to dtype('{out.dtype}')")
-        elif __backend__ == 'jax':
-            raise ValueError('JAX backend does not support the out parameter')
 
-    alpha_row, alpha_col = np.broadcast_to(alpha, (2,))
+    elif __backend__ == 'jax':
+        _mtp = prtools.jax._multi_dot_three
+        if out is not None:
+            raise ValueError('JAX backend does not support the out parameter')
 
     f = np.asarray(f)
-    m, n = f.shape
 
-    if shape is None:
-        shape = (m, n)
-    M, N = shape
+    out_shape, axes = _cook_nd_args(f, shape, axes)
+    in_shape = np.take(np.asarray(f.shape), axes)
 
+    m, n = in_shape
+    M, N = out_shape
+
+    alpha_row, alpha_col = np.broadcast_to(alpha, (2,))
     shift_row, shift_col = np.broadcast_to(shift, (2,))
     offset_row, offset_col = np.broadcast_to(offset, (2,))
 
-    E1, E2 = _dft2_matrices(m, n, M, N, alpha_row, alpha_col, shift_row, shift_col,
-                            offset_row, offset_col, forward)
-
-    F = np.dot(E1.dot(f), E2, out=out)
+    E1, E2 = _dft2_matrices(m, n, M, N, alpha_row, alpha_col, shift_row,
+                            shift_col, offset_row, offset_col, forward)
+    F = _mtp(E1, f, E2, axes, out)
 
     if unitary:
         F = np.multiply(F, np.sqrt(np.abs(alpha_row * alpha_col)), out=F)
@@ -117,28 +124,21 @@ def _dftcore(f, alpha, shape, shift, offset, unitary, out, forward):
     return F
 
 
-def _dft2_matrices(m, n, M, N, alphar, alphac, shiftr, shiftc, offsetr, offsetc, forward):
+def _dft2_matrices(m, n, M, N, alphar, alphac, shiftr, shiftc, offsetr,
+                   offsetc, forward):
     if forward:
-        sign = -1
+        c = -1j
     else:
-        sign = 1
+        c = 1j
     R, S, U, V = _dft2_coords(m, n, M, N)
-    E1 = np.exp(sign*2.0 * 1j * np.pi * alphar * np.outer(R+offsetr, U-shiftr)).T
-    E2 = np.exp(sign*2.0 * 1j * np.pi * alphac * np.outer(S+offsetc, V-shiftc))
-    return E1, E2
-
-
-def _idft2_matrices(m, n, M, N, alphar, alphac, shiftr, shiftc, offsetr, offsetc):
-    R, S, U, V = _dft2_coords(m, n, M, N)
-    E1 = np.exp(2.0 * 1j * np.pi * alphar * np.outer(R+offsetr, U-shiftr)).T
-    E2 = np.exp(2.0 * 1j * np.pi * alphac * np.outer(S+offsetc, V-shiftc))
+    E1 = np.exp(2.0 * c * np.pi * alphar * np.outer(R+offsetr, U-shiftr)).T
+    E2 = np.exp(2.0 * c * np.pi * alphac * np.outer(S+offsetc, V-shiftc))
     return E1, E2
 
 
 def _dft2_coords(m, n, M, N):
     # R and S are (r,c) coordinates in the (m x n) input plane f
     # V and U are (r,c) coordinates in the (M x N) output plane F
-
     R = np.arange(m) - np.floor(m/2.0)
     S = np.arange(n) - np.floor(n/2.0)
     U = np.arange(M) - np.floor(M/2.0)
@@ -147,7 +147,51 @@ def _dft2_coords(m, n, M, N):
     return R, S, U, V
 
 
-def idft2(F, alpha, shape=None, shift=(0,0), offset=(0,0), unitary=True, out=None):
+def _cook_nd_args(a, s=None, axes=None):
+    # slightly modified version of numpy's function of the same name
+    if s is None:
+        if axes is None:
+            if a.ndim == 2:
+                s = list(a.shape)
+            elif a.ndim == 3:
+                s = list(a.shape[1:3])
+            else:
+                raise ValueError("Array must have ndim == 2 or 3")
+        else:
+            s = np.take(a.shape, axes)
+    s = list(s)
+    if axes is None:
+        axes = list(range(-len(s), 0))
+    if len(s) != len(axes):
+        raise ValueError("Shape and axes have different lengths.")
+    return s, axes
+
+
+def _multi_dot_three(a, b, c, axes, out):
+    # compute the matrix triple product
+    #
+    # a few notes:
+    # * this method is similar to np.linalg.multi_dot although it is less
+    #   general - here we only consider the matrix triple product used as
+    #   a part of prtools.dft2
+    # * because we use np.matmul instead of np.linalg.multi_dot, we can
+    #   take advantage of broadcasting a and c when b.ndim = 3. This
+    #   eliminates a for loop in the code
+    # * the implementation used here supports b with ndim in (2, 3)
+    #   iterating over any of the 3 axes when b.ndim == 3
+    # * the implementation used here is actually slightly faster than
+    #   an equivalent call to np.linalg.multi_dot when b.ndim == 2
+    # * np.linalg.multi_dot chooses the fastest multiplication order from
+    #   [(ab)c, a(bc)] depending on the shapes of a, b, and c. There is no
+    #   difference when computing the dft because both a and c are square
+    #   matrices
+    ab = np.matmul(a, b, axes=[(0, 1), axes, axes])
+    out = np.matmul(ab, c, axes=[axes, (0, 1), axes], out=out)
+    return out
+
+
+def idft2(F, alpha, shape=None, shift=(0, 0), offset=(0, 0), axes=(-2, -1),
+          unitary=True, out=None):
     r"""Compute the 2-dimensional inverse discrete Fourier Transform.
 
     The IDFT is defined in one dimension as
@@ -165,9 +209,9 @@ def idft2(F, alpha, shape=None, shift=(0,0), offset=(0,0), unitary=True, out=Non
     F : array_like
         2D array to Fourier Transform
     alpha : float or array_like
-        Input plane sampling interval (frequency). If :attr:`alpha` is an array,
-        ``alpha[1]`` represents row-wise sampling and ``alpha[2]`` represents
-        column-wise sampling. If :attr:`alpha` is a scalar,
+        Input plane sampling interval (frequency). If :attr:`alpha` is an
+        array, ``alpha[1]`` represents row-wise sampling and ``alpha[2]``
+        represents column-wise sampling. If :attr:`alpha` is a scalar,
         ``alpha[1] = alpha[2] = alpha`` represents uniform sampling across the
         rows and columns of the input plane.
     shape : int or array_like, optional
@@ -175,14 +219,24 @@ def idft2(F, alpha, shape=None, shift=(0,0), offset=(0,0), unitary=True, out=Non
         ``F.shape = (shape[0], shape[1])``. If :attr:`shape` is a scalar,
         ``F.shape = (shape, shape)``. Default is ``F.shape``
     shift : array_like, optional
-        Number of pixels in (x,y) to shift the DC pixel in the output plane with
-        the origin centrally located in the plane. Default is `[0,0]`.
+        Number of pixels in (x,y) to shift the DC pixel in the output plane
+        with the origin centrally located in the plane. Default is `[0,0]`.
+    offset : array_like, optional
+        Number of pixels in (r,c) that the input plane is shifted relative to
+        the origin. Default is ``(0, 0)``.
+    axes : (2,) array_like of ints, optional
+        Axes over which to compute the DFT. If not given, the last two axes are
+        used.
     unitary : bool, optional
         Normalization flag. If ``True``, a normalization is performed on the
         output such that the DFT operation is unitary and energy is conserved
         through the Fourier transform operation (Parseval's theorem). In this
         way, the energy in in a limited-area DFT is a fraction of the total
         energy corresponding to the limited area. Default is ``True``.
+    out : ndarray or None
+        A location into which the result is stored. If provided, ``out.shape ==
+        shape`` and ``out.dtype == np.complex``. If not provided or None, a
+        freshly-allocated array is returned.
 
     Returns
     -------
@@ -205,13 +259,15 @@ def idft2(F, alpha, shape=None, shift=(0,0), offset=(0,0), unitary=True, out=Non
 
     * If the y-axis shift behavior is not what you are expecting, you most
       likely have your plotting axes flipped (matplotlib's default behavior is
-      to place [0,0] in the upper left corner of the axes). This may be resolved
-      by either flipping the sign of the y component of ``shift`` or by passing
-      ``origin = 'lower'`` to ``imshow()``.
+      to place [0,0] in the upper left corner of the axes). This may be
+      resolved by either flipping the sign of the y component of ``shift`` or
+      by passing ``origin = 'lower'`` to ``imshow()``.
 
     References
     ----------
-    [1] Soummer, et. al. Fast computation of Lyot-style coronagraph propagation (2007)
-    
+    [1] Soummer, et. al. Fast computation of Lyot-style coronagraph
+    propagation (2007)
+
     """
-    return _dftcore(F, alpha, shape, shift, offset,  unitary, out, forward=False)
+    return _dftcore(F, alpha, shape, shift, offset, axes, unitary,
+                    forward=False, out=out)

{prtools-1.3.2 → prtools-1.4.1}/prtools/jax.py

@@ -31,14 +31,22 @@ class JaxOptimizeResult:
     state: Any  #: Optimizer state
 
 
-def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
+def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=None,
+          fn_kwargs=None):
     """Minimize a scalar function of one or more variables using the L-BFGS
     algorithm
 
     Parameters
     ----------
-    fun : callable
-        The objective function to be minimied
+    fn : callable
+        The objective function to be minimized:
+
+        .. code:: python
+
+            fn(x, *fn_args, **fn_kwargs)
+
+        where ``x`` is a 1-D array with shape (n,) and ``fn_args`` and
+        ``fn_kwargs`` are optional positional and keyword arguments.
     x0 : jax.Array
         Initial guess
     gtol : float
@@ -53,6 +61,10 @@ def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
             callback(intermediate_result: JaxOptimizeResult)
 
         where ``intermediate_result`` is a :class:`JaxOptimizeResult`.
+    fn_args : iterable or None
+        Extra positional arguments passed to ``fn()``
+    fn_kwargs : dict or None
+        Extra keyword arguments passed to the ``fn()``
 
     Returns
     -------
@@ -71,14 +83,21 @@ def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
     if not any((gtol, maxiter)):
         raise ValueError('At least one termination tolerance must be specified.')
 
+    fn_args = [] if fn_args is None else fn_args
+    fn_kwargs = {} if fn_kwargs is None else fn_kwargs
+
     opt = optax.lbfgs()
-    value_and_grad_fun = optax.value_and_grad_from_state(fun)
-    
+    value_and_grad_fn = optax.value_and_grad_from_state(fn)
+
     def step(carry):
         params, state = carry
-        value, grad = value_and_grad_fun(params, state=state)
+        # NOTE: passing *args and **kwargs to value_and_grad_fun is very
+        # poorly documented in optax (as of v0.2.6 - 10/2025) but this
+        # seems to work for now
+        value, grad = value_and_grad_fn(params, *fn_args, state=state,
+                                        **fn_kwargs)
         updates, state = opt.update(
-            grad, state, params, value=value, grad=grad, value_fn=fun)
+            grad, state, params, value=value, grad=grad, value_fn=fn)
         if callback:
             res = JaxOptimizeResult(
                 n=otu.tree_get(state, 'count')-1,
@@ -107,8 +126,47 @@ def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
         grad=otu.tree_get(final_state, 'grad'),
         value=otu.tree_get(final_state, 'value'),
         state=final_state)
-    
+
     if callback:
         jax.debug.callback(callback, res)
 
     return res
+
+
+def _multi_dot_three(a, b, c, axes, out):
+    # compute the matrix triple product
+    #
+    # a few notes:
+    # * while numpy-based implementation of this method is based on
+    #   np.matmul, jax.numpy.matmul doesn't implement the axes argument so
+    #   we have to use jax.numpy.linalg.multi_dot instead
+    # * the implementation used here supports b with ndim in (2, 3)
+    #   iterating over any of the 3 axes when b.ndim == 3
+    # * jax.vmap handles the case when b.ndim == 3 compared with the numpy
+    #   equivalent of this function which does everything within the
+    #   confines of matmul using the axes argument
+    if b.ndim == 2:
+        return jax.numpy.linalg.multi_dot((a, b, c))
+    else:
+        iter_axis = _iter_axis(axes)
+        return jax.vmap(_multi_dot, in_axes=[None, iter_axis, None], out_axes=iter_axis)(a, b, c)
+
+def _multi_dot(a, b, c):
+        # wrapper function to support vmap call signature
+        return jax.numpy.linalg.multi_dot((a, b, c))
+
+
+def _iter_axis(axes):
+    # pure Python to avoid dealing with JAX array mutability issues
+    mask = [0, 1, 2]
+    for ax in axes:
+        mask[ax] = None
+    return [ax for ax in mask if ax is not None][0]
+
+
+def _ndrebin(a, f):
+    return jax.vmap(_rebin, in_axes=[0, None])(a, f)
+
+
+def _rebin(a, f):
+        return jax.numpy.reshape(a, (a.shape[0]//f, f, a.shape[1]//f, f)).sum(-1).sum(1)

{prtools-1.3.2 → prtools-1.4.1}/prtools/zernike.py

@@ -4,7 +4,7 @@ import prtools
 from prtools.backend import numpy as np
 
 
-def zernike(mask, index, normalize=True, rho=None, theta=None):
+def zernike(mask, index, normalize=True, order='noll', rho=None, theta=None):
     """Compute the circular Zernike polynomial for a given mask.
 
     Parameters
@@ -17,6 +17,8 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
+    order : {'noll', 'fringe', 'ansi'}, optional
+        Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional
         Radial coordinates of the mask array. :attr:`rho` should be 0 at the
         origin and 1 at the edge of the circle.
@@ -44,9 +46,15 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
 
     """
     mask = np.asarray(mask)
-    #out = np.zeros_like(mask)
-    #mask_slice = util.boundary_slice(mask, pad=0)
-    #mask = mask[mask_slice]
+
+    if order == 'noll':
+        n, m = noll_to_nm(index)
+    elif order == 'fringe':
+        n, m = fringe_to_nm(index)
+    elif order =='ansi':
+        n, m = ansi_to_nm(index)
+    else:
+        raise ValueError('Unknown order', order)
     
     if rho is None:
         rho, theta = zernike_coordinates(mask)
@@ -54,8 +62,6 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
         if theta is None:
             raise ValueError("Both rho and theta must be specified")
 
-    m, n = zernike_index(index)
-
     if m == 0:
         if n == 0:
             Z = mask
@@ -97,26 +103,24 @@ def R(m, n, rho):
         return R
     
 
-def zernike_compose(mask, coeffs, normalize=True, rho=None, theta=None):
+def zernike_compose(mask, coeffs, normalize=True, order='noll', rho=None, theta=None):
     """Create an OPD based on the supplied Zernike coefficients.
 
     Parameters
     ----------
     mask : array_like
         Binary mask defining the extent to compute the Zernike polynomial over.
-
     coeffs : array_like
         List of coefficients corresponding to Zernike indices (Noll ordering)
         used to create the OPD.
-
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
-
+    order : {'noll', 'fringe', 'ansi'}, optional
+        Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional
         Radial coordinates of the mask array. :attr:`rho` should be 0 at the
         origin and 1 at the edge of the circle.
-
     theta : array_like, optional
         Angular coordinates of the mask array in radians.
 
@@ -159,11 +163,12 @@ def zernike_compose(mask, coeffs, normalize=True, rho=None, theta=None):
     opd = np.zeros(mask.shape)
 
     for index, coeff in enumerate(coeffs):
-        opd += coeff * zernike(mask, index+1, normalize, rho, theta)
+        opd += coeff * zernike(mask, index+1, normalize, order, rho, theta)
 
     return opd
 
-def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=None):
+def zernike_basis(mask, modes, vectorize=False, normalize=True, order='noll',
+                  rho=None, theta=None):
     """Compute a Zernike basis set for a given mask.
 
     Parameters
@@ -180,6 +185,8 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
+    order : {'noll', 'fringe', 'ansi'}, optional
+        Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional
         Radial coordinates of the mask array. :attr:`rho` should be 0 at the
         origin and 1 at the edge of the circle.
@@ -203,7 +210,7 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
 
     basis = []
     for index, mode in enumerate(modes):
-        basis.append(zernike(mask, mode, normalize, rho, theta))
+        basis.append(zernike(mask, mode, normalize, order, rho, theta))
     basis = np.asarray(basis)
 
     if vectorize:
@@ -211,9 +218,10 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
         return basis.reshape(basis.shape[0], -1)
     else:
         return basis
-    
 
-def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
+
+def zernike_fit(opd, mask, modes, normalize=True, order='noll', rho=None,
+                theta=None):
     """Fit a Zernike basis set to an OPD.
 
     Parameters
@@ -227,6 +235,8 @@ def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
+    order : {'noll', 'fringe', 'ansi'}, optional
+        Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional
         Radial coordinates of the mask array. :attr:`rho` should be 0 at the
         origin and 1 at the edge of the circle.
@@ -263,14 +273,14 @@ def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
     opd = np.asarray(opd)
     mask = np.asarray(mask)
 
-    basis = zernike_basis(mask, modes, True, normalize, rho, theta)
+    basis = zernike_basis(mask, modes, True, normalize, order, rho, theta)
 
     basis = np.linalg.pinv(basis)
 
     return np.einsum('ij,i->j', basis, opd.ravel())
 
 
-def zernike_remove(opd, mask, modes, rho=None, theta=None):
+def zernike_remove(opd, mask, modes, order='noll', rho=None, theta=None):
     """Fit and remove a Zernike basis set from an OPD.
 
     Parameters
@@ -281,6 +291,8 @@ def zernike_remove(opd, mask, modes, rho=None, theta=None):
         Binary mask defining the extent to compute the Zernike basis over.
     modes : array_like
         List of modes (Noll ordering) to remove.
+    order : {'noll', 'fringe', 'ansi'}, optional
+        Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional
         Radial coordinates of the mask array. :attr:`rho` should be 0 at the
         origin and 1 at the edge of the circle.
@@ -303,35 +315,34 @@ def zernike_remove(opd, mask, modes, rho=None, theta=None):
     opd = np.asarray(opd)
     mask = np.asarray(mask)
 
-    coeffs = zernike_fit(opd, mask, modes, rho, theta)
-    fit_opd = zernike_compose(mask, coeffs, rho, theta)
+    coeffs = zernike_fit(opd, mask, modes, True, order, rho, theta)
+    fit_opd = zernike_compose(mask, coeffs, True, order, rho, theta)
 
     residual = opd - fit_opd
 
     return residual
 
 
-def zernike_index(j):
-    """Convert a 1-D Noll index to a 2-D radiul and azimuthal index.
+def noll_to_nm(j):
+    """Convert Noll index to (n, m) index.
 
     Parameters
     ----------
     j : int
-        Noll Zernike index.
+        Noll Zernike index
 
     Returns
     -------
     n : int
-        Radial Zernike index.
+        Radial degree.
     m : int
-        Azimuthal Zernike index.
-
+        Azimuthal degree.
     """
     if j < 1:
-        raise ValueError('Zernike index j must be a positive integer')
+        raise ValueError('Noll index must be greater than zero')
 
     # find the row j is in
-    n = int(np.ceil((-1 + np.sqrt(1 + 8*j)) / 2) - 1)
+    n = int(np.ceil((-1 + np.sqrt(1 + 8 * j)) / 2) - 1)
     if n == 0:
         m = 0
     else:
@@ -349,13 +360,63 @@ def zernike_index(j):
         else:
             row_m = [0]
 
-        for i in range(int(np.floor(n / 2))):
+        for _ in range(n // 2):
             row_m.append(row_m[-1] + 2)
             row_m.append(row_m[-1])
 
         m = row_m[r] * sign
 
-    return m, n
+    return int(n), int(m)
+
+
+def ansi_to_nm(j):
+    """Convert ANSI index to (n, m) index.
+
+    Parameters
+    ----------
+    j : int
+        ANSI index
+
+    Returns
+    -------
+    n : int
+        Radial degree
+    m : int
+        Azimuthal degree
+    """
+    if j < 0:
+        raise ValueError('ANSI index must be zero or larger')
+    
+    # https://en.wikipedia.org/wiki/Zernike_polynomials
+    n = np.floor((np.sqrt(8 * j + 1) - 1) / 2)
+    m = 2*j - n*(n + 2)
+    return int(n), int(m)
+
+
+def fringe_to_nm(j):
+    """Convert Fringe index to (n, m) index.
+
+    Parameters
+    ----------
+    j : int
+        Fringe index
+
+    Returns
+    -------
+    n : int
+        Radial degree
+    m : int
+        Azimuthal degree
+    """
+    if j < 1:
+        raise ValueError('Fringe index must be greater than zero')
+    
+    # https://github.com/arielmission-space/PAOS/blob/main/paos/classes/zernike.py
+    m_n = 2 * (np.ceil(np.sqrt(j)) - 1)
+    g_s = (m_n / 2) ** 2 + 1
+    n = m_n / 2 + np.floor((j - g_s) / 2)
+    m = (m_n - n) * (1 - np.mod(j - g_s, 2) * 2)
+    return int(n), int(m)
 
 
 def zernike_coordinates(mask, shift=None, angle=0, indexing='ij'):

{prtools-1.3.2 → prtools-1.4.1/prtools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prtools
-Version: 1.3.2
+Version: 1.4.1
 Summary: Utility functions for image-based phase retrieval
 Author: Andy Kee
 License: MIT License

{prtools-1.3.2 → prtools-1.4.1}/prtools.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ README.rst
 pyproject.toml
 prtools/__init__.py
 prtools/array.py
+prtools/convolve.py
 prtools/cost.py
 prtools/fourier.py
 prtools/jax.py
@@ -23,4 +24,5 @@ prtools/backend/_base.py
 prtools/backend/_jax.py
 prtools/backend/_numpy.py
 tests/test_array.py
+tests/test_convolve.py
 tests/test_fourier.py

{prtools-1.3.2 → prtools-1.4.1}/tests/test_array.py

@@ -65,3 +65,17 @@ def test_boundary(backend):
     x = np.zeros((10, 10))
     x[3:7, 2:8] = 1
     assert np.array_equal(prtools.boundary(x), (3, 6, 2, 7))
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_rebin(backend):
+    prtools.use(backend)
+    x = np.ones((10, 10))
+    assert np.array_equal(prtools.rebin(x, 2), 4*np.ones((5, 5)))
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_ndrebin(backend):
+    prtools.use(backend)
+    x = np.ones((3, 10, 10))
+    assert np.array_equal(prtools.rebin(x, 2), 4*np.ones((3, 5, 5)))

prtools-1.4.1/tests/test_convolve.py

@@ -0,0 +1,75 @@
+import numpy as np
+import pytest
+
+import prtools
+from .import BACKENDS
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+@pytest.mark.parametrize('shape', ((255,255), (256,256)))
+def test_gauss_kernel(backend, shape):
+    prtools.use(backend)
+
+    # note we need to use a sufficiently large sigma here to avoid wrapping
+    # in the FFT-computed version of the kernel
+    sigma = 5
+    #shape = (256,256)
+    G = prtools.gauss_kernel(shape, sigma, fftshift=True)
+
+    r = np.arange(-shape[0]//2,shape[0]//2)
+    c = np.arange(-shape[1]//2,shape[1]//2)
+    g = prtools.gauss(r, c, sigma)
+    G_fft = np.abs(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(g))))
+    G_fft = G_fft / np.max(G_fft)
+    
+    assert np.allclose(G, G_fft, atol=1e-6)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+@pytest.mark.parametrize('shape', ((255,255), (256,256)))
+def test_gauss_kernel_pixelscale(backend, shape):
+    prtools.use(backend)
+
+    r = np.linspace(-10, 10, shape[0])
+    c = np.linspace(-10, 10, shape[1])
+    nr, nc = len(r), len(c)
+    dr, dc = r[1] - r[0], c[1] - c[0]
+
+    G = prtools.gauss_kernel((nr, nc), sigma=1, pixelscale=(dr, dc),
+                             fftshift=True)
+
+    g = prtools.gauss(r, c, sigma=1)
+    G_fft = np.abs(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(g))))
+    G_fft = G_fft / np.max(G_fft)
+    
+    assert np.allclose(G, G_fft, atol=1e-6)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+@pytest.mark.parametrize('shape', ((255,255), (256,256)))
+def test_gauss_kernel_oversample(backend, shape):
+    prtools.use(backend)
+
+    sigma = 5
+    oversample = 3
+
+    G = prtools.gauss_kernel(shape, sigma=sigma, fftshift=False)
+    Go = prtools.gauss_kernel((shape[0]*oversample, shape[1]*oversample),
+                              sigma=sigma, oversample=oversample,
+                              fftshift=False)
+    
+    assert np.allclose(G[0:100,0:100], Go[0:100,0:100])
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+@pytest.mark.parametrize('shape', ((255,255), (256,256)))
+def test_pixel_kernel_oversample(backend, shape):
+    prtools.use(backend)
+
+    oversample = 3
+
+    P = prtools.pixel_kernel(shape, fftshift=False)
+    Po = prtools.pixel_kernel((shape[0]*oversample, shape[1]*oversample),
+                              oversample=oversample, fftshift=False)
+    
+    assert np.allclose(P[0:100,0:100], Po[0:100,0:100])

{prtools-1.3.2 → prtools-1.4.1}/tests/test_fourier.py

@@ -6,53 +6,43 @@ from .import BACKENDS
 
 
 @pytest.mark.parametrize('backend', BACKENDS)
-def test_dft2_even(backend):
+@pytest.mark.parametrize('shape', ((10,10), (11,11), (10,11)))
+def test_dft2(backend, shape):
     prtools.use(backend)
-    n = 10
-    f = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    m, n = shape
+    f = np.random.rand(m, n) + 1j * np.random.rand(m, n)
 
-    F_dft = prtools.dft2(f, 1/n, unitary=False)
+    F_dft = prtools.dft2(f, [1/m, 1/n], unitary=False)
     F_fft = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(f)))
 
     assert np.allclose(F_dft, F_fft, atol=1e-5)
 
 
 @pytest.mark.parametrize('backend', BACKENDS)
-def test_idft2_even(backend):
+@pytest.mark.parametrize('axes', ((0,1), (1,2), (0,2)))
+def test_dft2_cube(backend, axes):
     prtools.use(backend)
-    n = 10
-    F = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    shape = (10, 11, 12)
+    m, n = np.take(shape, axes)
+    f = np.random.uniform(size=shape) + 1j * np.random.uniform(size=shape)
 
-    f_dft = prtools.idft2(F, 1/n, unitary=False)
-    f_fft = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(F)))
+    F_dft = prtools.dft2(f, [1/m, 1/n], axes=axes, unitary=False)
+    F_fft = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(f), axes=axes))
 
-    assert np.allclose(f_dft, f_fft, atol=1e-5)
+    assert np.allclose(F_dft, F_fft, atol=5e-5)
 
 
 @pytest.mark.parametrize('backend', BACKENDS)
-def test_dft2_odd(backend):
+@pytest.mark.parametrize('shape', ((10,10), (11,11), (10,11)))
+def test_idft2(backend, shape):
     prtools.use(backend)
-    n = 11
-    f = np.random.rand(n, n) + 1j * np.random.rand(n, n)
-
-    F_dft = prtools.dft2(f, 1/n, unitary=False)
-    F_fft = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(f)))
-    
-    print(F_dft.dtype)
-    print(F_fft.dtype)
-    assert np.allclose(F_dft, F_fft, atol=1e-5)
-
-
-@pytest.mark.parametrize('backend', BACKENDS)
-def test_idft2_odd(backend):
-    prtools.use(backend)
-    n = 11
-    F = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    m, n = shape
+    f = np.random.rand(m, n) + 1j * np.random.rand(m, n)
 
-    f_dft = prtools.idft2(F, 1/n, unitary=False)
-    f_fft = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(F)))
+    F_dft = prtools.idft2(f, [1/m, 1/n], unitary=False)
+    F_fft = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(f)))
 
-    assert np.allclose(f_dft, f_fft, atol=1e-5)
+    assert np.allclose(F_dft, F_fft, atol=1e-5)
 
 
 @pytest.mark.parametrize('backend', BACKENDS)