""" CPU solvers ====================================== """ import scipy import numpy from ..src._helper import helper def cDiff(x, d_indx): """ Compute image gradient, which needs the results of indxmap_diff(Nd) :param x: The image array :param d_indx: The index of the shifted image :type x: numpy.float array, matrix size = Nd :type d_indx: int32 :returns: x_diff: Image gradient determined by d_indx :rtype: x_diff: numpy.complex64 """ x_diff=numpy.asarray(x.copy(),order='C') x_diff.flat = x_diff.ravel()[d_indx] - x.ravel() return x_diff def _create_kspace_sampling_density(nufft): """ Compute k-space sampling density """ y = numpy.ones(nufft.st['M'],dtype = numpy.complex64) # w = numpy.abs(nufft.xx2k(nufft.adjoint(y))) if nufft.parallel_flag is 1: w = numpy.abs( nufft.xx2k(nufft.adjoint(y)))[..., 0]#**2) )) else: w = numpy.abs( nufft.xx2k(nufft.adjoint(y)))#**2) )) nufft.st['w'] = w#self.nufftobj.vec2k(w) RTR=nufft.st['w'] # see __init__() in class "nufft" # print('RTR.shape = ', RTR.shape) return RTR def L1TVOLS(nufft, y, maxiter, rho ): # main function of solver """ L1-total variation regularized ordinary least square """ mu = 1.0 LMBD = rho*mu def AHA(x): x2 = nufft.selfadjoint(x) return x2 def AH(y): x2 = nufft.adjoint(y.reshape(nufft.st['M'], order='C')) return x2 uker = mu*_create_kspace_sampling_density(nufft) # print('uker.shape', uker.shape) uker = uker - LMBD* helper.create_laplacian_kernel(nufft) # print('uker.shape', uker.shape) # import matplotlib.pyplot # matplotlib.pyplot.imshow(abs(uker)) # matplotlib.pyplot.show() # print('y.shape=', y.shape) AHy = AH(y) # print('AHy.shape = ', AHy.shape) xkp1 = numpy.zeros_like(AHy) AHyk = numpy.zeros_like(AHy) # self._allo_split_variables() zz= [] bb = [] dd = [] d_indx, dt_indx = helper.indxmap_diff(nufft.st['Nd']) z=numpy.zeros(nufft.st['Nd'], dtype = nufft.dtype, order='C') ndims = len(nufft.st['Nd']) s_tmp = [] for pp in range(0,ndims): s_tmp += [0, ] for jj in range( 0, ndims): # n_dims + 1 for wavelets zz += [z.copy(),] bb += [z.copy(),] dd += [z.copy(),] # zf = z.copy() # bf = z.copy() # df = z.copy() n_dims = len(nufft.st['Nd'])#numpy.size(uf.shape) for outer in numpy.arange(0, maxiter): # for inner in numpy.arange(0,nInner): # solve Ku = rhs rhs = mu*AHyk# + df - bf) # right hand side for pp in range(0,ndims): # diff1 = dd[pp] - bb[pp] # diff1 = numpy.roll( diff1, 1, axis = pp) - diff1 rhs += LMBD*(cDiff(dd[pp] - bb[pp], dt_indx[pp])) # del diff1 # LMBD*(cDiff(dd[1] - bb[1], dt_indx[1])) ) # Note K = F' uker F # so K-1 ~ F xkp1 = nufft.k2xx_one2one( (nufft.xx2k_one2one(rhs)+1e-7) / (uker+1e-7)) # self._update_d(xkp1) for pp in range(0,ndims): # zz[pp] = numpy.roll( xkp1, -1, axis = pp) - xkp1 zz[pp] = cDiff(xkp1, d_indx[pp]) # zz[0] = cDiff(xkp1, d_indx[0]) # zz[1] = cDiff(xkp1, d_indx[1]) zf = AHA(xkp1) -AHy ''' soft-thresholding the edges ''' s = numpy.zeros_like(zz[pp]) # complex for pp in range(0,ndims): s_tmp[pp] = zz[pp] + bb[pp] s_r = numpy.hypot(s_tmp[pp].real, s_tmp[pp].imag) s = numpy.hypot(s_r, s.real) # s = numpy.sqrt(s**2 + s_tmp[pp]**2) # else: # s = s_tmp[pp] # s1 = zz[0] + bb[0] # s2 = zz[1] + bb[1] # s = sum((s_tmp[pp])**2 for pp in range(0,n_dims)) # s = s1**2 + s2**2 s += 1e-5 threshold_value = 1/LMBD r =(s > threshold_value)*(s-threshold_value)/s#numpy.maximum(s - threshold_value , 0.0)/s for pp in range(0,ndims): dd[pp] = s_tmp[pp]*r # dd[0] = s1*r # dd[1] = s2*r # df = zf+bf # # threshold_value=1.0/mu # # df.real =0.0+ (df.real>threshold_value)*(df.real - threshold_value) +(df.real<= - threshold_value)*(df.real+threshold_value) # df.imag = 0.0+(df.imag>threshold_value)*(df.imag - threshold_value) +(df.imag<= - threshold_value)*(df.imag+threshold_value) # df = sy # end of shrinkage for pp in range(0,ndims): bb[pp] += zz[pp] - dd[pp] # bb[0] += zz[0] - dd[0] # bb[1] += zz[1] - dd[1] # bf += zf - df # self._update_b() # update b based on the current u AHyk = AHyk - zf # Linearized Bregman iteration f^k+1 = f^k + f - Au # print(outer) return xkp1 #(u,u_stack) def _pipe_density(nufft, maxiter): ''' Create the density function by iterative solution Generate pHp matrix ''' # W = pipe_density(self.st['p']) # sampling density function W = numpy.ones(nufft.st['M'],dtype=nufft.dtype) # V1= self.st['p'].getH() # VVH = V.dot(V.getH()) for pp in range(0,maxiter): # E = self.st['p'].dot(V1.dot(W)) E = nufft.forward(nufft.adjoint(W)) W = (W/E) return W def solve(nufft, y, solver=None, *args, **kwargs): """ Solve NUFFT. The current version supports solvers = 'cg' or 'L1TVOLS' or 'L1TVLAD'. :param nufft: NUFFT_cpu object :param y: (M,) array, non-uniform data :return: x: image """ if ('cgs' == solver) or ('qmr' ==solver) or ('minres'==solver): raise TypeError(solver +' requires real symmetric matrix') if None == solver: # solver = 'cg' print("solver must be one of the following solvers; \newline dc, cg, bicg, bicgstab, gmres, lgmres, L1OLS, L1LAD") if 'dc' == solver: """ Density compensation method nufft.st['W'] will be computed if doesn't exist If nufft.st['W'] exist then x2 = nufft.adjoint(nufft.st['W']*y) input: y: (M,) array output: x2: Nd array """ print(solver) # try: # # x = nufft.adjoint(nufft.st['W']*y) # # except: W = _pipe_density( nufft, *args, **kwargs) x = nufft.adjoint(W*y) return x elif ('lsmr'==solver) or ('lsqr'==solver): """ Assymetric matrix A Minimize L2 norm of |y-Ax|_2 or |y-Ax|_2+|x|_2 Very stable """ # A = nufft.sp def sp(k): k2 = k.reshape(nufft.Kd, order='C') return nufft.k2y(k2).ravel() def spH(y): y2 = y.reshape(nufft.st['M'], order='C') return nufft.y2k(y2).ravel() # return nufft.spH.dot(nufft.sp.dot(x)) # print('shape', (nufft.st['M']*nufft.batch, nufft.Kdprod*nufft.batch)) # print('spH ') A = scipy.sparse.linalg.LinearOperator((nufft.st['M']*nufft.batch, nufft.Kdprod*nufft.batch), matvec = sp, rmatvec = spH, ) methods={'lsqr':scipy.sparse.linalg.lsqr, 'lsmr':scipy.sparse.linalg.lsmr,} k2 = methods[solver](A, y.flatten(), *args, **kwargs)#,show=True) xx = nufft.k2xx(nufft.vec2k(k2[0])) x= xx/nufft.sn return x#, k2[1:] elif 'L1TVOLS' == solver: return L1TVOLS(nufft, y, *args, **kwargs) # elif 'L1TVLAD' == solver: # return L1TVLAD(nufft, y, *args, **kwargs) else: """ Hermitian matrix A cg: conjugate gradient bicgstab: biconjugate gradient stablizing bicg: biconjugate gradient gmres: lgmres: """ # A = nufft.spHsp#nufft.st['p'].getH().dot(nufft.st['p']) def spHsp(x): k = x.reshape(nufft.Kd, order='C') return nufft.k2y2k(k).ravel() # return nufft.spH.dot(nufft.sp.dot(x)) A = scipy.sparse.linalg.LinearOperator((nufft.Kdprod*nufft.batch, nufft.Kdprod*nufft.batch), matvec = spHsp, rmatvec = spHsp, ) methods={'cg':scipy.sparse.linalg.cg, 'bicgstab':scipy.sparse.linalg.bicgstab, 'bicg':scipy.sparse.linalg.bicg, # 'cgs':scipy.sparse.linalg.cgs, 'gmres':scipy.sparse.linalg.gmres, 'lgmres':scipy.sparse.linalg.lgmres, # 'minres':scipy.sparse.linalg.minres, # 'qmr':scipy.sparse.linalg.qmr, } k2 = methods[solver](A, nufft.y2k(y).ravel(), *args, **kwargs)#,show=True) xx = nufft.k2xx(k2[0].reshape(nufft.Kd)) x= xx/nufft.sn return x# , k2[1:]