[BUG]: Allow for different measures with optional `subdomain_data` in variational forms

Summarize the issue

Consider the following variational form for a diffusion problem with a source term f defined on some subdomain.

dx_ft = ufl.dx(domain=mesh, subdomain_data=facet_tags)
dx = ufl.dx
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)

du_dt = u - u_n
theta = 0.5
u_mid = theta * u + (1 - theta) * u_n

G = du_dt * v * dx + dt * ufl.dot(ufl.grad(u_mid), ufl.grad(v)) * dx  - dt * f * v * dx_ft(S1_marker)
a, L = ufl.system(G)

Here I have intentionally created two measures where one of the measures contain subdomain_data. When solving this problem (see full example below) it yields an AssertionError which requires that the subdomain_data should be the same for all dx (see https://github.com/FEniCS/dolfinx/blob/3e7c708770274ecb27f03c37286f9f1de58e9d32/python/dolfinx/fem/forms.py#L245-L247).

The full backtrace is here

Traceback (most recent call last):
  File "/home/shared/mwe.py", line 65, in <module>
    problem = LinearProblem(
              ^^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/petsc.py", line 789, in __init__
    self._L = _create_form(
              ^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 337, in form
    return _create_form(form)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 331, in _create_form
    return _form(form)
           ^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 249, in _form
    assert all([d is data[0] for d in data])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
Exception ignored in: <function LinearProblem.__del__ at 0xffffa602c7c0>
Traceback (most recent call last):
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/petsc.py", line 831, in __del__
    self._solver.destroy()
    ^^^^^^^^^^^^
AttributeError: 'LinearProblem' object has no attribute '_solver'

In this case subdomain_data would be None for two elements and equal to facet_tags for the final element.

I think in scenarios where all elements are None and one is not (e.g equal to facet_tags), then I belive it is more natural to add facet_tags as subdomain_data to all the elements.

I would also like to note that in this example, it is easy to fix this problem, but in more complicated code you might build up your variational form in different parts of the code where some parts don't have access to the MeshTags, in which case this is more relevant.

Finally, if you think this behaviour should be kept, then I would encourage to add a better error message.

I am happy to help out by submitting a PR (just need to decide on the right path moving forward).

How to reproduce the bug

import ufl
import numpy as np

from petsc4py import PETSc
from mpi4py import MPI

from dolfinx.fem.petsc import LinearProblem
import dolfinx

# Define mesh
comm = MPI.COMM_WORLD
N = 20
mesh = dolfinx.mesh.create_unit_square(comm, N, N, dolfinx.cpp.mesh.CellType.triangle)
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))

tol = 1.0e-10
L = 0.3

def S1_subdomain(x):
    return np.logical_and(x[0] <= L + tol, x[1] <= L + tol)

S1_marker = 1
tdim = mesh.topology.dim
facets = dolfinx.mesh.locate_entities(mesh, tdim, S1_subdomain)
facet_tags = dolfinx.mesh.meshtags(
    mesh,
    tdim,
    facets,
    np.full(len(facets), S1_marker, dtype=np.int32),
)

f = dolfinx.fem.Constant(mesh, 1.0)

# Solution at previous time step
u_n = dolfinx.fem.Function(V)
u_n.name = "u_n"
# Solution
uh = dolfinx.fem.Function(V)
uh.name = "u"

t = 0
T = 1.0
num_steps = 50
dt = T / num_steps

dx_ft = ufl.dx(domain=mesh, subdomain_data=facet_tags)
dx = ufl.dx
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)

du_dt = u - u_n
theta = 0.5
u_mid = theta * u + (1 - theta) * u_n

G = (
    du_dt * v * dx
    + dt * ufl.dot(ufl.grad(u_mid), ufl.grad(v)) * dx
    - dt * f * v * dx_ft(S1_marker)
)
a, L = ufl.system(G)

problem = LinearProblem(
    a,
    L,
    u=uh,
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
    },
)

problem.A.zeroEntries()
dolfinx.fem.petsc.assemble_matrix(problem.A, problem.a)
problem.A.assemble()

vtx = dolfinx.io.VTXWriter(comm, "diffusion.bp", [uh], engine="BP4")
vtx.write(0.0)

for n in range(num_steps):
    with problem.b.localForm() as b_loc:
        b_loc.set(0)
    dolfinx.fem.petsc.assemble_vector(problem.b, problem.L)
    problem.b.ghostUpdate(
        addv=PETSc.InsertMode.ADD,
        mode=PETSc.ScatterMode.REVERSE,
    )

    problem.solver.solve(problem.b, uh.x.petsc_vec)
    u_n.x.array[:] = uh.x.array
    vtx.write(n + 1)

Minimal Example (Python)

No response

Output (Python)

No response

Version

main branch

DOLFINx git commit

No response

Installation

No response

Additional information

No response

FEniCS / dolfinx