def assemble_vector(form: ufl.form.Form,
constraint: MultiPointConstraint,
b: _PETSc.Vec = None) -> _PETSc.Vec:
"""
Assemble a linear form into vector b with corresponding multi point constraint
Parameters
----------
form
The linear form
constraint
The multi point constraint
b
PETSc vector to assemble into (optional)
Returns
-------
PETSc.Vec
The assembled linear form
"""
if b is None:
b = _cpp.la.petsc.create_vector(
constraint.function_space.dofmap.index_map,
constraint.function_space.dofmap.index_map_bs)
t = Timer("~MPC: Assemble vector (C++)")
with b.localForm() as b_local:
b_local.set(0.0)
dolfinx_mpc.cpp.mpc.assemble_vector(b_local, form,
constraint._cpp_object)
t.stop()
return b
def apply_lifting(b: _PETSc.Vec,
form: List[_fem.FormMetaClass],
bcs: List[List[_fem.DirichletBCMetaClass]],
constraint: MultiPointConstraint,
x0: List[_PETSc.Vec] = [],
scale: float = 1.0):
"""
Apply lifting to vector b, i.e.
.. math::
b <- b - scale * K^T (A_j (g_j - x0_j))
Parameters
----------
b
PETSc vector to assemble into
form
The linear form
bcs
List of Dirichlet boundary conditions
constraint
The multi point constraint
x0
List of vectors
scale
Scaling for lifting
Returns
-------
PETSc.Vec
The assembled linear form
"""
t = Timer("~MPC: Apply lifting (C++)")
with contextlib.ExitStack() as stack:
x0 = [stack.enter_context(x.localForm()) for x in x0]
x0_r = [x.array_r for x in x0]
b_local = stack.enter_context(b.localForm())
dolfinx_mpc.cpp.mpc.apply_lifting(b_local.array_w, form, bcs, x0_r,
scale, constraint._cpp_object)
t.stop()
def mesh_3D_dolfin(theta=0,
ct=CellType.tetrahedron,
ext="tetrahedron",
num_refinements=0,
N0=5):
timer = Timer("Create mesh")
def find_plane_function(p0, p1, p2):
"""
Find plane function given three points:
http://www.nabla.hr/CG-LinesPlanesIn3DA3.htm
"""
v1 = np.array(p1) - np.array(p0)
v2 = np.array(p2) - np.array(p0)
n = np.cross(v1, v2)
D = -(n[0] * p0[0] + n[1] * p0[1] + n[2] * p0[2])
return lambda x: np.isclose(0, np.dot(n, x) + D)
def over_plane(p0, p1, p2):
"""
Returns function that checks if a point is over a plane defined
by the points p0, p1 and p2.
"""
v1 = np.array(p1) - np.array(p0)
v2 = np.array(p2) - np.array(p0)
n = np.cross(v1, v2)
D = -(n[0] * p0[0] + n[1] * p0[1] + n[2] * p0[2])
return lambda x: n[0] * x[0] + n[1] * x[1] + D > -n[2] * x[2]
tmp_mesh_name = "tmp_mesh.xdmf"
r_matrix = rotation_matrix([1 / np.sqrt(2), 1 / np.sqrt(2), 0], -theta)
if MPI.COMM_WORLD.rank == 0:
# Create two coarse meshes and merge them
mesh0 = create_unit_cube(MPI.COMM_SELF, N0, N0, N0, ct)
mesh0.geometry.x[:, 2] += 1
mesh1 = create_unit_cube(MPI.COMM_SELF, 2 * N0, 2 * N0, 2 * N0, ct)
tdim0 = mesh0.topology.dim
num_cells0 = mesh0.topology.index_map(tdim0).size_local
cells0 = entities_to_geometry(
mesh0, tdim0,
np.arange(num_cells0, dtype=np.int32).reshape((-1, 1)), False)
tdim1 = mesh1.topology.dim
num_cells1 = mesh1.topology.index_map(tdim1).size_local
cells1 = entities_to_geometry(
mesh1, tdim1,
np.arange(num_cells1, dtype=np.int32).reshape((-1, 1)), False)
cells1 += mesh0.geometry.x.shape[0]
# Concatenate points and cells
points = np.vstack([mesh0.geometry.x, mesh1.geometry.x])
cells = np.vstack([cells0, cells1])
cell = Cell(ext, geometric_dimension=points.shape[1])
domain = Mesh(VectorElement("Lagrange", cell, 1))
# Rotate mesh
points = np.dot(r_matrix, points.T).T
mesh = create_mesh(MPI.COMM_SELF, cells, points, domain)
with XDMFFile(MPI.COMM_SELF, tmp_mesh_name, "w") as xdmf:
xdmf.write_mesh(mesh)
MPI.COMM_WORLD.barrier()
with XDMFFile(MPI.COMM_WORLD, tmp_mesh_name, "r") as xdmf:
mesh = xdmf.read_mesh()
# Refine coarse mesh
for i in range(num_refinements):
mesh.topology.create_entities(mesh.topology.dim - 2)
mesh = refine(mesh, redistribute=True)
tdim = mesh.topology.dim
fdim = tdim - 1
# Find information about facets to be used in meshtags
bottom_points = np.dot(
r_matrix,
np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]]).T)
bottom = find_plane_function(bottom_points[:, 0], bottom_points[:, 1],
bottom_points[:, 2])
bottom_facets = locate_entities_boundary(mesh, fdim, bottom)
top_points = np.dot(
r_matrix,
np.array([[0, 0, 2], [1, 0, 2], [0, 1, 2], [1, 1, 2]]).T)
top = find_plane_function(top_points[:, 0], top_points[:, 1],
top_points[:, 2])
top_facets = locate_entities_boundary(mesh, fdim, top)
# Determine interface facets
if_points = np.dot(
r_matrix,
np.array([[0, 0, 1], [1, 0, 1], [0, 1, 1], [1, 1, 1]]).T)
interface = find_plane_function(if_points[:, 0], if_points[:, 1],
if_points[:, 2])
i_facets = locate_entities_boundary(mesh, fdim, interface)
mesh.topology.create_connectivity(fdim, tdim)
top_interface = []
bottom_interface = []
facet_to_cell = mesh.topology.connectivity(fdim, tdim)
num_cells = mesh.topology.index_map(tdim).size_local
# Find top and bottom interface facets
cell_midpoints = compute_midpoints(mesh, tdim, range(num_cells))
top_cube = over_plane(if_points[:, 0], if_points[:, 1], if_points[:, 2])
for facet in i_facets:
i_cells = facet_to_cell.links(facet)
assert (len(i_cells == 1))
i_cell = i_cells[0]
if top_cube(cell_midpoints[i_cell]):
top_interface.append(facet)
else:
bottom_interface.append(facet)
# Create cell tags
num_cells = mesh.topology.index_map(tdim).size_local
cell_midpoints = compute_midpoints(mesh, tdim, range(num_cells))
top_cube_marker = 2
indices = []
values = []
for cell_index in range(num_cells):
if top_cube(cell_midpoints[cell_index]):
indices.append(cell_index)
values.append(top_cube_marker)
ct = meshtags(mesh, tdim, np.array(indices, dtype=np.intc),
np.array(values, dtype=np.intc))
# Create meshtags for facet data
markers = {
3: top_facets,
4: bottom_interface,
9: top_interface,
5: bottom_facets
} # , 6: left_facets, 7: right_facets}
indices = np.array([], dtype=np.intc)
values = np.array([], dtype=np.intc)
for key in markers.keys():
indices = np.append(indices, markers[key])
values = np.append(values,
np.full(len(markers[key]), key, dtype=np.intc))
sorted_indices = np.argsort(indices)
mt = meshtags(mesh, fdim, indices[sorted_indices], values[sorted_indices])
mt.name = "facet_tags"
fname = f"meshes/mesh_{ext}_{theta:.2f}.xdmf"
with XDMFFile(MPI.COMM_WORLD, fname, "w") as o_f:
o_f.write_mesh(mesh)
o_f.write_meshtags(ct)
o_f.write_meshtags(mt)
timer.stop()
def assemble_vector(form: _forms, constraint: MultiPointConstraint, b: _PETSc.Vec = None) -> _PETSc.Vec:
"""
Assemble a compiled DOLFINx form into vector b.
Parameters
----------
form
The complied linear form
constraint
The multi point constraint
b
PETSc vector to assemble into (optional)
"""
_log.log(_log.LogLevel.INFO, "Assemble MPC vector")
timer_vector = Timer("~MPC: Assemble vector (numba)")
# Unpack Function space data
V = form.function_spaces[0]
pos = V.mesh.geometry.dofmap.offsets
x_dofs = V.mesh.geometry.dofmap.array
x = V.mesh.geometry.x
dofs = V.dofmap.list().array
block_size = V.dofmap.index_map_bs
# Data from multipointconstraint
coefficients = constraint.coefficients()[0]
masters_adj = constraint.masters
c_to_s_adj = constraint.cell_to_slaves
cell_to_slave = c_to_s_adj.array
c_to_s_off = c_to_s_adj.offsets
is_slave = constraint.is_slave
mpc_data = (masters_adj.array, coefficients, masters_adj.offsets, cell_to_slave, c_to_s_off, is_slave)
slave_cells = extract_slave_cells(c_to_s_off)
# Get index map and ghost info
if b is None:
index_map = constraint.function_space.dofmap.index_map
vector = _cpp.la.petsc.create_vector(index_map, block_size)
else:
vector = b
# Pack constants and coefficients
form_coeffs = _cpp.fem.pack_coefficients(form)
form_consts = _cpp.fem.pack_constants(form)
tdim = V.mesh.topology.dim
num_dofs_per_element = V.dofmap.dof_layout.num_dofs
# Assemble vector with all entries
with vector.localForm() as b_local:
_cpp.fem.assemble_vector(b_local.array_w, form, form_consts, form_coeffs)
# Check if we need facet permutations
# FIXME: access apply_dof_transformations here
e0 = form.function_spaces[0].element
needs_transformation_data = e0.needs_dof_transformations or form.needs_facet_permutations
cell_perms = numpy.array([], dtype=numpy.uint32)
if needs_transformation_data:
V.mesh.topology.create_entity_permutations()
cell_perms = V.mesh.topology.get_cell_permutation_info()
if e0.needs_dof_transformations:
raise NotImplementedError("Dof transformations not implemented")
# Assemble over cells
subdomain_ids = form.integral_ids(_fem.IntegralType.cell)
num_cell_integrals = len(subdomain_ids)
is_complex = numpy.issubdtype(_PETSc.ScalarType, numpy.complexfloating)
nptype = "complex128" if is_complex else "float64"
ufcx_form = form.ufcx_form
if num_cell_integrals > 0:
V.mesh.topology.create_entity_permutations()
for i, id in enumerate(subdomain_ids):
cell_kernel = getattr(ufcx_form.integrals(_fem.IntegralType.cell)[i], f"tabulate_tensor_{nptype}")
active_cells = form.domains(_fem.IntegralType.cell, id)
coeffs_i = form_coeffs[(_fem.IntegralType.cell, id)]
with vector.localForm() as b:
assemble_cells(numpy.asarray(b), cell_kernel, active_cells[numpy.isin(active_cells, slave_cells)],
(pos, x_dofs, x), coeffs_i, form_consts,
cell_perms, dofs, block_size, num_dofs_per_element, mpc_data)
# Assemble exterior facet integrals
subdomain_ids = form.integral_ids(_fem.IntegralType.exterior_facet)
num_exterior_integrals = len(subdomain_ids)
if num_exterior_integrals > 0:
V.mesh.topology.create_entities(tdim - 1)
V.mesh.topology.create_connectivity(tdim - 1, tdim)
# Get facet permutations if required
facet_perms = numpy.array([], dtype=numpy.uint8)
if form.needs_facet_permutations:
facet_perms = V.mesh.topology.get_facet_permutations()
perm = (cell_perms, form.needs_facet_permutations, facet_perms)
for i, id in enumerate(subdomain_ids):
facet_kernel = getattr(ufcx_form.integrals(_fem.IntegralType.exterior_facet)[i],
f"tabulate_tensor_{nptype}")
coeffs_i = form_coeffs[(_fem.IntegralType.exterior_facet, id)]
facets = form.domains(_fem.IntegralType.exterior_facet, id)
facet_info = pack_slave_facet_info(facets, slave_cells)
num_facets_per_cell = len(V.mesh.topology.connectivity(tdim, tdim - 1).links(0))
with vector.localForm() as b:
assemble_exterior_slave_facets(numpy.asarray(b), facet_kernel, facet_info, (pos, x_dofs, x),
coeffs_i, form_consts, perm,
dofs, block_size, num_dofs_per_element, mpc_data, num_facets_per_cell)
timer_vector.stop()
return vector
def assemble_matrix(form: _forms, constraint: MultiPointConstraint,
bcs: List[_bcs] = None, diagval: _PETSc.ScalarType = 1.,
A: _PETSc.Mat = None):
"""
Assembles a compiled DOLFINx form with given a multi point constraint and possible
Dirichlet boundary conditions.
NOTE: Strong Dirichlet conditions cannot be on master dofs.
Parameters
----------
form
The compiled bilinear form
constraint
The multi point constraint
bcs
List of Dirichlet boundary conditions
diagval
Value to set on the diagonal of the matrix(Default 1)
A
PETSc matrix to assemble into (optional)
"""
timer_matrix = Timer("~MPC: Assemble matrix (numba)")
V = constraint.function_space
dofmap = V.dofmap
dofs = dofmap.list.array
# Pack MPC data for numba kernels
coefficients = constraint.coefficients()[0]
masters_adj = constraint.masters
c_to_s_adj = constraint.cell_to_slaves
cell_to_slave = c_to_s_adj.array
c_to_s_off = c_to_s_adj.offsets
is_slave = constraint.is_slave
mpc_data = (masters_adj.array, coefficients, masters_adj.offsets, cell_to_slave, c_to_s_off, is_slave)
slave_cells = extract_slave_cells(c_to_s_off)
# Create 1D bc indicator for matrix assembly
num_dofs_local = (dofmap.index_map.size_local + dofmap.index_map.num_ghosts) * dofmap.index_map_bs
is_bc = numpy.zeros(num_dofs_local, dtype=bool)
bcs = [] if bcs is None else bcs
if len(bcs) > 0:
for bc in bcs:
is_bc[bc.dof_indices()[0]] = True
# Get data from mesh
pos = V.mesh.geometry.dofmap.offsets
x_dofs = V.mesh.geometry.dofmap.array
x = V.mesh.geometry.x
# Pack constants and coefficients
form_coeffs = _cpp.fem.pack_coefficients(form)
form_consts = _cpp.fem.pack_constants(form)
# Create sparsity pattern and matrix if not supplied
if A is None:
pattern = create_sparsity_pattern(form, constraint)
pattern.assemble()
A = _cpp.la.petsc.create_matrix(V.mesh.comm, pattern)
A.zeroEntries()
# Assemble the matrix with all entries
_cpp.fem.petsc.assemble_matrix(A, form, form_consts, form_coeffs, bcs, False)
# General assembly data
block_size = dofmap.dof_layout.block_size
num_dofs_per_element = dofmap.dof_layout.num_dofs
tdim = V.mesh.topology.dim
# Assemble over cells
subdomain_ids = form.integral_ids(_fem.IntegralType.cell)
num_cell_integrals = len(subdomain_ids)
e0 = form.function_spaces[0].element
e1 = form.function_spaces[1].element
# Get dof transformations
needs_transformation_data = e0.needs_dof_transformations or e1.needs_dof_transformations or \
form.needs_facet_permutations
cell_perms = numpy.array([], dtype=numpy.uint32)
if needs_transformation_data:
V.mesh.topology.create_entity_permutations()
cell_perms = V.mesh.topology.get_cell_permutation_info()
# NOTE: Here we need to add the apply_dof_transformation and apply_dof_transformation transpose functions
# to support more exotic elements
if e0.needs_dof_transformations or e1.needs_dof_transformations:
raise NotImplementedError("Dof transformations not implemented")
is_complex = numpy.issubdtype(_PETSc.ScalarType, numpy.complexfloating)
nptype = "complex128" if is_complex else "float64"
ufcx_form = form.ufcx_form
if num_cell_integrals > 0:
V.mesh.topology.create_entity_permutations()
for i, id in enumerate(subdomain_ids):
coeffs_i = form_coeffs[(_fem.IntegralType.cell, id)]
cell_kernel = getattr(ufcx_form.integrals(_fem.IntegralType.cell)[i], f"tabulate_tensor_{nptype}")
active_cells = form.domains(_fem.IntegralType.cell, id)
assemble_slave_cells(A.handle, cell_kernel, active_cells[numpy.isin(active_cells, slave_cells)],
(pos, x_dofs, x), coeffs_i, form_consts, cell_perms, dofs,
block_size, num_dofs_per_element, mpc_data, is_bc)
# Assemble over exterior facets
subdomain_ids = form.integral_ids(_fem.IntegralType.exterior_facet)
num_exterior_integrals = len(subdomain_ids)
if num_exterior_integrals > 0:
V.mesh.topology.create_entities(tdim - 1)
V.mesh.topology.create_connectivity(tdim - 1, tdim)
# Get facet permutations if required
facet_perms = numpy.array([], dtype=numpy.uint8)
if form.needs_facet_permutations:
facet_perms = V.mesh.topology.get_facet_permutations()
perm = (cell_perms, form.needs_facet_permutations, facet_perms)
for i, id in enumerate(subdomain_ids):
facet_kernel = getattr(ufcx_form.integrals(_fem.IntegralType.exterior_facet)
[i], f"tabulate_tensor_{nptype}")
facets = form.domains(_fem.IntegralType.exterior_facet, id)
coeffs_i = form_coeffs[(_fem.IntegralType.exterior_facet, id)]
facet_info = pack_slave_facet_info(facets, slave_cells)
num_facets_per_cell = len(V.mesh.topology.connectivity(tdim, tdim - 1).links(0))
assemble_exterior_slave_facets(A.handle, facet_kernel, (pos, x_dofs, x), coeffs_i, form_consts,
perm, dofs, block_size, num_dofs_per_element, facet_info, mpc_data, is_bc,
num_facets_per_cell)
# Add mpc entries on diagonal
slaves = constraint.slaves
num_local_slaves = constraint.num_local_slaves
add_diagonal(A.handle, slaves[:num_local_slaves], diagval)
# Add one on diagonal for diriclet bc and slave dofs
# NOTE: In the future one could use a constant in the dirichletbc
if form.function_spaces[0] is form.function_spaces[1]:
A.assemblyBegin(_PETSc.Mat.AssemblyType.FLUSH)
A.assemblyEnd(_PETSc.Mat.AssemblyType.FLUSH)
_cpp.fem.petsc.insert_diagonal(A, form.function_spaces[0], bcs, diagval)
A.assemble()
timer_matrix.stop()
return A