prtools 2.0.0__tar.gz → 2.0.2__tar.gz

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

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

{prtools-2.0.0 → prtools-2.0.2}/prtools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '2.0.0'
+__version__ = '2.0.2'
 
 
 from prtools._array_api import use
@@ -8,6 +8,7 @@ from prtools.array import (
     pad,
     crop,
     boundary,
+    sample,
     rebin,
     rescale,
     normpow,
@@ -55,6 +56,7 @@ from prtools.shape import (
     circle,
     hexagon,
     rectangle,
+    ellipse,
     spider,
     sin,
     waffle

{prtools-2.0.0 → prtools-2.0.2}/prtools/array.py

@@ -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.0 → prtools-2.0.2}/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.0.0 → prtools-2.0.2}/prtools/shape.py

@@ -219,6 +219,68 @@ def rectangle(shape, width, height, shift=(0, 0), angle=0, antialias=True,
     return rect
 
 
+def ellipse(shape, xradius, yradius, shift=(0,0), angle=0, antialias=True,
+            indexing='ij'):
+    """Draw an ellipse
+
+    Parameters
+    ----------
+    shape : array_like
+        Size of output in pixels (nrows, ncols)
+    xradius : float
+        Semi-axis of the ellipse in pixels in the x (horizontal) direction.
+    yradius : float
+        Semi-axis of the ellipse in pixels in the y (vertical) direction.
+    shift : (2,) array_like, optional
+        How far to shift center in float (rows, cols). Default is (0, 0).
+    angle : float, optional
+        Rotation of the ellipse in degrees counterclockwise from horizontal.
+        Default is 0.
+    antialias : bool, optional
+        If True (default), the shape edges are antialiased.
+    indexing : {'ij', 'xy'}, optional
+        Matrix ('ij', default) or cartesian ('xy') indexing of output.
+
+    Returns
+    -------
+    ndarray
+
+    Examples
+    --------
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> ell = prtools.ellipse(shape=(256, 256), xradius=100, yradius=50)
+        >>> plt.imshow(ell, cmap='gray')
+
+    .. plot::
+        :include-source:
+        :context: reset
+        :scale: 50
+
+        >>> ell = prtools.ellipse(shape=(256, 256), xradius=100, yradius=50, angle=30)
+        >>> plt.imshow(ell, cmap='gray')
+
+    """
+    a, b = yradius, xradius
+    rr, cc = mesh(shape, shift, angle, indexing=indexing)
+
+    rho = np.sqrt((rr / a)**2 + (cc / b)**2)
+    grad = np.sqrt((rr / a**2)**2 + (cc / b**2)**2)
+
+    with np.errstate(divide='ignore', invalid='ignore'):
+        dist = (1 - rho) * rho / grad
+    dist = np.where(grad == 0, 1.0, dist)  # center is always interior
+    ell = np.clip(dist + 0.5, 0.0, 1.0)
+
+    if not antialias:
+        ell[ell > 0] = 1
+
+    return ell
+
+
 def spider(shape, width, angle=0, shift=(0, 0), antialias=True, indexing='ij'):
     """Draw a spider
 
@@ -261,10 +323,6 @@ def spider(shape, width, angle=0, shift=(0, 0), antialias=True, indexing='ij'):
     return 1 - rectangle(shape, len, width, shift, angle, antialias, indexing)
 
 
-def ellipse():
-    pass
-
-
 def gauss(shape, sigma, shift=(0, 0), indexing='ij'):
     """Generate a 2D Gaussian function.
 

{prtools-2.0.0 → prtools-2.0.2}/prtools/zernike.py

@@ -14,7 +14,7 @@ def zernike(mask, index, normalize=True, order='noll', rho=None, theta=None):
         Binary mask defining the extent to compute the Zernike polynomial
         over.
     index : int
-        Noll Zernike index as defined in [1]
+        Zernike mode index
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
@@ -112,8 +112,8 @@ def zernike_compose(mask, coeffs, normalize=True, order='noll', rho=None, theta=
     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.
+        List of coefficients corresponding to Zernike indices 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.
@@ -177,7 +177,7 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, order='noll',
     mask : array_like
         Binary mask defining the extent to compute the Zernike polynomial over.
     modes : array_like
-        List of modes (Noll ordering) to return.
+        List of modes to return.
     vectorize : bool, optional
         If True, the output is returned as a
         ``(length(modes), modes.shape[0]*modes.shape[1])`` array If False
@@ -228,11 +228,11 @@ def zernike_fit(opd, mask, modes, normalize=True, order='noll', rho=None,
     Parameters
     ----------
     opd : array_like
-        OPD to fit.
+        OPD to fit. NaN values in ``opd`` are excluded from the fit.
     mask : array_like
         Binary mask defining the extent to compute the Zernike basis over.
     modes : array_like
-        List of modes (Noll ordering) to fit.
+        List of modes to fit.
     normalize : bool, optional
         If True (default), the output is normalized according to [1]. If False,
         the output value ranges [-1, 1] over the mask.
@@ -271,14 +271,16 @@ def zernike_fit(opd, mask, modes, normalize=True, order='noll', rho=None,
     [1] Noll, RJ. Zernike polynomials and atmospheric turbulence. J Opt Soc Am 66, 207-211  (1976).
 
     """
-    opd = np.asarray(opd)
+    opd = np.asarray(opd).ravel()
     mask = np.asarray(mask)
-
     basis = zernike_basis(mask, modes, True, normalize, order, rho, theta)
+    
+    # valid data includes all non-NaN values inside the mask
+    valid = np.isfinite(opd) & mask.ravel().astype(bool)
+    
+    coeff, *_ = np.linalg.lstsq(basis[:, valid].T, opd[valid], rcond=None)
 
-    basis = np.linalg.pinv(basis)
-
-    return np.einsum('ij,i->j', basis, opd.ravel())
+    return coeff
 
 
 def zernike_remove(opd, mask, modes, order='noll', rho=None, theta=None):
@@ -291,7 +293,7 @@ def zernike_remove(opd, mask, modes, order='noll', rho=None, theta=None):
     mask : array_like
         Binary mask defining the extent to compute the Zernike basis over.
     modes : array_like
-        List of modes (Noll ordering) to remove.
+        List of modes to remove.
     order : {'noll', 'fringe', 'ansi'}, optional
         Zernike ordering scheme. Default is 'noll'.
     rho : array_like, optional

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

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

{prtools-2.0.0 → prtools-2.0.2}/tests/test_array.py

@@ -109,4 +109,45 @@ def test_register(xp):
 
     s = prtools.register(x2, x1, oversample=2)
 
-    assert(xp.allclose(xp.asarray(s), xp.asarray((-1,-1))))
+    assert xp.allclose(xp.asarray(s), xp.asarray((-1, -1)))
+
+
+@pytest.mark.parametrize('xp', BACKENDS)
+@pytest.mark.parametrize('shape, oversample, offset, result', [
+    ((2, 2), 1, 0, [[0, 1], [2, 3]]),
+    ((2, 2), 2, 0, [3]),
+    ((2, 2), 2, (-1, 0), [1]),
+    ((2, 2), 2, (0, -1), [2]),
+    ((2, 2), 2, (-1, -1), [0]),
+    ((3, 3), 1, 0, [[0, 1, 2], [3, 4, 5], [6, 7, 8]]),
+    ((3, 3), 3, 0, [4]),
+    ((3, 3), 3, (-1, 0), [1]),
+    ((3, 3), 3, (1, 0), [7]),
+    ((3, 3), 3, (0, -1), [3]),
+    ((3, 3), 3, (0, 1), [5]),
+    ((3, 3), 3, (-1, -1), [0]),
+    ((3, 3), 3, (-1, 1), [2]),
+    ((3, 3), 3, (1, -1), [6]),
+    ((3, 3), 3, (1, 1), [8]),
+    ((4, 4), 1, 0, [[0, 1, 2, 3], [4, 5, 6, 7],
+                    [8, 9, 10, 11], [12, 13, 14, 15]]),
+    ((4, 4), 2, 0, [[5, 7], [13, 15]]),
+    ])
+def test_sample(xp, shape, oversample, offset, result):
+    x = np.arange(np.prod(shape)).reshape(shape)
+    assert np.all(prtools.sample(x, oversample, offset) == result)
+
+@pytest.mark.parametrize('xp', BACKENDS)
+@pytest.mark.parametrize('shape, oversample, offset', [
+    ((2, 2), 1, 1),
+    ((2, 2), 1, -2),
+    ((2, 2), 2, 1),
+    ((3, 3), 3, -2),
+    ((3, 3), 3, 2),
+    ((4, 4), 2, 1),
+    ((4, 4), 2, -2)
+    ])
+def test_sample_value_error(xp, shape, oversample, offset):
+    x = np.arange(np.prod(shape)).reshape(shape)
+    with pytest.raises(ValueError):
+        prtools.sample(x, oversample, offset)