prtools 2.0.2__tar.gz → 2.1.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prtools
-Version: 2.0.2
+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.2 → prtools-2.1.0}/prtools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '2.0.2'
+__version__ = '2.1.0'
 
 
 from prtools._array_api import use
@@ -33,6 +33,8 @@ from prtools.fourier import dft2, idft2
 
 import prtools.jax
 
+from prtools.mesh import mesh
+
 from prtools.misc import (
     calcpsf,
     translation_defocus,
@@ -86,5 +88,6 @@ from prtools.zernike import (
     zernike_basis,
     zernike_fit,
     zernike_remove,
-    zernike_coordinates
+    zernike_coordinates,
+    zernike_coeff_legacy
 )

{prtools-2.0.2 → 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:
@@ -381,7 +381,7 @@ def sample(x, oversample, offset=0, axes=(-2, -1), indexing='ij'):
     axes = tuple(int(a) % x.ndim for a in axes)
 
     if indexing == 'xy':
-        offset = (-offset[1], offset[0])
+        offset = (offset[1], offset[0])
         oversample = (oversample[1], oversample[0])
 
     idx = [slice(None)] * x.ndim

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

@@ -64,7 +64,7 @@ def fftconv(array, kernel, axes=(-2, -1), normalize_kernel=True,
       (``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``.
+      ``sum(kernel) == 1``.
     """
     xp = array_namespace(array)
 

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.2 → 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.2 → 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