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 SplitSum(Functional):
""" F(x1,x2,...) = F1(x1) + F2(x2) + ... """
def __init__(self, fcts, conj=None):
Functional.__init__(self)
self.x = Variable(*[(F.x.size, ) for F in fcts])
self.fcts = fcts
self.conj = SplitSum([F.conj for F in fcts],
conj=self) if conj is None else conj
def __call__(self, x, grad=False):
X = self.x.vars(x)
results = [F(xi, grad=grad) for F, xi in zip(self.fcts, X)]
if grad:
val = sum([res[0][0] for res in results])
infeas = sum([res[0][1] for res in results])
dF = np.concatenate([res[1] for res in results])
return (val, infeas), dF
else:
val = sum([res[0] for res in results])
infeas = max([res[1] for res in results])
return (val, infeas)
def prox(self, tau):
if type(tau) is np.ndarray:
tau = self.x.vars(tau)
prox_ops = []
for F, Ftau in zip(self.fcts, tau):
prox_ops.append(F.prox(Ftau))
return SplitOp(prox_ops)
else:
return SplitOp([F.prox(tau) for F in self.fcts])
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 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 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 L1Norms(Functional):
""" F(x) = lbd * \sum_i |x[i,:,:]|
Supported norms are 'frobenius' and 'nuclear'
"""
def __init__(self, N, M, lbd, matrixnorm="frobenius", conj=None):
Functional.__init__(self)
assert matrixnorm in ['frobenius', 'nuclear']
self.x = Variable((N,) + M)
self.lbd = lbd
self.matrixnorm = matrixnorm
conjnorm = 'spectral' if matrixnorm == 'nuclear' else 'frobenius'
self.conj = L1NormsConj(N, M, lbd, conjnorm, conj=self) if conj is None else conj
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 = self.lbd*self._xnorms.sum()
infeas = 0
result = (val, infeas)
if grad:
assert self.matrixnorm == 'frobenius'
""" dF = 0 if x == 0 else x/|x| """
dF = x.copy()
where0 = (self._xnorms == 0)
wheren0 = ~where0
dF[where0,:,:] = 0
dF[wheren0,:,:] /= self._xnorms[wheren0]
result = result, dF.ravel()
return result
class MaxFct(Functional):
""" \sum_ik b[k]*max(0, x[i,k] - f[i,k]) """
def __init__(self, data, vol=None, mask=None, conj=None):
Functional.__init__(self)
self.f = np.atleast_2d(data)
self.x = Variable(self.f.shape)
self.vol = np.ones(data.shape[1]) if vol is None else vol
self.mask = np.ones(data.shape[0],
dtype=bool) if mask is None else mask
if conj is None:
self.conj = MaxFctConj(data, weights=vol, mask=mask, conj=self)
else:
self.conj = conj
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
posval = np.fmax(0, (x - self.f)[self.mask, :])
val = np.einsum('ik,k->', posval, self.vol)
infeas = 0
result = (val, infeas)
if grad:
df = np.zeros_like(x)
posval[posval > 0] = 1.0
df[self.mask, :] = np.einsum('ik,k->', posval, self.vol)
return result, df
else:
return result
class MaxFctConj(Functional):
""" sum_i <x[i,:],f[i,:]> + \delta_{0 <= b[k]*x[i,k] <= 1} """
def __init__(self, data, weights=None, mask=None, conj=None):
Functional.__init__(self)
self.f = np.atleast_2d(data)
self.x = Variable(self.f.shape)
self.weights = np.ones(data.shape[1]) if weights is None else weights
self.mask = np.ones(data.shape[0],
dtype=bool) if mask is None else mask
if conj is None:
self.conj = MaxFct(data, vol=vol, mask=mask, conj=self)
else:
self.conj = conj
self._prox = IntervProj(self.weights, self.f, mask=mask)
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
val = np.einsum('ik,ik->', x[self.mask, :], self.f[self.mask, :])
infeas = norm(np.fmin(0, x[self.mask, :]), ord=np.inf)
bx = np.einsum('ik,k->ik', x[self.mask, :], self.weights)
infeas += norm(np.fmin(0, 1.0 - bx), ord=np.inf)
result = (val, infeas)
if grad:
df = np.zeros_like(x)
df[self.mask, :] = self.f[self.mask, :]
result = result, df
return result
def prox(self, tau):
self._prox.a = tau
return self._prox
class HuberPerspective(Functional):
""" \sum_i -x[i,-1]*f(-x[i,:-1]/x[i,-1]) if x[i,-1] < 0
\sum_i lbd*|x[i,:-1]| if x[i,-1] == 0
+inf if x[i,-1] > 0
f(x) := | lbd*(0.5/alph*|x|^2), if |x| < alph,
| lbd*(|x| - alph/2), if |x| > alph.
"""
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 = TruncQuadEpiInd(N, M, lbd=dlbd, alph=dalph, conj=self)
else:
self.conj = conj
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
lbd, alph = self.lbd, self.alph
x1, x2 = x[:, :-1], -x[:, -1]
x1norm = np.linalg.norm(x1, axis=-1)
qmsk = (x1norm < alph * x2)
lmsk = (~qmsk) & (x2 > -1e-4) # relaxed positivity condition
val = 0.5 / alph * (x1norm[qmsk]**2 / x2[qmsk]).sum()
val += (x1norm[lmsk] - x2[lmsk] * alph / 2).sum()
val *= lbd
infeas = np.linalg.norm(np.fmin(0, x2), ord=np.inf)
return (val, infeas)
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 QuadEpiSupp(Functional):
""" \sum_i -x[i,-1]*f_i(-x[i,:-1]/x[i,-1]) if x[i,-1] < 0
and inf if x[i,-1] >= 0
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 = QuadEpiInd(N, M, a=da, b=db, c=dc, conj=self)
else:
self.conj = conj
def __call__(self, x, grad=False):
assert not grad
a, b, c = self.a, self.b, self.c
x = self.x.vars(x)[0]
msk = x[:, -1] < -1e-8
x1, x2 = x[msk, :-1], -x[msk, -1]
val = (0.5 * a * x1**2 / x2[:, None] +
b[msk] * x1).sum(axis=1) + x2 * c[msk]
val = val.sum()
if np.all(msk):
infeas = 0
else:
infeas = np.linalg.norm(x[~msk, -1], ord=np.inf)
return (val, infeas)
class BndSSDConj(Functional):
""" 0.5*<x,x>_b + sum[max(f1*x,f2*x)]
same as 0.5*<x,x>_b + sum[(a + sign(x)*b)*x]
where a = (f1+f2)/2 and b = (f2-f1)/2
"""
def __init__(self, f1, f2, vol=None, mask=None, conj=None):
Functional.__init__(self)
self.x = Variable(f1.shape)
self.a = 0.5 * (f1 + f2)
self.b = 0.5 * (f2 - f1)
self.vol = np.ones(f1.shape[1]) if vol is None else vol
self.mask = np.ones(f1.shape[0], dtype=bool) if mask is None else mask
if conj is None:
cj_vol = 1.0 / self.vol
self.conj = BndSSD(f1, f2, vol=cj_vol, mask=mask, conj=self)
else:
self.conj = conj
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
x_msk = x[self.mask, :]
a_msk, b_msk = self.a[self.mask, :], self.b[self.mask, :]
val = 0.5*np.einsum('ik,k->', x_msk**2, self.vol) \
+ np.einsum('ik,ik->', a_msk + np.sign(x_msk)*b_msk, x_msk)
infeas = 0
return (val, infeas)
class SSD(Functional):
""" 0.5*<x-f, x-f>_b + shift
where b is the volume element
"""
def __init__(self, data, vol=None, shift=0, mask=None, conj=None):
Functional.__init__(self)
self.x = Variable(data.shape)
self.f = np.atleast_2d(data)
self.shift = shift
self.vol = np.ones(data.shape[1]) if vol is None else vol
self.mask = np.ones(data.shape[0],
dtype=bool) if mask is None else mask
if conj is None:
cj_vol = 1.0 / self.vol
cj_data = np.zeros_like(self.f)
cj_data[self.mask, :] = np.einsum('ik,k->ik', self.f[self.mask, :],
-self.vol)
cj_shift = -0.5 * np.einsum('ik,k->', cj_data**2, cj_vol)
cj_shift -= self.shift
self.conj = SSD(cj_data,
shift=cj_shift,
vol=cj_vol,
mask=mask,
conj=self)
else:
self.conj = conj
prox_shift = np.zeros_like(self.f)
prox_shift[self.mask, :] = self.f[self.mask, :]
prox_shift = prox_shift.ravel()
self._prox = ShiftScaleOp(self.x.size, prox_shift, 0.5, 1.0)
def __call__(self, x, grad=False):
x = self.x.vars(x)[0]
val = 0.5 * np.einsum('ik,k->',
(x - self.f)[self.mask, :]**2, self.vol)
val += self.shift
infeas = 0
result = (val, infeas)
if grad:
df = np.zeros_like(x)
df[self.mask, :] = np.einsum('ik,k->ik',
(x - self.f)[self.mask, :], self.vol)
return result, df.ravel()
else:
return result
def prox(self, tau):
msk = self.mask
tauvol = np.zeros_like(self.f)
tauvol[msk, :] = (tau * np.ones(self.f.size)).reshape(
self.f.shape)[msk, :]
tauvol[msk, :] = np.einsum('ik,k->ik', tauvol[msk, :], self.vol)
self._prox.a = 1.0 / (1.0 + tauvol.ravel())
self._prox.b = tauvol.ravel()
if hasattr(self._prox, 'gpuvars'):
self._prox.gpuvars['a'][:] = self._prox.a
self._prox.gpuvars['b'][:] = self._prox.b
return self._prox
class ProxBndSSD(Operator):
""" y = max(x + tau*f1, min(x + tau*f2, alpha*x))/alpha """
def __init__(self, f1, f2, alpha=2.0, tau=1.0):
Operator.__init__(self)
self.x = Variable(f1.shape)
self.y = Variable(f1.shape)
self.f1 = f1
self.f2 = f2
self.alpha = alpha
self.tau = tau
def prepare_gpu(self, type_t="double"):
np_dtype = np.float64 if type_t == "double" else np.float32
taufact = "tau[i]*" if type(self.tau) is np.ndarray else "tau[0]*"
self.gpuvars = {
'f1': gpuarray.to_gpu(self.f1),
'f2': gpuarray.to_gpu(self.f2),
'alpha': gpuarray.to_gpu(self.alpha),
'tau': gpuarray.to_gpu(np.asarray(self.tau, dtype=np_dtype))
}
headstr = "{0} *x, {0} *y, "
headstr += "{0} *f1, {0} *f2, {0} *alpha, {0} *tau"
self._kernel = ElementwiseKernel(headstr.format(type_t),
("y[i] = fmax(x[i] + {0}f1[i], "\
+ "fmin(x[i] + {0}f2[i], "\
+ "alpha[i]*x[i]))/alpha[i]").format(taufact))
def _call_gpu(self, x, y=None, add=False, jacobian=False):
assert not jacobian
assert not add
g = self.gpuvars
y = x if y is None else y
self._kernel(x, y, g['f1'], g['f2'], g['alpha'], g['tau'])
def _call_cpu(self, x, y=None, add=False, jacobian=False):
assert not jacobian
assert not add
x = self.x.vars(x)[0]
y = self.y.vars(y)[0] if y is not None else x
np.fmax(x + self.tau * self.f1,
np.fmin(x + self.tau * self.f2, self.alpha * x),
out=y)
y /= self.alpha
class IntervProj(Operator):
""" y[i,k] = proj_[0,b[k]](x[i,k] - a[i,k]*shift[i,k]) """
def __init__(self, b, shift, a=1, mask=None):
Operator.__init__(self)
assert b.size == shift.shape[1]
self.x = Variable(shift.shape)
self.y = Variable(shift.shape)
self.b = b
self.a = a
self.shift = shift
self.mask = np.ones(shift.shape[0],
dtype=bool) if mask is None else mask
def prepare_gpu(self, type_t="double"):
# don't multiply with a if a is 1 (not 1.0!)
afact = "" if self.a is 1 else "a[0]*"
if type(self.a) is np.ndarray:
afact = "a[i]*"
np_dtype = np.float64 if type_t == "double" else np.float32
self.gpuvars = {
'b': gpuarray.to_gpu(self.b),
'shift': gpuarray.to_gpu(self.shift),
'a': gpuarray.to_gpu(np.asarray(self.a, dtype=np_dtype))
}
headstr = "{0} *x, {0} *y, {0} *shift, {0} *a, {0} *b".format(type_t)
self._kernel = ElementwiseKernel(headstr,
"y[i] = fmax(0.0, fmin(b[i%{}], x[i] - {}shift[i]))"\
.format(self.b.size, afact))
def _call_gpu(self, x, y=None, add=False, jacobian=False):
assert not jacobian
assert not add
g = self.gpuvars
y = x if y is None else y
self._kernel(x, y, g['shift'], g['a'], g['b'])
def _call_cpu(self, x, y=None, add=False, jacobian=False):
assert not jacobian
assert not add
x = self.x.vars(x)[0]
y = self.y.vars(y)[0] if y is not None else x
np.fmax(0.0, np.fmin(self.b[None, :], x - self.a * self.shift), out=y)
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 BndSSD(Functional):
""" 0.5*|max(0, f1 - x)|^2 + 0.5*|max(0, x - f2)|^2 """
def __init__(self, f1, f2, vol=None, mask=None, conj=None):
Functional.__init__(self)
self.x = Variable(f1.shape)
self.f1 = f1
self.f2 = f2
self.vol = np.ones(f1.shape[1]) if vol is None else vol
self.mask = np.ones(f1.shape[0], dtype=bool) if mask is None else mask
if conj is None:
cj_vol = 1.0 / self.vol
self.conj = BndSSDConj(f1, f2, vol=cj_vol, mask=mask, conj=self)
else:
self.conj = conj
f1_msk, f2_msk = [np.zeros_like(a) for a in [f1, f2]]
f1_msk[self.mask, :] = self.vol[None, :] * f1[self.mask, :]
f2_msk[self.mask, :] = self.vol[None, :] * f2[self.mask, :]
self._prox = ProxBndSSD(f1_msk, f2_msk)
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
posval1 = np.fmax(0.0, (self.f1 - x)[self.mask, :])
posval2 = np.fmax(0.0, (x - self.f2)[self.mask, :])
val = 0.5*np.einsum('ik,k->', posval1**2, self.vol) \
+ 0.5*np.einsum('ik,k->', posval2**2, self.vol)
infeas = 0
return (val, infeas)
def prox(self, tau):
msk = self.mask
tauvol = np.zeros_like(self.f1)
tauvol[msk, :] = (tau * np.ones(self.f1.size)).reshape(
self.f1.shape)[msk, :]
tauvol[msk, :] = np.einsum('ik,k->ik', tauvol[msk, :], self.vol)
tauvol += 1.0
self._prox.alpha = tauvol
self._prox.tau = tau
return self._prox
class ConstrainOp(Operator):
""" y[mask,:] = const, else identity """
def __init__(self, mask, const):
Operator.__init__(self)
self.x = Variable(const.shape)
self.y = self.x
self.mask = mask
self.const = const
scale = np.ones_like(const)
scale[mask, :] = 0.0
self._jacobian = ScaleOp(self.x.size, scale.ravel())
def prepare_gpu(self, type_t="double"):
self.const_gpu = gpuarray.to_gpu(self.const)
self.mask_gpu = gpuarray.to_gpu(self.mask.astype(np.int8))
N, M = self.x[0]['shape']
headstr = "%s *x, %s *y, %s *c, char *mask" % ((type_t, ) * 3)
self._kernel = ElementwiseKernel(
headstr, "y[i] = (mask[i/{}]) ? c[i] : x[i]".format(M))
self._kernel_add = ElementwiseKernel(
headstr, "y[i] += (mask[i/{}]) ? c[i] : x[i]".format(M))
self._jacobian.prepare_gpu(type_t=type_t)
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, self.const_gpu, self.mask_gpu)
else: self._kernel(x, y, self.const_gpu, self.mask_gpu)
if jacobian:
return self._jacobian
def _call_cpu(self, x, y=None, add=False, jacobian=False):
if y is not None:
y[:] = x
x = y
x = self.x.vars(x)[0]
x[self.mask, :] = self.const[self.mask, :]
if jacobian:
return self._jacobian
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 PosSSDConj(Functional):
""" 0.5*|max(0, x + f)|^2 - 0.5*|f|^2 """
def __init__(self, data, vol=None, mask=None, conj=None):
Functional.__init__(self)
self.x = Variable(data.shape)
self.f = np.atleast_2d(data)
self.shift = -0.5*np.einsum('ik,k->', data**2, vol)
self.vol = np.ones(data.shape[1]) if vol is None else vol
self.mask = np.ones(data.shape[0], dtype=bool) if mask is None else mask
if conj is None:
cj_vol = 1.0/self.vol
cj_data = np.zeros_like(self.f)
cj_data[self.mask,:] = np.einsum('ik,k->ik', self.f[self.mask,:], self.vol)
self.conj = PosSSD(cj_data, vol=cj_vol, mask=mask, conj=self)
else:
self.conj = conj
def __call__(self, x, grad=False):
assert not grad
x = self.x.vars(x)[0]
val = 0.5*np.einsum('ik,k->', np.fmax(0.0, (x + self.f)[self.mask,:])**2, self.vol)
val += self.shift
infeas = 0
return (val, infeas)
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(SublabelModel):
name = "quadratic"
def __init__(self,
*args,
lbd=5.0,
alph=np.inf,
fdscheme="centered",
**kwargs):
SublabelModel.__init__(self, *args, **kwargs)
self.lbd = lbd
self.alph = alph
self.fdscheme = fdscheme
logging.info("Init model '%s' (lambda=%.2e, alpha=%.2e, fdscheme=%s)" \
% (self.name, self.lbd, self.alph, self.fdscheme))
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
xvars = [('u', (N_image, L_labels)),
('w12', (M_tris, N_image, s_gamma + 1)),
('w', (M_tris, N_image, d_image, s_gamma))]
yvars = [
('p', (N_image, d_image, L_labels)),
('q', (N_image, L_labels)),
('v12', (M_tris, N_image, s_gamma + 1)),
('v3', (N_image, )),
('g12', (M_tris, N_image, d_image * s_gamma + 1)),
]
self.x = Variable(*xvars)
self.y = Variable(*yvars)
def setup_solver(self, *args):
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
d_image = self.data.d_image
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
Id_w2 = np.zeros((s_gamma + 1, d_image * s_gamma + 1), order='C')
Id_w2[-1, -1] = 1.0
Adext = np.zeros((M_tris, d_image * s_gamma, d_image * s_gamma + 1),
order='C')
Adext[:, :, :-1] = np.kron(np.eye(d_image), self.data.Ad)
self.linblocks.update({
'Grad':
GradientOp(imagedims, L_labels, scheme=self.fdscheme),
'PB':
IndexedMultAdj(L_labels, d_image * N_image, self.data.P,
self.data.B),
'Adext':
MatrixMultRBatched(N_image, Adext),
'Id_w2':
MatrixMultR(M_tris * N_image, Id_w2),
})
SublabelModel.setup_solver(self, *args)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
PAbOp = self.linblocks['PAb']
S_u_k = self.linblocks['S']
GradOp = self.linblocks['Grad']
PBLinOp = self.linblocks['PB']
AdMult = self.linblocks['Adext']
Id_w2 = self.linblocks['Id_w2']
if self.alph < np.inf:
etahat = HuberPerspective(M_tris * N_image,
s_gamma * d_image,
lbd=self.lbd,
alph=self.alph)
else:
etahat = QuadEpiSupp(M_tris * N_image,
s_gamma * d_image,
a=self.lbd)
Id_u = IdentityOp(x['u']['size'])
Id_w12 = IdentityOp(x['w12']['size'])
if self.data.constraints is not None:
constrmask, constru = self.data.constraints
constru_lifted = self.data.mfd.embed_barycentric(constru)[1]
Gu = ConstrainFct(constrmask, constru_lifted)
else:
Gu = PositivityFct(x['u']['size'])
self.pdhg_G = SplitSum([
Gu, # \delta_{u >= 0} or constraints
ZeroFct(x['w12']['size']), # 0
ZeroFct(x['w']['size']), # 0
])
self.pdhg_F = SplitSum([
IndicatorFct(y['p']['size']), # \delta_{p = 0}
IndicatorFct(y['q']['size']), # \delta_{q = 0}
self.epifct, # \max_{v \in epi(rho*)} <v12,v>
IndicatorFct(y['v3']['size'], c1=1), # \delta_{v3^i = 1}
etahat, # 0.5*lbd*\sum_ji |g1[j,i]|^2/|g2[j,i]|
])
self.pdhg_linop = BlockOp([
[GradOp, 0, PBLinOp], # p = Du - P'B'w
[Id_u, PAbOp, 0], # q = u - P'Ab'w12
[0, Id_w12, 0], # v12 = w12
[S_u_k, 0, 0], # v3^i = sum_k u[i,k]
[0, Id_w2, AdMult], # g12 = (Ad'w, w2)
])
class QuadEpiProj(Operator):
""" T(z)[i] = proj[epi(f_i)](z[i])
f_i(x) := 0.5*alph*|x|^2 + <b[i],x> + c[i]
= 0.5*alph*|x - shift1[i]|^2 + shift2[i]
shift1 := -b/alph
shift2 := -(0.5/alph*|b|^2 - c)
Orthogonal projections onto a (translated) paraboloid which requires root
finding for cubic polynomials.
Optionally, the paraboloid is truncated at |x1 - shift1| < lbd.
"""
def __init__(self, N, M, alph=1.0, shift=None, b=None, c=None, lbd=np.inf):
Operator.__init__(self)
assert lbd > 0
assert alph > 0
self.N, self.M = N, M
self.x = Variable((self.N, self.M + 1))
self.y = self.x
self.lbd = lbd
self.alph = alph
self.shift = shift
if not (b is None and c is None):
assert self.shift is None
self.b = np.zeros((N, M)) if b is None else b
self.c = np.zeros((N, )) if c is None else c
self.shift = np.zeros((N, M + 1))
self.shift[:, :-1] = -self.b / self.alph
self.shift[:, -1] = -(0.5 / self.alph *
(self.b**2).sum(axis=-1) - self.c)
def prepare_gpu(self, type_t="double"):
constvars = {
'QUAD_EPI_PROJ': 1,
'alph': np.float64(self.alph),
'N': self.N,
'M': self.M,
'TYPE_T': type_t,
}
if self.lbd < np.inf:
constvars['lbd'] = np.float64(self.lbd)
constvars['USE_LBD'] = 1
if self.shift is not None:
constvars['shift'] = self.shift
constvars['USE_SHIFT'] = 1
for f in ['sqrt', 'cbrt', 'acos', 'cos', 'fabs']:
constvars[f.upper()] = f if type_t == "double" else (f + "f")
files = [resource_stream('opymize.operators', 'proj.cu')]
templates = [("quadepiproj", "P", (self.N, 1, 1), (200, 1, 1))]
self._kernel = prepare_kernels(files, templates,
constvars)['quadepiproj']
def _call_gpu(self, x, y=None, add=False, jacobian=False):
assert not add
assert not jacobian
if y is not None:
y[:] = x.copy()
x = y
self._kernel(x)
def _call_cpu(self, x, y=None, add=False, jacobian=False):
assert not add
assert not jacobian
if y is None:
y = x
else:
y[:] = x
x, y = self.x.vars(x)[0], self.y.vars(y)[0]
alph, shift, lbd = self.alph, self.shift, self.lbd
if shift is not None: y -= shift
xnorms = np.linalg.norm(y[:, :-1], axis=-1)
msk_c0 = (xnorms == 0)
if lbd < np.inf:
msk_c1 = -xnorms / (
lbd * alph) + 0.5 * alph * lbd**2 + 1 / alph > y[:, -1]
msk_c2 = 0.5 * alph * lbd**2 > y[:, -1]
msk = (xnorms > lbd) & (~msk_c2)
y[msk, :-1] *= lbd / xnorms[msk, None]
msk = msk_c2 & (~msk_c1)
y[msk, :-1] *= lbd / xnorms[msk, None]
y[msk, -1] = 0.5 * alph * lbd**2
else:
msk_c1 = np.ones(y.shape[:-1], dtype=bool)
msk = msk_c1 & (~msk_c0) & (0.5 * alph * xnorms**2 > y[:, -1])
a_2 = 2.0 / alph**2
a = a_2 * (1 - y[msk, -1] * alph)
b = -a_2 * xnorms[msk]
ynorms = solve_reduced_monic_cubic(a, b)
y[msk, :-1] *= (ynorms / xnorms[msk])[:, None]
y[msk, -1] = 0.5 * alph * ynorms**2
y[msk_c0 & (y[:, -1] < 0), -1] = 0.0
if shift is not None: y += shift
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_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(SublabelModel):
name = "rof"
def __init__(self, *args, lbd=1.0, regularizer="tv", alph=np.inf,
fdscheme="centered", **kwargs):
SublabelModel.__init__(self, *args, **kwargs)
self.lbd = lbd
self.regularizer = regularizer
self.alph = alph
self.fdscheme = fdscheme
logging.info("Init model '%s' (%s regularizer, lambda=%.2e, "
"alpha=%.2e, fdscheme=%s)" \
% (self.name, self.regularizer, self.lbd,
self.alph, self.fdscheme))
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
xvars = [('u', (N_image, L_labels)),
('w12', (M_tris, N_image, s_gamma+1)),
('w', (M_tris, N_image, d_image, s_gamma))]
yvars = [('p', (N_image, d_image, L_labels)),
('q', (N_image, L_labels)),
('v12a', (M_tris, N_image, s_gamma+1)),
('v12b', (M_tris, N_image, s_gamma+1)),
('v3', (N_image,)),]
if self.regularizer == "tv":
yvars.append(('g', (M_tris, N_image, d_image, s_gamma)))
elif self.regularizer == "quadratic":
yvars.append(('g12', (M_tris, N_image, d_image*s_gamma+1)))
self.x = Variable(*xvars)
self.y = Variable(*yvars)
def setup_solver(self, *args):
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
Id_w2 = np.zeros((s_gamma+1,d_image*s_gamma+1), order='C')
Id_w2[-1,-1] = 1.0
Adext = np.zeros((M_tris,s_gamma,d_image*s_gamma+1), order='C')
Adext[:,:,:-1] = np.tile(self.data.Ad, (1,1,d_image))
# Ab (M_tris, s_gamma+1, s_gamma+1)
Ab_mats = -np.ones((M_tris, s_gamma+1, s_gamma+1),
dtype=np.float64, order='C')
Ab_mats[:,:,0:-1] = self.data.T[self.data.P]
Ab_mats[:] = np.linalg.inv(Ab_mats)
self.linblocks.update({
'PAbTri': IndexedMultAdj(L_labels, N_image, self.data.P, Ab_mats),
'Grad': GradientOp(imagedims, L_labels, scheme=self.fdscheme),
'PB': IndexedMultAdj(L_labels, d_image*N_image, self.data.P, self.data.B),
'Ad': MatrixMultRBatched(N_image*d_image, self.data.Ad),
'Adext': MatrixMultRBatched(N_image, Adext),
'Id_w2': MatrixMultR(M_tris*N_image, Id_w2),
})
SublabelModel.setup_solver(self, *args)
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
PAbOp = self.linblocks['PAbTri']
S_u_k = self.linblocks['S']
GradOp = self.linblocks['Grad']
PBLinOp = self.linblocks['PB']
AdMult = self.linblocks['Ad']
shift = np.tile(self.data.data_b, (M_tris,1,1)).reshape((-1, s_gamma))
c = 0.5*(shift**2).sum(axis=-1)
epifct1 = QuadEpiSupp(M_tris*N_image, s_gamma, b=-shift, c=c)
epifct2 = EpigraphSupp(np.ones((N_image, L_labels), dtype=bool),
[[np.arange(s_gamma+1)[None]]*M_tris]*N_image,
self.data.P, self.data.T, np.zeros((N_image, L_labels)))
Id_u = IdentityOp(x['u']['size'])
Id_w12 = IdentityOp(x['w12']['size'])
if self.data.constraints is not None:
constrmask, constru = self.data.constraints
constru_lifted = self.data.mfd.embed_barycentric(constru)[1]
Gu = ConstrainFct(constrmask, constru_lifted)
else:
Gu = PositivityFct(x['u']['size'])
self.pdhg_G = SplitSum([
Gu, # \delta_{u >= 0} or constraints
ZeroFct(x['w12']['size']), # 0
ZeroFct(x['w']['size']), # 0
])
F_summands = [
IndicatorFct(y['p']['size']), # \delta_{p = 0}
IndicatorFct(y['q']['size']), # \delta_{q = 0}
epifct1, # -0.5*v2a*|v1a/v2a + b|^2
epifct2, # \max_{v \in Delta} <v12b,v>
IndicatorFct(y['v3']['size'], c1=1), # \delta_{v3^i = 1}
]
op_blocks = [
[GradOp, 0, PBLinOp], # p = Du - P'B'w
[ Id_u, PAbOp, 0], # q = u - P'Ab'w12
[ 0, Id_w12, 0], # v12a = w12
[ 0, Id_w12, 0], # v12b = w12
[ S_u_k, 0, 0], # v3^i = sum_k u[i,k]
]
if self.regularizer == "tv":
l1norms = L1Norms(M_tris*N_image, (d_image, s_gamma), self.lbd, "nuclear")
F_summands.append(l1norms) # lbd*\sum_ji |g[j,i,:,:]|_nuc
op_blocks.append([ 0, 0, AdMult]) # g = A'w
elif self.regularizer == "quadratic":
if self.alph < np.inf:
etahat = HuberPerspective(M_tris*N_image, s_gamma*d_image,
lbd=self.lbd, alph=self.alph)
else:
etahat = QuadEpiSupp(M_tris*N_image, s_gamma*d_image, a=self.lbd)
F_summands.append(etahat) # 0.5*lbd*\sum_ji |g1[j,i]|^2/|g2[j,i]|
AdMult = self.linblocks['Adext']
Id_w2 = self.linblocks['Id_w2']
op_blocks.append([ 0, Id_w2, AdMult]) # g12 = (Ad'w, w2)
self.pdhg_F = SplitSum(F_summands)
self.pdhg_linop = BlockOp(op_blocks)
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(SublabelModel):
name = "tv"
def __init__(self, *args, lbd=1.0, fdscheme="centered", **kwargs):
SublabelModel.__init__(self, *args, **kwargs)
self.lbd = lbd
self.fdscheme = fdscheme
logging.info("Init model '%s' (lambda=%.2e,fdscheme=%s)" \
% (self.name, self.lbd, self.fdscheme))
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
xvars = [('u', (N_image, L_labels)),
('w12', (M_tris, N_image, s_gamma + 1)),
('w', (M_tris, N_image, d_image, s_gamma))]
yvars = [('p', (N_image, d_image, L_labels)),
('q', (N_image, L_labels)),
('v12', (M_tris, N_image, s_gamma + 1)), ('v3', (N_image, )),
('g', (M_tris, N_image, d_image, s_gamma))]
if self.data.R.shape[-1] == s_gamma + 1:
del xvars[1]
del yvars[2]
self.x = Variable(*xvars)
self.y = Variable(*yvars)
def setup_solver(self, *args):
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
d_image = self.data.d_image
self.linblocks.update({
'Grad':
GradientOp(imagedims, L_labels, scheme=self.fdscheme),
'PB':
IndexedMultAdj(L_labels, d_image * N_image, self.data.P,
self.data.B),
'Ad':
MatrixMultRBatched(N_image * d_image, self.data.Ad),
})
SublabelModel.setup_solver(self, *args)
def setup_dataterm_blocks(self):
if hasattr(self, 'epifct') or hasattr(self, 'rho'):
return
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
P = self.data.P
if self.data.R.shape[-1] == s_gamma + 1:
logging.info("Setup for model without sublabels...")
R = self.data.R.reshape(M_tris, N_image, s_gamma + 1)
self.rho = np.zeros((N_image, L_labels), order='C')
for j in range(M_tris):
for m in range(s_gamma + 1):
self.rho[:, P[j, m]] = R[j, :, m]
self.rho = self.rho.ravel()
else:
logging.info("Setup for sublabel-accurate model...")
self.epifct = EpigraphSupp(self.data.Rbase, self.data.Rfaces,
self.data.Qbary, self.data.Sbary,
self.data.R)
# Ab (M_tris, s_gamma+1, s_gamma+1)
Ab = np.zeros((M_tris, s_gamma + 1, s_gamma + 1),
dtype=np.float64,
order='C')
Ab[:] = np.eye(s_gamma + 1)[None]
Ab[..., -1] = -1
self.linblocks['PAb'] = IndexedMultAdj(L_labels, N_image, P, Ab)
self.linblocks.update({
'S':
MatrixMultR(N_image, np.ones((L_labels, 1), order='C')),
})
def setup_solver_pdhg(self):
x, y = self.x.vars(named=True), self.y.vars(named=True)
imagedims = self.data.imagedims
N_image = self.data.N_image
L_labels = self.data.L_labels
M_tris = self.data.M_tris
s_gamma = self.data.s_gamma
d_image = self.data.d_image
S_u_k = self.linblocks['S']
GradOp = self.linblocks['Grad']
PBLinOp = self.linblocks['PB']
AdMult = self.linblocks['Ad']
l1norms = L1Norms(M_tris * N_image, (d_image, s_gamma), self.lbd,
"nuclear")
Id_u = IdentityOp(x['u']['size'])
if self.data.constraints is not None:
constrmask, constru = self.data.constraints
constru_lifted = self.data.mfd.embed_barycentric(constru)[1]
Gu = ConstrainFct(constrmask, constru_lifted)
else:
Gu = PositivityFct(x['u']['size'])
if hasattr(self, 'epifct'):
PAbOp = self.linblocks['PAb']
Id_w12 = IdentityOp(x['w12']['size'])
G_summands = [
Gu, # \delta_{u >= 0} or constraints
ZeroFct(x['w12']['size']), # 0
ZeroFct(x['w']['size']), # 0
]
F_summands = [
IndicatorFct(y['p']['size']), # \delta_{p = 0}
IndicatorFct(y['q']['size']), # \delta_{q = 0}
self.epifct, # \max_{v \in epi(rho*)} <v12,v>
IndicatorFct(y['v3']['size'], c1=1), # \delta_{v3^i = 1}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
]
op_blocks = [
[GradOp, 0, PBLinOp], # p = Du - P'B'w
[Id_u, PAbOp, 0], # q = u - P'Ab'w12
[0, Id_w12, 0], # v12 = w12
[S_u_k, 0, 0], # v3^i = sum_k u[i,k]
[0, 0, AdMult], # g = A'w
]
else:
G_summands = [
Gu, # \delta_{u >= 0} or constraints
ZeroFct(x['w']['size']), # 0
]
F_summands = [
IndicatorFct(y['p']['size']), # \delta_{p = 0}
AffineFct(y['q']['size'], c=self.rho), # <q,rho>
IndicatorFct(y['v3']['size'], c1=1), # \delta_{v3^i = 1}
l1norms, # lbd*\sum_ji |g[j,i,:,:]|_nuc
]
op_blocks = [
[GradOp, PBLinOp], # p = Du - P'B'w
[Id_u, 0], # q = u
[S_u_k, 0], # v3^i = sum_k u[i,k]
[0, AdMult], # g = A'w
]
self.pdhg_G = SplitSum(G_summands)
self.pdhg_F = SplitSum(F_summands)
self.pdhg_linop = BlockOp(op_blocks)
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)