update to newer apgd

indigo/backends/backend.py

@@ -350,6 +350,7 @@ def FFT(self, shape, dtype, **kwargs):
         s = np.ones(n, order='F', dtype=dtype) / np.sqrt(n)
         S = self.Diag(s, name='scale')
         F = self.UnscaledFFT(shape, dtype, **kwargs)
+        return S,F # FIXME
         return S*F
 
     def FFTc(self, ft_shape, dtype, normalize=True, **kwargs):
@@ -363,9 +364,11 @@ def FFTc(self, ft_shape, dtype, normalize=True, **kwargs):
         mod = np.exp(1j * 2.0 * np.pi * mod).astype(dtype)
         M = self.Diag(mod, name='mod')
         if normalize:
-            F = self.FFT(ft_shape, dtype=dtype, **kwargs)
+            U, F = self.FFT(ft_shape, dtype=dtype, **kwargs)
+            return M,U,F,M # FIXME
         else:
             F = self.UnscaledFFT(ft_shape, dtype=dtype, **kwargs)
+            return M,F,M # FIXME
         return M*F*M
 
     def Zpad(self, M, N, mode='center', dtype=np.dtype('complex64'), **kwargs):
@@ -403,7 +406,7 @@ def Interp(self, N, coord, width, table, dtype=np.dtype('complex64'), **kwargs):
     def NUFFT(self, M, N, coord, width=3, n=128, oversamp=None, dtype=np.dtype('complex64'), **kwargs):
         assert len(M) == 3
         assert len(N) == 3
-        assert M[1:] == coord.shape[1:]
+        assert M[1:] == coord.shape[1:], (M, coord.shape)
 
         # target 448 x 270 x 640
         #   448 x 270 x 640   mkl-batch: 170.83 ms, 237.51 gflop/s  back-to-back: 121.76 ms, 333.23 gflop/s
@@ -430,7 +433,7 @@ def NUFFT(self, M, N, coord, width=3, n=128, oversamp=None, dtype=np.dtype('comp
         oN = tuple(int(on) for on in oN)
 
         Z = self.Zpad(oN, N, dtype=dtype, name='zpad')
-        F = self.FFTc(oN, dtype=dtype, name='fft')
+        M1,U,F,M2 = self.FFTc(oN, dtype=dtype, name='fft')
 
         beta = np.pi * np.sqrt(((width * 2. / omin) * (omin- 0.5)) ** 2 - 0.8)
         kb = signal.kaiser(2 * n + 1, beta)[n:]
@@ -439,6 +442,7 @@ def NUFFT(self, M, N, coord, width=3, n=128, oversamp=None, dtype=np.dtype('comp
         r = rolloff3(omin, width, beta, N)
         R = self.Diag(r, name='apod')
 
+        return G,M1,U,F,M2,Z,R # FIXME
         return G*F*Z*R
 
     def Convolution(self, kernel, normalize=True, name='noname'):
@@ -683,48 +687,45 @@ def cg(self, A, b_h, x_h, lamda=0.0, tol=1e-10, maxiter=100, team=None):
             log.info("cg reached maxiter")
         x.copy_to(x_h)
 
-    def apgd(self, gradf, proxg, alpha, x_h, maxiter=100, team=None):
-        '''Accelerated proximal gradient descent.
-        Solves for min_x f(x) + g(x)
-
-        Parameters
-        ----------
-        gradf : Gradient of f
-        proxg : Proximal of g
-        alpha : Step size
-        x0 : 1D array, initial solution
-        maxiter : int, optional
-        '''
-        x_k = self.copy_array(x_h)
-        y_k = x_k.copy()
-        y_k1 = x_k.copy()
-        x_k1 = x_k.copy()
+    def power(self, A, x_h, maxiter=10):
+        x = self.copy_array(x_h)
+        for it in range(maxiter):
+            A.eval(x, x)
+            s = np.sqrt(self.norm2(x))
+            self.scale(x, 1/s)
+        return s, x
 
-        gf = x_k.copy()
+    def apgd(self, gradf, proxg, alpha, x_h, maxiter=100, disp=None):
 
-        t_k = 1
+        gfx = self.zeros_like(x_h)
+        x = self.copy_array(x_h)
+        z = x.copy()
+        t = 1.0
 
-        for it in range(1,maxiter+1):
+        for it in range(maxiter):
             profile.extra['it'] = it
-
             with profile("iter"):
-                gradf(gf, y_k)
-                self.axpby(1, x_k, -alpha, gf)
 
-                proxg(x_k, alpha)
+                    x, z = z, x
+                    s = t
 
-                t_k1 = (1.0 + np.sqrt(1.0 + 4.0 * t_k**2)) / 2.0
+                    with profile("linop"):
+                        gradf(gfx, x)
 
-                t_ratio = (t_k - 1) / t_k1
-                self.axpby(0, y_k1, 1+t_ratio, x_k)
-                self.axpby(1, y_k1,  -t_ratio, x_k1)
+                    self.axpby(1, x, -alpha, gfx)  # x = x - alpha * gfx
 
-                x_k1.copy(x_k)
-                y_k.copy(y_k1)
+                    with profile("prox"):
+                        if proxg is not None:
+                            proxg(x, alpha)  # x = proxg(x)
 
-            log.info("iter %d", it)
+                    t = (1.0 + (1.0 + 4.0 * t**2)**0.5) / 2.0
+                    self.axpby((1.0 - s) / t, z, (s + t - 1.0) / t, x)
 
-        x_k.copy_to(x_h)
+                    if disp is not None:
+                        disp(x.to_host())
+
+        x.copy_to(x_h)
+        return x_h
 
     def max(self, val, arr):
         """ Computes elementwise maximum: arr[:] = max(arr, val). """

indigo/backends/mkl.py

@@ -66,7 +66,7 @@ def _zero(self):
             self._arr[:] = 0
 
         def __getitem__(self, slc):
-            d = self._arr[slc]
+            d = self._arr.reshape(self.shape, order='F')[slc]
             return self._backend.dndarray( self._backend, d.shape, d.dtype,
                 ld=self._leading_dim, own=False, data=d )
 

indigo/operators.py

@@ -26,13 +26,15 @@ def eval(self, y, x, alpha=1, beta=0, forward=True, left=True):
         """
         M, N = self.shape if forward else tuple(reversed(self.shape))
         if left: # left-multiply
+            xexp, yexp = x.size, y.size
             x = x.reshape( (N,-1) )
             y = y.reshape( (M,-1) )
-            assert x.shape[1] == y.shape[1], "Dimension mismatch"
+            xact, yact = x.size, y.size
+            assert x.shape[1] == y.shape[1], "Dimension mismatch: expected %d, got %d" % (x.shape[1], y.shape[1])
         else: # right-multiply
             x = x.reshape( (-1,M) )
             y = y.reshape( (-1,N) )
-            assert x.shape[0] == y.shape[0], "Dimension mismatch"
+            assert x.shape[0] == y.shape[0], "Dimension mismatch: expected %d, got %d" % (x.shape[0], y.shape[0])
         self._eval(y, x, alpha=alpha, beta=beta, forward=forward, left=left)
 
     @property
@@ -113,6 +115,15 @@ def has(self, *op_classes):
         from indigo.analyses import TreeHasOp
         return TreeHasOp(op_classes).search(self)
 
+    def save(self, fname):
+        from pickle import dump
+        dump(self, fname)
+
+    @staticmethod
+    def load(fname):
+        from pickle import load
+        return load(fname)
+
 
 class CompositeOperator(Operator):
     def __init__(self, backend, *children, **kwargs):
@@ -197,6 +208,7 @@ def __init__(self, backend, M, **kwargs):
         """
         super().__init__(backend, **kwargs)
         assert isinstance(M, spp.spmatrix)
+
         self._matrix = M
         self._matrix_d = None
 
@@ -434,6 +446,7 @@ def _eval_forward(self, y, x, alpha=1, beta=0, left=True):
         for C in self._children:
             h = C.shape[0]
             slc = slice( h_offset, h_offset+h )
+            y_slc = y[slc,:]
             C.eval( y[slc,:], x, alpha=alpha, beta=beta, forward=True, left=left)
             h_offset += h
 

indigo/test_compat.py

@@ -120,12 +120,11 @@ def test_compat_NUFFT(backend, X, Y, Z, RO, PS, K, oversamp, n, width):
     traj = indigo.util.rand64c( *t_dims ).real - 0.5
     kwargs = dict(oversamp=oversamp, width=width, n=n, dtype=x.dtype)
 
-    print(nc_dims, c_dims, traj.shape)
     G0 = pymr.linop.NUFFT(nc_dims, c_dims, traj, **kwargs)
     G1 = b.NUFFT(nc_dims[:3], c_dims[:3], traj, **kwargs)
 
     x_indigo = np.asfortranarray(x.reshape((-1,K), order='F'))
-    x_pmr = pymr.util.vec(x)
+    x_pmr = pymr.util.vec(x.copy())
     y_exp = G0 * x_pmr
     y_act = G1 * x_indigo
 

indigo/test_operators.py

@@ -3,10 +3,13 @@
 import scipy.sparse as spp
 import numpy.testing as npt
 from itertools import product
+from tempfile import TemporaryFile
 from scipy.signal import fftconvolve
 
 import indigo
+from indigo.operators import Operator
 from indigo.backends import available_backends
+
 BACKENDS = available_backends()
 
 @pytest.mark.parametrize("backend,M,N,K,density,alpha,beta",
@@ -646,7 +649,6 @@ def test_Kron_general(backend, L, Q, alpha, beta, eyeL, eyeR):
     product( BACKENDS, [23,45], [45,23], [1,2,3] ))
 def test_Convolution(backend, M, N, P):
     from scipy.signal import fftconvolve
-
     b = backend()
 
     #k = indigo.util.rand64c(M,N)
@@ -666,3 +668,55 @@ def test_Convolution(backend, M, N, P):
 
     pytest.xfail("under development")
     #npt.assert_allclose(y_act, y_exp, rtol=1e-5)
+
+
+@pytest.mark.parametrize("backend,M,N,K,density,alpha,beta",
+    product( BACKENDS, [3,4], [7,8], [1,8,9,17], [0.01,0.1,0.5,1], [0,.5,1], [0,.5,1] ))
+def test_pickle(backend, M, N, K, density, alpha, beta):
+    A_h = indigo.util.randM(M, N, density)
+    x = indigo.util.rand64c(N,K)
+    y = indigo.util.rand64c(M,K)
+
+    def test_b0():
+        b0 = backend()
+        A = b0.SpMatrix(A_h, name='A0')
+
+        from tempfile import NamedTemporaryFile
+        with NamedTemporaryFile(delete=False) as f:
+            A.save(f)
+
+        x_d = b0.copy_array(x)
+        y_d = b0.copy_array(y)
+        x_exp_d = b0.copy_array(x)
+        y_exp_d = b0.copy_array(y)
+
+        A.eval(y_exp_d, x_d, alpha=alpha, beta=beta)
+        A.H.eval(x_exp_d, y_d, alpha=alpha, beta=beta)
+
+        x_exp = x_exp_d.to_host()
+        y_exp = y_exp_d.to_host()
+
+        del b0
+
+        return f.name, x_exp, y_exp
+
+    def test_b1(fname):
+        with open(fname, 'rb') as f:
+            B = Operator.load(f)
+        b1 = B._backend
+
+        x_d = b1.copy_array(x)
+        y_d = b1.copy_array(y)
+        x_act_d = b1.copy_array(x)
+        y_act_d = b1.copy_array(y)
+
+        B.eval(y_act_d, x_d, alpha=alpha, beta=beta)
+        B.H.eval(x_act_d, y_d, alpha=alpha, beta=beta)
+
+        return x_act_d.to_host(), y_act_d.to_host()
+
+    saved, x_exp, y_exp = test_b0()
+    x_act, y_act = test_b1(saved)
+
+    npt.assert_allclose(y_act, y_exp, rtol=1e-5)
+    npt.assert_allclose(x_act, x_exp, rtol=1e-5)

indigo/transforms.py

@@ -136,7 +136,7 @@ def visit_Kron(self, node):
         node = self.generic_visit(node)
         L, R = node.children
         if isinstance(L, Eye):
-            L = L.realize()
+            return node._backend.BlockDiag([R]*L.shape[0]).realize()
         if isinstance(L, SpMatrix) and isinstance(R, SpMatrix):
             name = "({}(x){})".format(L._name, R._name)
             log.debug('realizing kron %s x %s', L._name, R._name)

requirements.txt

@@ -29,7 +29,6 @@ dependencies:
 - pip:
   - chardet==3.0.4
   - idna==2.6
-  - indigo==1.0.2
   - python-coveralls==2.9.1
   - pyyaml==3.12
   - requests==2.18.4