prtools 1.4.0__tar.gz → 1.4.2__tar.gz

{prtools-1.4.0/prtools.egg-info → prtools-1.4.2}/PKG-INFO

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

{prtools-1.4.0 → prtools-1.4.2}/prtools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '1.4.0'
+__version__ = '1.4.2'
 
 from prtools.backend import __backend__, use
 

{prtools-1.4.0 → prtools-1.4.2}/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.4.0 → prtools-1.4.2}/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.0 → prtools-1.4.2}/prtools/convolve.py

@@ -2,7 +2,7 @@ from prtools.backend import numpy as np
 
 
 def fftconv(array, kernel, normalize_kernel=True, fft_array=True,
-            fft_kernel=True):
+            fft_kernel=False, fftshift_kernel=False):
     r"""Convolve an array with a kernel using the FFT.
 
     The colvolution is computed as
@@ -14,20 +14,24 @@ def fftconv(array, kernel, normalize_kernel=True, fft_array=True,
     Parameters
     ----------
     array : array_like
-        Array to be convolved with `kernel`.
+        Array to be convolved with ``kernel``.
     kernel : array_like
-        Convolution kernel. Should have the same shape as `array`.
+        Convolution kernel. Should have the same shape as ``array``.
     normalize_kernel : bool, optional
-        If True (default), kernel is normalized so that `kernel =
-        kernel / np.sum(kernel)`.
+        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. If False, the array is assumed to be provided in the frequency
-        domain.
+        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 (default), the kernel is assumed to be provided in the spatial
-        domain. If False, the kernel is assumed to be provided in the frequency
+        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
     -------
@@ -40,7 +44,8 @@ def fftconv(array, kernel, normalize_kernel=True, fft_array=True,
     """
 
     a = np.fft.fft2(array) if fft_array else array
-    k = np.fft.fft2(array) if fft_kernel else kernel
+    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)
@@ -56,8 +61,9 @@ def gauss_blur(img, sigma, oversample=1):
     img : array_like
         Array to be blurred
     sigma : float or (2,) array_like
-        Standard deviation of Gaussian. Providing two values allows for
-        non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
+        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.
     
@@ -72,7 +78,10 @@ def gauss_blur(img, sigma, oversample=1):
         :context: reset
         :scale: 50
 
-        >>> img = prtools.circle((256,256), 100)
+        >>> 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')
@@ -82,20 +91,20 @@ def gauss_blur(img, sigma, oversample=1):
     """
     img = np.asarray(img)
     kernel = gauss_kernel(img.shape, sigma, oversample, fftshift=False)
-    return fftconv(img, kernel, normalize_kernel=True, fft_array=True,
-                   fft_kernel=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 on a discretely
-    sampled image using the FFT.
+    """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.
+        Oversampling factor of ``img``. Default is 1.
 
     Returns
     -------
@@ -106,36 +115,14 @@ def pixelate(img, oversample=1):
     -----
     To avoid the introduction of numerical artifacts, this function should be
     performed on data that is at least 2x oversampled.
-
-    Examples
-    --------
-    .. plot::
-        :include-source:
-        :context: reset
-        :scale: 50
-
-        >>> img = prtools.circle((17,17), 5)
-        >>> img_pixelate = prtools.pixelate(img)
-        >>> 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_pixelate, cmap='gray')
-        >>> ax[1].set_title('Pixelated image')
-
-    Note that the pixelation process preserves image power:
-
-    .. code:: pycon
-
-        >>> print(np.sum(img), np.sum(img_pixelate))
-        64.0, 64.0
     """
     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)
+                   fft_kernel=False, fftshift_kernel=False)
 
 
-def gauss(x1, x2, sigma, oversample=1, indexing='ij', normalize=False):
+def gauss(x1, x2, sigma, indexing='ij', normalize=False):
     """2D Gaussian function
 
     Parameters
@@ -145,10 +132,8 @@ def gauss(x1, x2, sigma, oversample=1, indexing='ij', normalize=False):
     sigma : float or (2,) array_like
         Standard deviation of Gaussian. Providing two values allows for
         non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
-    oversample : float, optional
-        Oversampling factor. Defailt is 1.
     indexing : {'ij', 'xy'}, optional
-        Matrix ('ij', default) or cartesian ('xy') indexing of output.
+        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.
@@ -170,7 +155,7 @@ def gauss(x1, x2, sigma, oversample=1, indexing='ij', normalize=False):
         >>> plt.imshow(g, cmap='gray')
     """
     xx1, xx2 = np.meshgrid(x1, x2, indexing=indexing)
-    sigma = np.broadcast_to(sigma, (2,)) / oversample
+    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))
@@ -178,14 +163,14 @@ def gauss(x1, x2, sigma, oversample=1, indexing='ij', normalize=False):
 
 
 def sinc(x1, x2, indexing='ij'):
-    r"""Normalized 2D sinc function
+    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.
+        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.
@@ -221,10 +206,10 @@ def sinc(x1, x2, indexing='ij'):
     return np.sinc(xx1) * np.sinc(xx2)
 
 
-def gauss_kernel(shape, sigma, oversample=1, fftshift=True):
+def gauss_kernel(shape, sigma, oversample=1, pixelscale=1.0, fftshift=False):
     """2D Gaussian filter kernel
 
-    This function returns the transfer function of a normalized 2D Gaussian
+    This function returns the Fourier transform of a normalized 2D Gaussian
     function.
 
     Parameters
@@ -232,12 +217,16 @@ def gauss_kernel(shape, sigma, oversample=1, fftshift=True):
     shape : int or tuple of int
         Shape of the kernel
     sigma : float
-        Standard deviation of Gaussian. Providing two values allows for
-        non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
+        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 (default), the kernel is FFT-shifted before it is returned.
+        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
     -------
@@ -251,35 +240,44 @@ def gauss_kernel(shape, sigma, oversample=1, fftshift=True):
 
     Examples
     --------
+    Create a Gaussian kernel with 
+
     .. plot::
         :include-source:
-        :context: reset
         :scale: 50
 
-        >>> blur_kernel = prtools.gauss_kernel((256,256), 1)
-        >>> plt.imshow(blur_kernel, cmap='gray')
+        >>> G = prtools.gauss_kernel((256,256), sigma=1, fftshift=True)
+        >>> plt.imshow(G, cmap='gray')
 
-    Note this is functionally equivalent to
+    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:
-        :context: reset
         :scale: 50
 
-        >>> r = np.fft.fftfreq(256)
-        >>> c = np.fft.fftfreq(256)
-        >>> sigma = 1
-        >>> blur_kernel = np.fft.fftshift(prtools.gauss(r, c, 1/(2*np.pi*sigma)))
-        >>> plt.imshow(blur_kernel, cmap='gray')
+        >>> 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(shape[0])
-    x2 = np.fft.fftfreq(shape[1])
+    x1 = np.fft.fftfreq(n=shape[0], d=pixelscale[0])
+    x2 = np.fft.fftfreq(n=shape[1], d=pixelscale[1])
 
-    k = gauss(x1, x2, 1/(2*np.pi*sigma), oversample, indexing='ij',
+    k = gauss(x1, x2, sigma=1/(2*np.pi*sigma*oversample), indexing='ij',
               normalize=False)
 
     if fftshift:
@@ -288,7 +286,7 @@ def gauss_kernel(shape, sigma, oversample=1, fftshift=True):
         return k
 
 
-def pixel_kernel(shape, oversample=1, fftshift=True):
+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
@@ -300,8 +298,12 @@ def pixel_kernel(shape, oversample=1, fftshift=True):
         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 (default), the kernel is FFT-shifted before it is returned.
+        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
     -------
@@ -320,25 +322,14 @@ def pixel_kernel(shape, oversample=1, fftshift=True):
         :context: reset
         :scale: 50
 
-        >>> pixel_mtf = prtools.pixel_kernel((256,256), 1)
-        >>> plt.imshow(pixel_mtf, cmap='gray', vmin=0)
-
-    Note this is functionally equivalent to
-
-    .. plot::
-        :include-source:
-        :context: reset
-        :scale: 50
-
-        >>> r = np.fft.fftfreq(256)
-        >>> c = np.fft.fftfreq(256)
-        >>> pixel_mtf = np.fft.fftshift(np.abs(prtools.sinc(r, c)))
+        >>> 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(shape[0]) / oversample
-    x2 = np.fft.fftfreq(shape[1]) / oversample
+    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'))
 

{prtools-1.4.0 → prtools-1.4.2}/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
@@ -76,36 +80,36 @@ def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
     [1] Soummer, et. al. Fast computation of Lyot-style coronagraph
     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
+    (M, N), axes = _cook_nd_args(f, shape, axes)
+    m, n = (f.shape[axes[0]], f.shape[axes[1]])
 
+    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 +121,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 +144,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(np.asarray(a.shape), np.asarray(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 +206,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 +216,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 +256,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.4.0 → prtools-1.4.2}/prtools/jax.py

@@ -31,8 +31,7 @@ class JaxOptimizeResult:
     state: Any  #: Optimizer state
 
 
-def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=None, 
-          fn_kwargs=None):
+def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_kwargs=None):
     """Minimize a scalar function of one or more variables using the L-BFGS
     algorithm
 
@@ -43,10 +42,10 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=None,
 
         .. 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.
+            fn(x, **fn_kwargs)
+
+        where ``x`` is a 1-D array with shape (n,) and ``fn_kwargs`` are
+        optional keyword arguments.
     x0 : jax.Array
         Initial guess
     gtol : float
@@ -61,10 +60,9 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=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()``
+        Extra keyword arguments passed to the objective function as
+        ``fn(x, **fn_kwargs)``
 
     Returns
     -------
@@ -83,21 +81,21 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=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_fn = optax.value_and_grad_from_state(fn)
-    
+
     def step(carry):
         params, state = carry
         # 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)
+        # seems to work for now.
+        value, grad = value_and_grad_fn(params, state=state, **fn_kwargs)
+        # NOTE: opt.update() doesn't have any mechanism for accepting *args so
+        # we're limited to using **kwargs here
         updates, state = opt.update(
-            grad, state, params, value=value, grad=grad, value_fn=fn)
+            grad, state, params, value=value, grad=grad, value_fn=fn, **fn_kwargs)
         if callback:
             res = JaxOptimizeResult(
                 n=otu.tree_get(state, 'count')-1,
@@ -126,8 +124,48 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=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.4.0 → prtools-1.4.2/prtools.egg-info}/PKG-INFO

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

{prtools-1.4.0 → prtools-1.4.2}/prtools.egg-info/SOURCES.txt

@@ -24,4 +24,6 @@ prtools/backend/_base.py
 prtools/backend/_jax.py
 prtools/backend/_numpy.py
 tests/test_array.py
-tests/test_fourier.py
+tests/test_convolve.py
+tests/test_fourier.py
+tests/test_jax.py

{prtools-1.4.0 → prtools-1.4.2}/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.2/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.4.0 → prtools-1.4.2}/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)

prtools-1.4.2/tests/test_jax.py

@@ -0,0 +1,18 @@
+import numpy as np
+import jax.numpy as jnp
+
+import prtools
+
+def test_lbfgs_fn_kwargs():
+    prtools.use('jax')
+
+    def fn(x, a, b):
+        return (x[0] + a)**2 + (x[1] + b)**2
+    
+    a = 2
+    b = -5
+    x0 = jnp.zeros(2)
+
+    res = prtools.jax.lbfgs(fn, x0, gtol=1e-5, maxiter=25, 
+                            fn_kwargs={'a': a, 'b': b})
+    assert(np.allclose(res.x, (-a, -b)))