prtools 1.3.0__tar.gz → 1.3.2__tar.gz

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

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

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

@@ -1,6 +1,6 @@
-__version__ = '1.3.0'
+__version__ = '1.3.2'
 
-from prtools._backend import __backend__, use
+from prtools.backend import __backend__, use
 
 from prtools.array import (
     centroid,
@@ -19,6 +19,8 @@ from prtools.cost import sserror
 
 from prtools.fourier import dft2, idft2
 
+import prtools.jax
+
 from prtools.misc import (
     calcpsf,
     translation_defocus,

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

@@ -1,8 +1,8 @@
 import warnings
 
 from prtools import __backend__
-from prtools._backend import numpy as np
-from prtools._backend import scipy
+from prtools.backend import numpy as np
+from prtools.backend import scipy
 
 
 def centroid(a, where=None, kind='absolute', indexing='ij'):
@@ -265,8 +265,9 @@ def boundary(x, threshold=0):
     rows = np.any(x, axis=1)
     cols = np.any(x, axis=0)
 
-    rmin, rmax = np.where(rows)[0][[0, -1]]
-    cmin, cmax = np.where(cols)[0][[0, -1]]
+    idx = np.array((0, -1))  # https://github.com/andykee/prtools/issues/6
+    rmin, rmax = np.where(rows)[0][idx]
+    cmin, cmax = np.where(cols)[0][idx]
 
     return rmin, rmax, cmin, cmax
 

prtools-1.3.0/prtools/_backend.py → prtools-1.3.2/prtools/backend/__init__.py

@@ -1,12 +1,22 @@
+import importlib
 import sys
 import types
 
-from prtools.backend import registry
-from prtools.backend._base import _BackendLibrary
+from ._base import BackendLibrary
 
 
-class _BackendName:
-    """Helper class for storing the backend name - needed since strings are
+def backend_available(name):
+    """Return True if backend library is available"""
+    try:
+        importlib.import_module(name)
+    except ImportError:
+        return False
+    else:
+        return True
+        
+
+class BackendName:
+    """Helper class for changing the backend name - needed since strings are
     immutable.
     """
 
@@ -23,10 +33,32 @@ class _BackendName:
         return self.name == other
 
 
-class _Backend(types.ModuleType):
-    __backend__ = _BackendName(None)
-    numpy = _BackendLibrary(None)
-    scipy = _BackendLibrary(None)
+class BackendRegistry:
+
+    _BUILTIN_BACKEND = ['numpy', 'jax']
+
+    def _get_backend_module(self, backend):
+        return importlib.import_module(f'._{backend}', 'prtools.backend')
+
+    def load_numpy(self, backend):
+        backend = self._get_backend_module(backend)
+        return backend.Numpy()
+
+    def load_scipy(self, backend):
+        return self._get_backend_module(backend).Scipy()
+
+
+registry = BackendRegistry()
+
+class Backend(types.ModuleType):
+    
+    numpy = BackendLibrary(None)
+    scipy = BackendLibrary(None)
+
+    __backend__ = BackendName(None)
+
+    JAX_AVAILABLE = backend_available('jax') & backend_available('optax')
+    NUMPY_AVAILABLE = backend_available('numpy')
 
     @classmethod
     def use(cls, backend):
@@ -67,7 +99,7 @@ class _Backend(types.ModuleType):
 
 
 # use NumPy by default
-_Backend.use('numpy')
+Backend.use('numpy')
 
 # https://stackoverflow.com/a/72911884
-sys.modules[__name__].__class__ = _Backend
+sys.modules[__name__].__class__ = Backend

prtools-1.3.2/prtools/backend/_base.py

@@ -0,0 +1,7 @@
+
+class BackendLibrary:
+    def __init__(self, module):
+        self.module = module
+
+    def __getattr__(self, name):
+        return getattr(self.module, name)

prtools-1.3.2/prtools/backend/_jax.py

@@ -0,0 +1,49 @@
+import importlib
+
+from ._base import BackendLibrary
+
+
+class Numpy(BackendLibrary):
+    def __init__(self):
+        super().__init__(importlib.import_module('jax.numpy'))
+
+    def broadcast_to(self, array, shape):
+        # jax numpy.broadcast_to expects an array input
+        array = self.module.asarray(array)
+        return self.module.broadcast_to(array, shape)
+
+    def divide(self, a, b, out=None):
+        # jax.numpy.divide doesn't support the `out` parameter so we
+        # ignore it
+        return self.module.divide(a, b)
+
+    def dot(self, a, b, out=None):
+        # jax.numpy.dot doesn't support the `out` parameter so we ignore it
+        return self.module.dot(a, b)
+
+    def floor(self, x, *args, **kwargs):
+        # jax numpy.floor expects a scalar or array input. It also doesn't
+        # support the `out` parameter
+        kwargs.pop('out', None)
+        x = self.module.asarray(x)
+        return self.module.floor(x, *args, **kwargs)
+
+    def max(self, a, *args, **kwargs):
+        # jax numpy.max expects an array input
+        array = self.module.asarray(a)
+        return self.module.max(array, *args, **kwargs)
+
+    def multiply(self, a, b, out=None):
+        # jax.numpy.multiply doesn't support the `out` parameter so we
+        # ignore it
+        return self.module.multiply(a, b)
+
+    def sum(self, a, *args, **kwargs):
+        kwargs.pop('out', None)
+        a = self.module.asarray(a)
+        return self.module.sum(a, *args, **kwargs)
+
+
+class Scipy(BackendLibrary):
+    def __init__(self):
+        super().__init__(importlib.import_module('jax.scipy'))

prtools-1.3.2/prtools/backend/_numpy.py

@@ -0,0 +1,13 @@
+from ._base import BackendLibrary
+
+
+class Numpy(BackendLibrary):
+    def __init__(self):
+        import numpy
+        super().__init__(numpy)
+
+
+class Scipy(BackendLibrary):
+    def __init__(self):
+        import scipy
+        super().__init__(scipy)

{prtools-1.3.0 → prtools-1.3.2}/prtools/cost.py

@@ -1,4 +1,4 @@
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def sserror(data, est, mask=None, gain_bias_invariant=False, ghat=None):

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

@@ -1,7 +1,5 @@
-import functools
-
 from prtools import __backend__
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
@@ -137,7 +135,6 @@ def _idft2_matrices(m, n, M, N, alphar, alphac, shiftr, shiftc, offsetr, offsetc
     return E1, E2
 
 
-@functools.lru_cache(maxsize=32)
 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

prtools-1.3.2/prtools/jax.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass
+from typing import Any
+
+from prtools import __backend__
+from prtools.backend import JAX_AVAILABLE
+
+if JAX_AVAILABLE:
+    import jax
+    import optax
+    import optax.tree_utils as otu
+
+
+def register_dataclass(cls):
+    if JAX_AVAILABLE:
+        data_fields = ['x', 'n', 'grad', 'value', 'state']
+        meta_fields = []
+        cls = jax.tree_util.register_dataclass(cls,
+                                               data_fields=data_fields,
+                                               meta_fields=meta_fields)
+    return cls
+
+
+@register_dataclass
+@dataclass
+class JaxOptimizeResult:
+    """Represents the optimization result."""
+    x: Any  #: The solution of the optimization
+    n: Any  #: Number of iterations performed by the optimizer
+    value: Any  #: Value of objective function at x
+    grad: Any  #: Value of objective function gradient at x
+    state: Any  #: Optimizer state
+
+
+def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
+    """Minimize a scalar function of one or more variables using the L-BFGS
+    algorithm
+
+    Parameters
+    ----------
+    fun : callable
+        The objective function to be minimied
+    x0 : jax.Array
+        Initial guess
+    gtol : float
+        Iteration stops when ``l2_norm(grad) <= gtol``
+    maxiter : int
+        Maximum number of iterations
+    callback : callable, optional
+        A callable called after each iteration with the signature
+
+        .. code:: python
+
+            callback(intermediate_result: JaxOptimizeResult)
+
+        where ``intermediate_result`` is a :class:`JaxOptimizeResult`.
+
+    Returns
+    -------
+    res: JaxOptimizeResult
+        The optimization result. See :class:`JaxOptimizeResult` for a
+        description of attributes.
+
+    """
+    if not JAX_AVAILABLE:
+        message = "jax and optax must be installed to use method `lbfgs`."
+        raise ModuleNotFoundError(message)
+
+    if __backend__ != 'jax':
+        raise RuntimeError('JAX backend must be selected')
+
+    if not any((gtol, maxiter)):
+        raise ValueError('At least one termination tolerance must be specified.')
+
+    opt = optax.lbfgs()
+    value_and_grad_fun = optax.value_and_grad_from_state(fun)
+    
+    def step(carry):
+        params, state = carry
+        value, grad = value_and_grad_fun(params, state=state)
+        updates, state = opt.update(
+            grad, state, params, value=value, grad=grad, value_fn=fun)
+        if callback:
+            res = JaxOptimizeResult(
+                n=otu.tree_get(state, 'count')-1,
+                x=params,
+                grad=grad,
+                value=value,
+                state=state)
+            jax.debug.callback(callback, res)
+        params = optax.apply_updates(params, updates)
+        return params, state
+
+    def continuing_criterion(carry):
+        _, state = carry
+        iter_num = otu.tree_get(state, 'count')
+        grad = otu.tree_get(state, 'grad')
+        grad_norm = otu.tree_l2_norm(grad)
+        return (iter_num == 0) | ((iter_num < maxiter) & (grad_norm >= gtol))
+
+    init_carry = (x0, opt.init(x0))
+    final_params, final_state = jax.lax.while_loop(
+        continuing_criterion, step, init_carry)
+
+    res = JaxOptimizeResult(
+        n=otu.tree_get(final_state, 'count'),
+        x=final_params,
+        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

{prtools-1.3.0 → prtools-1.3.2}/prtools/misc.py

@@ -1,5 +1,5 @@
 import prtools
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def calcpsf(amp, opd, wavelength, sampling, shape, oversample=2,
@@ -18,7 +18,7 @@ def calcpsf(amp, opd, wavelength, sampling, shape, oversample=2,
         Propagation sampling term defined as
 
         .. math::
-            \mbox{sampling} = \\frac{dx \ du}{z}
+            \mbox{sampling} = \frac{dx \ du}{z}
 
         where *dx* is the pupil plane spatial sampling, *du* is the image
         plane spatial sampling, and *z* is the propagation distance. If a

{prtools-1.3.0 → prtools-1.3.2}/prtools/segmented.py

@@ -1,7 +1,7 @@
 import collections
 
 import prtools
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def hex_segments(rings, seg_radius, seg_gap, rotate=False, antialias=True,

{prtools-1.3.0 → prtools-1.3.2}/prtools/shape.py

@@ -1,5 +1,5 @@
 from prtools import __backend__
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def mesh(shape, shift=(0, 0), angle=0, indexing='ij'):

{prtools-1.3.0 → prtools-1.3.2}/prtools/stats.py

@@ -1,5 +1,5 @@
 import prtools
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def ee(a, energy=0.8, center=None):

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

@@ -1,7 +1,7 @@
 from math import factorial
 
 import prtools
-from prtools._backend import numpy as np
+from prtools.backend import numpy as np
 
 
 def zernike(mask, index, normalize=True, rho=None, theta=None):

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

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

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

@@ -3,10 +3,10 @@ NOTICE
 README.rst
 pyproject.toml
 prtools/__init__.py
-prtools/_backend.py
 prtools/array.py
 prtools/cost.py
 prtools/fourier.py
+prtools/jax.py
 prtools/misc.py
 prtools/segmented.py
 prtools/shape.py
@@ -18,4 +18,9 @@ prtools.egg-info/SOURCES.txt
 prtools.egg-info/dependency_links.txt
 prtools.egg-info/requires.txt
 prtools.egg-info/top_level.txt
-tests/test_array.py
+prtools/backend/__init__.py
+prtools/backend/_base.py
+prtools/backend/_jax.py
+prtools/backend/_numpy.py
+tests/test_array.py
+tests/test_fourier.py

{prtools-1.3.0 → prtools-1.3.2}/pyproject.toml

@@ -36,8 +36,5 @@ Documentation = "https://andykee.github.io/prtools/"
 [project.optional-dependencies]
 jax = ["jax", "optax"]
 
-[tool.setuptools]
-packages = ["prtools"]
-
 [tool.setuptools.dynamic]
 version = {attr = "prtools.__version__"} 

prtools-1.3.2/tests/test_array.py

@@ -0,0 +1,67 @@
+import warnings
+
+import numpy as np
+import pytest
+
+import prtools
+from .import BACKENDS
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_centroid(backend):
+    prtools.use(backend)
+    x = np.zeros((5, 5))
+    x[2, 2] = 1
+    assert np.array_equal(prtools.centroid(x), [2, 2])
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_centroid_where(backend):
+    prtools.use(backend)
+    x = np.zeros((5, 5))
+    x[2, 2] = 1
+    x[1, 1] = 1
+    mask = np.ones_like(x)
+    mask[1, 1] = 0
+    assert np.array_equal(prtools.centroid(x, where=mask), [2, 2])
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_centroid_nan(backend):
+    prtools.use(backend)
+    x = np.zeros((5, 5))
+    x[2, 2] = 1
+    x[2, 3] = np.nan
+    assert np.array_equal(prtools.centroid(x), [2, 2])
+
+
+x, _ = np.meshgrid(range(10), range(10))
+x[2, 2] = 100
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_medfix(backend):
+    prtools.use(backend)
+    m = np.zeros_like(x)
+    m[2, 2] = 1
+    y = prtools.medfix(x, mask=m, kernel=(3, 3))
+    assert y[2, 2] == 2
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_medfix_bigmask(backend):
+    prtools.use(backend)
+    m = np.zeros_like(x)
+    m[2:6, 2:6] = 1
+    with warnings.catch_warnings():
+        warnings.simplefilter('error')
+        y = prtools.medfix(x, mask=m, kernel=(3, 3))
+    assert np.all(np.isnan(y[3:5, 3:5]))
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_boundary(backend):
+    prtools.use(backend)
+    x = np.zeros((10, 10))
+    x[3:7, 2:8] = 1
+    assert np.array_equal(prtools.boundary(x), (3, 6, 2, 7))

prtools-1.3.2/tests/test_fourier.py

@@ -0,0 +1,95 @@
+import numpy as np
+import pytest
+
+import prtools
+from .import BACKENDS
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_dft2_even(backend):
+    prtools.use(backend)
+    n = 10
+    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)))
+
+    assert np.allclose(F_dft, F_fft, atol=1e-5)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_idft2_even(backend):
+    prtools.use(backend)
+    n = 10
+    F = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+
+    f_dft = prtools.idft2(F, 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)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_dft2_odd(backend):
+    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)
+
+    f_dft = prtools.idft2(F, 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)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_dft2_unitary(backend):
+    prtools.use(backend)
+    n = 10
+    f = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    f_power = np.sum(np.abs(f)**2)
+
+    F = prtools.dft2(f, 1/n, unitary=True)
+    F_power = np.sum(np.abs(F)**2)
+
+    assert np.allclose(f_power, F_power)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_dft2_shift(backend):
+    prtools.use(backend)
+    n = 100
+    shift = np.round(np.random.uniform(low=-25, high=25, size=2))
+    f = np.ones((n, n))
+
+    F = prtools.dft2(f, 1/n, shift=shift)
+
+    (xc, yc) = (np.floor(n/2), np.floor(n/2))
+    (x, y) = prtools.centroid(np.abs(F)**2)
+    observed_shift = (x-xc, y-yc)
+    assert np.array_equal(shift, observed_shift)
+
+
+@pytest.mark.parametrize('backend', BACKENDS)
+def test_dft2_rectangle(backend):
+    prtools.use(backend)
+    m, n = 10, 11
+    f = np.random.rand(m, n) + 1j * np.random.rand(m, n)
+
+    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)

prtools-1.3.0/tests/test_array.py

@@ -1,45 +0,0 @@
-import warnings
-
-import numpy as np
-
-import prtools
-
-
-def test_centroid():
-    x = np.zeros((5,5))
-    x[2,2] = 1
-    assert(np.array_equal(prtools.centroid(x), [2,2]))
-
-
-def test_centroid_where():
-    x = np.zeros((5,5))
-    x[2,2] = 1
-    x[1,1] = 1
-    mask = np.ones_like(x)
-    mask[1,1] = 0
-    assert(np.array_equal(prtools.centroid(x, where=mask), [2,2]))
-
-
-def test_centroid_nan():
-    x = np.zeros((5,5))
-    x[2,2] = 1
-    x[2,3] = np.nan
-    assert(np.array_equal(prtools.centroid(x), [2,2]))
-
-
-x, _ = np.meshgrid(range(10), range(10))
-x[2,2] = 100
-
-def test_medfix():
-    m = np.zeros_like(x)
-    m[2,2] = 1
-    y = prtools.medfix(x, mask=m, kernel=(3,3))
-    assert(y[2,2] == 2)
-
-def test_medfix_bigmask():
-    m = np.zeros_like(x)
-    m[2:6, 2:6] = 1
-    with warnings.catch_warnings():
-        warnings.simplefilter('error')
-        y = prtools.medfix(x, mask=m, kernel=(3,3))
-    assert(np.all(np.isnan(y[3:5,3:5])))