Test case for dolfinx_mpc periodic condition bug

main.py

from mpi4py import MPI
from petsc4py import PETSc
from dolfinx.mesh import CellType, create_unit_square
from dolfinx.fem import Function, FunctionSpace
import ufl
from dolfinx import fem, io
from ufl import dx, grad, inner, lhs, rhs
from dolfinx import cpp as _cpp
import numpy as np

from dolfinx_mpc import MultiPointConstraint
from dolfinx_mpc.assemble_matrix import assemble_matrix, create_sparsity_pattern
from dolfinx_mpc.assemble_vector import apply_lifting, assemble_vector

msh = create_unit_square(MPI.COMM_WORLD, 20, 20, CellType.triangle)

P1 = ufl.FiniteElement("Lagrange", msh.ufl_cell(), 1)

space = FunctionSpace(msh, P1)

comm = msh.comm

D = 1.0
dt = 0.0002

def u_func(x):
    mask = np.logical_and( np.logical_and( x[0] < 0.5001, x[0] >= -0.1 ), np.logical_and( x[1] < 0.5001, x[1] >= -0.1 ))
    mask2 = np.logical_not( mask )
    res = np.ndarray(x.shape[1])
    res[mask] = 1.0
    res[mask2] = 0.0
    return res

def periodic_boundary(x):
    return np.isclose(x[0], 1.0)

def periodic_relation(x):
    out_x = np.zeros(x.shape)
    out_x[0] = 1.0 - x[0]
    # out_x[0] = x[0] - 5.0e-5
    out_x[1] = x[1]
    out_x[2] = x[2]
    return out_x

mpc = MultiPointConstraint(space)
mpc.create_periodic_constraint_geometrical(space, periodic_boundary, periodic_relation, [], 1)
mpc.finalize()

# print(mpc.masters)
# print(mpc.num_local_slaves)

u_trial = ufl.TrialFunction(mpc.function_space)
u_test = ufl.TestFunction(mpc.function_space)

u = Function(mpc.function_space)
u0 = Function(mpc.function_space)
u.interpolate(u_func)
u0.interpolate(u_func)

form = u_trial / dt * u_test * dx - u0 / dt * u_test * dx + D * inner( grad(u_trial), grad(u_test) ) * dx

fem_form_lhs = fem.form(lhs(form))
fem_form_rhs = fem.form(rhs(form))

pattern = create_sparsity_pattern(fem_form_lhs, mpc)
pattern.assemble()

A = _cpp.la.petsc.create_matrix(comm, pattern)
b = _cpp.la.petsc.create_vector(mpc.function_space.dofmap.index_map,
                                mpc.function_space.dofmap.index_map_bs)

opts = PETSc.Options()
ksp = PETSc.KSP().create(comm)
ksp.setOperators(A)
opts["ksp_error_if_not_converged"] = True
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType(PETSc.Mat.SolverType.MUMPS)
ksp.setFromOptions()

u_file = io.XDMFFile(MPI.COMM_WORLD, "u.xdmf", "w")
u_file.write_mesh(msh)

u_file.close()

for i in range(20):
    # print(self.mpc.masters)
    A.zeroEntries()
    with b.localForm() as b_local:
        b_local.set(0.0)
    assemble_matrix( fem_form_lhs, mpc, bcs=[], A=A )
    A.assemble()
    assemble_vector( fem_form_rhs, mpc, b = b )
    apply_lifting(b, [fem_form_lhs], bcs=[[]], constraint=mpc, x0=[u.vector], scale=1.0)
    b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    fem.petsc.set_bc(b, [], x0=u.vector, scale=1.0)
    ksp.setOperators(A)
    ksp.solve(b, u.vector)
    u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    mpc.backsubstitution(u.vector)
    u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    u0.x.array[:] = u.x.array
    u_file = io.XDMFFile(MPI.COMM_WORLD, "u.xdmf", "a")
    u_file.write_function(u, i)
    u_file.close()