def par_xstep(i): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}_{G_i}`, one of the disjoint problems of optimizing :math:`\mathbf{x}`. Parameters ---------- i : int Index of grouping to update """ global mp_X global mp_DX YU0f = sl.rfftn(mp_Y0[[i]] - mp_U0[[i]], mp_Nv, mp_axisN) YU1f = sl.rfftn( mp_Y1[mp_grp[i]:mp_grp[i + 1]] - 1 / mp_alpha * mp_U1[mp_grp[i]:mp_grp[i + 1]], mp_Nv, mp_axisN) if mp_Cd == 1: b = np.conj(mp_Df[mp_grp[i]:mp_grp[i + 1]]) * YU0f + mp_alpha**2 * YU1f Xf = sl.solvedbi_sm(mp_Df[mp_grp[i]:mp_grp[i + 1]], mp_alpha**2, b, mp_cache[i], axis=mp_axisM) else: b = sl.inner(np.conj(mp_Df[mp_grp[i]:mp_grp[i + 1]]), YU0f, axis=mp_C) + mp_alpha**2 * YU1f Xf = sl.solvemdbi_ism(mp_Df[mp_grp[i]:mp_grp[i + 1]], mp_alpha**2, b, mp_axisM, mp_axisC) mp_X[mp_grp[i]:mp_grp[i + 1]] = sl.irfftn(Xf, mp_Nv, mp_axisN) mp_DX[i] = sl.irfftn( sl.inner(mp_Df[mp_grp[i]:mp_grp[i + 1]], Xf, mp_axisM), mp_Nv, mp_axisN)
def xstep_check(self, b): r"""Check the minimisation of the Augmented Lagrangian with respect to :math:`\mathbf{x}` by method `xstep` defined in derived classes. This method should be called at the end of any `xstep` method. """ if self.opt['LinSolveCheck']: Zop = lambda x: sl.inner(self.Zf, x, axis=self.cri.axisM) ZHop = lambda x: sl.inner(np.conj(self.Zf), x, axis=self.cri.axisK) ax = ZHop(Zop(self.Xf)) + self.rho*self.Xf self.xrrs = sl.rrs(ax, b) else: self.xrrs = None
def xstep(self): r"""Minimise Augmented Lagrangian with respect to block vector :math:`\mathbf{x} = \left( \begin{array}{ccc} \mathbf{x}_0^T & \mathbf{x}_1^T & \ldots \end{array} \right)^T\;`. """ # This test reflects empirical evidence that two slightly # different implementations are faster for single or # multi-channel data. This kludge is intended to be temporary. if self.cri.Cd > 1: for i in range(self.Nb): self.xistep(i) else: self.YU[:] = self.Y[..., np.newaxis] - self.U b = np.swapaxes(self.ZSf[..., np.newaxis], self.cri.axisK, -1) \ + self.rho*sl.rfftn(self.YU, None, self.cri.axisN) for i in range(self.Nb): self.Xf[..., i] = sl.solvedbi_sm( self.Zf[..., [i], :], self.rho, b[..., i], axis=self.cri.axisM) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) if self.opt['LinSolveCheck']: ZSfs = np.sum(self.ZSf, axis=self.cri.axisK, keepdims=True) YU = np.sum(self.Y[..., np.newaxis] - self.U, axis=-1) b = ZSfs + self.rho*sl.rfftn(YU, None, self.cri.axisN) Xf = self.swapaxes(self.Xf) Zop = lambda x: sl.inner(self.Zf, x, axis=self.cri.axisM) ZHop = lambda x: np.conj(self.Zf) * x ax = np.sum(ZHop(Zop(Xf)) + self.rho*Xf, axis=self.cri.axisK, keepdims=True) self.xrrs = sl.rrs(ax, b) else: self.xrrs = None
def obfn_dfd(self): r"""Compute data fidelity term :math:`(1/2) \| \sum_m \mathbf{d}_m * \mathbf{x}_m - \mathbf{s} \|_2^2`. """ Ef = sl.inner(self.Zf, self.obfn_fvarf(), axis=self.cri.axisM) \ - self.Sf return sl.rfl2norm2(Ef, self.S.shape, axis=self.cri.axisN) / 2.0
def obfn_dfd(self): r"""Compute data fidelity term :math:`(1/2) \| W \left( \sum_m \mathbf{d}_m * \mathbf{x}_m - \mathbf{s} \right) \|_2^2`. """ XF = sl.rfftn(self.obfn_fvar(), mp_Nv, mp_axisN) DX = np.moveaxis( sl.irfftn(sl.inner(mp_Df, XF, mp_axisM), mp_Nv, mp_axisN), mp_axisM, self.cri.axisM) return np.sum((self.W * (DX - self.S))**2) / 2.0
def obfn_dfd(self): r"""Compute data fidelity term :math:`(1/2) \| W \left( \sum_m \mathbf{d}_m * \mathbf{x}_m - \mathbf{s} \right) \|_2^2`. """ Ef = sl.inner(self.Zf, self.obfn_fvarf(), axis=self.cri.axisM) \ - self.Sf return (np.linalg.norm(self.W * sl.irfftn(Ef, self.cri.Nv, self.cri.axisN))**2) / 2.0
def cbpdnmd_xstep(k): """Do the X step of the cbpdn stage. The only parameter is the slice index `k` and there are no return values; all inputs and outputs are from and to global variables. """ YU0 = mp_Z_Y0[k] + mp_S[k] - mp_Z_U0[k] YU1 = mp_Z_Y1[k] - mp_Z_U1[k] if mp_cri.Cd == 1: b = np.conj(mp_Df) * spl.rfftn(YU0, None, mp_cri.axisN) + \ spl.rfftn(YU1, None, mp_cri.axisN) Xf = spl.solvedbi_sm(mp_Df, 1.0, b, axis=mp_cri.axisM) else: b = spl.inner(np.conj(mp_Df), spl.rfftn(YU0, None, mp_cri.axisN), axis=mp_cri.axisC) + \ spl.rfftn(YU1, None, mp_cri.axisN) Xf = spl.solvemdbi_ism(mp_Df, 1.0, b, mp_cri.axisM, mp_cri.axisC) mp_Z_X[k] = spl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN) mp_DX[k] = spl.irfftn(spl.inner(mp_Df, Xf), mp_cri.Nv, mp_cri.axisN)
def reconstruct(self, D=None, X=None): """Reconstruct representation.""" if D is None: D = self.getdict(crop=False) if X is None: X = self.getcoef() Df = sl.rfftn(D, self.xstep.cri.Nv, self.xstep.cri.axisN) Xf = sl.rfftn(X, self.xstep.cri.Nv, self.xstep.cri.axisN) DXf = sl.inner(Df, Xf, axis=self.xstep.cri.axisM) return sl.irfftn(DXf, self.xstep.cri.Nv, self.xstep.cri.axisN)
def cnst_A0T(self, Y0): r"""Compute :math:`A_0^T \mathbf{y}_0` component of :math:`A^T \mathbf{y}` (see :meth:`.ADMMTwoBlockCnstrnt.cnst_AT`). """ # This calculation involves non-negligible computational cost. It # should be possible to disable relevant diagnostic information # (dual residual) to avoid this cost. Y0f = sl.rfftn(Y0, None, self.cri.axisN) return sl.irfftn(sl.inner(np.conj(self.Zf), Y0f, axis=self.cri.axisK), self.cri.Nv, self.cri.axisN)
def cnst_A0(self, X, Xf=None): r"""Compute :math:`A_0 \mathbf{x}` component of ADMM problem constraint. """ # This calculation involves non-negligible computational cost # when Xf is None (i.e. the function is not being applied to # self.X). if Xf is None: Xf = sl.rfftn(X, None, self.cri.axisN) return sl.irfftn(sl.inner(self.Zf, Xf, axis=self.cri.axisM), self.cri.Nv, self.cri.axisN)
def eval_grad(self): """Compute gradient in Fourier domain.""" # Compute X D - S Ryf = self.eval_Rf(self.Yf) gradf = sl.inner(np.conj(self.Zf), Ryf, axis=self.cri.axisK) # Multiple channel signal, single channel dictionary if self.cri.C > 1 and self.cri.Cd == 1: gradf = np.sum(gradf, axis=self.cri.axisC, keepdims=True) return gradf
def ccmodmd_xstep(k): """Do the X step of the ccmod stage. The only parameter is the slice index `k` and there are no return values; all inputs and outputs are from and to global variables. """ YU0 = mp_D_Y0 - mp_D_U0[k] YU1 = mp_D_Y1[k] + mp_S[k] - mp_D_U1[k] b = spl.rfftn(YU0, None, mp_cri.axisN) + \ np.conj(mp_Zf[k]) * spl.rfftn(YU1, None, mp_cri.axisN) Xf = spl.solvedbi_sm(mp_Zf[k], 1.0, b, axis=mp_cri.axisM) mp_D_X[k] = spl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN) mp_DX[k] = spl.irfftn(spl.inner(Xf, mp_Zf[k]), mp_cri.Nv, mp_cri.axisN)
def evaluate(self): """Evaluate functional value of previous iteration.""" X = mp_Z_Y Xf = mp_Zf Df = mp_Df Sf = mp_Sf Ef = spl.inner(Df[np.newaxis, ...], Xf, axis=self.xstep.cri.axisM + 1) - Sf Ef = np.swapaxes(Ef, 0, self.xstep.cri.axisK + 1)[0] dfd = spl.rfl2norm2(Ef, self.xstep.S.shape, axis=self.xstep.cri.axisN) / 2.0 rl1 = np.sum(np.abs(X)) obj = dfd + self.xstep.lmbda * rl1 return (obj, dfd, rl1)
def xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`. """ self.YU[:] = self.Y - self.U self.block_sep0(self.YU)[:] += self.S YUf = sl.rfftn(self.YU, None, self.cri.axisN) b = sl.inner(np.conj(self.Zf), self.block_sep0(YUf), axis=self.cri.axisK) + self.block_sep1(YUf) self.Xf[:] = sl.solvemdbi_ism(self.Zf, 1.0, b, self.cri.axisM, self.cri.axisK) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) self.xstep_check(b)
def cbpdn_setdict(): """Set the dictionary for the cbpdn stage. There are no parameters or return values because all inputs and outputs are from and to global variables. """ global mp_DSf # Set working dictionary for cbpdn step and compute DFT of dictionary # D and of D^T S mp_Df[:] = spl.rfftn(mp_D_Y, mp_cri.Nv, mp_cri.axisN) if mp_cri.Cd == 1: mp_DSf[:] = np.conj(mp_Df) * mp_Sf else: mp_DSf[:] = spl.inner(np.conj(mp_Df[np.newaxis, ...]), mp_Sf, axis=mp_cri.axisC + 1)
def setcoef(self, Z): """Set coefficient array.""" # If the dictionary has a single channel but the input (and # therefore also the coefficient map array) has multiple # channels, the channel index and multiple image index have # the same behaviour in the dictionary update equation: the # simplest way to handle this is to just reshape so that the # channels also appear on the multiple image index. if self.cri.Cd == 1 and self.cri.C > 1: Z = Z.reshape(self.cri.Nv + (1,) + (self.cri.Cx*self.cri.K,) + (self.cri.M,)) self.Z = np.asarray(Z, dtype=self.dtype) self.Zf = sl.rfftn(self.Z, self.cri.Nv, self.cri.axisN) # Compute X^H S self.ZSf = sl.inner(np.conj(self.Zf), self.Sf, self.cri.axisK)
def relax_AX(self): """The parent class method that this method overrides only implements the relaxation step for the variables of the baseline consensus algorithm. This method calls the overridden method and then implements the relaxation step for the additional variables required for the mask decoupling modification to the baseline algorithm. """ super(ConvCnstrMODMaskDcpl_Consensus, self).relax_AX() self.AX1nr = sl.irfftn(sl.inner(self.Zf, self.swapaxes(self.Xf), axis=self.cri.axisM), self.cri.Nv, self.cri.axisN) if self.rlx == 1.0: self.AX1 = self.AX1nr else: alpha = self.rlx self.AX1 = alpha*self.AX1nr + (1-alpha)*(self.Y1 + self.S)
def xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`. """ self.cgit = None self.YU[:] = self.Y - self.U self.block_sep0(self.YU)[:] += self.S YUf = sl.rfftn(self.YU, None, self.cri.axisN) b = sl.inner(np.conj(self.Zf), self.block_sep0(YUf), axis=self.cri.axisK) + self.block_sep1(YUf) self.Xf[:], cgit = sl.solvemdbi_cg( self.Zf, 1.0, b, self.cri.axisM, self.cri.axisK, self.opt['CG', 'StopTol'], self.opt['CG', 'MaxIter'], self.Xf) self.cgit = cgit self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) self.xstep_check(b)
def eval_grad(self): """Compute gradient in Fourier domain.""" # Compute X D - S self.Ryf[:] = self.eval_Rf(self.Yf) # Map to spatial domain to multiply by mask Ry = sl.irfftn(self.Ryf, self.cri.Nv, self.cri.axisN) # Multiply by mask self.WRy[:] = (self.W**2) * Ry # Map back to frequency domain WRyf = sl.rfftn(self.WRy, self.cri.Nv, self.cri.axisN) gradf = sl.inner(np.conj(self.Zf), WRyf, axis=self.cri.axisK) # Multiple channel signal, single channel dictionary if self.cri.C > 1 and self.cri.Cd == 1: gradf = np.sum(gradf, axis=self.cri.axisC, keepdims=True) return gradf
def evaluate(self): """Evaluate functional value of previous iteration.""" if self.opt['AccurateDFid']: DX = self.reconstruct() W = self.dstep.W S = self.dstep.S else: W = mp_W S = mp_S Xf = mp_Zf Df = mp_Df DX = spl.irfftn( spl.inner(Df[np.newaxis, ...], Xf, axis=self.xstep.cri.axisM + 1), self.xstep.cri.Nv, np.array(self.xstep.cri.axisN) + 1) dfd = (np.linalg.norm(W * (DX - S))**2) / 2.0 rl1 = np.sum(np.abs(self.getcoef())) obj = dfd + self.xstep.lmbda * rl1 return (obj, dfd, rl1)
def evaluate(self): """Evaluate functional value of previous iteration.""" if self.opt['AccurateDFid']: if self.dmethod == 'fista': D = self.dstep.getdict(crop=False) else: D = self.dstep.var_y() if self.xmethod == 'fista': X = self.xstep.getcoef() else: X = self.xstep.var_y() Df = sl.rfftn(D, self.xstep.cri.Nv, self.xstep.cri.axisN) Xf = sl.rfftn(X, self.xstep.cri.Nv, self.xstep.cri.axisN) Sf = self.xstep.Sf Ef = sl.inner(Df, Xf, axis=self.xstep.cri.axisM) - Sf dfd = sl.rfl2norm2( Ef, self.xstep.S.shape, axis=self.xstep.cri.axisN) / 2.0 rl1 = np.sum(np.abs(X)) return dict(DFid=dfd, RegL1=rl1, ObjFun=dfd + self.xstep.lmbda * rl1) else: return None
def eval_Rf(self, Vf): """Evaluate smooth term in Vf.""" return sl.inner(self.Df, Vf, axis=self.cri.axisM) - self.Sf