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):
        local_con = con_facet_cell.links(ifacet)
        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))
    a = ufl.inner(ufl.grad(u), ufl.grad(v))*dx
    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)
    u.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
    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):
        local_con = con_facet_cell.links(ifacet)
        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))
    a = ufl.dot(ufl.grad(x), ufl.grad(v))*dx
    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)
    x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
    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()