def test_nest_matrix(pushpop_parameters): # Create Matrices and insert into nest A00 = PETScMatrix() A01 = PETScMatrix() A10 = PETScMatrix() mesh = UnitSquareMesh(12, 12) V = FunctionSpace(mesh, "Lagrange", 2) Q = FunctionSpace(mesh, "Lagrange", 1) u, v = TrialFunction(V), TestFunction(V) p, q = TrialFunction(Q), TestFunction(Q) assemble(u * v * dx, tensor=A00) assemble(p * v * dx, tensor=A01) assemble(u * q * dx, tensor=A10) AA = PETScNestMatrix([A00, A01, A10, None]) # Create compatible RHS Vectors and insert into nest u = PETScVector() p = PETScVector() x = PETScVector() A00.init_vector(u, 1) A01.init_vector(p, 1) AA.init_vectors(x, [u, p])
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)