Periodic BCs for DG in dolfinx

I wanted to apply a bi-periodic boundary condition for NS equation. I have used these two tutorials in writing my code.
Tutorials:
https://github.com/jorgensd/dolfinx_mpc/blob/v0.6.1.post1/python/demos/demo_periodic_gep.py
https://docs.fenicsproject.org/dolfinx/main/python/demos/demo_navier-stokes.html

Here is my code:

from mpi4py import MPI
from petsc4py import PETSc
import h5py
import dolfinx_mpc
# +
import numpy as np
from dolfinx import default_real_type, fem, io, mesh
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
from dolfinx.fem import Function, dirichletbc, form, functionspace, locate_dofs_geometrical
from dolfinx.fem.petsc import set_bc
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_unit_cube, locate_entities_boundary, meshtags
from dolfinx_mpc import MultiPointConstraint, apply_lifting, assemble_matrix, assemble_vector
from dolfinx_mpc.utils import log_info
from ufl import (
    CellDiameter,
    FacetNormal,
    TestFunction,
    TrialFunction,
    avg,
    conditional,
    div,
    dot,
    dS,
    ds,
    dx,
    grad,
    gt,
    inner,
    outer,
)

if np.issubdtype(PETSc.ScalarType, np.complexfloating):  # type: ignore
    print("Demo should only be executed with DOLFINx real mode")
    exit(0)
# -


# We also define some helper functions that will be used later


# +
def norm_L2(comm, v):
    """Compute the L2(Ω)-norm of v"""
    return np.sqrt(comm.allreduce(fem.assemble_scalar(fem.form(inner(v, v) * dx)), op=MPI.SUM))


def domain_average(msh, v):
    """Compute the average of a function over the domain"""
    vol = msh.comm.allreduce(
        fem.assemble_scalar(fem.form(fem.Constant(msh, default_real_type(1.0)) * dx)), op=MPI.SUM
    )
    return (1 / vol) * msh.comm.allreduce(fem.assemble_scalar(fem.form(v * dx)), op=MPI.SUM)


def u_e_expr(x):
    """Expression for the exact velocity solution to Kovasznay flow"""
    return np.vstack(
        (
            -np.cos(2*np.pi*x[0])*np.sin(2*np.pi*x[1]),
            np.sin(2*np.pi*x[0])*np.cos(2*np.pi*x[1]),
            # 1
            # - np.exp((Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0])
            # * np.cos(2 * np.pi * x[1]),
            # (Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2))
            # / (2 * np.pi)
            # * np.exp((Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0])
            # * np.sin(2 * np.pi * x[1]),
        )
    )


def p_e_expr(x):
    """Expression for the exact pressure solution to Kovasznay flow"""
    #return (0.0) * (1 - np.exp(2 * (Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0]))
    return np.zeros_like(x[0]) 


def f_expr(x):
    """Expression for the applied force"""
    return np.vstack((np.zeros_like(x[0]), np.zeros_like(x[0])))


def boundary_marker(x):
    return (
        np.isclose(x[0], 0.0)
        | np.isclose(x[0], 1.0)
        | np.isclose(x[1], 0.0)
        | np.isclose(x[1], 1.0)
    )

### Periodic boundary condition
# -

# We define some simulation parameters
n = 30
num_time_steps = 100
t_end = 1
#Re = 25  # Reynolds Number
k = 1  # Polynomial degree

# Next, we create a mesh and the required functions spaces over it.
# Since the velocity uses an $H(\text{div})$-conforming function
# space, we also create a vector valued discontinuous Lagrange space to
# interpolate into for artifact free visualisation.

# +
msh = mesh.create_unit_square(MPI.COMM_WORLD, n, n)

# Function spaces for the velocity and for the pressure
V = fem.functionspace(msh, ("Raviart-Thomas", k + 1))
Q = fem.functionspace(msh, ("Discontinuous Lagrange", k))

# Funcion space for visualising the velocity field
gdim = msh.geometry.dim
W = fem.functionspace(msh, ("Discontinuous Lagrange", k + 1, (gdim,)))

# Enforcing periodic boundary condition
#### Periodic boundary condition
fdim = msh.topology.dim - 1
bcs = []
pbc_directions = []
pbc_slave_tags = []
pbc_is_slave = []
pbc_is_master = []
pbc_meshtags = []
pbc_slave_to_master_maps = []

def generate_pbc_slave_to_master_map(i):
    def pbc_slave_to_master_map(x):
        out_x = x.copy()
        out_x[i] = x[i] - 1
        return out_x
    return pbc_slave_to_master_map

def generate_pbc_is_slave(i):
    return lambda x: np.isclose(x[i], 1)

def generate_pbc_is_master(i):
    return lambda x: np.isclose(x[i], 0)

# Parse boundary conditions
for i, bc_type in enumerate("periodic"):
    if bc_type == "dirichlet":
        u_bc = fem.Function(V)
        u_bc.x.array[:] = 0

        def dirichletboundary(x):
            return np.logical_or(np.isclose(x[i], 0), np.isclose(x[i], 1))
        facets = locate_entities_boundary(mesh, fdim, dirichletboundary)
        topological_dofs = fem.locate_dofs_topological(V, fdim, facets)
        bcs.append(fem.dirichletbc(u_bc, topological_dofs))
    elif bc_type == "periodic":
        pbc_directions.append(i)
        pbc_slave_tags.append(i + 2)
        pbc_is_slave.append(generate_pbc_is_slave(i))
        pbc_is_master.append(generate_pbc_is_master(i))
        pbc_slave_to_master_maps.append(generate_pbc_slave_to_master_map(i))

        facets = locate_entities_boundary(msh, fdim, pbc_is_slave[-1])
        arg_sort = np.argsort(facets)
        pbc_meshtags.append(meshtags(msh,
                                      fdim,
                                      facets[arg_sort],
                                      np.full(len(facets), pbc_slave_tags[-1], dtype=np.int32)))

# Create MultiPointConstraint object
mpc = MultiPointConstraint(V)
N_pbc = len(pbc_directions)
for i in range(N_pbc):
    if N_pbc > 1:
        def pbc_slave_to_master_map(x):
            out_x = pbc_slave_to_master_maps[i](x)
            idx = pbc_is_slave[(i + 1) % N_pbc](x)
            out_x[pbc_directions[i]][idx] = np.nan
            return out_x
    else:
        pbc_slave_to_master_map = pbc_slave_to_master_maps[i]

    mpc.create_periodic_constraint_topological(V, pbc_meshtags[i],
                                                pbc_slave_tags[i],
                                                pbc_slave_to_master_map,
                                                bcs)
if len(pbc_directions) > 1:
    # Map intersection(slaves_x, slaves_y) to intersection(masters_x, masters_y),
    # i.e. map the slave dof at (1, 1) to the master dof at (0, 0)
    def pbc_slave_to_master_map(x):
        out_x = x.copy()
        out_x[0] = x[0] - 1
        out_x[1] = x[1] - 1
        idx = np.logical_and(pbc_is_slave[0](x), pbc_is_slave[1](x))
        out_x[0][~idx] = np.nan
        out_x[1][~idx] = np.nan
        return out_x
    mpc.create_periodic_constraint_topological(V, pbc_meshtags[1],
                                                pbc_slave_tags[1],
                                                pbc_slave_to_master_map,
                                                bcs)
mpc.finalize()

###

# Define trial and test functions

u, v = TrialFunction(V), TestFunction(V)
p, q = TrialFunction(Q), TestFunction(Q)

delta_t = fem.Constant(msh, default_real_type(t_end / num_time_steps))
alpha = fem.Constant(msh, default_real_type(6.0 * k**2))

h = CellDiameter(msh)
n = FacetNormal(msh)


def jump(phi, n):
    return outer(phi("+"), n("+")) + outer(phi("-"), n("-"))

# Defining RHS and LHS
a_01 = -inner(p, div(v)) * dx
a_10 = -inner(div(u), q) * dx

f = fem.Function(W)
L_0 = inner(f, v) * dx 
L_1 = inner(fem.Constant(msh, default_real_type(0.0)), q) * dx
L = fem.form([L_0, L_1])

u_h = fem.Function(V)
u_h.interpolate(u_e_expr)


# Defining solution functions
p_h = fem.Function(Q)
p_h.name = "p"
offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs

u_vis = fem.Function(W)
u_vis.name = "u"
u_vis.interpolate(u_h)

# Write initial condition to file
t = 0.0
try:
    u_file = io.VTXWriter(msh.comm, "u.bp", u_vis)
    p_file = io.VTXWriter(msh.comm, "p.bp", p_h)
    u_file.write(t)
    p_file.write(t)
except AttributeError:
    print("File output requires ADIOS2.")

# Create function to store solution and previous time step
u_n = fem.Function(V)
u_n.x.array[:] = u_h.x.array
# -

# Now we add the time stepping and convective terms

# +
lmbda = conditional(gt(dot(u_n, n), 0), 1, 0)
u_uw = lmbda("+") * u("+") + lmbda("-") * u("-")
a_00 = (
    inner(u / delta_t, v) * dx
    - inner(u, div(outer(v, u_n))) * dx
    + inner((dot(u_n, n))("+") * u_uw, v("+")) * dS
    + inner((dot(u_n, n))("-") * u_uw, v("-")) * dS
    + inner(dot(u_n, n) * lmbda * u, v) * ds
)
a = fem.form([[a_00, a_01], [a_10, None]])

L_0 += inner(u_n / delta_t, v) * dx - inner(dot(u_n, n) * (1 - lmbda) * u_h, v) * ds
L = fem.form([L_0, L_1])

A = assemble_matrix_block(a, bcs=bcs)
A.assemble()
b = assemble_vector_block(L, a, bcs=bcs)

# # Create and configure solver
ksp = PETSc.KSP().create(msh.comm)  # type: ignore
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
opts = PETSc.Options()  # type: ignore
opts["mat_mumps_icntl_14"] = 80  # Increase MUMPS working memory
opts["mat_mumps_icntl_24"] = 1  # Option to support solving a singular matrix (pressure nullspace)
opts["mat_mumps_icntl_25"] = 0  # Option to support solving a singular matrix (pressure nullspace)
opts["ksp_error_if_not_converged"] = 1
ksp.setFromOptions()
x = A.createVecRight()

# Time stepping loop
for n in range(num_time_steps):
    t += delta_t.value

    A.zeroEntries()
    fem.petsc.assemble_matrix_block(A, a, bcs=bcs)  # type: ignore
    A.assemble()

    with b.localForm() as b_loc:
        b_loc.set(0)
    fem.petsc.assemble_vector_block(b, L, a, bcs=bcs)  # type: ignore

    # Compute solution
    ksp.solve(b, x)

    u_h.x.array[:offset] = x.array_r[:offset]
    u_h.x.scatter_forward()
    p_h.x.array[: (len(x.array_r) - offset)] = x.array_r[offset:]
    p_h.x.scatter_forward()
    p_h.x.array[:] -= domain_average(msh, p_h)

    u_vis.interpolate(u_h)

    # Write to file
    try:
        u_file.write(t)
        p_file.write(t)
    except NameError:
        pass

    # Update u_n
    u_n.x.array[:] = u_h.x.array

try:
    u_file.close()
    p_file.close()
except NameError:
    pass
# -

# Now we compare the computed solution to the exact solution

# +
# Function spaces for exact velocity and pressure
V_e = fem.functionspace(msh, ("Lagrange", k + 3, (gdim,)))
Q_e = fem.functionspace(msh, ("Lagrange", k + 2))

u_e = fem.Function(V_e)
u_e.interpolate(u_e_expr)

p_e = fem.Function(Q_e)
p_e.interpolate(p_e_expr)

# Compute errors
e_u = norm_L2(msh.comm, u_h - u_e)
e_div_u = norm_L2(msh.comm, div(u_h))

# This scheme conserves mass exactly, so check this
assert np.isclose(e_div_u, 0.0, atol=float(1.0e5 * np.finfo(default_real_type).eps))
p_e_avg = domain_average(msh, p_e)
e_p = norm_L2(msh.comm, p_h - (p_e - p_e_avg))

if msh.comm.rank == 0:
    print(f"e_u = {e_u}")
    print(f"e_div_u = {e_div_u}")
    print(f"e_p = {e_p}")
# -

I’m not sure if I have implemented periodic BCs correctly into the solver.

I didn’t go through the entire code, but I think you should be careful as you are using a DG space for the pressure. The DOFs of Q “live” within the cell so you probably won’t find them by just type

def generate_pbc_is_master(i):
    return lambda x: np.isclose(x[i], 0)

In addition, you only give periodic constraints for the Functionspace V

    mpc.create_periodic_constraint_topological(V, pbc_meshtags[1],
                                                pbc_slave_tags[1],
                                                pbc_slave_to_master_map,
                                                bcs)

, but no Dirichlet bc for Q.

Thanks for noticing. Even with these corrections, I still need to implement mpc into the solver which is impossible at the moment due to the mpc’s inability to measure dS