prtools 1.4.1__tar.gz → 1.4.2__tar.gz

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

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

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

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

{prtools-1.4.1 → prtools-1.4.2}/prtools/backend/_jax.py

@@ -44,11 +44,11 @@ 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)
+    #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):

{prtools-1.4.1 → prtools-1.4.2}/prtools/fourier.py

@@ -100,11 +100,8 @@ def _dftcore(f, alpha, shape, shift, offset, axes, unitary, forward, out):
 
     f = np.asarray(f)
 
-    out_shape, axes = _cook_nd_args(f, shape, axes)
-    in_shape = np.take(np.asarray(f.shape), axes)
-
-    m, n = in_shape
-    M, N = out_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,))
@@ -158,7 +155,7 @@ def _cook_nd_args(a, s=None, axes=None):
             else:
                 raise ValueError("Array must have ndim == 2 or 3")
         else:
-            s = np.take(a.shape, axes)
+            s = np.take(np.asarray(a.shape), np.asarray(axes))
     s = list(s)
     if axes is None:
         axes = list(range(-len(s), 0))

{prtools-1.4.1 → 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)
+            fn(x, **fn_kwargs)
 
-        where ``x`` is a 1-D array with shape (n,) and ``fn_args`` and
-        ``fn_kwargs`` are optional positional and keyword arguments.
+        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,7 +81,6 @@ 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()
@@ -93,11 +90,12 @@ def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=None,
         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,
@@ -151,9 +149,10 @@ def _multi_dot_three(a, b, c, axes, out):
         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))
+    # wrapper function to support vmap call signature
+    return jax.numpy.linalg.multi_dot((a, b, c))
 
 
 def _iter_axis(axes):
@@ -169,4 +168,4 @@ def _ndrebin(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)
+        return jax.numpy.reshape(a, (a.shape[0]//f, f, a.shape[1]//f, f)).sum(-1).sum(1)

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

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

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

@@ -25,4 +25,5 @@ prtools/backend/_jax.py
 prtools/backend/_numpy.py
 tests/test_array.py
 tests/test_convolve.py
-tests/test_fourier.py
+tests/test_fourier.py
+tests/test_jax.py

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)))