def get_diagonal(A, d, solve_mode=True):
"""
Compute the diagonal of the square operator A
or its inverse A^{-1} (if solve_mode=True).
"""
ej, xj = Vector(), Vector()
if hasattr(A, "init_vector"):
A.init_vector(ej, 1)
A.init_vector(xj, 0)
else:
A.get_operator().init_vector(ej, 1)
A.get_operator().init_vector(xj, 0)
ncol = ej.size()
da = np.zeros(ncol, dtype=ej.array().dtype)
for j in range(ncol):
ej[j] = 1.
if solve_mode:
A.solve(xj, ej)
else:
A.mult(ej, xj)
da[j] = xj[j]
ej[j] = 0.
d.set_local(da)
def get_diagonal(A, d):
"""
Compute the diagonal of the square operator A.
Use Solver2Operator if A^-1 is needed.
"""
ej, xj = Vector(), Vector()
A.init_vector(ej, 1)
A.init_vector(xj, 0)
g_size = ej.size()
d.zero()
for gid in xrange(g_size):
owns_gid = ej.owns_index(gid)
if owns_gid:
SetToOwnedGid(ej, gid, 1.)
ej.apply("insert")
A.mult(ej, xj)
if owns_gid:
val = GetFromOwnedGid(xj, gid)
SetToOwnedGid(d, gid, val)
SetToOwnedGid(ej, gid, 0.)
ej.apply("insert")
d.apply("insert")
def get_diagonal(A, d):
"""
Compute the diagonal of the square operator :math:`A`.
Use :code:`Solver2Operator` if :math:`A^{-1}` is needed.
"""
ej, xj = Vector(d.mpi_comm()), Vector(d.mpi_comm())
A.init_vector(ej, 1)
A.init_vector(xj, 0)
g_size = ej.size()
d.zero()
for gid in range(g_size):
owns_gid = ej.owns_index(gid)
if owns_gid:
SetToOwnedGid(ej, gid, 1.)
ej.apply("insert")
A.mult(ej, xj)
if owns_gid:
val = GetFromOwnedGid(xj, gid)
SetToOwnedGid(d, gid, val)
SetToOwnedGid(ej, gid, 0.)
ej.apply("insert")
d.apply("insert")
def get_diagonal(A, d):
"""
Compute the diagonal of the square operator :math:`A`.
Use :code:`Solver2Operator` if :math:`A^{-1}` is needed.
"""
ej, xj = Vector(d.mpi_comm()), Vector(d.mpi_comm())
A.init_vector(ej,1)
A.init_vector(xj,0)
g_size = ej.size()
d.zero()
for gid in range(g_size):
owns_gid = ej.owns_index(gid)
if owns_gid:
SetToOwnedGid(ej, gid, 1.)
ej.apply("insert")
A.mult(ej,xj)
if owns_gid:
val = GetFromOwnedGid(xj, gid)
SetToOwnedGid(d, gid, val)
SetToOwnedGid(ej, gid, 0.)
ej.apply("insert")
d.apply("insert")
class TVPD():
""" Total variation using primal-dual Newton """
def __init__(self, parameters):
"""
TV regularization in primal-dual format
Input parameters:
* k = regularization parameter
* eps = regularization constant (see above)
* rescaledradiusdual = radius of dual set
* exact = use full TV (bool)
* PCGN = use GN Hessian to precondition (bool);
not recommended for performance but can help avoid num.instability
* print (bool)
"""
self.parameters = {}
self.parameters['k'] = 1.0
self.parameters['eps'] = 1e-2
self.parameters['rescaledradiusdual'] = 1.0
self.parameters['exact'] = False
self.parameters['PCGN'] = False
self.parameters['print'] = False
self.parameters['correctcost'] = True
self.parameters['amg'] = 'default'
assert parameters.has_key('Vm')
self.parameters.update(parameters)
self.Vm = self.parameters['Vm']
k = self.parameters['k']
eps = self.parameters['eps']
exact = self.parameters['exact']
amg = self.parameters['amg']
self.m = Function(self.Vm)
testm = TestFunction(self.Vm)
trialm = TrialFunction(self.Vm)
# WARNING: should not be changed.
# As it is, code only works with DG0
if self.parameters.has_key('Vw'):
Vw = self.parameters['Vw']
else:
Vw = FunctionSpace(self.Vm.mesh(), 'DG', 0)
self.wx = Function(Vw)
self.wxrs = Function(Vw) # re-scaled dual variable
self.wxhat = Function(Vw)
self.gwx = Function(Vw)
self.wy = Function(Vw)
self.wyrs = Function(Vw) # re-scaled dual variable
self.wyhat = Function(Vw)
self.gwy = Function(Vw)
self.wxsq = Vector()
self.wysq = Vector()
self.normw = Vector()
self.factorw = Vector()
testw = TestFunction(Vw)
trialw = TrialFunction(Vw)
normm = inner(nabla_grad(self.m), nabla_grad(self.m))
TVnormsq = normm + Constant(eps)
TVnorm = sqrt(TVnormsq)
if self.parameters['correctcost']:
meshtmp = UnitSquareMesh(self.Vm.mesh().mpi_comm(), 10, 10)
Vtmp = FunctionSpace(meshtmp, 'CG', 1)
x = SpatialCoordinate(meshtmp)
correctioncost = 1. / assemble(sqrt(4.0 * x[0] * x[0]) * dx)
else:
correctioncost = 1.0
self.wkformcost = Constant(k * correctioncost) * TVnorm * dx
if exact:
sys.exit(1)
# self.w = nabla_grad(self.m)/TVnorm # full Hessian
# self.Htvw = inner(Constant(k) * nabla_grad(testm), self.w) * dx
self.misfitwx = inner(testw, self.wx * TVnorm - self.m.dx(0)) * dx
self.misfitwy = inner(testw, self.wy * TVnorm - self.m.dx(1)) * dx
self.Htvx = assemble(inner(Constant(k) * testm.dx(0), trialw) * dx)
self.Htvy = assemble(inner(Constant(k) * testm.dx(1), trialw) * dx)
self.massw = inner(TVnorm * testw, trialw) * dx
mpicomm = self.Vm.mesh().mpi_comm()
invMwMat, VDM, VDM = setupPETScmatrix(Vw, Vw, 'aij', mpicomm)
for ii in VDM.dofs():
invMwMat[ii, ii] = 1.0
invMwMat.assemblyBegin()
invMwMat.assemblyEnd()
self.invMwMat = PETScMatrix(invMwMat)
self.invMwd = Vector()
self.invMwMat.init_vector(self.invMwd, 0)
self.invMwMat.init_vector(self.wxsq, 0)
self.invMwMat.init_vector(self.wysq, 0)
self.invMwMat.init_vector(self.normw, 0)
self.invMwMat.init_vector(self.factorw, 0)
u = Function(Vw)
uflrank = len(u.ufl_shape)
if uflrank == 0:
ones = ("1.0")
elif uflrank == 1:
ones = (("1.0", "1.0"))
else:
sys.exit(1)
u = interpolate(Constant(ones), Vw)
self.one = u.vector()
self.wkformAx = inner(testw, trialm.dx(0) - \
self.wx * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAxrs = inner(testw, trialm.dx(0) - \
self.wxrs * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAy = inner(testw, trialm.dx(1) - \
self.wy * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAyrs = inner(testw, trialm.dx(1) - \
self.wyrs * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
kovsq = Constant(k) / TVnorm
self.wkformGNhess = kovsq * inner(nabla_grad(trialm),
nabla_grad(testm)) * dx
factM = 1e-2 * k
M = assemble(inner(testm, trialm) * dx)
self.sMass = M * factM
self.Msolver = PETScKrylovSolver('cg', 'jacobi')
self.Msolver.parameters["maximum_iterations"] = 2000
self.Msolver.parameters["relative_tolerance"] = 1e-24
self.Msolver.parameters["absolute_tolerance"] = 1e-24
self.Msolver.parameters["error_on_nonconvergence"] = True
self.Msolver.parameters["nonzero_initial_guess"] = False
self.Msolver.set_operator(M)
if amg == 'default': self.amgprecond = amg_solver()
else: self.amgprecond = amg
if self.parameters['print']:
print '[TVPD] TV regularization -- primal-dual method',
if self.parameters['PCGN']:
print ' -- PCGN',
print ' -- k={}, eps={}'.format(k, eps)
print '[TVPD] preconditioner = {}'.format(self.amgprecond)
print '[TVPD] correction cost with factor={}'.format(
correctioncost)
def isTV(self):
return True
def isPD(self):
return True
def cost(self, m):
""" evaluate the cost functional at m """
setfct(self.m, m)
return assemble(self.wkformcost)
def costvect(self, m_in):
return self.cost(m_in)
def _assemble_invMw(self):
""" Assemble inverse of matrix Mw,
weighted mass matrix in dual space """
# WARNING: only works if Mw is diagonal (e.g, DG0)
Mw = assemble(self.massw)
Mwd = get_diagonal(Mw)
as_backend_type(self.invMwd).vec().pointwiseDivide(\
as_backend_type(self.one).vec(),\
as_backend_type(Mwd).vec())
self.invMwMat.set_diagonal(self.invMwd)
def grad(self, m):
""" compute the gradient at m """
setfct(self.m, m)
self._assemble_invMw()
self.gwx.vector().zero()
self.gwx.vector().axpy(1.0, assemble(self.misfitwx))
normgwx = norm(self.gwx.vector())
self.gwy.vector().zero()
self.gwy.vector().axpy(1.0, assemble(self.misfitwy))
normgwy = norm(self.gwy.vector())
if self.parameters['print']:
print '[TVPD] |gw|={}'.format(np.sqrt(normgwx**2 + normgwy**2))
return self.Htvx*(self.wx.vector() - self.invMwd*self.gwx.vector()) \
+ self.Htvy*(self.wy.vector() - self.invMwd*self.gwy.vector())
#return assemble(self.Htvw) - self.Htv*(self.invMwd*self.gw.vector())
def gradvect(self, m_in):
return self.grad(m_in)
def assemble_hessian(self, m):
""" build Hessian matrix at given point m """
setfct(self.m, m)
self._assemble_invMw()
self.Ax = assemble(self.wkformAx)
Hxasym = MatMatMult(self.Htvx, MatMatMult(self.invMwMat, self.Ax))
Hx = (Hxasym + Transpose(Hxasym)) * 0.5
Axrs = assemble(self.wkformAxrs)
Hxrsasym = MatMatMult(self.Htvx, MatMatMult(self.invMwMat, Axrs))
Hxrs = (Hxrsasym + Transpose(Hxrsasym)) * 0.5
self.Ay = assemble(self.wkformAy)
Hyasym = MatMatMult(self.Htvy, MatMatMult(self.invMwMat, self.Ay))
Hy = (Hyasym + Transpose(Hyasym)) * 0.5
Ayrs = assemble(self.wkformAyrs)
Hyrsasym = MatMatMult(self.Htvy, MatMatMult(self.invMwMat, Ayrs))
Hyrs = (Hyrsasym + Transpose(Hyrsasym)) * 0.5
self.H = Hx + Hy
self.Hrs = Hxrs + Hyrs
PCGN = self.parameters['PCGN']
if PCGN:
HGN = assemble(self.wkformGNhess)
self.precond = HGN + self.sMass
else:
self.precond = self.Hrs + self.sMass
def hessian(self, mhat):
return self.Hrs * mhat
def compute_what(self, mhat):
""" Compute update direction for what, given mhat """
self.wxhat.vector().zero()
self.wxhat.vector().axpy(
1.0, self.invMwd * (self.Ax * mhat - self.gwx.vector()))
normwxhat = norm(self.wxhat.vector())
self.wyhat.vector().zero()
self.wyhat.vector().axpy(
1.0, self.invMwd * (self.Ay * mhat - self.gwy.vector()))
normwyhat = norm(self.wyhat.vector())
if self.parameters['print']:
print '[TVPD] |what|={}'.format(
np.sqrt(normwxhat**2 + normwyhat**2))
def update_w(self, mhat, alphaLS, compute_what=True):
""" update dual variable in direction what
and update re-scaled version """
if compute_what: self.compute_what(mhat)
self.wx.vector().axpy(alphaLS, self.wxhat.vector())
self.wy.vector().axpy(alphaLS, self.wyhat.vector())
# rescaledradiusdual=1.0: checked empirically to be max radius acceptable
rescaledradiusdual = self.parameters['rescaledradiusdual']
# wx**2
as_backend_type(self.wxsq).vec().pointwiseMult(\
as_backend_type(self.wx.vector()).vec(),\
as_backend_type(self.wx.vector()).vec())
# wy**2
as_backend_type(self.wysq).vec().pointwiseMult(\
as_backend_type(self.wy.vector()).vec(),\
as_backend_type(self.wy.vector()).vec())
# |w|
self.normw = self.wxsq + self.wysq
as_backend_type(self.normw).vec().sqrtabs()
# |w|/r
as_backend_type(self.normw).vec().pointwiseDivide(\
as_backend_type(self.normw).vec(),\
as_backend_type(self.one*rescaledradiusdual).vec())
# max(1.0, |w|/r)
# as_backend_type(self.factorw).vec().pointwiseMax(\
# as_backend_type(self.one).vec(),\
# as_backend_type(self.normw).vec())
count = pointwiseMaxCount(self.factorw, self.normw, 1.0)
# rescale wx and wy
as_backend_type(self.wxrs.vector()).vec().pointwiseDivide(\
as_backend_type(self.wx.vector()).vec(),\
as_backend_type(self.factorw).vec())
as_backend_type(self.wyrs.vector()).vec().pointwiseDivide(\
as_backend_type(self.wy.vector()).vec(),\
as_backend_type(self.factorw).vec())
minf = self.factorw.min()
maxf = self.factorw.max()
if self.parameters['print']:
# print 'min(factorw)={}, max(factorw)={}'.format(minf, maxf)
print '[TVPD] perc. dual entries rescaled={:.2f} %, min(factorw)={}, max(factorw)={}'.format(\
100.*float(count)/self.factorw.size(), minf, maxf)
def getprecond(self):
""" Precondition by inverting the TV Hessian """
solver = PETScKrylovSolver('cg', self.amgprecond)
solver.parameters["maximum_iterations"] = 2000
solver.parameters["relative_tolerance"] = 1e-24
solver.parameters["absolute_tolerance"] = 1e-24
solver.parameters["error_on_nonconvergence"] = True
solver.parameters["nonzero_initial_guess"] = False
# used to compare iterative application of preconditioner
# with exact application of preconditioner:
#solver = PETScLUSolver("petsc")
#solver.parameters['symmetric'] = True
#solver.parameters['reuse_factorization'] = True
solver.set_operator(self.precond)
return solver
def init_vector(self, u, dim):
self.sMass.init_vector(u, dim)
