Thermal contact resistance problem

rtc.py

import dolfinx.fem.petsc
from mpi4py import MPI

import dolfinx
import numpy as np
import ufl
from petsc4py import PETSc

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)
    submesh.topology.create_connectivity(entity_tag.dim, submesh.topology.dim)
    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)

    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


N = 50
M = 2 * N
mesh = dolfinx.mesh.create_unit_square(
    MPI.COMM_WORLD, M, M, dolfinx.mesh.CellType.triangle, ghost_mode=dolfinx.mesh.GhostMode.shared_facet
)


# def inner_square(x, tol=1e-13):
#     return (0.3 - tol <= x[0]) & (x[0] <= 0.8 + tol) & (x[1] <= 0.6 + tol) & (0.4 - tol <= x[1])


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


def right_square(x, tol=1e-13):
    return x[0] >= 0.5 - tol


tdim = mesh.topology.dim
mesh.topology.create_connectivity(tdim - 1, tdim)
cell_map = mesh.topology.index_map(tdim)
num_cells = cell_map.size_local + cell_map.num_ghosts
all_cells = np.arange(num_cells, dtype=np.int32)
marker = np.full(num_cells, 2, dtype=np.int32)
marker[dolfinx.mesh.locate_entities(mesh, tdim, left_square)] = 1
marker[dolfinx.mesh.locate_entities(mesh, tdim, right_square)] = 2
cell_tags = dolfinx.mesh.meshtags(
    mesh, tdim, np.arange(num_cells, dtype=np.int32), marker)

fdim = mesh.topology.dim - 1
facet_map = mesh.topology.index_map(fdim)
num_facets = facet_map.size_local + facet_map.num_ghosts
all_facets = np.arange(num_facets, dtype=np.int32)
facet_marker = np.ones(num_facets, dtype=np.int32)
left_marker = 2
right_marker = 3
facet_marker[dolfinx.mesh.locate_entities_boundary(
    mesh, tdim-1, lambda x: np.isclose(x[0], 0))] = left_marker
facet_marker[dolfinx.mesh.locate_entities_boundary(
    mesh, tdim-1, lambda x: np.isclose(x[0], 1))] = right_marker
facet_tags = dolfinx.mesh.meshtags(
    mesh, tdim-1, np.arange(num_facets, dtype=np.int32), facet_marker)


all_tags = np.sort([2, 1])  # Should be the same on every process

submeshes = []
sub_to_parent_maps = []
sub_facet_tags = []
sub_to_parent_facet_maps = []
for tag in all_tags:
    sm, s2pe, s2pv, _ = dolfinx.mesh.create_submesh(
        mesh, cell_tags.dim, cell_tags.find(tag))
    submeshes.append(sm)
    sub_to_parent_maps.append(s2pe)
    sft, sfm = transfer_meshtags_to_submesh(mesh, facet_tags, sm, s2pv, s2pe)
    sub_facet_tags.append(sft)
    sub_to_parent_facet_maps.append(sfm)

interface_marker = np.full(num_facets, -2, dtype=np.int32)
counter = 0
for i in range(len(submeshes)):
    for j in range(i+1, len(submeshes)):

        interface = np.intersect1d(
            sub_to_parent_facet_maps[i], sub_to_parent_facet_maps[j])
        interface_marker[interface] = counter
        counter += 1
ft_2 = dolfinx.mesh.meshtags(mesh, fdim, all_facets, interface_marker)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "debug.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_meshtags(ft_2, mesh.geometry)

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

Ts = [dolfinx.fem.functionspace(submesh, ("Lagrange", 1))
      for submesh in submeshes]
uhs = [dolfinx.fem.Function(Ti, name=f"T{i}") for i, Ti in enumerate(Ts)]
W = ufl.MixedFunctionSpace(*Ts)
us = ufl.TrialFunctions(W)
vs = ufl.TestFunctions(W)

# Add standard terms to variational form
dx = ufl.Measure("dx", domain=mesh, subdomain_data=cell_tags)
a = sum([ufl.inner(ufl.grad(us[i]), ufl.grad(vs[i]))*dx(all_tags[i])
         for i in range(len(submeshes))])


L = sum([dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0))
        * vs[i]*dx(all_tags[i]) for i in range(len(submeshes))])


entity_maps = {}
for i, submesh in enumerate(submeshes):
    parent_map = np.full(num_cells, -1, dtype=np.int32)
    parent_map[sub_to_parent_maps[i]] = np.arange(
        len(sub_to_parent_maps[i]), dtype=np.int32)
    entity_maps[submesh] = parent_map


# Compute integration entities and consistently switch the order
# Additionally pad entity maps for sparsity pattern
integral_data = []
counter = 0
for i in range(len(submeshes)):
    for j in range(i+1, len(submeshes)):
        integration_data = dolfinx.fem.compute_integration_domains(
            dolfinx.fem.IntegralType.interior_facet, mesh.topology, ft_2.find(counter), ft_2.dim)
        ordered_integration_data = integration_data.reshape(-1, 4).copy()
        switch = cell_tags.values[ordered_integration_data[:, 0]
                                  ] < cell_tags.values[ordered_integration_data[:, 2]]
        # Order restriction on one side
        if True in switch:
            ordered_integration_data[switch, :] = ordered_integration_data[switch][
                :, [2, 3, 0, 1]
            ]
        integral_data.append((counter, ordered_integration_data.flatten()))
        counter += 1
        # Pad entity maps for sparsity pattern
        parent_cells_plus = ordered_integration_data[:, 0]
        parent_cells_minus = ordered_integration_data[:, 2]
        entity_maps[submeshes[i]
                    ][parent_cells_plus] = entity_maps[submeshes[i]][parent_cells_minus]
        entity_maps[submeshes[j]
                    ][parent_cells_minus] = entity_maps[submeshes[j]][parent_cells_plus]


dS = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)

RTC = dolfinx.fem.Constant(mesh, 0.5)
counter = 0
for i in range(len(submeshes)):
    for j in range(i+1, len(submeshes)):
        a -= (us[i]("-") - us[j]("+"))/RTC * \
            (vs[i]("-") - vs[j]("+"))*dS(counter)
        counter += 1


left_bc_val = 0
right_bc_val = 1

# Create bcs
bcs = []
for i in range(len(submeshes)):
    sub_dofs = dolfinx.fem.locate_dofs_topological(
        Ts[i], fdim, sub_facet_tags[i].find(left_marker))
    bc_left = dolfinx.fem.dirichletbc(
        dolfinx.default_scalar_type(left_bc_val), sub_dofs, Ts[i])

    sub_dofs_right = dolfinx.fem.locate_dofs_topological(
        Ts[i], fdim, sub_facet_tags[i].find(right_marker))
    bc_right = dolfinx.fem.dirichletbc(
        dolfinx.default_scalar_type(right_bc_val), sub_dofs_right, Ts[i])
    bcs.append(bc_left)
    bcs.append(bc_right)


a = dolfinx.fem.form(ufl.extract_blocks(a), entity_maps=entity_maps)
L = dolfinx.fem.form(ufl.extract_blocks(L), entity_maps=entity_maps)
A = dolfinx.fem.petsc.create_matrix_block(a)
b = dolfinx.fem.petsc.create_vector_block(L)

A.zeroEntries()
with b.localForm() as loc:
    loc.set(0)
dolfinx.fem.petsc.assemble_matrix_block(A, a, bcs=bcs)
dolfinx.fem.petsc.assemble_vector_block(b, L, a, bcs=bcs)
A.assemble()
b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES,
              PETSc.ScatterMode.FORWARD)
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(PETSc.Mat.SolverType.MUMPS)
w_blocked = dolfinx.fem.petsc.create_vector_block(L)
ksp.solve(b, w_blocked)
w_blocked.ghostUpdate(PETSc.InsertMode.INSERT_VALUES,
                      PETSc.ScatterMode.FORWARD)
assert (reason := ksp.getConvergedReason()
        ) > 0, f"Solver did not converge: {reason}"
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)
for i in range(len(submeshes)):
    uhs[i].x.array[:] = local_values[i]
    with dolfinx.io.VTXWriter(mesh.comm, f"T{i}.bp", [uhs[i]]) as bp:
        bp.write(0.0)

rtc_dg.py

from mpi4py import MPI

import dolfinx
import ufl
import numpy as np

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.triangle)


V = dolfinx.fem.functionspace(mesh, ("DG", 1))


u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)

Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
RTC = dolfinx.fem.Function(Q)
RTC.x.array[:] = 0.5

def interface(x):
    return np.isclose(x[0], 0.5)


mesh.topology.create_connectivity(mesh.topology.dim-1, mesh.topology.dim)
interface_facets = dolfinx.mesh.locate_entities(mesh, mesh.topology.dim-1,interface)
dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim-1, interface_facets)

def left(x):
    return np.isclose(x[0], 0.0)

def right(x):
    return np.isclose(x[0], 1.0)

left_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim-1, left)
right_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim-1, right)
BCSPACE = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
uD = dolfinx.fem.Function(BCSPACE)
left_dofs = dolfinx.fem.locate_dofs_topological(BCSPACE, mesh.topology.dim-1, left_facets)
right_dofs = dolfinx.fem.locate_dofs_topological(BCSPACE, mesh.topology.dim-1, right_facets)
uD.x.array[left_dofs] = 0.0
uD.x.array[right_dofs] = 1.0

num_facets = mesh.topology.index_map(mesh.topology.dim-1).size_local + mesh.topology.index_map(mesh.topology.dim-1).num_ghosts
facets_values = np.ones(num_facets, dtype=np.int32)
dir_val = 2
interface_val = 3
facets_values[left_facets] = dir_val
facets_values[right_facets] = dir_val
facets_values[interface_facets] = interface_val

ft = dolfinx.mesh.meshtags(mesh, mesh.topology.dim-1, np.arange(num_facets, dtype=np.int32),
                           facets_values)

ds = ufl.Measure("ds", domain=mesh, subdomain_data=ft, subdomain_id=dir_val)
dS = ufl.Measure("dS", domain=mesh, subdomain_data=ft)
dSi = dS(interface_val)
dSo = dS(1)

h = 2*ufl.Circumradius(mesh)
alpha = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(10))

n = ufl.FacetNormal(mesh)
F = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
F -= ufl.inner(ufl.avg(ufl.grad(u)), ufl.jump(v, n)) * dSo
F -= ufl.inner(ufl.avg(ufl.grad(v)), ufl.jump(u, n)) * dSo
F += alpha/ufl.avg(h) *ufl.inner(ufl.jump(u), ufl.jump(v)) * dSo
F += - ufl.inner(n, ufl.grad(v)) * u * ds + alpha / h * ufl.inner(u, v) * ds
x = ufl.SpatialCoordinate(mesh)
f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0))# 10 * ufl.cos(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
F -= ufl.inner(f, v) * ufl.dx
F -= - ufl.inner(n, ufl.grad(v)) * uD * ds + alpha / h * ufl.inner(uD, v) * ds

F -= ufl.jump(u)/RTC*ufl.jump(v)*dSi

a, L = ufl.system(F)

uh = dolfinx.fem.Function(V)
import dolfinx.fem.petsc
problem = dolfinx.fem.petsc.LinearProblem(a, L, u=uh, bcs=[],
                                          petsc_options={"ksp_type": "preonly", "pc_type": "lu",
                                                         "pc_factor_mat_solver_type": "mumps"})

problem.solve()

with dolfinx.io.VTXWriter(mesh.comm, "uh.bp", [uh]) as bp:
    bp.write(0.0)