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 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 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 = 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 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 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 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[:] = spl.rfftn(mp_D_Y0, mp_cri.Nv, mp_cri.axisN)
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] = spl.rfftn(mp_Z_Y1[k], mp_cri.Nv, mp_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 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 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 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 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 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_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 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 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 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
: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 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 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 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 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 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 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)
