Unable to assemble matrix for coupled fields with support on parent mesh and submesh

Hello all,
I am trying to solve an equation with two fields, one defined on a domain in 2D (mesh) and the second one defined on a subdomain in 2D also (submesh). The two fields are coupled in the subdomain. I am having trouble to define the variational form and to assemble the corresponding matrix for the coupling terms, when the two fields appear in the form. Essentially, I am trying to do the following:

import numpy as np
import ufl
from dolfinx import fem
import dolfinx.mesh
import dolfinx
from mpi4py import MPI

# Create a unit square mesh
domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim)
submesh_entities = dolfinx.mesh.locate_entities(domain, dim=2, marker=lambda x: (x[0] < 0.5) )
submesh, sub_to_parent, vertex_map, geom_map = dolfinx.mesh.create_submesh(domain, domain.topology.dim, submesh_entities)


fn_space_field1 = fem.functionspace(domain, ("Lagrange", 1))
fn_space_field2 = fem.functionspace(submesh, ("Lagrange", 1))

fn_space = ufl.MixedFunctionSpace(fn_space_field1, fn_space_field2)

field1_trial, field2_trial = ufl.TrialFunctions(fn_space)
field1_test,  field2_test  = ufl.TestFunctions(fn_space)

dx_all = ufl.Measure("dx", domain=domain)
dx_sub = ufl.Measure("dx", domain=submesh)


F0 = ufl.inner(field1_trial, field2_test) * dx_sub 
F0_compiled = dolfinx.fem.form(F0)

M0 = dolfinx.fem.petsc.assemble_matrix(F0_compiled, bcs=[])
M0.assemble()

This code throws an error:

RuntimeError: Incompatible mesh. entity_maps must be provided.

I have tried to provide the entity_maps, following many posts on submeshes, and went through some examples referenced there, but somehow, I am unable to make it work.

I am using fenicsx, version 0.9.

Any help would be appreciated.

Hello all,
After many trials and error, and re-reading the following example:

https://jsdokken.com/FEniCS-workshop/src/multiphysics/coupling.html

I was able to find a solution to my problem. Here is the solution:

import numpy as np
import ufl
from dolfinx import fem
import dolfinx.mesh
import dolfinx
import dolfinx.fem.petsc
from mpi4py import MPI


# Create a unit square mesh
domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)


# Define cell tags on the mesh, for the two different subdomains

def Omega1(x, tol=1e-14):
    return x[0] <= 0.5 + tol

def Omega2(x, tol=1e-14):
    return 0.5-tol <= x[0]

tdim = domain.topology.dim
cell_map = domain.topology.index_map(tdim)
num_cells_local = cell_map.size_local + cell_map.num_ghosts
marker = np.empty(num_cells_local, dtype=np.int32)
omega1_marker = 1
omega2_marker = 2

marker[dolfinx.mesh.locate_entities(domain, tdim, Omega1)] = omega1_marker
marker[dolfinx.mesh.locate_entities(domain, tdim, Omega2)] = omega2_marker

cell_tags = dolfinx.mesh.meshtags(domain, domain.topology.dim,
                                    np.arange(num_cells_local, dtype=np.int32),
                                    marker)

omega1_subdomain, omega1_cell_map, omega1_vertex_map, _ = dolfinx.mesh.create_submesh(domain, cell_tags.dim, cell_tags.find(omega1_marker))
omega2_subdomain, omega2_cell_map, omega2_vertex_map, _ = dolfinx.mesh.create_submesh(domain, cell_tags.dim, cell_tags.find(omega2_marker))


# Invert the cell maps
mesh_to_omega1_entity = np.full(num_cells_local, -1, dtype=np.int32)
mesh_to_omega1_entity[omega1_cell_map] = np.arange(len(omega1_cell_map), dtype=np.int32)

mesh_to_omega2_entity = np.full(num_cells_local, -1, dtype=np.int32)
mesh_to_omega2_entity[omega2_cell_map] = np.arange(len(omega2_cell_map), dtype=np.int32)

entity_maps = {}
entity_maps[omega1_subdomain] = mesh_to_omega1_entity
entity_maps[omega2_subdomain] = mesh_to_omega2_entity

# Create function space and fields
fn_space_field0 = fem.functionspace(domain, ("Lagrange", 1))
fn_space_field1 = fem.functionspace(omega1_subdomain, ("Lagrange", 1))
fn_space_field2 = fem.functionspace(omega2_subdomain, ("Lagrange", 1))

fn_space = ufl.MixedFunctionSpace(fn_space_field0, fn_space_field1, fn_space_field2)

field0_trial, field1_trial, field2_trial = ufl.TrialFunctions(fn_space)
field0_test,  field1_test,  field2_test  = ufl.TestFunctions(fn_space)

# Create measure on subdomains
dx_all = ufl.Measure("dx", domain=domain)
dx = ufl.Measure("dx", domain=domain, subdomain_data=cell_tags)
dx_omega1 = dx(omega1_marker)
dx_omega2 = dx(omega2_marker)

# Create the variational form
F0  = ufl.inner(field0_trial, field1_test) * dx_omega1
F0 += ufl.inner(field1_trial, field0_test) * dx_omega1
F0 += ufl.inner(field1_trial, field1_test) * dx_omega1
F0 += ufl.inner(field0_trial, field2_test) * dx_omega2
F0 += ufl.inner(field2_trial, field0_test) * dx_omega2
F0 += ufl.inner(field2_trial, field2_test) * dx_omega2
F0 += ufl.inner(field0_trial, field0_test) * dx_all

F0_blocked = ufl.extract_blocks(F0)

F0_compiled = dolfinx.fem.form(F0_blocked, entity_maps=entity_maps)

M0 = dolfinx.fem.petsc.assemble_matrix_block(F0_compiled, bcs=[])
M0.assemble()

Now, the code is running and the matrix is assembling. Let me know if you have any feedback on this solution and if this is the proper way to proceed.