Multiphysics: Solving PDEs on subdomains

Multiphysics: Solving PDEs on subdomains#

So far we have considered problems where the PDE is defined over the entire domain. However, in many cases this is not accurate. An example of this is fluid-structure interaction, where the fluid and solid domains are coupled. In this case, the PDEs are defined over different subdomains, and the coupling is done at the interface between the subdomains. In this section, we will show how to solve PDEs on subdomains using FEniCS.

We will consider a simple problem where we have a domain \(\Omega\) that is divided into two subdomains \(\Omega_1\) and \(\Omega_2\). In each of these domains we want to solve a PDE (that is not coupled to the other domain).

We will consider the following PDEs:

\[\begin{split} \begin{align*} - \nabla \cdot (\kappa \nabla T) &= f \quad \text{in } \Omega_1, \\ \kappa \nabla T \cdot \mathbf{n} &= g \quad \text{on } \Gamma, \\ T &= g \quad \text{ on } \partial \Omega_1\setminus\Gamma, \\ - \nabla \cdot ( \nabla \mathbf{u}) - \nabla \bar p &= \mathbf{f} \quad \text{in } \Omega_2 \\ \nabla \cdot \mathbf{u} &= 0 \quad \text{in } \Omega_2 \\ \mathbf{u} &= \mathbf{h} \text{ on } \partial {\Omega_{2,D}} \\ \nabla \mathbf{u} \cdot \mathbf{n} + \bar p \mathbf{n} &= \mathbf{0} \quad \text{on } \partial_{\Omega_{2, N}} \end{align*} \end{split}\]

We start by creating our mesh, a unit square, that we will split into two compartments

\[\begin{split} \begin{align*} \Omega_1 &= [0.7, 1] \times [0, 1]\\ \Omega_2 &= [0, 0.7] \times [0, 1]\\ \end{align*} \end{split}\]

from mpi4py import MPI
import dolfinx
import numpy as np

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 40, 40)

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

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

tdim = mesh.topology.dim
cell_map = mesh.topology.index_map(tdim)
num_cells_local = cell_map.size_local + cell_map.num_ghosts
marker = np.empty(num_cells_local, dtype=np.int32)
heat_marker = 1
stokes_marker = 3

marker[dolfinx.mesh.locate_entities(mesh, tdim, Omega0)] = heat_marker
marker[dolfinx.mesh.locate_entities(mesh, tdim, Omega1)] = stokes_marker

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

This section will contain alot of figures, to illustrate the different steps. Expand to see the code for creating a plotter of meshes and meshtags.

[   0    1    2 ... 3197 3198 3199] [1 1 1 ... 3 3 3]

Next we define the boundaries for our sub-domains

[   0    3    4    6   12   14   23   25   37   39   54   56   74   76
 99  123  125  152  154  184  186  219  221  256  257  259  263
300  307  342  344  354  389  391  404  439  441  457  492  494
548  550  572  607  609  634  669  671  699  734  736  767  802
838  873  875  912  947  949  989 1024 1026 1069 1104 1106 1152
1189 1238 1273 1275 1327 1362 1364 1419 1454 1456 1514 1549 1551
1647 1649 1713 1748 1750 1817 1852 1854 1924 1959 1961 2034 2069
2147 2182 2184 2263 2298 2300 2382 2417 2419 2504 2538 2539 2541
2656 2658 2742 2770 2772 2856 2881 2883 2967 2989 2991 3075 3094
3180 3196 3198 3282 3295 3297 3381 3391 3393 3477 3484 3486 3570
3576 3660 3661 3663 3745 3747 3826 3828 3904 3906 3979 3981 4051
4120 4122 4186 4188 4249 4251 4309 4311 4366 4368 4420 4422 4471
4519 4521 4564 4566 4606 4608 4645 4647 4681 4683 4714 4716 4744
4771 4773 4795 4797 4816 4818 4834 4836 4849 4851 4861 4863 4870
4876 4878 4879] [4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3
4 3 2 4 3 2 4 3 2 4 3 2 4 3 2 4 3 2 4 3 2 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4
3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4
3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 3 4 4 4 4 4
4 4 4 1 4 1 4 1 4 1 4 1 4 1 4 1 4 1 4 4 4 4 4 4 2 4 2 4 2 4 2 4 2 4 2 4
4 4 4 4 4 4 4 4 4 4 4 4 4 4]

Next, we want to extract a sub-mesh containing only the the cellsthat will be used in the Stokes problem. We do this with dolfinx.mesh.create_submesh

stokes_mesh, stokes_cell_map, stokes_vertex_map, _ = dolfinx.mesh.create_submesh(mesh, cell_tags.dim, cell_tags.find(stokes_marker))

Creating a submesh

dolfinx.mesh.create_submesh takes in three inputs:

The mesh we want to extract a mesh from
The dimension of the entities we want to make the mesh from. This can be any number \([0, tdim]\). We define a submesh consisting of a subset of cells from the input mesh is a submesh of co-dimension 0, while a submesh consisting fo facets from the input mesh is a submesh of co-dimension 1
A list of integers defining the entities (index local to process, and including ghosted entities)

The function returns four objects

The submesh (as a dolfinx.mesh.Mesh)
A map from each entity cell in the submesh to the corresponding entity in the input mesh
A map from each vertex in the submesh to the corresponding vertex in the input mesh topology
A map from each node in the submesh to the corresponding node in the input mesh geometry

Transfering meshtags to the sub-mesh#

There are many ways on could use the created sub-mesh. We will start by treating it as a normal mesh. For this, we would like to transfer the cell_tags and facet_tags defined on the whole mesh. We do this with the following function

Show code cell source Hide code cell source

import numpy.typing as npt
def transfer_meshtags_to_submesh(
    mesh: dolfinx.mesh.Mesh,
    entity_tag: dolfinx.mesh.MeshTags,
    submesh: dolfinx.mesh.Mesh,
    sub_vertex_to_parent: npt.NDArray[np.int32],
    sub_cell_to_parent: npt.NDArray[np.int32]
)-> tuple[dolfinx.mesh.MeshTags, npt.NDArray[np.int32]]:
    """
    Transfer a meshtag from a parent mesh to a sub-mesh.

    Args:
        mesh: Mesh containing the meshtags
        entity_tag: The meshtags object to transfer
        submesh: The submesh to transfer the `entity_tag` to
        sub_to_vertex_map: Map from each vertex in `submesh` to the corresponding
            vertex in the `mesh`
        sub_cell_to_parent: Map from each cell in the `submesh` to the corresponding
            entity in the `mesh`
    Returns:
        The entity tag defined on the submesh, and a map from the entities in the
        `submesh` to the entities in the `mesh`.

    """

    tdim = mesh.topology.dim
    cell_imap = mesh.topology.index_map(tdim)
    num_cells = cell_imap.size_local + cell_imap.num_ghosts
    mesh_to_submesh = np.full(num_cells, -1)
    mesh_to_submesh[sub_cell_to_parent] = np.arange(
        len(sub_cell_to_parent), dtype=np.int32
    )
    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)

    submesh.topology.create_connectivity(entity_tag.dim, 0)

    num_child_entities = (
        submesh.topology.index_map(entity_tag.dim).size_local
        + submesh.topology.index_map(entity_tag.dim).num_ghosts
    )
    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)

    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
    c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)

    child_markers = np.full(num_child_entities, 0, dtype=np.int32)

    mesh.topology.create_connectivity(entity_tag.dim, 0)
    mesh.topology.create_connectivity(entity_tag.dim, mesh.topology.dim)
    p_f_to_v = mesh.topology.connectivity(entity_tag.dim, 0)
    p_f_to_c = mesh.topology.connectivity(entity_tag.dim, mesh.topology.dim)
    sub_to_parent_entity_map = np.full(num_child_entities, -1, dtype=np.int32)
    for facet, value in zip(entity_tag.indices, entity_tag.values):
        facet_found = False
        for cell in p_f_to_c.links(facet):
            if facet_found:
                break
            if (child_cell := mesh_to_submesh[cell]) != -1:
                for child_facet in c_c_to_e.links(child_cell):
                    child_vertices = c_e_to_v.links(child_facet)
                    child_vertices_as_parent = sub_vertex_to_parent[child_vertices]
                    is_facet = np.isin(
                        child_vertices_as_parent, p_f_to_v.links(facet)
                    ).all()
                    if is_facet:
                        child_markers[child_facet] = value
                        facet_found = True
                        sub_to_parent_entity_map[child_facet] = facet
    tags = dolfinx.mesh.meshtags(
        submesh,
        entity_tag.dim,
        np.arange(num_child_entities, dtype=np.int32),
        child_markers,
    )
    tags.name = entity_tag.name
    return tags, sub_to_parent_entity_map

stokes_facet_tags, stokes_facet_map = transfer_meshtags_to_submesh(mesh, facet_tags, stokes_mesh,
                                                                   stokes_vertex_map, stokes_cell_map)

[   0    1    4    6   12   14   23   25   37   39   54   56   74   76
 99  123  125  152  154  184  186  219  221  257  259  298  300
344  389  391  439  441  492  494  548  550  607  609  669  671
736  802  804  873  875  947  949 1024 1026 1104 1106 1187 1189
1274 1276 1359 1361 1444 1446 1529 1531 1614 1616 1699 1701 1784
1869 1871 1954 1956 2039 2041 2124 2126 2209 2211 2293 2295 2374
2452 2454 2527 2529 2599 2601 2668 2670 2734 2736 2797 2799 2857
2914 2916 2968 2970 3019 3021 3067 3069 3112 3114 3154 3156 3193
3229 3231 3262 3264 3292 3294 3319 3321 3343 3345 3364 3366 3382
3397 3399 3409 3411 3418 3420 3424 3426 3427] [3 4 3 4 3 4 3 4 3 2 3 2 3 2 3 2 3 2 3 2 3 2 3 2 3 4 3 4 3 4 3 4 3 4 3 4 3
3 4 3 4 3 4 3 4 3 4 3 4 3 4 3 4 3 4 3 4 4 3 4 3 4 3 4 3 4 3 4 3 4 3 4 3
3 4 3 4 3 4 4 4 4 4 4 4 4 1 4 1 4 1 4 1 4 1 4 1 4 1 4 1 4 4 4 4 4 4 2 4
4 2 4 2 4 2 4 2 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4]

Solve Stokes problem on submesh#

We can now solve the Stokes problem on this sub-mesh

Solve the Stokes problem on the submesh

See the section Mixed problems for hints on how to set up relevant boundary conditions and the variational form

Expand the below dropdown to see the solution

Show code cell source Hide code cell source

import ufl
import basix
import scipy

# Define variational form on submesh
el_u = basix.ufl.element("Lagrange", stokes_mesh.basix_cell(), 3, shape=(stokes_mesh.geometry.dim,))
el_p = basix.ufl.element("Lagrange", stokes_mesh.basix_cell(), 2)
el_mixed = basix.ufl.mixed_element([el_u, el_p])
W = dolfinx.fem.functionspace(stokes_mesh, el_mixed)
u, p = ufl.TrialFunctions(W)
v, q = ufl.TestFunctions(W)
wh = dolfinx.fem.Function(W)
uh, ph = ufl.split(wh)
g = dolfinx.fem.Constant(stokes_mesh, dolfinx.default_scalar_type((0, 0)))
f = dolfinx.fem.Constant(stokes_mesh, dolfinx.default_scalar_type((0, 0)))
F = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
F += ufl.inner(p, ufl.div(v)) * ufl.dx
F += ufl.inner(ufl.div(u), q) * ufl.dx
F -= ufl.inner(f, v) * ufl.dx
a, L = ufl.system(F)

# Create boundary conditions
W0 = W.sub(0)
V, V_to_W0 = W0.collapse()
u_wall = dolfinx.fem.Function(V)
u_wall.x.array[:] = 0
stokes_walls = np.union1d(stokes_facet_tags.find(outer_marker), stokes_facet_tags.find(interface_marker))
dofs_wall = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, stokes_walls)
bc_wall = dolfinx.fem.dirichletbc(u_wall, dofs_wall, W0)
u_inlet = dolfinx.fem.Function(V)
u_inlet.interpolate(lambda x: (0.5*x[1], 0*x[0]))
dofs_inlet = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, stokes_facet_tags.find(inlet_marker))
bc_inlet = dolfinx.fem.dirichletbc(u_inlet, dofs_inlet, W0)
bcs = [bc_wall, bc_inlet]

# Compile form and assemble system
a_compiled = dolfinx.fem.form(a)
L_compiled = dolfinx.fem.form(L)
A = dolfinx.fem.create_matrix(a_compiled)
b = dolfinx.fem.create_vector(L_compiled)
A_scipy = A.to_scipy()
dolfinx.fem.assemble_matrix(A, a_compiled, bcs=bcs)
dolfinx.fem.assemble_vector(b.array, L_compiled)
dolfinx.fem.apply_lifting(b.array, [a_compiled], [bcs])
b.scatter_reverse(dolfinx.la.InsertMode.add)
[bc.set(b.array) for bc in bcs]

# Solve with SPLU
import scipy.sparse
A_inv = scipy.sparse.linalg.splu(A_scipy)
wh = dolfinx.fem.Function(W)
wh.x.array[:] = A_inv.solve(b.array)

Solving with integration over full mesh#

As we have seen above, we can create a sub-mesh and use it as one would use any other mesh. However, as we want to use this mesh in multi-physics problems, we want to exploit the relation to the parent domain.

In this section we will illustrate this for the Poisson problem defined at the top. We will define \(\kappa\) on the whole domain \(\Omega\), and not on the sub-mesh.

K = dolfinx.fem.functionspace(mesh, ("DG", 0))
kappa = dolfinx.fem.Function(K)
kappa.x.array[:] = 25
subset_cells = dolfinx.mesh.locate_entities(mesh, tdim, lambda x: x[1] > 0.3)
kappa.interpolate(lambda x: 12*x[0]+x[1], cells0 = subset_cells)

We create the submesh and the function space for the temperature on the submesh

heat_mesh, heat_cell_map, heat_vertex_map, _ = dolfinx.mesh.create_submesh(mesh, cell_tags.dim, cell_tags.find(heat_marker))
heat_facet_tags , _= transfer_meshtags_to_submesh(mesh, facet_tags, heat_mesh,
                                                                   heat_vertex_map, heat_cell_map)

However, in this case \(\kappa\) lives on the parent mesh. There are two ways of handling this:

Interpolate \(\kappa\) onto the heat_mesh
Integrate over a restricted section of the parent mesh

Interpolate to submesh#

We can interpolate data onto a submesh.

K_sub = dolfinx.fem.functionspace(heat_mesh, ("DG", 0))
kappa_sub = dolfinx.fem.Function(K_sub)

We do this by supplying to lists to interpolate. As we have seen before cells0 relate to what cells in the incoming space (K_sub) we want to interpolate data onto. Now, we also need to pass the information about which cells in mesh that relates to cells0. We can retrieve all this information from the heat_cell_map

kappa_sub.interpolate(kappa, cells0=heat_cell_map, cells1=np.arange(len(heat_cell_map)))

However, as we have already seen how to deal with problems that are pure related to quantities on a submesh, we will move to option number 2.

Integration with function from parent and sub-mesh#

The first thing we have to decide on, is which of the domains mesh and heat_mesh that we want to associate with the integration. We will use mesh as it more easily generalizes to co-dimension 1 problems.

dx = ufl.Measure("dx", domain=mesh, subdomain_data=cell_tags)

Note that there is only a sub-set of cells in the dx measure that is also in heat_mesh. Therefore we restrict the integral

dx_heat = dx(heat_marker)

Next, we set up the problem as done previously, with test-functions and trial-functions on the heat_mesh

T = dolfinx.fem.functionspace(heat_mesh, ("Lagrange", 1))
t = ufl.TrialFunction(T)
dt = ufl.TestFunction(T)

x = ufl.SpatialCoordinate(mesh)
f = 10*ufl.sin(5*ufl.pi*x[1]) + x[0]**3
F_heat = ufl.inner(kappa*ufl.grad(t), ufl.grad(dt))*dx_heat - ufl.inner(f, dt)*dx_heat

Warning

As we choose mesh as the integration domain, all spatial quantities such as ufl.SpatialCoordinate and ufl.FacetNormal should be defined wrt. this domain.

We create the boundary condition as before

T_bndry = dolfinx.fem.Function(T)
T_bndry.x.array[:] = 0
heat_mesh.topology.create_connectivity(heat_mesh.topology.dim-1, heat_mesh.topology.dim)
heat_bc_dofs = dolfinx.fem.locate_dofs_topological(T, heat_mesh.topology.dim-1, heat_facet_tags.find(outer_marker))
bc_heat = dolfinx.fem.dirichletbc(T_bndry, heat_bc_dofs)
bcs_heat = [bc_heat]

Note that so far very little have been different from just solving on the submesh. However, now an important distinction is introduced, namely the entity map

Entity maps

Entity maps are an input to the dolfinx.fem.form function, and should be on the following form: The maps is a dictionary, where each key is a dolfinx.mesh.Mesh that is part of the form, but is not the chosen integration domain. For each key, we have a map that maps an entity from the integrating mesh to an entity in the mesh that is the key. Note that in the case where we use the full mesh as integration domain, this is the inverse map of the one returned by dolfinx.mesh.create_submesh.

We therefore create this inverse map. As this map is a sparse map (not all cells in the full mesh is in the submesh), we use -1 to indicate that a cell is not part of the submesh.

mesh_to_heat_entity = np.full(num_cells_local, -1, dtype=np.int32)
mesh_to_heat_entity[heat_cell_map] = np.arange(len(heat_cell_map), dtype=np.int32)
entity_maps = {heat_mesh: mesh_to_heat_entity}

We can now compile our forms

a_heat, L_heat = dolfinx.fem.form(ufl.system(F_heat), entity_maps=entity_maps)

Now we can solve assemble the system

A_heat = dolfinx.fem.assemble_matrix(a_heat, bcs=bcs_heat)
A_heat_scipy = A_heat.to_scipy()

b_heat = dolfinx.fem.assemble_vector(L_heat)
dolfinx.fem.apply_lifting(b.array, [a_heat], [bcs_heat])
b_heat.scatter_reverse(dolfinx.la.InsertMode.add)
[bc.set(b_heat.array) for bc in bcs_heat]

[None]

Why did we not pass matrices to assemble vector/matrix?

In all the previous examples we have seen, we have been explicitly creating the matrices and passing them to assemble. There are many good reasons for this if one calls assemble multiple times in a program: As the matrix is sparse, we want to pre-compute the sparsity pattern of the matrix before inserting data into it. dolfinx.fem.create_matrix estimates the sparsity pattern based on the variational form and creates the appropriate CSR matrix. It is cheaper to zero out the initial contributions in the matrix than creating a new one.

We can now solve the system

A_heat_inv = scipy.sparse.linalg.splu(A_heat_scipy)

th = dolfinx.fem.Function(T, name="Temperature")
th.x.array[:] = A_heat_inv.solve(b_heat.array)

Combined assembly#

As we aim to solve multiphysics problems in a monolitic way, we will go through the basic steps of setting up such as system. We define the function spaces on the sub-meshes as done previously, however, instead of creating a basix.ufl.mixed_element, we will create individual function spaces for uandp`.

T = dolfinx.fem.functionspace(heat_mesh, ("Lagrange", 1))
V = dolfinx.fem.functionspace(stokes_mesh, ("Lagrange", 2, (stokes_mesh.topology.dim, )))
Q = dolfinx.fem.functionspace(stokes_mesh, ("Lagrange", 1))

Next, we use ufl.MixedFunctionSpace to create a representation of the monolitic problem.

W = ufl.MixedFunctionSpace(T, V, Q)

We can now create test and trial functions

t, u, p = ufl.TrialFunctions(W)
dt, du, dp = ufl.TestFunctions(W)

We define the relevant integration meshes. For consistency, we will use the mesh as the integration domain for both problems.

dx = ufl.Measure("dx", domain=mesh, subdomain_data=cell_tags)
dx_thermal = dx(heat_marker)
dx_fluid = dx(stokes_marker)

Instead of creating separate definitions of each form, we use a single variational form

g = dolfinx.fem.Constant(stokes_mesh, dolfinx.default_scalar_type((0, 0)))
f_fluid = dolfinx.fem.Constant(stokes_mesh, dolfinx.default_scalar_type((0, 0)))
a = ufl.inner(ufl.grad(u), ufl.grad(du)) * dx_fluid
a += ufl.inner(p, ufl.div(du)) * dx_fluid
a += ufl.inner(ufl.div(u), dp) * dx_fluid
L = ufl.inner(f_fluid, du) * dx_fluid
L += ufl.inner(dolfinx.fem.Constant(stokes_mesh, 0.0), dp) * dx_fluid
x = ufl.SpatialCoordinate(mesh)
f_heat = 10*ufl.sin(5*ufl.pi*x[1]) + x[0]**3
a += ufl.inner(kappa*ufl.grad(t), ufl.grad(dt))*dx_heat
L += ufl.inner(f, dt)*dx_thermal

We can extract the matrix block structure by calling ufl.extract_blocks

a_blocked = ufl.extract_blocks(a)
L_blocked = ufl.extract_blocks(L)

A_(0, 0)={ ({ A | A_{i_{39}} = (grad(v_1^0))[i_{39}] * f }) : (grad(v_0^0)) } * dx(<Mesh #0>[1], {})
A_(0, 1)=None
A_(0, 2)=None
A_(1, 0)=None
A_(1, 1)={ conj(((grad(v_0^1)) : (grad(v_1^1)))) } * dx(<Mesh #0>[3], {})
A_(1, 2)={ v_1^2 * (conj((div(v_0^1)))) } * dx(<Mesh #0>[3], {})
A_(2, 0)=None
A_(2, 1)={ (conj((v_0^2))) * (div(v_1^1)) } * dx(<Mesh #0>[3], {})
A_(2, 2)=None
L_0 { (10 * sin(15.707963267948966 * x[1]) + x[0] ** 3) * (conj((v_0^0))) } * dx(<Mesh #0>[1], {})
L_1 { (c_3) : (v_0^1) } * dx(<Mesh #0>[3], {})
L_2 { c_4 * (conj((v_0^2))) } * dx(<Mesh #0>[3], {})

Extract blocks

For a bi-linear form, the output of extract_blocks is an nested list of size M, where each sub-list is also of length M, representing the MxM blocked matrix. For extract blocks on a linear form, we need to supply an integral for each the M components. If we fail to do so, we will get a smaller blocked vector. Therefore we added a integral \(\int_{\Omega_s} 0 * \delta p ~\mathrm{d}x\) to ensure that we get the right block structure.

We create the appropriate inverse maps and compile the forms

mesh_to_heat_entity = np.full(num_cells_local, -1, dtype=np.int32)
mesh_to_heat_entity[heat_cell_map] = np.arange(len(heat_cell_map), dtype=np.int32)
mesh_to_stokes_entity = np.full(num_cells_local, -1, dtype=np.int32)
mesh_to_stokes_entity[stokes_cell_map] = np.arange(len(stokes_cell_map), dtype=np.int32)
entity_maps = {heat_mesh: mesh_to_heat_entity, stokes_mesh: mesh_to_stokes_entity}
a_blocked_compiled = dolfinx.fem.form(a_blocked, entity_maps=entity_maps)
L_blocked_compiled = dolfinx.fem.form(L_blocked, entity_maps=entity_maps)

We create the boundary condition as before

T_bndry = dolfinx.fem.Function(T)
T_bndry.x.array[:] = 0
heat_mesh.topology.create_connectivity(heat_mesh.topology.dim-1, heat_mesh.topology.dim)
heat_bc_dofs = dolfinx.fem.locate_dofs_topological(T, heat_mesh.topology.dim-1, heat_facet_tags.find(outer_marker))
bc_heat = dolfinx.fem.dirichletbc(T_bndry, heat_bc_dofs)

stokes_walls = np.union1d(stokes_facet_tags.find(outer_marker), stokes_facet_tags.find(interface_marker))
dofs_wall = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, stokes_walls)
u_wall = dolfinx.fem.Function(V)
u_wall.x.array[:] = 0
bc_wall = dolfinx.fem.dirichletbc(u_wall, dofs_wall)
u_inlet = dolfinx.fem.Function(V)
u_inlet.interpolate(lambda x: (0.5*x[1], 0*x[0]))
dofs_inlet = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, stokes_facet_tags.find(inlet_marker))
bc_inlet = dolfinx.fem.dirichletbc(u_inlet, dofs_inlet)

bcs = [bc_heat, bc_wall, bc_inlet]

We could now use scipy.sparse.vstack and [scipy.sparse.hstack]https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.hstack.html to create the full matrix in a scipy compatible format, and solve as we have shown before.

However, DOLFINx provides wrapper to assemble this into PETSc blocked and PETSc nest matrices. We will illustrate how to do this with PETSc blocked matrices.

from petsc4py import PETSc
import dolfinx.fem.petsc

A = dolfinx.fem.petsc.create_matrix_block(a_blocked_compiled)
A.zeroEntries()
dolfinx.fem.petsc.assemble_matrix_block(A, a_blocked_compiled, bcs=bcs)
A.assemble()

b = dolfinx.fem.petsc.create_vector_block(L_blocked_compiled)
dolfinx.fem.petsc.assemble_vector_block(b, L_blocked_compiled, a_blocked_compiled, bcs=bcs)
b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)

Assembly of blocked matrices and vectors

We observe that assembling blocked matrices looks like how we assemble native matrices. However, we observe that the lifting procedure is embedded in assemble_block_vector, as this code gets a bit compilcated for block systems.

We create the KSP object and solve the system

ksp = PETSc.KSP().create(mesh.comm)
ksp.setOperators(A)
ksp.setType(PETSc.KSP.Type.PREONLY)
ksp.getPC().setType(PETSc.PC.Type.LU)
ksp.getPC().setFactorSolverType("mumps")

w_blocked = dolfinx.fem.petsc.create_vector_block(L_blocked_compiled)

ksp.solve(b, w_blocked)
assert ksp.getConvergedReason() > 0, "Solve failed"

We extract the individual functions from the blocked solution vector and visualize them

blocked_maps = [(space.dofmap.index_map, space.dofmap.index_map_bs) for space in W.ufl_sub_spaces()]
local_values = dolfinx.cpp.la.petsc.get_local_vectors(w_blocked, blocked_maps)

Th = dolfinx.fem.Function(T, name="Temperature")
uh = dolfinx.fem.Function(V, name="Velocity")
ph = dolfinx.fem.Function(Q, name="Pressure")
Th.x.array[:] = local_values[0]
uh.x.array[:] = local_values[1]
ph.x.array[:] = local_values[2]