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(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, 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 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 = sl.rfftn(YU0, None, mp_cri.axisN) + \ np.conj(mp_Zf[k]) * sl.rfftn(YU1, None, mp_cri.axisN) Xf = sl.solvedbi_sm(mp_Zf[k], 1.0, b, axis=mp_cri.axisM) mp_D_X[k] = sl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN) mp_DX[k] = sl.irfftn(sl.inner(Xf, mp_Zf[k]), mp_cri.Nv, mp_cri.axisN)
def setdict(self, D=None, B=None): """Set dictionary array.""" if D is not None: self.D = np.asarray(D, dtype=self.dtype) if B is not None: self.B = np.asarray(B, dtype=self.dtype) if B is not None or not hasattr(self, 'Gamma'): self.Gamma, self.Q = np.linalg.eigh(self.B.T.dot(self.B)) self.Gamma = np.abs(self.Gamma) if D is not None or not hasattr(self, 'Df'): self.Df = sl.rfftn(self.D, self.cri.Nv, self.cri.axisN) # Fold square root of Gamma into the dictionary array to enable # use of the solvedbi_sm solver shpg = [1] * len(self.cri.shpD) shpg[self.cri.axisC] = self.Gamma.shape[0] Gamma2 = np.sqrt(self.Gamma).reshape(shpg) self.gDf = Gamma2 * self.Df if self.opt['HighMemSolve']: self.c = sl.solvedbd_sm_c( self.gDf, np.conj(self.gDf), (self.mu / self.rho) * self.GHGf + 1.0, self.cri.axisM) else: self.c = None
def xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`.""" self.YU[:] = self.Y - self.U YUf = sl.rfftn(self.YU, None, self.cri.axisN) # The sum is over the extra axis indexing spatial gradient # operators G_i, *not* over axisM b = self.DSf + self.rho*(YUf[..., -1] + self.Wtv * np.sum( np.conj(self.Gf) * YUf[..., 0:-1], axis=-1)) if self.cri.Cd == 1: self.Xf[:] = sl.solvedbi_sm( self.Df, self.rho*self.GHGf + self.rho, b, self.c, self.cri.axisM) else: self.Xf[:] = sl.solvemdbi_ism( self.Df, self.rho*self.GHGf + self.rho, b, self.cri.axisM, self.cri.axisC) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) if self.opt['LinSolveCheck']: Dop = lambda x: sl.inner(self.Df, x, axis=self.cri.axisM) if self.cri.Cd == 1: DHop = lambda x: np.conj(self.Df) * x else: DHop = lambda x: sl.inner(np.conj(self.Df), x, axis=self.cri.axisC) ax = DHop(Dop(self.Xf)) + (self.rho*self.GHGf + self.rho)*self.Xf self.xrrs = sl.rrs(ax, b) else: self.xrrs = None
def obfn_fvarf(self): """Variable to be evaluated in computing data fidelity term, depending on 'fEvalX' option value. """ return self.Xf if self.opt['fEvalX'] else \ sl.rfftn(self.Y, None, self.cri.axisN)
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 Zf = sl.rfftn(self.YU, None, self.cri.axisN) Z0f = self.block_sep0(Zf) Z1f = self.block_sep1(Zf) DZ0f = np.conj(self.Df) * Z0f DZ0fBQ = sl.dot(self.B.dot(self.Q).T, DZ0f, axis=self.cri.axisC) Z1fQ = sl.dot(self.Q.T, Z1f, axis=self.cri.axisC) b = DZ0fBQ + Z1fQ Xh = sl.solvedbd_sm(self.gDf, (self.mu / self.rho) * self.GHGf + 1.0, b, self.c, axis=self.cri.axisM) self.Xf[:] = sl.dot(self.Q, Xh, axis=self.cri.axisC) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) if self.opt['LinSolveCheck']: DDXf = np.conj(self.Df) * sl.inner(self.Df, self.Xf, axis=self.cri.axisM) DDXfBB = sl.dot(self.B.T.dot(self.B), DDXf, axis=self.cri.axisC) ax = self.rho * (DDXfBB + self.Xf) + \ self.mu * self.GHGf * self.Xf b = self.rho * (sl.dot(self.B.T, DZ0f, axis=self.cri.axisC) + Z1f) self.xrrs = sl.rrs(ax, b) else: self.xrrs = None
def cbpdnmd_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. """ # Set working dictionary for cbpdn step and compute DFT of dictionary D mp_Df[:] = sl.rfftn(mp_D_Y0, mp_cri.Nv, mp_cri.axisN)
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 cnst_AT(self, X): r"""Compute :math:`A^T \mathbf{x}` where :math:`A \mathbf{x}` is a component of ADMM problem constraint. In this case :math:`A^T \mathbf{x} = (G_r^T \;\; G_c^T) \mathbf{x}`. """ Xf = sl.rfftn(X, axes=self.axes) return np.sum(sl.irfftn(np.conj(self.Gf)*Xf, self.axsz, axes=self.axes), axis=self.Y.ndim-1)
def cnst_A1T(self, Y1): r"""Compute :math:`A_1^T \mathbf{y}_1` component of :math:`A^T \mathbf{y}`. In this case :math:`A_1^T \mathbf{y}_1 = (\Gamma_0^T \;\; \Gamma_1^T \;\; \ldots) \mathbf{y}_1`. """ Y1f = sl.rfftn(Y1, None, axes=self.cri.axisN) return sl.irfftn(np.conj(self.GDf) * Y1f, 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 ccmodmd_setcoef(k): """Set the coefficient maps for 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. """ # Set working coefficient maps for ccmod step and compute DFT of # coefficient maps Z mp_Zf[k] = sl.rfftn(mp_Z_Y1[k], mp_cri.Nv, mp_cri.axisN)
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) * sl.rfftn(YU0, None, mp_cri.axisN) + \ sl.rfftn(YU1, None, mp_cri.axisN) Xf = sl.solvedbi_sm(mp_Df, 1.0, b, axis=mp_cri.axisM) else: b = sl.inner(np.conj(mp_Df), sl.rfftn(YU0, None, mp_cri.axisN), axis=mp_cri.axisC) + \ sl.rfftn(YU1, None, mp_cri.axisN) Xf = sl.solvemdbi_ism(mp_Df, 1.0, b, mp_cri.axisM, mp_cri.axisC) mp_Z_X[k] = sl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN) mp_DX[k] = sl.irfftn(sl.inner(mp_Df, Xf), mp_cri.Nv, mp_cri.axisN)
def cnst_A(self, X, Xf=None): r"""Compute :math:`A \mathbf{x}` component of ADMM problem constraint. In this case :math:`A \mathbf{x} = (G_r^T \;\; G_c^T)^T \mathbf{x}`. """ if Xf is None: Xf = sl.rfftn(X, axes=self.axes) return sl.irfftn(self.Gf*Xf[..., np.newaxis], self.axsz, axes=self.axes)
def cnst_A0T(self, X): r"""Compute :math:`A_0^T \mathbf{x}` where :math:`A_0 \mathbf{x}` is a component of the ADMM problem constraint. In this case :math:`A_0^T \mathbf{x} = (\Gamma_0^T \;\; \Gamma_1^T \;\; \ldots ) \mathbf{x}`. """ Xf = sl.rfftn(X, axes=self.cri.axisN) return self.Wtv[..., np.newaxis] * sl.irfftn( np.conj(self.Gf) * Xf[..., 0:-1], self.cri.Nv, axes=self.cri.axisN)
def cnst_A0(self, X, Xf=None): r"""Compute :math:`A_0 \mathbf{x}` component of ADMM problem constraint. In this case :math:`A_0 \mathbf{x} = (\Gamma_0^T \;\; \Gamma_1^T \;\; \ldots )^T \mathbf{x}`. """ if Xf is None: Xf = sl.rfftn(X, axes=self.cri.axisN) return self.Wtv[..., np.newaxis] * sl.irfftn( self.Gf * Xf[..., np.newaxis], self.cri.Nv, axes=self.cri.axisN)
def cnst_A0(self, X, Xf=None): r"""Compute :math:`A_0 \mathbf{x}` component of ADMM problem constraint. """ if Xf is None: Xf = sl.rfftn(X, None, self.cri.axisN) return sl.irfftn( sl.dot(self.B, sl.inner(self.Df, Xf, axis=self.cri.axisM), axis=self.cri.axisC), self.cri.Nv, self.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 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 * sl.rfftn(YU, None, mp_cri.axisN) Xf = sl.solvedbi_sm(mp_Zf[k], mp_drho, b, axis=mp_cri.axisM) mp_D_X[k] = sl.irfftn(Xf, mp_cri.Nv, mp_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 cnst_AT(self, X): r"""Compute :math:`A^T \mathbf{x}` where :math:`A \mathbf{x}` is a component of ADMM problem constraint. In this case :math:`A^T \mathbf{x} = (G_r^T \;\; G_c^T) \mathbf{x}`. """ Xf = sl.rfftn(X, axes=self.axes) return np.sum(sl.irfftn(np.conj(self.Gf) * Xf, self.axsz, axes=self.axes), axis=self.Y.ndim - 1)
def cnst_A(self, X, Xf=None): r"""Compute :math:`A \mathbf{x}` component of ADMM problem constraint. In this case :math:`A \mathbf{x} = (G_r^T \;\; G_c^T)^T \mathbf{x}`. """ if Xf is None: Xf = sl.rfftn(X, axes=self.axes) return sl.irfftn(self.Gf * Xf[..., np.newaxis], self.axsz, axes=self.axes)
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 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_A1(self, X, Xf=None): r"""Compute :math:`A_1 \mathbf{x}` component of ADMM problem constraint. In this case :math:`A_1 \mathbf{x} = (\Gamma_0^T \;\; \Gamma_1^T \;\; \ldots )^T \mathbf{x}`. """ if Xf is None: Xf = sl.rfftn(X, axes=self.cri.axisN) return sl.irfftn( sl.inner(self.GDf, Xf[..., np.newaxis], axis=self.cri.axisM), self.cri.Nv, self.cri.axisN)
def proximal_step(self, gradf=None): """ Compute proximal update (gradient descent + constraint). """ if gradf is None: gradf = self.eval_gradf() Vf = self.Yf - (1. / self.L) * gradf V = sl.irfftn(Vf, self.cri.Nv, self.cri.axisN) self.X = self.eval_proxop(V) self.Xf = sl.rfftn(self.X, None, 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 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.dot(self.B.T, np.conj(self.Df) * Y0f, axis=self.cri.axisC), self.cri.Nv, self.cri.axisN)
def __init__(self, D, S, lmbda, mu=0.0, opt=None, dimK=None, dimN=2): if opt is None: opt = ConvBPDNRecTV.Options() # Infer problem dimensions and set relevant attributes of self self.cri = cbpdn.ConvRepIndexing(D, S, dimK=dimK, dimN=dimN) # Call parent class __init__ Nx = np.product(self.cri.shpX) yshape = list(self.cri.shpX) yshape[self.cri.axisM] += len(self.cri.axisN) * self.cri.Cd super(ConvBPDNRecTV, self).__init__(Nx, yshape, yshape, S.dtype, opt) # Set l1 term scaling and weight array self.lmbda = self.dtype.type(lmbda) self.Wl1 = np.asarray(opt['L1Weight'], dtype=self.dtype) self.mu = self.dtype.type(mu) if hasattr(opt['TVWeight'], 'ndim') and opt['TVWeight'].ndim > 0: self.Wtv = np.asarray(opt['TVWeight'].reshape((1,)*(dimN+2) + opt['TVWeight'].shape), dtype=self.dtype) else: # Wtv is a scalar: no need to change shape self.Wtv = self.dtype.type(opt['TVWeight']) # Set penalty parameter self.set_attr('rho', opt['rho'], dval=(50.0*self.lmbda + 1.0), dtype=self.dtype) # Set rho_xi attribute self.set_attr('rho_xi', opt['AutoRho','RsdlTarget'], dval=1.0, dtype=self.dtype) # Reshape D and S to standard layout self.D = np.asarray(D.reshape(self.cri.shpD), dtype=self.dtype) self.S = np.asarray(S.reshape(self.cri.shpS), dtype=self.dtype) # Compute signal in DFT domain self.Sf = sl.rfftn(self.S, None, self.cri.axisN) self.Gf, GHGf = sl.GradientFilters(self.cri.dimN+3, self.cri.axisN, self.cri.Nv, dtype=self.dtype) # Initialise byte-aligned arrays for pyfftw self.YU = sl.pyfftw_empty_aligned(self.Y.shape, dtype=self.dtype) xfshp = list(self.cri.shpX) xfshp[dimN-1] = xfshp[dimN-1]//2 + 1 self.Xf = sl.pyfftw_empty_aligned(xfshp, dtype=sl.complex_dtype(self.dtype)) self.setdict()
def setcoef(self, Z): """Set coefficient array.""" # This method largely replicates the method from parent class # ConvCnstrMOD_Consensus that it overrides. The inherited # method is overridden to avoid the superfluous computation of # self.ZSf in that method, which is not required for the # modified algorithm with mask decoupling 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)
def xstep(self, S, W, lmbda, dimK): """Solve CSC problem for training data `S`.""" if self.opt['CUDA_CBPDN']: Z = cucbpdn.cbpdnmsk(self.D.squeeze(), S[..., 0], W.squeeze(), lmbda, self.opt['CBPDN']) Z = Z.reshape(self.cri.Nv + (1, 1, self.cri.M,)) self.Z[:] = np.asarray(Z, dtype=self.dtype) self.Zf = sl.rfftn(self.Z, self.cri.Nv, self.cri.axisN) self.Sf = sl.rfftn(S.reshape(self.cri.shpS), self.cri.Nv, self.cri.axisN) self.xstep_itstat = None else: # Create X update object (external representation is expected!) xstep = cbpdn.ConvBPDNMaskDcpl(self.D.squeeze(), S, lmbda, W, self.opt['CBPDN'], dimK=dimK, dimN=self.cri.dimN) xstep.solve() self.Sf = sl.rfftn(S.reshape(self.cri.shpS), self.cri.Nv, self.cri.axisN) self.setcoef(xstep.getcoef()) self.xstep_itstat = xstep.itstat[-1] if xstep.itstat else None
def xstep(self): """Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`.""" b = self.AHSf + np.sum(np.conj(self.GAf) * sl.rfftn(self.Y-self.U, axes=self.axes), axis=self.Y.ndim-1) self.Xf = b / (self.AHAf + self.GHGf) self.X = sl.irfftn(self.Xf, None, axes=self.axes) if self.opt['LinSolveCheck']: ax = (self.AHAf + self.GHGf)*self.Xf self.xrrs = sl.rrs(ax, b) else: self.xrrs = None
def setdict(self, D=None): """Set dictionary array.""" if D is not None: self.D = np.asarray(D, dtype=self.dtype) self.Df = sl.rfftn(self.D, self.cri.Nv, self.cri.axisN) self.GDf = self.Gf * (self.Wtv * self.Df)[..., np.newaxis] # Compute D^H S self.DSf = np.conj(self.Df) * self.Sf if self.cri.Cd > 1: self.DSf = np.sum(self.DSf, axis=self.cri.axisC, keepdims=True)
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 * sl.rfftn(YU, None, mp_cri.axisN) if mp_cri.Cd == 1: Xf = sl.solvedbi_sm(mp_Df, mp_xrho, b, axis=mp_cri.axisM) else: Xf = sl.solvemdbi_ism(mp_Df, mp_xrho, b, mp_cri.axisM, mp_cri.axisC) mp_Z_X[k] = sl.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 setcoef(self, Z): """Set coefficient array.""" # This method largely replicates the method from parent class # ConvCnstrMOD_Consensus that it overrides. The inherited # method is overridden to avoid the superfluous computation of # self.ZSf in that method, which is not required for the # modified algorithm with mask decoupling 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)
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.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 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)
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 setdict(self, D=None): """Set dictionary array.""" if D is not None: self.D = np.asarray(D, dtype=self.dtype) self.Df = sl.rfftn(self.D, self.cri.Nv, self.cri.axisN) # Compute D^H S self.DSf = np.conj(self.Df) * self.Sf if self.cri.Cd > 1: self.DSf = np.sum(self.DSf, axis=self.cri.axisC, keepdims=True) if self.opt['HighMemSolve'] and self.cri.Cd == 1: self.c = sl.solvedbi_sm_c(self.Df, np.conj(self.Df), self.rho*self.GHGf + self.rho, self.cri.axisM) else: self.c = None
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 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 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[:] = sl.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[:] = sl.inner(np.conj(mp_Df[np.newaxis, ...]), mp_Sf, axis=mp_cri.axisC + 1)
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 xstep(self): """The xstep of the baseline consensus class from which this class is derived is re-used to implement the xstep of the modified algorithm by replacing ``self.ZSf``, which is constant in the baseline algorithm, with a quantity derived from the additional variables ``self.Y1`` and ``self.U1``. It is also necessary to set the penalty parameter to unity for the duration of the x step. """ self.YU1[:] = self.Y1 - self.U1 self.ZSf = np.conj( self.Zf) * (self.Sf + sl.rfftn(self.YU1, None, self.cri.axisN)) rho = self.rho self.rho = 1.0 super(ConvCnstrMODMaskDcpl_Consensus, self).xstep() self.rho = rho
def setcoef(self, A): """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: A = A.reshape(self.cri.Nv + (1, ) + (self.cri.Cx * self.cri.K, ) + (self.cri.M, )) self.A = np.asarray(A, dtype=self.dtype) self.Af = sl.rfftn(self.A, self.cri.Nv, self.cri.axisN) # Compute X^H S self.ASf = np.sum(np.conj(self.Af) * self.Sf, self.cri.axisK, keepdims=True)
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 eval_gradf(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 WRy = (self.W**2) * Ry # Map back to frequency domain WRyf = sl.rfftn(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 xstep(self): r"""Minimise Augmented Lagrangian with respect to :math:`\mathbf{x}`.""" self.YU[:] = self.Y - self.U Zf = sl.rfftn(self.YU, None, self.cri.axisN) ZfQ = sl.dot(self.Q.T, Zf, axis=self.cri.axisC) b = self.DSfBQ + self.rho * ZfQ Xh = sl.solvedbi_sm(self.gDf, self.rho, b, self.c, axis=self.cri.axisM) self.Xf[:] = sl.dot(self.Q, Xh, axis=self.cri.axisC) self.X = sl.irfftn(self.Xf, self.cri.Nv, self.cri.axisN) if self.opt['LinSolveCheck']: DDXf = np.conj(self.Df) * sl.inner( self.Df, self.Xf, axis=self.cri.axisM) DDXfBB = sl.dot(self.B.T.dot(self.B), DDXf, axis=self.cri.axisC) ax = DDXfBB + self.rho * self.Xf b = sl.dot(self.B.T, self.DSf, axis=self.cri.axisC) + \ self.rho * Zf self.xrrs = sl.rrs(ax, b) else: self.xrrs = None