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 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 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 reconstruct(self, X=None): """Reconstruct representation.""" if X is None: X = self.X Xf = sl.rfftn(X, None, self.cri.axisN) Sf = np.sum(self.Df * Xf, axis=self.cri.axisM) return sl.irfftn(Sf, self.cri.Nv, self.cri.axisN)
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 (\sum_m \mathbf{d}_m * \mathbf{x}_{m} - \mathbf{s}) \|_2^2` """ Ef = self.eval_Rf(self.Xf) E = sl.irfftn(Ef, self.cri.Nv, self.cri.axisN) return (np.linalg.norm(self.W * E)**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 ccmod_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. """ YU = mp_D_Y - mp_D_U[k] b = mp_ZSf[k] + mp_drho * spl.rfftn(YU, None, mp_cri.axisN) Xf = spl.solvedbi_sm(mp_Zf[k], mp_drho, b, axis=mp_cri.axisM) mp_D_X[k] = spl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN)
def reconstruct(self, D=None): """Reconstruct representation.""" if D is None: Df = self.Xf else: Df = sl.rfftn(D, None, self.cri.axisN) Sf = np.sum(self.Zf * Df, axis=self.cri.axisM) return sl.irfftn(Sf, self.cri.Nv, self.cri.axisN)
def xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`. """ self.YU[:] = self.Y - self.U b = self.ZSf + self.rho*sl.rfftn(self.YU, None, self.cri.axisN) self.Xf[:] = sl.solvemdbi_ism(self.Zf, self.rho, b, self.cri.axisM, self.cri.axisK) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) self.xstep_check(b)
def compute_residuals(self): """Compute residuals and stopping thresholds. The parent class method is overridden to ensure that the residual calculations include the additional variables introduced in the modification to the baseline algorithm. """ # The full primary residual is straightforward to compute from # the primary residuals for the baseline algorithm and for the # additional variables r0 = self.rsdl_r(self.AXnr, self.Y) r1 = self.AX1nr - self.Y1 - self.S r = np.sqrt(np.sum(r0**2) + np.sum(r1**2)) # The full dual residual is more complicated to compute than the # full primary residual ATU = self.swapaxes(self.U) + sl.irfftn( np.conj(self.Zf) * sl.rfftn(self.U1, self.cri.Nv, self.cri.axisN), self.cri.Nv, self.cri.axisN) s = self.rho * np.linalg.norm(ATU) # The normalisation factor for the full primal residual is also not # straightforward nAX = np.sqrt(np.linalg.norm(self.AXnr)**2 + np.linalg.norm(self.AX1nr)**2) nY = np.sqrt(np.linalg.norm(self.Y)**2 + np.linalg.norm(self.Y1)**2) rn = max(nAX, nY, np.linalg.norm(self.S)) # The normalisation factor for the full dual residual is # straightforward to compute sn = self.rho * np.sqrt(np.linalg.norm(self.U)**2 + np.linalg.norm(self.U1)**2) # Final residual values and stopping tolerances depend on # whether standard or normalised residuals are specified via the # options object if self.opt['AutoRho', 'StdResiduals']: epri = np.sqrt(self.Nc)*self.opt['AbsStopTol'] + \ rn*self.opt['RelStopTol'] edua = np.sqrt(self.Nx)*self.opt['AbsStopTol'] + \ sn*self.opt['RelStopTol'] else: if rn == 0.0: rn = 1.0 if sn == 0.0: sn = 1.0 r /= rn s /= sn epri = np.sqrt(self.Nc)*self.opt['AbsStopTol']/rn + \ self.opt['RelStopTol'] edua = np.sqrt(self.Nx)*self.opt['AbsStopTol']/sn + \ self.opt['RelStopTol'] return r, s, epri, edua
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_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 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 cbpdn_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. """ YU = mp_Z_Y[k] - mp_Z_U[k] b = mp_DSf[k] + mp_xrho * spl.rfftn(YU, None, mp_cri.axisN) if mp_cri.Cd == 1: Xf = spl.solvedbi_sm(mp_Df, mp_xrho, b, axis=mp_cri.axisM) else: Xf = spl.solvemdbi_ism(mp_Df, mp_xrho, b, mp_cri.axisM, mp_cri.axisC) mp_Z_X[k] = spl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN)
def xistep(self, i): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}` component :math:`\mathbf{x}_i`. """ self.YU[:] = self.Y - self.U[..., i] b = np.take(self.ZSf, [i], axis=self.cri.axisK) + \ self.rho*sl.rfftn(self.YU, None, self.cri.axisN) self.Xf[..., i] = sl.solvedbi_sm(np.take( self.Zf, [i], axis=self.cri.axisK), self.rho, b, axis=self.cri.axisM) self.X[..., i] = sl.irfftn(self.Xf[..., i], self.cri.Nv, self.cri.axisN)
def xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`. """ self.cgit = None self.YU[:] = self.Y - self.U b = self.ZSf + self.rho*sl.rfftn(self.YU, None, self.cri.axisN) self.Xf[:], cgit = sl.solvemdbi_cg(self.Zf, self.rho, 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 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 obfn_f(self, Xf=None): r"""Compute data fidelity term :math:`(1/2) \| W (\sum_m \mathbf{d}_m * \mathbf{x}_{m} - \mathbf{s}) \|_2^2`. This is used for backtracking. Since the backtracking is computed in the DFT, it is important to preserve the DFT scaling. """ if Xf is None: Xf = self.Xf Rf = self.eval_Rf(Xf) R = sl.irfftn(Rf, self.cri.Nv, self.cri.axisN) WRf = sl.rfftn(self.W * R, self.cri.Nv, self.cri.axisN) return 0.5 * np.linalg.norm(WRf.flatten(), 2)**2
def proximal_step(self, gradf=None): """Compute proximal update (gradient descent + constraint). Variables are mapped back and forth between input and frequency domains. """ if gradf is None: gradf = self.eval_grad() self.Vf[:] = self.Yf - (1. / self.L) * gradf V = sl.irfftn(self.Vf, self.cri.Nv, self.cri.axisN) self.X[:] = self.eval_proxop(V) self.Xf = sl.rfftn(self.X, None, self.cri.axisN) return gradf
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 D X - 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 = np.conj(self.Df) * WRyf # Multiple channel signal, multiple channel dictionary if 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)