Solve on subdomain using multiphenicsx

I am trying to solve a PDE on a subdomain without creating a submesh because of the potential overhead. For this task I am using the multiphenicsx library. I assemble my matrix using multiphenicsx.fem.petsc.assemble_matrix with a degree of freedom restriction map and upon calling ksp.setOperators(matrix), I get a null pointer error. Here is a MWE where I try to solve a Laplace equation on the left side of the unit square:

from dolfinx import io, fem, mesh, cpp
import ufl
import numpy as np
from mpi4py import MPI
import multiphenicsx
import multiphenicsx.fem.petsc
import petsc4py.PETSc

def exact_sol(x):
return 2 -(x[0]**2 + x[1]**2)
def rhs():
return 4

def main():
# BACKGROUND MESH
nels_side = 2
bg_mesh  = mesh.create_unit_square(MPI.COMM_WORLD, nels_side, nels_side, mesh.CellType.quadrilateral)
cdim = bg_mesh.topology.dim
fdim = cdim - 1

# SOLUTION SPACES
Vbg   = fem.functionspace(bg_mesh, ("Lagrange", 1),)
dg0bg = fem.functionspace(bg_mesh, ("Discontinuous Lagrange", 0),)
Vactive = Vbg.clone()
active_dofs        = fem.locate_dofs_geometrical(Vactive, lambda x : x[0] <= 0.5)
active_els         = np.array(fem.locate_dofs_geometrical(dg0bg, lambda x : x[0] <= 0.5), dtype=np.int32)
restriction        = multiphenicsx.fem.DofMapRestriction(Vactive.dofmap, active_dofs)
active_els_tag     = mesh.meshtags(bg_mesh, cdim,
active_els,
np.ones(active_els.shape, dtype=np.int32) )

# LOCATE ACTIVE BOUNDARY
bfacets = []
bg_mesh.topology.create_connectivity(fdim, cdim)
con_facet_cell = bg_mesh.topology.connectivity(1, 2)
for ifacet in range(con_facet_cell.num_nodes):
incident_active_els = 0
for el in local_con:
if el in active_els:
incident_active_els += 1
if (incident_active_els==1):
bfacets.append(ifacet)
bdofs = fem.locate_dofs_topological(Vactive, fdim, bfacets,)

# BC
u_bc = fem.Function(Vactive)
u_bc.interpolate(exact_sol)
bc = fem.dirichletbc(u_bc, bdofs)

# FORMS
dx = ufl.Measure("dx")(subdomain_data=active_els_tag)
(u, v) = (ufl.TrialFunction(Vactive),ufl.TestFunction(Vactive))
l = rhs()*v*dx
a_cpp = fem.form(a)
l_cpp = fem.form(l)
A = multiphenicsx.fem.petsc.assemble_matrix(a_cpp,
bcs=[bc],
restriction=(restriction, restriction))
A.assemble()
L = multiphenicsx.fem.petsc.assemble_vector(l_cpp,
restriction=restriction,)

# SOLVE
u = multiphenicsx.fem.petsc.create_vector(l_cpp, restriction=restriction)
ksp = petsc4py.PETSc.KSP()
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(L, u)
ksp.destroy()

if __name__=="__main__":
main()

I get:

petsc4py.PETSc.Error: error code 85
[0] KSPSetOperators() at /usr/local/petsc/src/ksp/ksp/interface/itcreate.c:509
[0] Null argument, when expecting valid pointer

using the latest versions of dolfinx and multiphenicsx on the fenicsx development docker container. I would appreciate any help from @francesco-ballarin or anyone else in finding a potential bug.

PS: I could not find dolfinx.mesh.create_submesh in the documentation website.

There is a missing ksp.create(COMM). The full serial script writes:

from dolfinx import io, fem, mesh, cpp
import ufl
import numpy as np
from mpi4py import MPI
import multiphenicsx
import multiphenicsx.fem.petsc
import dolfinx.fem.petsc
import petsc4py.PETSc
import pdb

def exact_sol(x):
return 2 -(x[0]**2 + x[1]**2)
def rhs():
return 4

def main():
# BACKGROUND MESH
nels = 8
bg_mesh  = mesh.create_unit_square(MPI.COMM_WORLD, nels, nels, mesh.CellType.quadrilateral)
cdim = bg_mesh.topology.dim
fdim = cdim - 1

# SOLUTION SPACES
Vbg   = fem.functionspace(bg_mesh, ("Lagrange", 1),)
dg0bg = fem.functionspace(bg_mesh, ("Discontinuous Lagrange", 0),)
Vactive = Vbg.clone()
active_dofs        = fem.locate_dofs_geometrical(Vactive, lambda x : x[0] <= 0.5)
active_els         = np.array(fem.locate_dofs_geometrical(dg0bg, lambda x : x[0] <= 0.5), dtype=np.int32)
restriction        = multiphenicsx.fem.DofMapRestriction(Vactive.dofmap, active_dofs)
active_els_tag     = mesh.meshtags(bg_mesh, cdim,
active_els,
np.ones(active_els.shape, dtype=np.int32) )

# LOCATE ACTIVE BOUNDARY
bfacets = []
bg_mesh.topology.create_connectivity(fdim, cdim)
con_facet_cell = bg_mesh.topology.connectivity(1, 2)
for ifacet in range(con_facet_cell.num_nodes):
incident_active_els = 0
for el in local_con:
if el in active_els:
incident_active_els += 1
if (incident_active_els==1):
bfacets.append(ifacet)
bdofs = fem.locate_dofs_topological(Vactive, fdim, bfacets,)

# BC
u_bc = fem.Function(Vactive)
u_bc.interpolate(exact_sol)
bc = fem.dirichletbc(u_bc, bdofs)

# FORMS
dx = ufl.Measure("dx")(subdomain_data=active_els_tag)
(x, v) = (ufl.TrialFunction(Vactive),ufl.TestFunction(Vactive))
l = rhs()*v*dx
a_cpp = fem.form(a)
l_cpp = fem.form(l)
A = multiphenicsx.fem.petsc.assemble_matrix(a_cpp,
bcs=[bc],
restriction=(restriction, restriction))
B = dolfinx.fem.petsc.assemble_matrix(a_cpp, bcs=[bc],)
A.assemble()
B.assemble()
L = multiphenicsx.fem.petsc.assemble_vector(l_cpp,
restriction=restriction,)
multiphenicsx.fem.petsc.apply_lifting(L, [a_cpp], [[bc]], restriction=restriction,)# I think [[bc]] is awkward
multiphenicsx.fem.petsc.set_bc(L,[bc],restriction=restriction)

# SOLVE
x = multiphenicsx.fem.petsc.create_vector(l_cpp, restriction=restriction)
ksp = petsc4py.PETSc.KSP()
ksp.create(bg_mesh.comm)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(L, x)
ksp.destroy()

# GET FUNCTION ON SUBDOMAIN
usub = fem.Function(Vactive, name="uh")
with usub.vector.localForm() as usub_vector_local, \
multiphenicsx.fem.petsc.VecSubVectorWrapper(x, Vactive.dofmap, restriction) as x_wrapper:
usub_vector_local[:] = x_wrapper
x.destroy()
uexact = fem.Function(Vactive, name="uex")
uexact.interpolate( exact_sol )
# WRITE
with io.VTKFile(bg_mesh.comm, "out/res.pvd", "w") as ofile:
ofile.write_function([usub, uexact])

if __name__=="__main__":
main()