prtools 2.0.1__tar.gz → 2.1.0__tar.gz

{prtools-2.0.1/prtools.egg-info → prtools-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prtools
-Version: 2.0.1
+Version: 2.1.0
 Summary: Utility functions for image-based phase retrieval
 Author: Andy Kee
 License: MIT License
@@ -45,7 +45,7 @@ Requires-Dist: scipy
 Requires-Dist: array_api_compat
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
-Requires-Dist: optax; extra == "jax"
+Requires-Dist: optax>=0.2.5; extra == "jax"
 Dynamic: license-file
 
 prtools

{prtools-2.0.1 → prtools-2.1.0}/prtools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '2.0.1'
+__version__ = '2.1.0'
 
 
 from prtools._array_api import use
@@ -8,6 +8,7 @@ from prtools.array import (
     pad,
     crop,
     boundary,
+    sample,
     rebin,
     rescale,
     normpow,
@@ -32,6 +33,8 @@ from prtools.fourier import dft2, idft2
 
 import prtools.jax
 
+from prtools.mesh import mesh
+
 from prtools.misc import (
     calcpsf,
     translation_defocus,
@@ -85,5 +88,6 @@ from prtools.zernike import (
     zernike_basis,
     zernike_fit,
     zernike_remove,
-    zernike_coordinates
+    zernike_coordinates,
+    zernike_coeff_legacy
 )

{prtools-2.0.1 → prtools-2.1.0}/prtools/array.py

@@ -75,7 +75,7 @@ def centroid(x, where=None, mode='absolute', indexing='ij'):
     c = xp.dot(cc[where].ravel(), anorm.ravel())
 
     if mode == 'center':
-        rc, cc = np.array(x.shape)/2
+        rc, cc = np.floor(np.array(x.shape)/2)
         r = r - rc
         c = c - cc
 
@@ -261,7 +261,7 @@ def crop(x, shape, shift=None, mode='absolute', indexing='ij'):
     if shift is not None:
         shift = np.asarray(shift, dtype=int)
         if indexing == 'xy':
-            shift = (-shift[1], shift[0])
+            shift = (shift[1], shift[0])
         if mode == 'absolute':
             shift = np.array(shift - np.array(x.shape)//2)
     else:
@@ -326,6 +326,82 @@ def boundary(x, threshold=0):
     return rmin, rmax, cmin, cmax
 
 
+def sample(x, oversample, offset=0, axes=(-2, -1), indexing='ij'):
+    """Extract samples from an oversampled array.
+
+    The sampled axes are tiled into contiguous blocks of ``oversample``
+    samples, with one sample being taken from each block. No averaging or
+    filtering is performed.
+
+    Parameters
+    ----------
+    x : array_like
+        Array to sample
+    oversample : int or tuple of ints
+        Decimation factor. Providing two values defines asymmetric
+        oversampling with ordering interpreted according to ``indexing``.
+    offset : int or tuple of ints, optional
+        Position of the sample with each block, measured from
+        ``oversample//2``. The sub-sampled position reachable by ``offset`` is
+        quantized to ``1/oversample``. Default is 0. If two values are
+        provided, their ordering is interpreted according to ``indexing``.
+    axes : (2,) sequence of ints, optional
+        Axes over which to sample. If not provided, the last two axes are
+        used.
+    indexing : {'ij', 'xy'}, optional
+        Matrix (‘ij’, default) or cartesian (‘xy’) indexing of ``x`` grid.
+
+    Returns
+    -------
+    x : ndarray
+        Each sampled axis has length ``x.shape[axis]//oversample``. If axis
+        length is not a multiple of ``oversample``, samples taken beyond the
+        last whole block are ignored.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> psf = prtools.calcpsf(amp, opd, wave, sampling, (32, 32),
+        ...                       oversample=8)
+        >>> k = prtools.pixel_kernel(psf.shape, oversample=8)
+        >>> native = prtools.sample(prtools.fftconv(psf, k), 8)
+        >>> native.shape
+        (32, 32)
+
+    """
+    xp = array_namespace(x)
+    x = xp.asarray(x)
+
+    if indexing not in ('ij', 'xy'):
+        raise ValueError("Valid values for indexing are 'xy' and 'ij'.")
+
+    oversample = tuple(int(v) for v in np.broadcast_to(oversample, (2,)))
+    offset = tuple(int(v) for v in np.broadcast_to(offset, (2,)))
+    axes = tuple(int(a) % x.ndim for a in axes)
+
+    if indexing == 'xy':
+        offset = (offset[1], offset[0])
+        oversample = (oversample[1], oversample[0])
+
+    idx = [slice(None)] * x.ndim
+
+    for ax, m, o in zip(axes, oversample, offset):
+
+        lo, hi = -(m // 2), m - 1 - m // 2
+        if not lo <= o <= hi:
+            raise ValueError(f'offset {o} lies outside one pixel for oversample '
+                             f'{m} on axis {ax}; require {lo} <= offset <= {hi}')
+
+        n = x.shape[ax] // m
+        start = m // 2 + o
+        stop = start + m * n
+
+        idx[ax] = slice(start, stop, m)
+
+    return x[tuple(idx)]
+
+
 def rebin(x, factor):
     """Rebin an array by an integer factor.
 

{prtools-2.0.1 → prtools-2.1.0}/prtools/convolve.py

@@ -7,7 +7,7 @@ def fftconv(array, kernel, axes=(-2, -1), 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
+    The convolution is computed as
 
     .. math::
 
@@ -18,12 +18,13 @@ def fftconv(array, kernel, axes=(-2, -1), normalize_kernel=True,
     array : array_like
         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``. May be
+        given in either the frequency or spatial domain, as indicated by
+        ``fft_kernel``. The kernel is applied exactly as supplied. See Notes
+        for normalization details.
     axes : sequence of ints, optional
         Axes over which to apply the convolution. If not given, the last two
         axes are used.
-    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
@@ -46,6 +47,24 @@ def fftconv(array, kernel, axes=(-2, -1), normalize_kernel=True,
     --------
     :func:`~prtools.gauss_blur`
     :func:`~prtools.pixelate`
+
+    Notes
+    -----
+    **Kernel normalization** This function does not normalize or scale the
+    supplied kernel in any way.
+
+    The sum of ``array`` is conserved when the kernel's transfer function is
+    unity at zero frequency. Depending on which domain the kernel is supplied
+    in, the kernel normaliztion approach is different:
+
+    * A frequency-domain kernel (``fft_kernel=False``) should be scaled such
+      that ``kernel[0,0] == 1``. Note that ``kernel[0,0]`` is the
+      zero-frequency term only for a kernel in standard (unshifted) FFT layout.
+      If the kernel origin lies at the center of the array
+      (``fftshift_kernel=True``), the DC term is at ``kernel[nr//2, nc//2]``.
+
+    * A spatial-domain kernal(``fft_kernel=True``) should have
+      ``sum(kernel) == 1``.
     """
     xp = array_namespace(array)
 
@@ -53,9 +72,6 @@ def fftconv(array, kernel, axes=(-2, -1), normalize_kernel=True,
     k = xp.fft.fftshift(kernel) if fftshift_kernel else kernel
     k = xp.fft.fft2(k) if fft_kernel else k
 
-    if normalize_kernel:
-        k = k / xp.sum(k)
-
     return xp.fft.ifft2(a*k, axes=axes).real
 
 
@@ -312,10 +328,10 @@ def gauss_kernel(shape, sigma, oversample=1, pixelscale=1.0, fftshift=False,
         return k
 
 
-def pixel_kernel(shape, oversample=1, pixelscale=1.0, fftshift=False, 
+def pixel_kernel(shape, oversample=1, pixelscale=1.0, fftshift=False,
                  xp=None):
     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.
 
@@ -361,7 +377,7 @@ def pixel_kernel(shape, oversample=1, pixelscale=1.0, fftshift=False,
     x1 = xp.fft.fftfreq(shape[0], d=pixelscale[0]) * oversample
     x2 = xp.fft.fftfreq(shape[1], d=pixelscale[1]) * oversample
 
-    k = xp.abs(sinc(x1, x2, indexing='ij', xp=xp))
+    k = sinc(x1, x2, indexing='ij', xp=xp)
 
     if fftshift:
         return xp.fft.fftshift(k)

prtools-2.1.0/prtools/cost.py

@@ -0,0 +1,120 @@
+from prtools._array_api import array_namespace
+
+
+def sserror(data, est, mask=None, gain_bias_invariant=False, ghat=None):
+    r"""Compute the normalized sum squared error between two arrays.
+
+    The error between measured data :math:`\tilde{s}` and estimated data
+    :math:`\hat{s}` is computed as the gain invariant error metric
+
+    .. math::
+
+        \Phi = \frac{1}{K}\sum_k
+        \frac{\sum_{x,y} w(x,y,k)\left[\alpha_k\hat{s}(x,y,k) -
+        \tilde{s}(x,y,k)\right]^2}{\sum_{x,y} w(x,y,k)\tilde{s}^2(x,y,k)}
+
+    for :math:`k = 1, 2, ... K` independent measurements, where
+    :math:`w` is a weighting function given by `mask` and each
+    :math:`\alpha_k` is chosen to minimize the error. The result is
+    independent of any relative scaling (gain) difference between the inputs.
+
+    If `gain_bias_invariant` is ``True``, the gain and bias invariant error
+    metric is computed instead
+
+    .. math::
+
+        \Phi = \frac{1}{K}\sum_k
+        \frac{\sum_{x,y} w(x,y,k)\left[\alpha_k\hat{g}(x,y,k) -
+        \tilde{g}(x,y,k)\right]^2}{\sum_{x,y} w(x,y,k)\tilde{s}^2(x,y,k)}
+
+    where :math:`\hat{g}` and :math:`\tilde{g}` are :math:`\hat{s}` and
+    :math:`\tilde{s}` with their weighted means subtracted and each
+    :math:`\alpha_k` is chosen to minimize the error. The result is
+    independent of any relative scaling (gain) and offset (bias) differences
+    between the inputs.
+
+    Parameters
+    ----------
+    data : array_like
+        Ideal or measured data with shape (M, N) or (K, M, N)
+    est : array_like
+        Estimated or modeled data with the same shape as `data`
+    mask : array_like, optional
+        Weighting function :math:`w` applied to the inputs. Entries equal to
+        1 (or True) are fully included in the error computation and entries
+        equal to 0 (or False) are ignored. Intermediate values partially
+        weight the corresponding elements. Mask must either have the same
+        shape as the inputs or be broadcastable to the same shape. If None
+        (default), all elements are equally weighted.
+    gain_bias_invariant : bool, optional
+        If True, the error is computed to be independent of any gain or bias
+        differences between the inputs. If False (default), the error is
+        computed to be independent of any gain differences between the inputs.
+    ghat : array_like, optional
+        Precomputed bias-subtracted estimate :math:`\hat{g}`. Only used if
+        `gain_bias_invariant` is True. If None (default), :math:`\hat{g}`
+        is computed from `est`.
+
+    Returns
+    -------
+    error : float
+        Normalized sum squared error
+
+    References
+    ----------
+    [1] S. Thurman and J. Fienup, "Phase retrieval with signal bias", J. Opt. Soc. Am. A/Vol. 26, No. 4 (2009)
+
+    [2] A. Jurling and J. Fienup, "Applications of algorithmic differentiation to phase retrieval algorithms", J. Opt. Soc. Am. A/Vol. 31, No. 7 (2014)
+
+    """
+    xp = array_namespace(data, est)
+
+    if gain_bias_invariant:
+        if ghat is None:
+            ghat = g(est, mask, xp)  # Eq. 12
+        mtil = g(data, mask, xp)  # Eq. 13
+        mhat = ghat
+    else:
+        mtil = data
+        mhat = est
+
+    alpha = _a(mtil, mhat, mask, xp)  # Eq. 4 or Eq. 14
+    alpha = alpha[..., xp.newaxis, xp.newaxis]
+
+    num = _wsum(xp.square(alpha * mhat - mtil), mask, xp)
+    den = _wsum(xp.square(data), mask, xp)
+
+    # Eq. 3 or Eq. 11
+    return xp.mean(num/den)
+
+
+def _wsum(x, mask, xp):
+    # Weighted sum over the last two axes
+    if mask is not None:
+        x = mask * x
+    return xp.sum(x, axis=(-2, -1))
+
+
+def g(data, mask, xp):
+    # This function implements Eqs. 12 and 13 in [1]
+    if mask is None:
+        numer = xp.sum(data, axis=(-2, -1))
+        denom = xp.prod(xp.asarray(data.shape[-2:]))
+    else:
+        numer = xp.sum(mask * data, axis=(-2, -1))
+        denom = xp.sum(mask, axis=(-2, -1))
+    x = numer/denom
+    if data.ndim == 3:
+        x = x[:, xp.newaxis, xp.newaxis]
+    return data - x
+
+
+def _a(gtil, ghat, mask, xp):
+    # Eq. 4 or Eq. 14 in [1]
+    try:
+        return _wsum(gtil * ghat, mask, xp)/_wsum(xp.square(ghat), mask, xp)
+    except FloatingPointError as e:
+        if xp.all(ghat == 0):
+            raise ZeroDivisionError('ghat must be nonzero')
+        else:
+            raise e

{prtools-2.0.1 → prtools-2.1.0}/prtools/jax.py

@@ -87,7 +87,7 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_kwargs=None):
         _, state = carry
         iter_num = otu.tree_get(state, 'count')
         grad = otu.tree_get(state, 'grad')
-        grad_norm = otu.tree_l2_norm(grad)
+        grad_norm = otu.tree_norm(grad)
         return (iter_num == 0) | ((iter_num < maxiter) & (grad_norm >= gtol))
 
     init_carry = (x0, opt.init(x0))

prtools-2.1.0/prtools/mesh.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+
+def mesh(shape, shift=(0, 0), angle=0, indexing='ij'):
+    """Generate a standard mesh.
+
+    Parameters
+    ----------
+    shape : array_like
+        Size of output in pixels (nrows, ncols)
+    shift : (2,) array_like, optional
+        Coordinate translation from center of containing array.
+        Default is (0, 0). Shift is applied according to the specified
+        ``indexing`` scheme.
+    angle : float, optional
+        Mesh rotation in degrees. When ``indexing = 'xy'`` the rotation is
+        measured from the +x axis toward the +y axis. When ``indexing = 'ij'``
+        the rotation is measured from the +i axis toward the +j axis. Note
+        that rotations occur about the shifted center if ``shift`` is nonzero.
+        Default is 0.
+    indexing : {'ij', 'xy'}, optional
+        Matrix ('ij', default) or Cartesian ('xy') indexing of mesh.
+
+    Returns
+    -------
+    list of ndarrays
+
+    """
+    shape = np.broadcast_to(np.asarray(shape), (2,))
+    if indexing == 'xy':
+        shape = np.flip(shape)
+    angle = np.radians(angle)
+
+    x1, x2 = np.meshgrid(np.arange(shape[0]) - np.floor(shape[0]/2) - shift[0],
+                         np.arange(shape[1]) - np.floor(shape[1]/2) - shift[1],
+                         indexing=indexing)
+
+    X1 = x1 * np.cos(angle) + x2 * np.sin(angle)
+    X2 = x1 * -np.sin(angle) + x2 * np.cos(angle)
+    return X1, X2

{prtools-2.0.1 → prtools-2.1.0}/prtools/segmented.py

@@ -5,7 +5,7 @@ import numpy as np
 import prtools
 
 
-def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,
+def hex_segments(rings, seg_radius, seg_gap, rotate=True, antialias=True,
                  flatten=False, pad=2, drop=(0,)):
     """Draw a segmented aperture made up of hexagonal segments
 
@@ -18,9 +18,8 @@ def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,
     seg_gap : float
         Spacing between adjacent segments in pixels
     rotate : bool, optional
-        Rotate segments so that flat sides are aligned with the Y direction 
-        instead of the default orientation which is aligned with the X 
-        direction.
+        If True (default), the segment mask is rotated so that flat sides are
+        aligned with the X direction instead of the Y direction.
     antialias : bool, optional
         If True (default), the segment edges are antialiased.
     flatten : bool, optional
@@ -29,8 +28,8 @@ def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,
     pad : int, optional
         Number of additional pixels to zero-pad output by. Default is 2.
     drop : list_like, optional
-        Segments to exclude from output. 0 represents the central segment. 
-    
+        Segments to exclude from output. 0 represents the central segment.
+
     Returns
     -------
     ndarray
@@ -63,20 +62,24 @@ def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,
     mask = []
 
     if 0 not in drop:
-        mask.append(prtools.hexagon(shape, seg_radius, shift=(0, 0), antialias=antialias, rotate=rotate))
+        mask.append(prtools.hexagon(shape, seg_radius, shift=(0, 0),
+                                    antialias=antialias, rotate=rotate,
+                                    indexing='ij'))
 
     seg = 1
     for ring in range(1, rings+1):
         for h in hex_ring(ring):
             r, c = hex_to_rc(h, seg_radius + seg_gap/2, rotate)
             if seg not in drop:
-                mask.append(prtools.hexagon(shape, seg_radius, shift=(r, c), antialias=antialias, rotate=rotate))
+                mask.append(prtools.hexagon(shape, seg_radius, shift=(r, c),
+                                            antialias=antialias, rotate=rotate,
+                                            indexing='ij'))
             seg += 1
-    
+
     mask = np.asarray(mask)
     if flatten:
         mask = np.sum(mask, axis=0)
-    
+
     return mask
 
 
@@ -84,6 +87,8 @@ def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,
 # exactly from http://www.redblobgames.com/grids/hexagons/
 
 _Hex = collections.namedtuple("Hex", ["q", "r", "s"])
+
+
 def Hex(q, r, s):
     assert not (round(q + r + s) != 0), "q + r + s must be 0"
     return _Hex(q, r, s)
@@ -93,7 +98,7 @@ def hex_add(a, b):
     return Hex(a.q + b.q, a.r + b.r, a.s + b.s)
 
 
-hex_directions = [Hex(1, 0, -1), Hex(1, -1, 0), Hex(0, -1, 1), 
+hex_directions = [Hex(1, 0, -1), Hex(1, -1, 0), Hex(0, -1, 1),
                   Hex(-1, 0, 1), Hex(-1, 1, 0), Hex(0, 1, -1)]
 
 
@@ -117,13 +122,13 @@ def hex_ring(radius):
 
 
 def hex_to_xy(hex, radius, rotate=False):
-    if rotate:
+    if not rotate:
         x = radius * (np.sqrt(3) * hex.q + np.sqrt(3)/2 * hex.r)
         y = radius * (3/2 * hex.r)
     else:
         x = radius * (3/2 * hex.q)
         y = radius * (np.sqrt(3)/2 * hex.q + np.sqrt(3) * hex.r)
-        
+
     return x, y