diff --git a/indigo/backends/backend.py b/indigo/backends/backend.py index 8e73601..7fe263c 100644 --- a/indigo/backends/backend.py +++ b/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). """ diff --git a/indigo/backends/mkl.py b/indigo/backends/mkl.py index 64a32e2..4e8a249 100644 --- a/indigo/backends/mkl.py +++ b/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 ) diff --git a/indigo/operators.py b/indigo/operators.py index 1f74a13..9e2ce75 100644 --- a/indigo/operators.py +++ b/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 diff --git a/indigo/test_compat.py b/indigo/test_compat.py index 99fc54d..9f3251b 100644 --- a/indigo/test_compat.py +++ b/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 diff --git a/indigo/test_operators.py b/indigo/test_operators.py index 3b41f03..148ea9e 100644 --- a/indigo/test_operators.py +++ b/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) diff --git a/indigo/transforms.py b/indigo/transforms.py index 7d614af..90ab808 100644 --- a/indigo/transforms.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index 6ee3805..2139ca8 100644 --- a/requirements.txt +++ b/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