class MaskedAffineFct(Functional):
""" F(x) = sum(c[mask,:]*x[mask,:]) + \delta_{x[not(mask),:] == 0} """
def __init__(self, mask, c, conj=None):
Functional.__init__(self)
self.x = Variable(c.shape)
self.mask = mask.astype(bool)
self.nmask = ~self.mask
self.c = c
if conj is None:
from opymize.functionals import MaskedIndicatorFct
self.conj = MaskedIndicatorFct(mask, c, conj=self)
else:
self.conj = conj
scale = self.x.vars(self.x.new())[0]
scale[self.mask, :] = 1.0
self._prox = ShiftScaleOp(self.x.size, self.c.ravel(), scale.ravel(),
-1)
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
val = np.einsum('ik,ik->', x[self.mask, :], self.c[self.mask, :])
infeas = 0.0 if np.all(self.mask) else norm(x[self.nmask, :],
ord=np.inf)
result = (val, infeas)
if grad:
dF = self.x.new()
dF[self.mask, :] = self.c[self.mask, :]
result = (result, dF.ravel)
return result
def prox(self, tau):
self._prox.b = -tau
if hasattr(self._prox, 'gpuvars'):
self._prox.gpuvars['b'][:] = np.atleast_1d(self._prox.b)
return self._prox
class NegProj(Operator):
""" T(x) = min(0, x), in the elementwise sense"""
def __init__(self, N):
Operator.__init__(self)
self.x = Variable(N)
self.y = Variable(N)
self._jacobian = NihilOp(N, keep=self.x.new(dtype=bool))
def prepare_gpu(self, type_t="double"):
self._kernel = ElementwiseKernel("%s *x, %s *y" % ((type_t, ) * 2),
"y[i] = (x[i] > 0) ? 0 : x[i]")
self._kernel_add = ElementwiseKernel("%s *x, %s *y" % ((type_t, ) * 2),
"y[i] += (x[i] > 0) ? 0 : x[i]")
def _call_gpu(self, x, y=None, add=False, jacobian=False):
y = x if y is None else y
if add:
self._kernel_add(x, y)
else:
self._kernel(x, y)
if jacobian:
self._jacobian.keep = (x.get() < 0)
return self._jacobian
def _call_cpu(self, x, y=None, add=False, jacobian=False):
y = x if y is None else y
if jacobian:
self._jacobian.keep = (x < 0)
if add:
y += np.fmin(0, x)
else:
np.fmin(0, x, out=y)
if jacobian: return self._jacobian
class ConstrainFct(Functional):
""" F(x) = 0 if x[mask,:]==c[mask,:] else infty
The mask is only applied to the first component of x
"""
def __init__(self, mask, c, conj=None):
Functional.__init__(self)
self.x = Variable(c.shape)
self.mask = mask
self.c = c
if conj is None:
from opymize.functionals import MaskedAffineFct
self.conj = MaskedAffineFct(mask, c, conj=self)
else:
self.conj = conj
self._prox = ConstrainOp(mask, c)
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
val = 0
infeas = norm(x[self.mask, :] - self.c[self.mask, :], ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau
return self._prox
class QuadEpiInd(Functional):
""" \sum_i \delta_{f_i(x[i,:-1]) \leq x[i,-1]}
f_i(x) := 0.5*a*|x|^2 + <b[i],x> + c[i]
"""
def __init__(self, N, M, a=1.0, b=None, c=None, conj=None):
Functional.__init__(self)
assert a > 0
self.x = Variable((N, M + 1))
self.a = a
self.b = np.zeros((N, M)) if b is None else b
self.c = np.zeros((N, )) if c is None else c
if conj is None:
da, db, dc = quad_dual_coefficients(self.a, self.b, self.c)
self.conj = QuadEpiSupp(N, M, a=da, b=db, c=dc, conj=self)
else:
self.conj = conj
self._prox = QuadEpiProj(N, M, alph=a, b=b, c=c)
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
fx = (0.5 * self.a * x[:, :-1]**2 +
self.b * x[:, :-1]).sum(axis=1) + self.c
dif = fx - x[:, -1]
val = 0
infeas = np.linalg.norm(np.fmax(0, dif), ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau!
return self._prox
class L1NormsConj(Functional):
""" F(x) = \sum_i \delta_{|x[i,:,:]| \leq lbd}
Supported norms are 'frobenius' and 'spectral'
"""
def __init__(self, N, M, lbd, matrixnorm="frobenius", conj=None):
Functional.__init__(self)
assert matrixnorm in ['frobenius', 'spectral']
self.x = Variable((N,) + M)
self.lbd = lbd
self.matrixnorm = matrixnorm
conjnorm = 'nuclear' if matrixnorm == 'spectral' else 'frobenius'
self.conj = L1Norms(N, M, lbd, conjnorm, conj=self) if conj is None else conj
self._prox = L1NormsProj(N, M, self.lbd, matrixnorm)
self._xnorms = np.zeros((N,), order='C')
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
norms(x, self._xnorms, self.matrixnorm)
val = 0
infeas = norm(np.fmax(0, self._xnorms - self.lbd), ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau!
return self._prox
class IndicatorFct(Functional):
""" F(x) = c2 if x==c1 else infty (use broadcasting in c1 if necessary) """
def __init__(self, N, c1=0, c2=0, conj=None):
Functional.__init__(self)
self.c1 = c1
self.c2 = c2
self.x = Variable(N)
from opymize.functionals import AffineFct
self.conj = AffineFct(N, b=-c2, c=c1,
conj=self) if conj is None else conj
self._prox = ConstOp(N, self.x.new() + c1)
def __call__(self, x, grad=False):
val = self.c2
infeas = norm(x - self.c1, ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau
return self._prox
class L12ProjJacobian(LinOp):
""" Jacobian of L1NormsProj for Frobenius norm """
def __init__(self, N, M, lbd):
# xnorms[i] = 1.0/|xbar[i,:,:]|_2
# exterior = (xbar > lbd)
LinOp.__init__(self)
self.x = Variable((N, M[0] * M[1]))
self.y = self.x
self.lbd = lbd
self.adjoint = self
self.extind = np.zeros(N, dtype=bool)
self.intind = np.zeros(N, dtype=bool)
self.xbar_normed = self.x.vars(self.x.new())[0]
self.lbd_norms = np.zeros(N)
def update(self, xbar, exterior, xnorms):
self.extind[:] = exterior
self.intind[:] = ~self.extind
self.xbar_normed[:] = xbar.reshape(self.xbar_normed.shape)
self.xbar_normed[exterior, :] *= xnorms[exterior, None]
self.lbd_norms[:] = self.lbd * xnorms
def _call_cpu(self, x, y=None, add=False):
x = self.x.vars(x)[0]
if add or y is None:
yy = self.y.vars(self.y.new())[0]
else:
yy = self.y.vars(y)[0]
# xn[i,k] = xbar[i,k]/|xbar[i,:]|
xn = self.xbar_normed
# y[norms <= lbd] = x
yy[self.intind, :] = x[self.intind, :]
# y[norms > lbd] = lbd/|xbar|*(x - <xn,x>*xn)
yy[self.extind, :] = xn[self.extind, :]
yy[self.extind, :] *= -np.einsum('ik,ik->i', xn[self.extind, :],
x[self.extind, :])[:, None]
yy[self.extind, :] += x[self.extind, :]
yy[self.extind, :] *= self.lbd_norms[self.extind, None]
if y is None:
if add: x += yy
else: x[:] = yy
elif add:
y += yy
class EpigraphInd(Functional):
""" F(x) = 0 if x[j,i] \in epi(f[i]_j*) for every j,i else infty
More precisely:
F(x) = 0 if <v[k],x[j,i,:-1]> - b[i,k] <= x[j,i,-1]
for any i,j,k with I[i,J[j]][k] == True
"""
def __init__(self, I, If, J, v, b, conj=None):
"""
Args:
I : ndarray of bools, shape (nfuns, npoints)
If : nfuns lists of nregions arrays, shape (nfaces,ndim+1) each
J : ndarray of ints, shape (nregions, nsubpoints)
v : ndarray of floats, shape (npoints, ndim)
b : ndarray of floats, shape (nfuns, npoints)
"""
Functional.__init__(self)
nfuns, npoints = I.shape
nregions, nsubpoints = J.shape
ndim = v.shape[1]
self.I, self.J, self.v, self.b = I, J, v, b
self.x = Variable((nregions, nfuns, ndim + 1))
if conj is None:
self.conj = EpigraphSupp(I, If, J, v, b, conj=self)
else:
self.conj = conj
self.A = self.conj.A
self.b = self.conj.b
self._prox = EpigraphProj(I, J, v, b, Ab=(self.A, self.b))
def __call__(self, x, grad=False):
val = 0
infeas = max(0, np.amax(self.A.dot(x) - self.b))
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau
return self._prox
class NegativityFct(Functional):
""" F(x) = 0 if x <= 0 else infty """
def __init__(self, N, conj=None):
Functional.__init__(self)
self.x = Variable(N)
self.conj = PositivityFct(N, conj=self) if conj is None else conj
self._prox = NegProj(N)
def __call__(self, x, grad=False):
val = 0
infeas = norm(np.fmax(0, x), ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau
return self._prox
class TruncQuadEpiInd(Functional):
""" \sum_i \delta_{|x| \leq lbd} + \delta_{f(x[i,:-1]) \leq x[i,-1]}
f(x) := 0.5*alph*|x|^2
"""
def __init__(self, N, M, lbd=1.0, alph=1.0, conj=None):
Functional.__init__(self)
assert lbd > 0
assert alph > 0
self.x = Variable((N, M + 1))
self.lbd = lbd
self.alph = alph
if conj is None:
dlbd, dalph = self.lbd, self.alph * self.lbd
self.conj = HuberPerspective(N, M, lbd=dlbd, alph=dalph, conj=self)
else:
self.conj = conj
self._prox = QuadEpiProj(N, M, lbd=self.lbd, alph=self.alph)
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
val = 0
lbd, alph = self.lbd, self.alph
x1norm = np.linalg.norm(x[:, :-1], axis=-1)
dif = x1norm - self.lbd
infeas = np.linalg.norm(np.fmax(0, dif), ord=np.inf)
fx = 0.5 * alph * np.fmin(self.lbd, x1norm)**2
dif = fx - x[:, -1]
infeas += np.linalg.norm(np.fmax(0, dif), ord=np.inf)
result = (val, infeas)
if grad:
result = result, self.x.new()
return result
def prox(self, tau):
# independent of tau!
return self._prox
class SemismoothNewtonSystem(LinOp):
""" Block matrix of the following form:
[[ I - K, tau*K*A^T ],
[ -sigma*H*A, I - H ]]
"""
def __init__(self, A, tau, sigma):
LinOp.__init__(self)
self.A = A
self.tau = tau
self.sigma = sigma
self.x = Variable((A.x.size,), (A.y.size,))
self.y = self.x
self.xtmp = self.x.new()
self.K = None
self.H = None
self.adjoint = SemismoothNewtonSystemAdjoint(self)
def _call_cpu(self, x, y=None, add=False):
assert y is not None
assert add is False
x1tmp, x2tmp = self.x.vars(self.xtmp)
x1, x2 = self.x.vars(x)
y1, y2 = self.y.vars(y)
self.A.adjoint(x2, x1tmp)
self.K(x1tmp, y1)
y1[:] = x1 + self.tau*y1
self.K(x1, x1tmp)
y1 -= x1tmp
self.A(x1, x2tmp)
self.H(x2tmp, y2)
y2[:] = x2 - self.sigma*y2
self.H(x2, x2tmp)
y2 -= x2tmp
class SemismoothNewton(object):
def __init__(self, g, f, A):
self.g = g
self.f = f
self.linop = A
self.xy = Variable((self.g.x.size,), (self.f.x.size,))
self.itervars = { 'xyk': self.xy.new() }
self.constvars = { 'tau': 1.0, 'sigma': 1.0 }
def obj_primal(self, x, Ax):
obj, infeas = self.g(x)
obj2, infeas2 = self.f(Ax)
return obj + obj2, infeas + infeas2
def obj_dual(self, ATy, y):
obj, infeas = self.g.conj(-ATy)
obj2, infeas2 = self.f.conj(y)
return -obj - obj2, infeas + infeas2
def res(self, xy, xyres, subdiff=False, norm=True):
c = self.constvars
x, y = self.xy.vars(xy)
xgrad, ygrad = self.xy.vars(xyres)
# xgrad = x - Prox[tau*g](x - tau * A^T * y)
self.linop.adjoint(y, xgrad)
xgrad[:] = x - c['tau']*xgrad
K = self.gprox(xgrad, jacobian=subdiff)
xgrad[:] = x - xgrad
# ygrad = y - Prox[sigma*fc](y + sigma * A * x)
self.linop(x, ygrad)
ygrad[:] = y + c['sigma']*ygrad
H = self.fconjprox(ygrad, jacobian=subdiff)
ygrad[:] = y - ygrad
if subdiff:
self.M.K = K
self.M.H = H
if norm:
return 0.5*np.einsum('i,i->', xyres, xyres)
def iteration_step(self, _iter, use_scipy=False):
i = self.itervars
c = self.constvars
# Set up Newton system
res_normsq = self.res(i['xyk'], i['xyres'], subdiff=True)
# Solve modified Newton system
logging.info('Start linsolve')
if use_scipy:
scipy_solve(self.M, i['xyres'], i['dk'], lbd=self.lbd,
tmpvars=(i['cg_rk'], i['cg_rkp1'], i['xytmp'], i['cg_dk']),
tol=res_normsq, solver=scipy.sparse.linalg.lsqr)
else:
CG_solve(self.M, i['xyres'], i['dk'], lbd=self.lbd,
tmpvars=(i['cg_rk'], i['cg_rkp1'], i['xytmp'], i['cg_dk']),
tol=res_normsq)
logging.info('Stop linsolve')
# Armijo backtracking
self.M(i['dk'], i['xytmp'])
p_gradf = np.einsum('i,i->', i['xyres'], i['xytmp'])
armijo_fun = lambda xy: self.res(xy, i['xyres'], subdiff=False)
new_normsq, alpha = armijo(armijo_fun, i['xyk'], i['xykp1'], i['dk'], \
p_gradf, res_normsq)
# Update, taken from:
# "Distributed Newton Methods for Deep Neural Networks"
# by C.-C. Wang et al. (arxiv: https://arxiv.org/abs/1802.00130)
Md_normsq = 0.5*np.einsum('i,i->', i['xytmp'], i['xytmp'])
rho = (new_normsq - res_normsq)/(alpha*p_gradf + alpha**2*Md_normsq)
if rho > 0.75:
self.lbd *= self.lbd_drop
elif rho < 0.25:
self.lbd *= self.lbd_boost
logging.debug("#{:6d}: alpha = {: 9.6g}, "\
"rho = {: 9.6g}, lbd = {: 9.6g}, normsq = {: 9.6g}"\
.format(_iter, alpha, rho, self.lbd, res_normsq))
i['xyk'][:] = i['xykp1']
def prepare_stepsizes(self):
i = self.itervars
c = self.constvars
step_factor = 1.0
logging.info("Estimating optimal step bound...")
op_norm, itn = normest(self.linop)
# round (floor) to 3 significant digits
bnd = truncate(1.0/op_norm**2, 3) # < 1/|K|^2
bnd *= 1.25 # boost
fact = step_factor # tau/sigma
c['sigma'] = np.sqrt(bnd/fact)
c['tau'] = bnd/c['sigma']
logging.info("Constant steps: %f (%f | %f)" % (bnd,c['sigma'],c['tau']))
self.gprox = self.g.prox(c['tau'])
self.fconjprox = self.f.conj.prox(c['sigma'])
self.M = SemismoothNewtonSystem(self.linop, c['tau'], c['sigma'])
self.lbd = 1.0
self.lbd_drop = 0.1
self.lbd_boost = 5.0
def solve(self, continue_at=None, granularity=50,
term_relgap=1e-5, term_infeas=None, term_maxiter=int(5e2)):
i = self.itervars
c = self.constvars
if continue_at is not None:
i['xyk'][:] = continue_at
i['xykp1'] = i['xyk'].copy()
i['xyres'] = i['xyk'].copy()
i['dk'] = i['xyk'].copy()
i['xytmp'] = i['xyk'].copy()
i['cg_dk'] = i['xyk'].copy()
i['cg_rk'] = i['xyk'].copy()
i['cg_rkp1'] = i['xyk'].copy()
xk, yk = self.xy.vars(i['xyk'])
xkp1, ykp1 = self.xy.vars(i['xykp1'])
self.prepare_stepsizes()
if term_infeas is None:
term_infeas = term_relgap
obj_p = obj_d = infeas_p = infeas_d = relgap = 0.
logging.info("Solving (steps<%d)..." % term_maxiter)
with GracefulInterruptHandler() as interrupt_hdl:
_iter = 0
while _iter < term_maxiter:
self.iteration_step(_iter)
_iter += 1
if interrupt_hdl.interrupted or _iter % granularity == 0:
if interrupt_hdl.interrupted:
print("Interrupt (SIGINT) at iter=%d" % _iter)
self.linop(xk, ykp1)
self.linop.adjoint(yk, xkp1)
obj_p, infeas_p = self.obj_primal(xk, ykp1)
obj_d, infeas_d = self.obj_dual(xkp1, yk)
# compute relative primal-dual gap
relgap = (obj_p - obj_d) / max(np.spacing(1), obj_d)
logging.info("#{:6d}: objp = {: 9.6g} ({: 9.6g}), " \
"objd = {: 9.6g} ({: 9.6g}), " \
"gap = {: 9.6g}, " \
"relgap = {: 9.6g} ".format(
_iter, obj_p, infeas_p,
obj_d, infeas_d,
obj_p - obj_d,
relgap
))
if np.abs(relgap) < term_relgap \
and max(infeas_p, infeas_d) < term_infeas:
break
if interrupt_hdl.interrupted:
break
return {
'objp': obj_p,
'objd': obj_d,
'infeasp': infeas_p,
'infeasd': infeas_d,
'relgap': relgap
}
@property
def state(self):
return self.xy.vars(self.itervars['xyk'])
class Model(ModelHARDI_SHM):
name = "sh_l_tvw2"
def __init__(self, *args, gradnorm="frobenius", **kwargs):
ModelHARDI_SHM.__init__(self, *args, **kwargs)
self.gradnorm = 'nuclear' if gradnorm == "spectral" else "frobenius"
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
self.x = Variable(
('u1', (n_image, l_labels)),
('u2', (n_image, l_labels)),
('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold)),
)
self.y = Variable(
('p', (n_image, d_image, l_shm)),
('g', (m_gradients, n_image, d_image, s_manifold)),
('q0', (n_image, )),
('q1', (n_image, l_labels)),
('q2', (n_image, l_labels)),
)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
u1k, u2k, vk, wk = self.x.vars(self.state[0])
u1k[:] = 1.0 / np.einsum('k->', c['b'])
vk[:, 0] = .5 / np.sqrt(np.pi)
c['G'] = np.zeros((m_gradients, s_manifold, l_shm), order='C')
c['G'][:] = sym_shm_sample_grad(c['Y'], self.data.b_sph.v_grad)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
dataterm = SSD(c['f'], vol=c['b'], mask=c['inpaint_nloc'])
self.pdhg_G = SplitSum([
PositivityFct(x['u1']['size']), # \delta_{u1 >= 0}
dataterm, # 0.5*<u2-f,u2-f>
ZeroFct(x['v']['size']), # 0
ZeroFct(x['w']['size']) # 0
])
GradOp = GradientOp(imagedims, l_shm)
GMult = TangledMatrixMultR(n_image * d_image, c['G'][:, :, None, :])
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
YMult = MatrixMultR(n_image, c['Y'], trans=True)
YMMult = MatrixMultR(n_image, c['YM'], trans=True)
m_u = ScaleOp(x['u1']['size'], -1)
m_w = ScaleOp(x['w']['size'], -1)
self.pdhg_linop = BlockOp([
[0, 0, GradOp, GMult], # p = Dv + G'w
[0, 0, 0, m_w], # g = -w
[bMult, 0, 0, 0], # q0 = <b,u1>
[m_u, 0, YMult, 0], # q1 = Yv - u1
[0, m_u, YMMult, 0] # q2 = YMv - u2
])
l1norms = L1Norms(m_gradients * n_image, (d_image, s_manifold),
c['lbd'], self.gradnorm)
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u1']['size']), # \delta_{q1 = 0}
IndicatorFct(x['u1']['size']) # \delta_{q2 = 0}
])
def setup_solver_cvx(self):
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
self.cvx_x = Variable(
('p', (l_shm, d_image, n_image)),
('g', (n_image, m_gradients, s_manifold, d_image)),
('q0', (n_image, )),
('q1', (l_labels, n_image)),
('q2', (l_labels, n_image)),
)
self.cvx_y = Variable(
('u1', (n_image, l_labels)),
('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold)),
('misc', (n_image * m_gradients, )),
)
p, g, q0, q1, q2 = [
cvxVariable(*a['shape']) for a in self.cvx_x.vars()
]
self.cvx_vars = p + sum(g, []) + [q0, q1, q2]
self.cvx_obj = cvx.Maximize(
0.5 * cvx.sum(cvx.diag(c['b']) * cvx.square(c['f'].T)) -
0.5 * cvx.sum(
cvx.diag(1.0 / c['b']) *
cvx.square(q2 + cvx.diag(c['b']) * c['f'].T)) - cvx.sum(q0))
div_op = sparse_div_op(imagedims)
self.cvx_dual = True
self.cvx_constr = []
# u1_constr
for i in range(n_image):
self.cvx_constr.append(c['b'] * q0[i] - q1[:, i] >= 0)
# v_constr
for i in range(n_image):
for k in range(l_shm):
Yk = cvx.vec(c['Y'][:, k])
self.cvx_constr.append(
Yk.T*(c['M'][k]*q2[:,i] + q1[:,i]) \
- cvxOp(div_op, p[k], i) == 0)
# w_constr
for j in range(m_gradients):
Gj = c['G'][j, :, :]
for i in range(n_image):
for t in range(d_image):
for l in range(s_manifold):
self.cvx_constr.append(
g[i][j][l,t] == sum([Gj[l,k]*p[k][t,i] \
for k in range(l_shm)]))
# additional inequality constraints
for i in range(n_image):
for j in range(m_gradients):
self.cvx_constr.append(cvx.norm(g[i][j], 2) <= c['lbd'])
class Model(ModelHARDI):
name = "n_w_tvw"
def __init__(self, *args, dataterm="W1", gradnorm="frobenius", **kwargs):
ModelHARDI.__init__(self, *args, **kwargs)
self.gradnorm = 'nuclear' if gradnorm == "spectral" else "frobenius"
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
xvars = [('u', (n_image, l_labels)),
('w', (m_gradients, n_image, d_image, s_manifold))]
yvars = [('p', (n_image, d_image, l_labels)),
('g', (m_gradients, n_image, d_image, s_manifold)),
('q', (n_image, ))]
self.dataterm = dataterm
if self.dataterm == "W1":
xvars.append(('w0', (m_gradients, n_image, s_manifold)))
yvars.append(('p0', (n_image, l_labels)))
yvars.append(('g0', (m_gradients, n_image, s_manifold)))
elif self.dataterm != "quadratic":
raise Exception("Dataterm unknown: %s" % self.dataterm)
self.x = Variable(*xvars)
self.y = Variable(*yvars)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
x = self.x.vars(self.state[0], named=True)
x['u'][:] = 1.0 / np.einsum('k->', c['b'])
f = self.data.odf
f_flat = f.reshape(-1, l_labels).T
f = np.array(f_flat.reshape((l_labels, ) + imagedims), order='C')
normalize_odf(f, c['b'])
self.f = np.array(f.reshape(l_labels, -1).T, order='C')
logging.info("HARDI setup ({l_labels} labels; " \
"img: {imagedims}; lambda={lbd:.3g}) ready.".format(
lbd=c['lbd'],
l_labels=c['l_labels'],
imagedims="x".join(map(str,c['imagedims']))))
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
GradOp = GradientOp(imagedims, l_labels, weights=c['b'])
PBLinOp = IndexedMultAdj(l_labels, d_image * n_image, c['P'], c['B'])
AMult = MatrixMultRBatched(n_image * d_image, c['A'])
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
l1norms = L1Norms(m_gradients * n_image, (d_image, s_manifold),
c['lbd'], self.gradnorm)
if self.dataterm == "W1":
self.pdhg_G = SplitSum([
PositivityFct(x['u']['size']), # \delta_{u >= 0}
ZeroFct(x['w']['size']), # 0
ZeroFct(x['w0']['size']), # 0
])
PBLinOp0 = IndexedMultAdj(l_labels, n_image, c['P'], c['B'])
AMult0 = MatrixMultRBatched(n_image, c['A'])
bMult0 = DiagMatrixMultR(n_image, c['b'])
self.pdhg_linop = BlockOp([
[GradOp, PBLinOp, 0], # p = diag(b)Du - P'B'w
[0, AMult, 0], # g = A'w
[bMult, 0, 0], # q = <b,u>
[bMult0, 0, PBLinOp0], # p0 = diag(b) u - P'B'w0
[0, 0, AMult0] # g0 = A'w0
])
diag_b_f = np.einsum('ik,k->ik', self.f, c['b'])
dataterm = ConstrainFct(c['inpaint_nloc'], diag_b_f)
l1norms0 = L1Norms(m_gradients * n_image, (1, s_manifold), 1.0,
"frobenius")
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q']['size'],
c1=c['b_precond']), # \delta_{q = 1}
dataterm, # \delta_{p0 = diag(b)f}
l1norms0, # \sum_ji |g0[j,i,:]|_2
])
elif self.dataterm == "quadratic":
dataterm = PosSSD(self.f, vol=c['b'], mask=c['inpaint_nloc'])
self.pdhg_G = SplitSum([
dataterm, # 0.5*<u-f,u-f>_b + \delta_{u >= 0}
ZeroFct(x['w']['size']), # 0
])
self.pdhg_linop = BlockOp([
[GradOp, PBLinOp], # p = diag(b)Du - P'B'w
[0, AMult], # g = A'w
[bMult, 0], # q = <b,u>
])
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q']['size'],
c1=c['b_precond']), # \delta_{q = 1}
])
class Model(ModelHARDI_SHM):
name = "sh_bndl2_tvw"
def __init__(self, *args, gradnorm="frobenius", conf_lvl=0.9, **kwargs):
ModelHARDI_SHM.__init__(self, *args, **kwargs)
self.gradnorm = 'nuclear' if gradnorm == "spectral" else "frobenius"
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
self.data.init_bounds(conf_lvl)
_, f1, f2 = self.data.bounds
c['f1'], c['f2'] = [np.array(a.T, order='C') for a in [f1, f2]]
self.x = Variable(
('u1', (n_image, l_labels)),
('u2', (n_image, l_labels)),
('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold)),
)
self.y = Variable(
('p', (n_image, d_image, l_labels)),
('g', (m_gradients, n_image, d_image, s_manifold)),
('q0', (n_image, )),
('q1', (n_image, l_labels)),
('q2', (n_image, l_labels)),
)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
u1k, u2k, vk, wk = self.x.vars(self.state[0])
u1k[:] = 1.0 / np.einsum('k->', c['b'])
vk[:, 0] = .5 / np.sqrt(np.pi)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
dataterm = BndSSD(c['f1'], c['f2'], vol=c['b'], mask=c['inpaint_nloc'])
self.pdhg_G = SplitSum([
PositivityFct(x['u1']['size']), # \delta_{u1 >= 0}
dataterm, # 0.5*|max(0, f1 - u2)|^2 + 0.5*|max(0, u2 - f2)|^2
ZeroFct(x['v']['size']), # 0
ZeroFct(x['w']['size']) # 0
])
GradOp = GradientOp(imagedims, l_labels, weights=c['b'])
PBLinOp = IndexedMultAdj(l_labels, d_image * n_image, c['P'], c['B'])
AMult = MatrixMultRBatched(n_image * d_image, c['A'])
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
YMult = MatrixMultR(n_image, c['Y'], trans=True)
YMMult = MatrixMultR(n_image, c['YM'], trans=True)
m_u = ScaleOp(x['u1']['size'], -1)
self.pdhg_linop = BlockOp([
[GradOp, 0, 0, PBLinOp], # p = diag(b)Du1 - P'B'w
[0, 0, 0, AMult], # g = A'w
[bMult, 0, 0, 0], # q0 = <b,u1>
[m_u, 0, YMult, 0], # q1 = Yv - u1
[0, m_u, YMMult, 0] # q2 = YMv - u2
])
l1norms = L1Norms(m_gradients * n_image, (d_image, s_manifold),
c['lbd'], self.gradnorm)
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u1']['size']), # \delta_{q1 = 0}
IndicatorFct(x['u1']['size']) # \delta_{q2 = 0}
])
def setup_solver_cvx(self):
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
self.cvx_x = Variable(
('p', (l_labels, d_image, n_image)),
('g', (n_image, m_gradients, s_manifold, d_image)),
('q0', (n_image, )),
('q1', (l_labels, n_image)),
('q2', (l_labels, n_image)),
)
self.cvx_y = Variable(
('u1', (n_image, l_labels)),
('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold)),
('misc', (n_image * m_gradients, )),
)
p, g, q0, q1, q2 = [
cvxVariable(*a['shape']) for a in self.cvx_x.vars()
]
self.cvx_vars = p + sum(g, []) + [q0, q1, q2]
fid_fun_dual = 0
for i in range(n_image):
for k in range(l_labels):
fid_fun_dual += -cvx.power(q2[k,i],2)/2 \
- cvx.maximum(q2[k,i]*c['f1'][i,k],
q2[k,i]*c['f2'][i,k])
self.cvx_obj = cvx.Maximize(fid_fun_dual - cvx.sum(q0))
div_op = sparse_div_op(imagedims)
self.cvx_dual = True
self.cvx_constr = []
# u1_constr
for i in range(n_image):
for k in range(l_labels):
self.cvx_constr.append(
c['b'][k] *
(q0[i] - cvxOp(div_op, p[k], i)) - q1[k, i] >= 0)
# v_constr
for i in range(n_image):
for k in range(l_shm):
Yk = cvx.vec(c['Y'][:, k])
self.cvx_constr.append(Yk.T *
(c['M'][k] * q2[:, i] + q1[:, i]) == 0)
# w_constr
for j in range(m_gradients):
Aj = c['A'][j, :, :]
Bj = c['B'][j, :, :]
Pj = c['P'][j, :]
for i in range(n_image):
for t in range(d_image):
for l in range(s_manifold):
self.cvx_constr.append(Aj * g[i][j][l, t] == sum(
[Bj[l, m] * p[Pj[m]][t, i] for m in range(3)]))
# additional inequality constraints
for i in range(n_image):
for j in range(m_gradients):
self.cvx_constr.append(cvx.norm(g[i][j], 2) <= c['lbd'])
class Model(ModelHARDI_SHM):
name = "sh_bndl1_tvc"
def __init__(self, *args, conf_lvl=0.9, **kwargs):
ModelHARDI_SHM.__init__(self, *args, **kwargs)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
self.data.init_bounds(conf_lvl)
_, f1, f2 = self.data.bounds
c['f1'], c['f2'] = [np.array(a.T, order='C') for a in [f1, f2]]
self.x = Variable(
('u1', (n_image, l_labels)),
('u2', (n_image, l_labels)),
('v', (n_image, l_shm)),
)
self.y = Variable(
('p', (n_image, d_image, l_shm)),
('q0', (n_image, )),
('q1', (n_image, l_labels)),
('q2', (n_image, l_labels)),
('q3', (n_image, l_labels)),
('q4', (n_image, l_labels)),
)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
u1k, u2k, vk = self.x.vars(self.state[0])
u1k[:] = 1.0 / np.einsum('k->', c['b'])
vk[:, 0] = .5 / np.sqrt(np.pi)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
self.pdhg_G = SplitSum([
PositivityFct(x['u1']['size']), # \delta_{u1 >= 0}
ZeroFct(x['u1']['size']), # 0
ZeroFct(x['v']['size']) # 0
])
GradOp = GradientOp(imagedims, l_shm)
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
YMult = MatrixMultR(n_image, c['Y'], trans=True)
YMMult = MatrixMultR(n_image, c['YM'], trans=True)
m_u = ScaleOp(x['u1']['size'], -1)
dbMult = DiagMatrixMultR(n_image, c['b'])
mdbMult = DiagMatrixMultR(n_image, -c['b'])
self.pdhg_linop = BlockOp([
[0, 0, GradOp], # p = Dv
[bMult, 0, 0], # q0 = <b,u1>
[m_u, 0, YMult], # q1 = Yv - u1
[0, m_u, YMMult], # q2 = YMv - u2
[0, mdbMult, 0], # q3 = -diag(b) u2
[0, dbMult, 0] # q4 = diag(b) u2
])
l1norms = L1Norms(n_image, (d_image, l_shm), c['lbd'], "frobenius")
LowerBoundFct = MaxFct(np.einsum('ik,k->ik', c['f1'], -c['b']))
UpperBoundFct = MaxFct(np.einsum('ik,k->ik', c['f2'], c['b']))
self.pdhg_F = SplitSum([
l1norms, # lbd*\sum_i |p[i,:,:]|_2
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u1']['size']), # \delta_{q1 = 0}
IndicatorFct(x['u1']['size']), # \delta_{q2 = 0}
LowerBoundFct, # |max(0, q3 + diag(b)f1)|_1
UpperBoundFct # |max(0, q4 - diag(b)f2)|_1
])
def setup_solver_cvx(self):
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
self.cvx_x = Variable(
('p', (l_shm, d_image, n_image)),
('q0', (n_image, )),
('q1', (l_labels, n_image)),
('q2', (l_labels, n_image)),
('q3', (l_labels, n_image)),
('q4', (l_labels, n_image)),
)
self.cvx_y = Variable(
('u1', (n_image, l_labels)),
('v', (n_image, l_shm)),
('misc', (n_image * l_labels * 5 + n_image, )),
)
p, q0, q1, q2, q3, q4 = [
cvxVariable(*a['shape']) for a in self.cvx_x.vars()
]
self.cvx_vars = p + [q0, q1, q2, q3, q4]
self.cvx_obj = cvx.Maximize(
cvx.vec(q3).T * cvx.vec(cvx.diag(c['b']) * c['f1'].T) -
cvx.vec(q4).T * cvx.vec(cvx.diag(c['b']) * c['f2'].T) -
cvx.sum(q0))
div_op = sparse_div_op(imagedims)
self.cvx_dual = True
self.cvx_constr = []
# u1_constr
for i in range(n_image):
self.cvx_constr.append(c['b'][:] * q0[i] - q1[:, i] >= 0)
# v_constr
for i in range(n_image):
for k in range(l_shm):
Yk = cvx.vec(c['Y'][:, k])
self.cvx_constr.append(
Yk.T*(c['M'][k]*q2[:,i] + q1[:,i]) \
- cvxOp(div_op, p[k], i) == 0)
# additional inequality constraints
for i in range(n_image):
for k in range(l_labels):
self.cvx_constr.append(0 <= q3[k, i])
self.cvx_constr.append(q3[k, i] <= 1)
self.cvx_constr.append(0 <= q4[k, i])
self.cvx_constr.append(q4[k, i] <= 1)
self.cvx_constr.append(q4[k, i] - q3[k, i] - q2[k, i] == 0)
for i in range(n_image):
self.cvx_constr.append(sum(cvx.sum_squares(p[k][:,i]) \
for k in range(l_shm)) <= c['lbd']**2)
class Model(ModelHARDI_SHM):
name = "sh_w_tvw"
def __init__(self, *args, dataterm="W1", gradnorm="frobenius", **kwargs):
ModelHARDI_SHM.__init__(self, *args, **kwargs)
self.gradnorm = 'nuclear' if gradnorm == "spectral" else "frobenius"
c = self.constvars
b_sph = self.data.b_sph
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
xvars = [('u', (n_image, l_labels)), ('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold))]
yvars = [('p', (n_image, d_image, l_labels)),
('g', (m_gradients, n_image, d_image, s_manifold)),
('q0', (n_image, )), ('q1', (n_image, l_labels))]
self.dataterm = dataterm
if self.dataterm == "W1":
xvars.append(('w0', (m_gradients, n_image, s_manifold)))
yvars.append(('p0', (n_image, l_labels)))
yvars.append(('g0', (m_gradients, n_image, s_manifold)))
elif self.dataterm != "quadratic":
raise Exception("Dataterm unknown: %s" % self.dataterm)
self.x = Variable(*xvars)
self.y = Variable(*yvars)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
x = self.x.vars(self.state[0], named=True)
x['u'][:] = 1.0 / np.einsum('k->', c['b'])
x['v'][:, 0] = .5 / np.sqrt(np.pi)
f = self.data.odf
f_flat = f.reshape(-1, l_labels).T
f = np.array(f_flat.reshape((l_labels, ) + imagedims), order='C')
normalize_odf(f, c['b'])
self.f = np.array(f.reshape(l_labels, -1).T, order='C')
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
GradOp = GradientOp(imagedims, l_labels, weights=c['b'])
PBLinOp = IndexedMultAdj(l_labels, d_image * n_image, c['P'], c['B'])
AMult = MatrixMultRBatched(n_image * d_image, c['A'])
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
YMult = MatrixMultR(n_image, c['Y'], trans=True)
m_u = ScaleOp(x['u']['size'], -1)
l1norms = L1Norms(m_gradients * n_image, (d_image, s_manifold),
c['lbd'], self.gradnorm)
if self.dataterm == "W1":
self.pdhg_G = SplitSum([
PositivityFct(x['u']['size']), # \delta_{u >= 0}
ZeroFct(x['v']['size']), # 0
ZeroFct(x['w']['size']), # 0
ZeroFct(x['w0']['size']), # 0
])
PBLinOp0 = IndexedMultAdj(l_labels, n_image, c['P'], c['B'])
AMult0 = MatrixMultRBatched(n_image, c['A'])
bMult0 = DiagMatrixMultR(n_image, c['b'])
self.pdhg_linop = BlockOp([
[GradOp, 0, PBLinOp, 0], # p = diag(b)Du - P'B'w
[0, 0, AMult, 0], # g = A'w
[bMult, 0, 0, 0], # q0 = <b,u>
[m_u, YMult, 0, 0], # q1 = Yv - u
[bMult0, 0, 0, PBLinOp0], # p0 = diag(b) u - P'B'w0
[0, 0, 0, AMult0] # g0 = A'w0
])
diag_b_f = np.einsum('ik,k->ik', self.f, c['b'])
dataterm = ConstrainFct(c['inpaint_nloc'], diag_b_f)
l1norms0 = L1Norms(m_gradients * n_image, (1, s_manifold), 1.0,
"frobenius")
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u']['size']), # \delta_{q1 = 0}
dataterm, # \delta_{p0 = diag(b)f}
l1norms0, # \sum_ji |g0[j,i,:]|_2
])
elif self.dataterm == "quadratic":
dataterm = PosSSD(self.f, vol=c['b'], mask=c['inpaint_nloc'])
self.pdhg_G = SplitSum([
dataterm, # 0.5*<u-f,u-f>_b + \delta_{u >= 0}
ZeroFct(x['v']['size']), # 0
ZeroFct(x['w']['size']), # 0
])
self.pdhg_linop = BlockOp([
[GradOp, 0, PBLinOp], # p = diag(b)Du - P'B'w
[0, 0, AMult], # g = A'w
[bMult, 0, 0], # q0 = <b,u>
[m_u, YMult, 0], # q1 = Yv - u
])
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u']['size']), # \delta_{q1 = 0}
])
def setup_solver_cvx(self):
if self.dataterm != "W1":
raise Exception("Only W1 dataterm is implemented in CVX")
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
m_gradients = c['m_gradients']
s_manifold = c['s_manifold']
l_shm = c['l_shm']
f_flat = self.data.odf.reshape(-1, l_labels).T
f = np.array(f_flat.reshape((l_labels, ) + imagedims), order='C')
normalize_odf(f, c['b'])
f_flat = f.reshape(l_labels, n_image)
self.cvx_x = Variable(
('p', (l_labels, d_image, n_image)),
('g', (n_image, m_gradients, s_manifold, d_image)),
('q0', (n_image, )),
('q1', (l_labels, n_image)),
('p0', (l_labels, n_image)),
('g0', (n_image, m_gradients, s_manifold)),
)
self.cvx_y = Variable(
('u', (n_image, l_labels)),
('v', (n_image, l_shm)),
('w', (m_gradients, n_image, d_image, s_manifold)),
('w0', (m_gradients, n_image, s_manifold)),
('misc', (n_image * m_gradients, )),
)
p, g, q0, q1, p0, g0 = [
cvxVariable(*a['shape']) for a in self.cvx_x.vars()
]
self.cvx_vars = p + sum(g, []) + [q0, q1, p0] + g0
self.cvx_obj = cvx.Maximize(-cvx.vec(f_flat).T *
cvx.vec(cvx.diag(c['b']) * p0) -
cvx.sum(q0))
div_op = sparse_div_op(imagedims)
self.cvx_dual = True
self.cvx_constr = []
# u_constr
for i in range(n_image):
for k in range(l_labels):
self.cvx_constr.append(
c['b'][k]*(q0[i] + p0[k,i] \
- cvxOp(div_op, p[k], i)) - q1[k,i] >= 0)
# v_constr
for i in range(n_image):
for k in range(l_shm):
Yk = cvx.vec(c['Y'][:, k])
self.cvx_constr.append(-Yk.T * q1[:, i] == 0)
# w_constr
for j in range(m_gradients):
Aj = c['A'][j, :, :]
Bj = c['B'][j, :, :]
Pj = c['P'][j, :]
for i in range(n_image):
for t in range(d_image):
for l in range(s_manifold):
self.cvx_constr.append(
Aj*g[i][j][l,t] == sum([Bj[l,m]*p[Pj[m]][t,i] \
for m in range(3)]))
# w0_constr
for j in range(m_gradients):
Aj = c['A'][j, :, :]
Bj = c['B'][j, :, :]
Pj = c['P'][j, :]
for i in range(n_image):
self.cvx_constr.append(Aj * g0[i][j, :].T == Bj * p0[Pj, i])
# additional inequality constraints
for i in range(n_image):
for j in range(m_gradients):
self.cvx_constr.append(cvx.norm(g[i][j], 2) <= c['lbd'])
self.cvx_constr.append(cvx.norm(g0[i][j, :], 2) <= 1.0)
class Model(ModelHARDI_SHM):
name = "sh_bndl2_tvc"
def __init__(self, *args, conf_lvl=0.9, **kwargs):
ModelHARDI_SHM.__init__(self, *args, **kwargs)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
self.data.init_bounds(conf_lvl)
_, f1, f2 = self.data.bounds
c['f1'], c['f2'] = [np.array(a.T, order='C') for a in [f1, f2]]
self.x = Variable(
('u1', (n_image, l_labels)),
('u2', (n_image, l_labels)),
('v', (n_image, l_shm)),
)
self.y = Variable(
('p', (n_image, d_image, l_shm)),
('q0', (n_image, )),
('q1', (n_image, l_labels)),
('q2', (n_image, l_labels)),
)
# start with a uniform distribution in each voxel
self.state = (self.x.new(), self.y.new())
u1k, u2k, vk = self.x.vars(self.state[0])
u1k[:] = 1.0 / np.einsum('k->', c['b'])
vk[:, 0] = .5 / np.sqrt(np.pi)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
dataterm = BndSSD(c['f1'], c['f2'], vol=c['b'], mask=c['inpaint_nloc'])
self.pdhg_G = SplitSum([
PositivityFct(x['u1']['size']), # \delta_{u1 >= 0}
dataterm, # 0.5*|max(0, f1 - u2)|^2 + 0.5*|max(0, u2 - f2)|^2
ZeroFct(x['v']['size']) # 0
])
GradOp = GradientOp(imagedims, l_shm)
bMult = MatrixMultR(n_image, c['b_precond'] * c['b'][:, None])
YMult = MatrixMultR(n_image, c['Y'], trans=True)
YMMult = MatrixMultR(n_image, c['YM'], trans=True)
m_u = ScaleOp(x['u1']['size'], -1)
self.pdhg_linop = BlockOp([
[0, 0, GradOp], # p = Dv
[bMult, 0, 0], # q0 = <b,u1>
[m_u, 0, YMult], # q1 = Yv - u1
[0, m_u, YMMult] # q2 = YMv - u2
])
l1norms = L1Norms(n_image, (d_image, l_shm), c['lbd'], "frobenius")
self.pdhg_F = SplitSum([
l1norms, # lbd*\sum_i |p[i,:,:]|_2
IndicatorFct(y['q0']['size'],
c1=c['b_precond']), # \delta_{q0 = 1}
IndicatorFct(x['u1']['size']), # \delta_{q1 = 0}
IndicatorFct(x['u1']['size']) # \delta_{q2 = 0}
])
def setup_solver_cvx(self):
c = self.constvars
imagedims = c['imagedims']
n_image = c['n_image']
d_image = c['d_image']
l_labels = c['l_labels']
l_shm = c['l_shm']
self.cvx_x = Variable(
('p', (l_shm, d_image, n_image)),
('q0', (n_image, )),
('q1', (l_labels, n_image)),
('q2', (l_labels, n_image)),
)
self.cvx_y = Variable(
('u1', (n_image, l_labels)),
('v', (n_image, l_shm)),
('misc', (n_image, )),
)
p, q0, q1, q2 = [cvxVariable(*a['shape']) for a in self.cvx_x.vars()]
self.cvx_vars = p + [q0, q1, q2]
fid_fun_dual = 0
for i in range(n_image):
for k in range(l_labels):
fid_fun_dual += -1.0/c['b'][k]*(cvx.power(q2[k,i],2)/2 \
+ cvx.maximum(q2[k,i]*c['b'][k]*c['f1'][i,k],
q2[k,i]*c['b'][k]*c['f2'][i,k]))
self.cvx_obj = cvx.Maximize(fid_fun_dual - cvx.sum(q0))
div_op = sparse_div_op(imagedims)
self.cvx_dual = True
self.cvx_constr = []
# u1_constr
for i in range(n_image):
self.cvx_constr.append(c['b'] * q0[i] - q1[:, i] >= 0)
# v_constr
for i in range(n_image):
for k in range(l_shm):
Yk = cvx.vec(c['Y'][:, k])
self.cvx_constr.append(
Yk.T*(c['M'][k]*q2[:,i] + q1[:,i]) \
- cvxOp(div_op, p[k], i) == 0)
# additional inequality constraints
for i in range(n_image):
self.cvx_constr.append(sum(cvx.sum_squares(p[k][:,i]) \
for k in range(l_shm)) <= c['lbd']**2)
