def test_assemble_functional_dx(mode):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
M = 1.0 * dx(domain=mesh)
value = dolfinx.fem.assemble_scalar(M)
value = mesh.mpi_comm().allreduce(value, op=MPI.SUM)
assert value == pytest.approx(1.0, 1e-12)
x = ufl.SpatialCoordinate(mesh)
M = x[0] * dx(domain=mesh)
value = dolfinx.fem.assemble_scalar(M)
value = mesh.mpi_comm().allreduce(value, op=MPI.SUM)
assert value == pytest.approx(0.5, 1e-12)
def test_lambda_assembler():
"""Tests assembly with a lambda function
"""
mesh = UnitSquareMesh(MPI.COMM_WORLD, 5, 5)
V = fem.FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
a = inner(u, v) * dx
# Initial assembly
a_form = fem.Form(a)
rdata = []
cdata = []
vdata = []
def mat_insert(rows, cols, vals):
vdata.append(vals)
rdata.append(numpy.repeat(rows, len(cols)))
cdata.append(numpy.tile(cols, len(rows)))
return 0
dolfinx.cpp.fem.assemble_matrix(mat_insert, a_form._cpp_object, [])
vdata = numpy.array(vdata).flatten()
cdata = numpy.array(cdata).flatten()
rdata = numpy.array(rdata).flatten()
mat = scipy.sparse.coo_matrix((vdata, (rdata, cdata)))
v = numpy.ones(mat.shape[1])
s = MPI.COMM_WORLD.allreduce(mat.dot(v).sum(), MPI.SUM)
assert numpy.isclose(s, 1.0)
def test_basic_assembly_constant(mode):
"""Tests assembly with Constant
The following test should be sensitive to order of flattening the
matrix-valued constant.
"""
mesh = UnitSquareMesh(MPI.COMM_WORLD, 5, 5, ghost_mode=mode)
V = fem.FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
c = fem.Constant(mesh, [[1.0, 2.0], [5.0, 3.0]])
a = inner(c[1, 0] * u, v) * dx + inner(c[1, 0] * u, v) * ds
L = inner(c[1, 0], v) * dx + inner(c[1, 0], v) * ds
# Initial assembly
A1 = dolfinx.fem.assemble_matrix(a)
A1.assemble()
b1 = dolfinx.fem.assemble_vector(L)
b1.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
c.value = [[1.0, 2.0], [3.0, 4.0]]
A2 = dolfinx.fem.assemble_matrix(a)
A2.assemble()
b2 = dolfinx.fem.assemble_vector(L)
b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
assert (A1 * 3.0 - A2 * 5.0).norm() == pytest.approx(0.0)
assert (b1 * 3.0 - b2 * 5.0).norm() == pytest.approx(0.0)
def test_assemble_derivatives():
"""This test checks the original_coefficient_positions, which may change
under differentiation (some coefficients and constants are
eliminated)"""
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12)
Q = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
u = dolfinx.Function(Q)
v = ufl.TestFunction(Q)
du = ufl.TrialFunction(Q)
b = dolfinx.Function(Q)
c1 = fem.Constant(mesh, [[1.0, 0.0], [3.0, 4.0]])
c2 = fem.Constant(mesh, 2.0)
with b.vector.localForm() as b_local:
b_local.set(2.0)
# derivative eliminates 'u' and 'c1'
L = ufl.inner(c1, c1) * v * dx + c2 * b * inner(u, v) * dx
a = derivative(L, u, du)
A1 = dolfinx.fem.assemble_matrix(a)
A1.assemble()
a = c2 * b * inner(du, v) * dx
A2 = dolfinx.fem.assemble_matrix(a)
A2.assemble()
assert (A1 - A2).norm() == pytest.approx(0.0, rel=1e-12, abs=1e-12)
def test_ghost_mesh_assembly(mode, dx, ds):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
V = FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
dx = dx(mesh)
ds = ds(mesh)
f = Function(V)
with f.vector.localForm() as f_local:
f_local.set(10.0)
a = inner(f * u, v) * dx + inner(u, v) * ds
L = inner(f, v) * dx + inner(2.0, v) * ds
# Initial assembly
A = fem.assemble_matrix(a)
A.assemble()
assert isinstance(A, PETSc.Mat)
b = fem.assemble_vector(L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
assert isinstance(b, PETSc.Vec)
# Check that the norms are the same for all three modes
normA = A.norm()
assert normA == pytest.approx(0.6713621455570528, rel=1.e-6, abs=1.e-12)
normb = b.norm()
assert normb == pytest.approx(1.582294032953906, rel=1.e-6, abs=1.e-12)
def test_basic_assembly(mode):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
V = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
f = dolfinx.Function(V)
with f.vector.localForm() as f_local:
f_local.set(10.0)
a = inner(f * u, v) * dx + inner(u, v) * ds
L = inner(f, v) * dx + inner(2.0, v) * ds
# Initial assembly
A = dolfinx.fem.assemble_matrix(a)
A.assemble()
assert isinstance(A, PETSc.Mat)
b = dolfinx.fem.assemble_vector(L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
assert isinstance(b, PETSc.Vec)
# Second assembly
normA = A.norm()
A.zeroEntries()
A = dolfinx.fem.assemble_matrix(A, a)
A.assemble()
assert isinstance(A, PETSc.Mat)
assert normA == pytest.approx(A.norm())
normb = b.norm()
with b.localForm() as b_local:
b_local.set(0.0)
b = dolfinx.fem.assemble_vector(b, L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
assert isinstance(b, PETSc.Vec)
assert normb == pytest.approx(b.norm())
# Vector re-assembly - no zeroing (but need to zero ghost entries)
with b.localForm() as b_local:
b_local.array[b.local_size:] = 0.0
dolfinx.fem.assemble_vector(b, L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
assert 2.0 * normb == pytest.approx(b.norm())
# Matrix re-assembly (no zeroing)
dolfinx.fem.assemble_matrix(A, a)
A.assemble()
assert 2.0 * normA == pytest.approx(A.norm())
def test_ghost_mesh_dS_assembly(mode, dS):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
V = FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
dS = dS(mesh)
a = inner(avg(u), avg(v)) * dS
# Initial assembly
A = fem.assemble_matrix(a)
A.assemble()
assert isinstance(A, PETSc.Mat)
# Check that the norms are the same for all three modes
normA = A.norm()
print(normA)
assert normA == pytest.approx(2.1834054713561906, rel=1.e-6, abs=1.e-12)
def test_assembly_bcs(mode):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
V = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
a = inner(u, v) * dx + inner(u, v) * ds
L = inner(1.0, v) * dx
def boundary(x):
return numpy.logical_or(x[0] < 1.0e-6, x[0] > 1.0 - 1.0e-6)
bdofsV = dolfinx.fem.locate_dofs_geometrical(V, boundary)
u_bc = dolfinx.fem.Function(V)
with u_bc.vector.localForm() as u_local:
u_local.set(1.0)
bc = dolfinx.fem.dirichletbc.DirichletBC(u_bc, bdofsV)
# Assemble and apply 'global' lifting of bcs
A = dolfinx.fem.assemble_matrix(a)
A.assemble()
b = dolfinx.fem.assemble_vector(L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
g = b.duplicate()
with g.localForm() as g_local:
g_local.set(0.0)
dolfinx.fem.set_bc(g, [bc])
f = b - A * g
dolfinx.fem.set_bc(f, [bc])
# Assemble vector and apply lifting of bcs during assembly
b = dolfinx.fem.assemble_vector(L)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
b_bc = dolfinx.fem.assemble_vector(L)
dolfinx.fem.apply_lifting(b_bc, [a], [[bc]])
b_bc.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.set_bc(b_bc, [bc])
assert (f - b_bc).norm() == pytest.approx(0.0, rel=1e-12, abs=1e-12)
def test_eigen_assembly(mode):
"""Compare assembly into scipy.CSR matrix with PETSc assembly"""
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
Q = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
u = ufl.TrialFunction(Q)
v = ufl.TestFunction(Q)
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
def boundary(x):
return numpy.logical_or(x[0] < 1.0e-6, x[0] > 1.0 - 1.0e-6)
bdofsQ = dolfinx.fem.locate_dofs_geometrical(Q, boundary)
u_bc = dolfinx.function.Function(Q)
with u_bc.vector.localForm() as u_local:
u_local.set(1.0)
bc = dolfinx.fem.dirichletbc.DirichletBC(u_bc, bdofsQ)
A1 = dolfinx.fem.assemble_matrix(a, [bc])
A1.assemble()
A2 = dolfinx.fem.assemble_csr_matrix(a, [bc])
assert numpy.isclose(A1.norm(), scipy.sparse.linalg.norm(A2))
def test_assemble_functional_ds(mode):
mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
M = 1.0 * ds(domain=mesh)
value = dolfinx.fem.assemble_scalar(M)
value = mesh.mpi_comm().allreduce(value, op=MPI.SUM)
assert value == pytest.approx(4.0, 1e-12)
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 = UnitSquareMesh(MPI.COMM_WORLD, 32, 31, ghost_mode=mode)
P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
V0 = dolfinx.fem.FunctionSpace(mesh, P)
V1 = V0.clone()
def boundary(x):
return numpy.logical_or(x[0] < 1.0e-6, x[0] > 1.0 - 1.0e-6)
# Locate facets on boundary
facetdim = mesh.topology.dim - 1
bndry_facets = dolfinx.mesh.locate_entities_boundary(
mesh, facetdim, boundary)
bdofsV0 = dolfinx.fem.locate_dofs_topological(V0, facetdim, bndry_facets)
bdofsV1 = dolfinx.fem.locate_dofs_topological(V1, facetdim, bndry_facets)
u_bc0 = dolfinx.fem.Function(V0)
u_bc0.vector.set(50.0)
u_bc0.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
u_bc1 = dolfinx.fem.Function(V1)
u_bc1.vector.set(20.0)
u_bc1.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
bcs = [
dolfinx.fem.dirichletbc.DirichletBC(u_bc0, bdofsV0),
dolfinx.fem.dirichletbc.DirichletBC(u_bc1, bdofsV1)
]
# 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 = dolfinx.Function(V0)
a00 = inner(u, v) * dx
a01 = zero * inner(p, v) * dx
a10 = zero * inner(u, q) * dx
a11 = inner(p, q) * dx
L0 = inner(f, v) * dx
L1 = inner(g, q) * dx
def monitor(ksp, its, rnorm):
pass
# print("Norm:", its, rnorm)
A0 = dolfinx.fem.assemble_matrix_block([[a00, a01], [a10, a11]], bcs)
b0 = dolfinx.fem.assemble_vector_block([L0, L1], [[a00, a01], [a10, a11]],
bcs)
A0.assemble()
A0norm = A0.norm()
b0norm = b0.norm()
x0 = A0.createVecLeft()
ksp = PETSc.KSP()
ksp.create(mesh.mpi_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 = dolfinx.fem.assemble_matrix_nest([[a00, a01], [a10, a11]],
bcs,
diagonal=1.0)
A1.assemble()
b1 = dolfinx.fem.assemble_vector_nest([L0, L1])
dolfinx.fem.apply_lifting_nest(b1, [[a00, a01], [a10, a11]], bcs)
for b_sub in b1.getNestSubVecs():
b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
bcs0 = dolfinx.cpp.fem.bcs_rows(
dolfinx.fem.assemble._create_cpp_form([L0, L1]), bcs)
dolfinx.fem.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.mpi_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 = dolfinx.fem.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
u0_bc = dolfinx.fem.Function(V0)
u0_bc.vector.set(50.0)
u0_bc.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
u1_bc = dolfinx.fem.Function(V1)
u1_bc.vector.set(20.0)
u1_bc.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
bdofsW0_V0 = dolfinx.fem.locate_dofs_topological((W.sub(0), V0), facetdim,
bndry_facets)
bdofsW1_V1 = dolfinx.fem.locate_dofs_topological((W.sub(1), V1), facetdim,
bndry_facets)
bcs = [
dolfinx.fem.dirichletbc.DirichletBC(u0_bc, bdofsW0_V0, W.sub(0)),
dolfinx.fem.dirichletbc.DirichletBC(u1_bc, bdofsW1_V1, W.sub(1))
]
A2 = dolfinx.fem.assemble_matrix(a, bcs)
A2.assemble()
b2 = dolfinx.fem.assemble_vector(L)
dolfinx.fem.apply_lifting(b2, [a], [bcs])
b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.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.mpi_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 = UnitSquareMesh(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 = dolfinx.fem.FunctionSpace(mesh, P0)
V1 = dolfinx.fem.FunctionSpace(mesh, P1)
def boundary(x):
return numpy.logical_or(x[0] < 1.0e-6, x[0] > 1.0 - 1.0e-6)
# Locate facets on boundary
facetdim = mesh.topology.dim - 1
bndry_facets = dolfinx.mesh.locate_entities_boundary(
mesh, facetdim, boundary)
bdofsV1 = dolfinx.fem.locate_dofs_topological(V1, facetdim, bndry_facets)
u_bc = dolfinx.fem.Function(V1)
with u_bc.vector.localForm() as u_local:
u_local.set(50.0)
bc = dolfinx.fem.dirichletbc.DirichletBC(u_bc, bdofsV1)
# 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 = dolfinx.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 = [[a00, a01], [a10, a11]]
L_block = [L0, L1]
# Monolithic blocked
A0 = dolfinx.fem.assemble_matrix_block(a_block, [bc])
A0.assemble()
b0 = dolfinx.fem.assemble_vector_block(L_block, a_block, [bc])
assert A0.getType() != "nest"
Anorm0 = A0.norm()
bnorm0 = b0.norm()
# Nested (MatNest)
A1 = dolfinx.fem.assemble_matrix_nest(a_block, [bc],
[["baij", "aij"], ["aij", ""]])
A1.assemble()
Anorm1 = nest_matrix_norm(A1)
assert Anorm0 == pytest.approx(Anorm1, 1.0e-12)
b1 = dolfinx.fem.assemble_vector_nest(L_block)
dolfinx.fem.apply_lifting_nest(b1, a_block, [bc])
for b_sub in b1.getNestSubVecs():
b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
bcs0 = dolfinx.cpp.fem.bcs_rows(
dolfinx.fem.assemble._create_cpp_form(L_block), [bc])
dolfinx.fem.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 = dolfinx.fem.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
bdofsW_V1 = dolfinx.fem.locate_dofs_topological(
(W.sub(1), V1), mesh.topology.dim - 1, bndry_facets)
bc = dolfinx.fem.dirichletbc.DirichletBC(u_bc, bdofsW_V1, W.sub(1))
A2 = dolfinx.fem.assemble_matrix(a, [bc])
A2.assemble()
b2 = dolfinx.fem.assemble_vector(L)
dolfinx.fem.apply_lifting(b2, [a], [[bc]])
b2.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.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)
x2 = b2.copy()
ksp = PETSc.KSP()
ksp.create(mesh.mpi_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)
@pytest.mark.parametrize("mesh", [
UnitSquareMesh(
MPI.COMM_WORLD, 12, 11, ghost_mode=dolfinx.cpp.mesh.GhostMode.none),
UnitSquareMesh(MPI.COMM_WORLD,
12,
11,
ghost_mode=dolfinx.cpp.mesh.GhostMode.shared_facet),
UnitCubeMesh(
MPI.COMM_WORLD, 3, 7, 3, ghost_mode=dolfinx.cpp.mesh.GhostMode.none),
UnitCubeMesh(MPI.COMM_WORLD,
3,
7,
3,
ghost_mode=dolfinx.cpp.mesh.GhostMode.shared_facet)
])
def test_assembly_solve_taylor_hood(mesh):
"""Assemble Stokes problem with Taylor-Hood elements and solve."""
P2 = fem.VectorFunctionSpace(mesh, ("Lagrange", 2))
def test_eigen_assembly(tempdir): # noqa: F811
"""Compare assembly into scipy.CSR matrix with PETSc assembly"""
def compile_eigen_csr_assembler_module():
cpp_code_header = f"""
<%
setup_pybind11(cfg)
cfg['include_dirs'] = {dolfinx_pc["include_dirs"] + [petsc4py.get_include()] + [str(pybind_inc())]}
cfg['compiler_args'] = {["-D" + dm for dm in dolfinx_pc["define_macros"]]}
cfg['compiler_args'] = ['-std=c++17']
cfg['libraries'] = {dolfinx_pc["libraries"]}
cfg['library_dirs'] = {dolfinx_pc["library_dirs"]}
%>
"""
cpp_code = """
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl.h>
#include <vector>
#include <Eigen/Sparse>
#include <petscsys.h>
#include <dolfinx/fem/assembler.h>
#include <dolfinx/fem/DirichletBC.h>
#include <dolfinx/fem/Form.h>
template<typename T>
Eigen::SparseMatrix<T, Eigen::RowMajor>
assemble_csr(const dolfinx::fem::Form<T>& a,
const std::vector<std::shared_ptr<const dolfinx::fem::DirichletBC<T>>>& bcs)
{
std::vector<Eigen::Triplet<T>> triplets;
const auto mat_add
= [&triplets](std::int32_t nrow, const std::int32_t* rows,
std::int32_t ncol, const std::int32_t* cols, const T* v)
{
for (int i = 0; i < nrow; ++i)
for (int j = 0; j < ncol; ++j)
triplets.emplace_back(rows[i], cols[j], v[i * ncol + j]);
return 0;
};
dolfinx::fem::assemble_matrix<T>(mat_add, a, bcs);
auto map0 = a.function_space(0)->dofmap()->index_map;
auto map1 = a.function_space(1)->dofmap()->index_map;
Eigen::SparseMatrix<T, Eigen::RowMajor> mat(
map0->block_size() * (map0->size_local() + map0->num_ghosts()),
map1->block_size() * (map1->size_local() + map1->num_ghosts()));
mat.setFromTriplets(triplets.begin(), triplets.end());
return mat;
}
PYBIND11_MODULE(eigen_csr, m)
{
m.def("assemble_matrix", &assemble_csr<PetscScalar>);
}
"""
path = pathlib.Path(tempdir)
open(pathlib.Path(tempdir, "eigen_csr.cpp"),
"w").write(cpp_code + cpp_code_header)
rel_path = path.relative_to(pathlib.Path(__file__).parent)
p = str(rel_path).replace("/", ".") + ".eigen_csr"
return cppimport.imp(p)
def assemble_csr_matrix(a, bcs):
"""Assemble bilinear form into an SciPy CSR matrix, in serial."""
module = compile_eigen_csr_assembler_module()
_a = dolfinx.fem.assemble._create_cpp_form(a)
A = module.assemble_matrix(_a, bcs)
if _a.function_spaces[0].id == _a.function_spaces[1].id:
for bc in bcs:
if _a.function_spaces[0].contains(bc.function_space):
bc_dofs = bc.dof_indices[:, 0]
A[bc_dofs, bc_dofs] = 1.0
return A
mesh = UnitSquareMesh(MPI.COMM_SELF, 12, 12)
Q = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
u = ufl.TrialFunction(Q)
v = ufl.TestFunction(Q)
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
bdofsQ = dolfinx.fem.locate_dofs_geometrical(
Q, lambda x: numpy.logical_or(x[0] < 1.0e-6, x[0] > 1.0 - 1.0e-6))
u_bc = dolfinx.function.Function(Q)
with u_bc.vector.localForm() as u_local:
u_local.set(1.0)
bc = dolfinx.fem.dirichletbc.DirichletBC(u_bc, bdofsQ)
A1 = dolfinx.fem.assemble_matrix(a, [bc])
A1.assemble()
A2 = assemble_csr_matrix(a, [bc])
assert numpy.isclose(A1.norm(), scipy.sparse.linalg.norm(A2))
