def F_block(self, snes, x, F): assert x.getType() != "nest" assert F.getType() != "nest" x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) with F.localForm() as f_local: f_local.set(0.0) offset = 0 x_array = x.getArray(readonly=True) for var in self.soln_vars: size_local = var.vector.getLocalSize() var.vector.array[:] = x_array[offset:offset + size_local] var.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) offset += size_local assemble_vector_block(F, self.L, self.a, bcs=self.bcs, x0=x, scale=-1.0)
def blocked_solve(): """Blocked (monolithic) solver""" A = assemble_matrix_block(form([[a00, a01], [a10, a11]]), bcs=[bc0, bc1]) A.assemble() P = assemble_matrix_block(form([[p00, p01], [p10, p11]]), bcs=[bc0, bc1]) P.assemble() b = assemble_vector_block(form([L0, L1]), form([[a00, a01], [a10, a11]]), bcs=[bc0, bc1]) ksp = PETSc.KSP() ksp.create(mesh.comm) ksp.setOperators(A, P) ksp.setType("minres") pc = ksp.getPC() pc.setType('lu') ksp.setTolerances(rtol=1.0e-8, max_it=50) ksp.setFromOptions() x = A.createVecRight() ksp.solve(b, x) assert ksp.getConvergedReason() > 0 return b.norm(), x.norm(), A.norm(), P.norm()
def solve_thc_system(J, F, intern_var0, intern_var1, w0, w1, bcs, rtol=1.e-6, atol=1.e-12, max_its=10): """Solve system for temperature-humidity-co2""" rank = MPI.COMM_WORLD.rank # Preallocate diagonal block matrices t0 = time() A = create_matrix_block(J) if rank == 0: logger.info(f"[Timer] Preallocation matrix time: {time() - t0:.3f}") t0 = time() t0 = time() b = create_vector_block(F) if rank == 0: logger.info(f"[Timer] Preallocation vector time: {time() - t0:.3f}") assemble_vector_block(b, F, J, bcs) resid_norm0 = b.norm() if rank == 0: print(f"Initial THC residual: {resid_norm0:.2e}") # Newton iteration for THC blocks for iter in range(max_its): if rank == 0: logger.info(f"Newton iteration for THC {iter}") if iter > 0: for bc in bcs: with bc.value.vector.localForm() as locvec: locvec.set(0.0) A.zeroEntries() with b.localForm() as local: local.set(0.0) size = A.getSize()[0] local_size = A.getLocalSize()[0] t0 = time() assemble_matrix_block(A, J, bcs) A.assemble() _time = time() - t0 if rank == 0: logger.info(f"[Timer] A1 size: {size}, local size: {local_size}") logger.info(f"[Timer] A1 assembly: {_time:.4f}") logger.info(f"[Timer] A1 assembly dofs/s: {size / _time:.1f}") t0 = time() with b.localForm() as local: local.set(0.0) assemble_vector_block(b, F, J, bcs) _time = time() - t0 if rank == 0: logger.info(f"[Timer] b1 assembly: {_time}") logger.info(f"[Timer] b1 assembly dofs/s: {size / _time}") t0 = time() ksp = PETSc.KSP() ksp.create(MPI.COMM_WORLD) ksp.setOptionsPrefix("thc") ksp.setFromOptions() x = A.createVecLeft() ksp.setOperators(A) ksp.solve(b, x) if rank == 0: its = ksp.its t1 = time() - t0 dofsps = int(size / t1) logger.info(f"[Timer] A1 converged in: {its}") logger.info(f"[Timer] A1 solve {t1:1.3}") logger.info(f"[Timer] A1 solver dofs/s: {dofsps}") # Update all subfunctions u = [w1["temp"], w1["phi"], w1["co2"]] offset = 0 for i in range(3): size_local = u[i].vector.getLocalSize() u[i].vector.array[:] += x.array_r[offset:offset + size_local] offset += size_local u[i].vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) with b.localForm() as local: local.set(0.0) assemble_vector_block(b, F, J, bcs) norm = b.norm() if resid_norm0 == 0.0: rel_resid_norm = 0.0 else: rel_resid_norm = norm / resid_norm0 if rank == 0: logger.info("---") logger.info(f"Abs. resid norm: {norm:.2e}") logger.info(f"Rel. resid norm: {rel_resid_norm:.2e}") logger.info("---") if rel_resid_norm < rtol or norm < atol: break
def test_assembly_solve_block(mode): """Solve a two-field mass-matrix like problem with block matrix approaches and test that solution is the same""" mesh = create_unit_square(MPI.COMM_WORLD, 32, 31, ghost_mode=mode) P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1) V0 = FunctionSpace(mesh, P) V1 = V0.clone() # Locate facets on boundary facetdim = mesh.topology.dim - 1 bndry_facets = locate_entities_boundary(mesh, facetdim, lambda x: np.logical_or(np.isclose(x[0], 0.0), np.isclose(x[0], 1.0))) bdofsV0 = locate_dofs_topological(V0, facetdim, bndry_facets) bdofsV1 = locate_dofs_topological(V1, facetdim, bndry_facets) u0_bc = PETSc.ScalarType(50.0) u1_bc = PETSc.ScalarType(20.0) bcs = [dirichletbc(u0_bc, bdofsV0, V0), dirichletbc(u1_bc, bdofsV1, V1)] # Variational problem u, p = ufl.TrialFunction(V0), ufl.TrialFunction(V1) v, q = ufl.TestFunction(V0), ufl.TestFunction(V1) f = 1.0 g = -3.0 zero = Function(V0) a00 = form(inner(u, v) * dx) a01 = form(zero * inner(p, v) * dx) a10 = form(zero * inner(u, q) * dx) a11 = form(inner(p, q) * dx) L0 = form(inner(f, v) * dx) L1 = form(inner(g, q) * dx) def monitor(ksp, its, rnorm): pass # print("Norm:", its, rnorm) A0 = assemble_matrix_block([[a00, a01], [a10, a11]], bcs=bcs) b0 = assemble_vector_block([L0, L1], [[a00, a01], [a10, a11]], bcs=bcs) A0.assemble() A0norm = A0.norm() b0norm = b0.norm() x0 = A0.createVecLeft() ksp = PETSc.KSP() ksp.create(mesh.comm) ksp.setOperators(A0) ksp.setMonitor(monitor) ksp.setType('cg') ksp.setTolerances(rtol=1.0e-14) ksp.setFromOptions() ksp.solve(b0, x0) x0norm = x0.norm() # Nested (MatNest) A1 = assemble_matrix_nest([[a00, a01], [a10, a11]], bcs=bcs, diagonal=1.0) A1.assemble() b1 = assemble_vector_nest([L0, L1]) apply_lifting_nest(b1, [[a00, a01], [a10, a11]], bcs=bcs) for b_sub in b1.getNestSubVecs(): b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) bcs0 = bcs_by_block([L0.function_spaces[0], L1.function_spaces[0]], bcs) set_bc_nest(b1, bcs0) b1.assemble() b1norm = b1.norm() assert b1norm == pytest.approx(b0norm, 1.0e-12) A1norm = nest_matrix_norm(A1) assert A0norm == pytest.approx(A1norm, 1.0e-12) x1 = b1.copy() ksp = PETSc.KSP() ksp.create(mesh.comm) ksp.setMonitor(monitor) ksp.setOperators(A1) ksp.setType('cg') ksp.setTolerances(rtol=1.0e-12) ksp.setFromOptions() ksp.solve(b1, x1) x1norm = x1.norm() assert x1norm == pytest.approx(x0norm, rel=1.0e-12) # Monolithic version E = P * P W = FunctionSpace(mesh, E) u0, u1 = ufl.TrialFunctions(W) v0, v1 = ufl.TestFunctions(W) a = inner(u0, v0) * dx + inner(u1, v1) * dx L = inner(f, v0) * ufl.dx + inner(g, v1) * dx a, L = form(a), form(L) bdofsW0_V0 = locate_dofs_topological(W.sub(0), facetdim, bndry_facets) bdofsW1_V1 = locate_dofs_topological(W.sub(1), facetdim, bndry_facets) bcs = [dirichletbc(u0_bc, bdofsW0_V0, W.sub(0)), dirichletbc(u1_bc, bdofsW1_V1, W.sub(1))] A2 = assemble_matrix(a, bcs=bcs) A2.assemble() b2 = assemble_vector(L) apply_lifting(b2, [a], [bcs]) b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(b2, bcs) A2norm = A2.norm() b2norm = b2.norm() assert A2norm == pytest.approx(A0norm, 1.0e-12) assert b2norm == pytest.approx(b0norm, 1.0e-12) x2 = b2.copy() ksp = PETSc.KSP() ksp.create(mesh.comm) ksp.setMonitor(monitor) ksp.setOperators(A2) ksp.setType('cg') ksp.getPC().setType('jacobi') ksp.setTolerances(rtol=1.0e-12) ksp.setFromOptions() ksp.solve(b2, x2) x2norm = x2.norm() assert x2norm == pytest.approx(x0norm, 1.0e-10)
def test_matrix_assembly_block(mode): """Test assembly of block matrices and vectors into (a) monolithic blocked structures, PETSc Nest structures, and monolithic structures""" mesh = create_unit_square(MPI.COMM_WORLD, 4, 8, ghost_mode=mode) p0, p1 = 1, 2 P0 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), p0) P1 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), p1) V0 = FunctionSpace(mesh, P0) V1 = FunctionSpace(mesh, P1) # Locate facets on boundary facetdim = mesh.topology.dim - 1 bndry_facets = locate_entities_boundary(mesh, facetdim, lambda x: np.logical_or(np.isclose(x[0], 0.0), np.isclose(x[0], 1.0))) bdofsV1 = locate_dofs_topological(V1, facetdim, bndry_facets) u_bc = PETSc.ScalarType(50.0) bc = dirichletbc(u_bc, bdofsV1, V1) # Define variational problem u, p = ufl.TrialFunction(V0), ufl.TrialFunction(V1) v, q = ufl.TestFunction(V0), ufl.TestFunction(V1) f = 1.0 g = -3.0 zero = Function(V0) a00 = inner(u, v) * dx a01 = inner(p, v) * dx a10 = inner(u, q) * dx a11 = inner(p, q) * dx L0 = zero * inner(f, v) * dx L1 = inner(g, q) * dx a_block = form([[a00, a01], [a10, a11]]) L_block = form([L0, L1]) # Monolithic blocked A0 = assemble_matrix_block(a_block, bcs=[bc]) A0.assemble() b0 = assemble_vector_block(L_block, a_block, bcs=[bc]) assert A0.getType() != "nest" Anorm0 = A0.norm() bnorm0 = b0.norm() # Nested (MatNest) A1 = assemble_matrix_nest(a_block, bcs=[bc], mat_types=[["baij", "aij"], ["aij", ""]]) A1.assemble() Anorm1 = nest_matrix_norm(A1) assert Anorm0 == pytest.approx(Anorm1, 1.0e-12) b1 = assemble_vector_nest(L_block) apply_lifting_nest(b1, a_block, bcs=[bc]) for b_sub in b1.getNestSubVecs(): b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) bcs0 = bcs_by_block([L.function_spaces[0] for L in L_block], [bc]) set_bc_nest(b1, bcs0) b1.assemble() bnorm1 = math.sqrt(sum([x.norm()**2 for x in b1.getNestSubVecs()])) assert bnorm0 == pytest.approx(bnorm1, 1.0e-12) # Monolithic version E = P0 * P1 W = FunctionSpace(mesh, E) u0, u1 = ufl.TrialFunctions(W) v0, v1 = ufl.TestFunctions(W) a = inner(u0, v0) * dx + inner(u1, v1) * dx + inner(u0, v1) * dx + inner( u1, v0) * dx L = zero * inner(f, v0) * ufl.dx + inner(g, v1) * dx a, L = form(a), form(L) bdofsW_V1 = locate_dofs_topological(W.sub(1), mesh.topology.dim - 1, bndry_facets) bc = dirichletbc(u_bc, bdofsW_V1, W.sub(1)) A2 = assemble_matrix(a, bcs=[bc]) A2.assemble() b2 = assemble_vector(L) apply_lifting(b2, [a], bcs=[[bc]]) b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(b2, [bc]) assert A2.getType() != "nest" assert A2.norm() == pytest.approx(Anorm0, 1.0e-9) assert b2.norm() == pytest.approx(bnorm0, 1.0e-9)
def test_matrix_assembly_block_nl(): """Test assembly of block matrices and vectors into (a) monolithic blocked structures, PETSc Nest structures, and monolithic structures in the nonlinear setting """ mesh = create_unit_square(MPI.COMM_WORLD, 4, 8) p0, p1 = 1, 2 P0 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), p0) P1 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), p1) V0 = FunctionSpace(mesh, P0) V1 = FunctionSpace(mesh, P1) def initial_guess_u(x): return np.sin(x[0]) * np.sin(x[1]) def initial_guess_p(x): return -x[0]**2 - x[1]**3 def bc_value(x): return np.cos(x[0]) * np.cos(x[1]) facetdim = mesh.topology.dim - 1 bndry_facets = locate_entities_boundary( mesh, facetdim, lambda x: np.logical_or(np.isclose(x[0], 0.0), np.isclose(x[0], 1.0))) u_bc = Function(V1) u_bc.interpolate(bc_value) bdofs = locate_dofs_topological(V1, facetdim, bndry_facets) bc = dirichletbc(u_bc, bdofs) # Define variational problem du, dp = ufl.TrialFunction(V0), ufl.TrialFunction(V1) u, p = Function(V0), Function(V1) v, q = ufl.TestFunction(V0), ufl.TestFunction(V1) u.interpolate(initial_guess_u) p.interpolate(initial_guess_p) f = 1.0 g = -3.0 F0 = inner(u, v) * dx + inner(p, v) * dx - inner(f, v) * dx F1 = inner(u, q) * dx + inner(p, q) * dx - inner(g, q) * dx a_block = form([[derivative(F0, u, du), derivative(F0, p, dp)], [derivative(F1, u, du), derivative(F1, p, dp)]]) L_block = form([F0, F1]) # Monolithic blocked x0 = create_vector_block(L_block) scatter_local_vectors(x0, [u.vector.array_r, p.vector.array_r], [(u.function_space.dofmap.index_map, u.function_space.dofmap.index_map_bs), (p.function_space.dofmap.index_map, p.function_space.dofmap.index_map_bs)]) x0.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # Ghosts are updated inside assemble_vector_block A0 = assemble_matrix_block(a_block, bcs=[bc]) b0 = assemble_vector_block(L_block, a_block, bcs=[bc], x0=x0, scale=-1.0) A0.assemble() assert A0.getType() != "nest" Anorm0 = A0.norm() bnorm0 = b0.norm() # Nested (MatNest) x1 = create_vector_nest(L_block) for x1_soln_pair in zip(x1.getNestSubVecs(), (u, p)): x1_sub, soln_sub = x1_soln_pair soln_sub.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) soln_sub.vector.copy(result=x1_sub) x1_sub.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) A1 = assemble_matrix_nest(a_block, bcs=[bc]) b1 = assemble_vector_nest(L_block) apply_lifting_nest(b1, a_block, bcs=[bc], x0=x1, scale=-1.0) for b_sub in b1.getNestSubVecs(): b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) bcs0 = bcs_by_block([L.function_spaces[0] for L in L_block], [bc]) set_bc_nest(b1, bcs0, x1, scale=-1.0) A1.assemble() assert A1.getType() == "nest" assert nest_matrix_norm(A1) == pytest.approx(Anorm0, 1.0e-12) assert b1.norm() == pytest.approx(bnorm0, 1.0e-12) # Monolithic version E = P0 * P1 W = FunctionSpace(mesh, E) dU = ufl.TrialFunction(W) U = Function(W) u0, u1 = ufl.split(U) v0, v1 = ufl.TestFunctions(W) U.sub(0).interpolate(initial_guess_u) U.sub(1).interpolate(initial_guess_p) F = inner(u0, v0) * dx + inner(u1, v0) * dx + inner(u0, v1) * dx + inner(u1, v1) * dx \ - inner(f, v0) * ufl.dx - inner(g, v1) * dx J = derivative(F, U, dU) F, J = form(F), form(J) bdofsW_V1 = locate_dofs_topological((W.sub(1), V1), facetdim, bndry_facets) bc = dirichletbc(u_bc, bdofsW_V1, W.sub(1)) A2 = assemble_matrix(J, bcs=[bc]) A2.assemble() b2 = assemble_vector(F) apply_lifting(b2, [J], bcs=[[bc]], x0=[U.vector], scale=-1.0) b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(b2, [bc], x0=U.vector, scale=-1.0) assert A2.getType() != "nest" assert A2.norm() == pytest.approx(Anorm0, 1.0e-12) assert b2.norm() == pytest.approx(bnorm0, 1.0e-12)