Facet normal and facet tangent vector approximation in dolfinx

normals_and_tangents.py

# Copyright (C) 2023 Jørgen S. Dokken and Halvor Herlyng
#
# SPDX-License-Identifier:    MIT

import ufl
import meshio

import numpy   as np
import dolfinx as dfx

from mpi4py   import MPI
from petsc4py import PETSc

from dolfinx.fem.petsc     import assemble_matrix, assemble_vector, apply_lifting, set_bc
from dolfinx.cpp.fem.petsc import insert_diagonal

def tangential_projection(u: ufl.Coefficient, n: ufl.FacetNormal) -> ufl.Coefficient:
    """
    See for instance:
    https://link.springer.com/content/pdf/10.1023/A:1022235512626.pdf
    """
    return (ufl.Identity(u.ufl_shape[0]) - ufl.outer(n, n)) * u

def facet_vector_approximation(V: dfx.fem.FunctionSpace,
                               mt: dfx.mesh.MeshTags | None = None,
                               mt_id: int | None = None,
                               tangent: bool = False,
                               interior: bool = False,
                               jit_options: dict | None = None,
                               form_compiler_options: dict | None = None) -> dfx.fem.Function:
    """
    Approximate the facet normal or tangent on a surface by projecting it into the function space for a set of facets.

    Args:
        V: The function space to project into. Needs to be defined on discontinuous elements.

        mt: The `dolfinx.mesh.MeshTags` containing facet markers.

        mt_id: The id for the facets in `mt` we want to represent the normal or tangent on.

        tangent: To approximate the tangent of the facets set this flag to `True`.

        jit_options: Parameters used in CFFI JIT compilation of C code generated by FFCx.
            See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py#L22-L37
            for all available parameters. Takes priority over all other parameter values.

        form_compiler_options: Parameters used in FFCx compilation of this form. Run `ffcx - -help` at
            the commandline to see all available options. Takes priority over all
            other parameter values, except for `scalar_type` which is determined by
            DOLFINx.
"""
    
    if not V.element.basix_element.discontinuous:
        raise RuntimeError("Discontinuous elements are required for a proper representation of the facet vectors.")
    jit_options = jit_options if jit_options is not None else {}
    form_compiler_options = form_compiler_options if form_compiler_options is not None else {}

    comm  = V.mesh.comm # MPI Communicator
    n     = ufl.FacetNormal(V.mesh) # UFL representation of mesh facet normal
    nh    = dfx.fem.Function(V)     # Function for the facet vector approximation
    u, v  = ufl.TrialFunction(V), ufl.TestFunction(V) # Trial and test functions

    # Define integral measure
    if interior:
        # Create interior facet integral measure
        dS = ufl.dS(domain=V.mesh) if mt is None else ufl.dS(domain=V.mesh, subdomain_data=mt, subdomain_id=mt_id)
    else:
        # Create exterior facet integral measure
        ds = ufl.ds(domain=V.mesh) if mt is None else ufl.ds(domain=V.mesh, subdomain_data=mt, subdomain_id=mt_id) 
    
    # If tangent==True, the right-hand side of the problem should be a tangential projection of the facet normal vector.
    if tangent:
        if V.mesh.geometry.dim == 1:
            raise ValueError("Tangent not defined for 1D problem")
        elif V.mesh.geometry.dim == 2:
            if interior:
                a = (ufl.inner(u('+'), v('+')) + ufl.inner(u('-'), v('-'))) * dS
                L = ufl.inner(ufl.as_vector([-n('+')[1], n('+')[0]]), v('+')) * dS \
                  + ufl.inner(ufl.as_vector([-n('-')[1], n('-')[0]]), v('-')) * dS
            else:
                a = ufl.inner(u, v) * ds
                L = ufl.inner(ufl.as_vector([-n[1], n[0]]), v) * ds
        else:
            c = dfx.fem.Constant(V.mesh, (1.0, 1.0, 1.0)) # Vector to tangentially project the facet normal vectors on
            if interior:
                a = (ufl.inner(u('+'), v('+')) + ufl.inner(u('-'), v('-'))) * dS
                L = ufl.inner(tangential_projection(c, n('+')), v('+')) * dS \
                  + ufl.inner(tangential_projection(c, n('-')), v('-')) * dS
            else:
                a = ufl.inner(u, v) * ds
                L = ufl.inner(tangential_projection(c, n), v) * ds

    # If tangent==false the right-hand side is simply the facet normal vector.
    else:
        if interior:
            a = (ufl.inner(u('+'), v('+')) + ufl.inner(u('-'), v('-'))) * dS
            L = (ufl.inner(n('+'), v('+')) + ufl.inner(n('-'), v('-'))) * dS
        else:
            a = ufl.inner(u, v) * ds
            L = ufl.inner(n, v) * ds

    # Find all boundary dofs, which are the dofs where we want to solve for the facet vector approximation.
    # Start by assembling test functions integrated over the boundary integral measure.
    ones = dfx.fem.Constant(V.mesh, dfx.default_scalar_type((1,) * V.mesh.geometry.dim)) # A vector of ones
    if interior:
        local_val = dfx.fem.form((ufl.dot(ones, v('+')) + ufl.dot(ones, v('-')))*dS)
    else:
        local_val = dfx.fem.form(ufl.dot(ones, ufl.TestFunction(V))*ds) 
    local_vec = dfx.fem.assemble_vector(local_val)

    # For the dofs that do not lie on the boundary of the mesh the assembled vector has value zero.
    # Extract these dofs and use them to deactivate the corresponding block in the linear system we will solve.
    bdry_dofs_zero_val  = np.flatnonzero(np.isclose(local_vec.array, 0))
    deac_blocks = np.unique(bdry_dofs_zero_val // V.dofmap.bs).astype(np.int32)

    # Create sparsity pattern by manipulating the blocks to be deactivated and set
    # a zero Dirichlet boundary condition for these dofs.
    bilinear_form = dfx.fem.form(a, jit_options=jit_options,
                                 form_compiler_options=form_compiler_options)
    pattern = dfx.fem.create_sparsity_pattern(bilinear_form)
    pattern.insert_diagonal(deac_blocks)
    pattern.finalize()
    u_0 = dfx.fem.Function(V)
    u_0.vector.set(0)
    bc_deac = dfx.fem.dirichletbc(u_0, deac_blocks)

    # Create the matrix
    A = dfx.cpp.la.petsc.create_matrix(comm, pattern)
    A.zeroEntries()

    # Assemble the matrix with all entries
    form_coeffs = dfx.cpp.fem.pack_coefficients(bilinear_form._cpp_object)
    form_consts = dfx.cpp.fem.pack_constants(bilinear_form._cpp_object)
    assemble_matrix(A, bilinear_form, constants=form_consts, coeffs=form_coeffs, bcs=[bc_deac])

    # Insert the diagonal with the deactivated blocks.
    if bilinear_form.function_spaces[0] is bilinear_form.function_spaces[1]:
        A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH)
        A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH)
        insert_diagonal(A=A, V=bilinear_form.function_spaces[0], bcs=[bc_deac._cpp_object], diagonal=1.0)
    A.assemble()

    # Assemble the linear form and the right-hand side vector.
    linear_form = dfx.fem.form(L, jit_options=jit_options,
                               form_compiler_options=form_compiler_options)
    b = assemble_vector(linear_form)


    # Apply lifting to the right-hand side vector and set boundary conditions.
    apply_lifting(b, [bilinear_form], [[bc_deac]])
    b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b, [bc_deac])

    # Setup a linear solver using the Conjugate Gradient method.
    solver = PETSc.KSP().create(MPI.COMM_WORLD)
    solver.setType("cg")
    solver.rtol = 1e-8
    solver.setOperators(A)

    # Solve the linear system and perform ghost update.
    solver.solve(b, nh.vector)
    nh.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    print(np.unique(nh.vector[:]))

    # Normalize the vectors to get the unit facet normal/tangent vector.
    nh_norm = ufl.sqrt(ufl.inner(nh, nh)) # Norm of facet vector
    if V.mesh.geometry.dim == 1:
        nh_norm_vec = ufl.as_vector((nh[0]/nh_norm))
    elif V.mesh.geometry.dim == 2:
        nh_norm_vec = ufl.as_vector((nh[0]/nh_norm, nh[1]/nh_norm))
    elif V.mesh.geometry.dim == 3:
        nh_norm_vec = ufl.as_vector((nh[0]/nh_norm, nh[1]/nh_norm, nh[2]/nh_norm))

    nh_normalized = dfx.fem.Expression(nh_norm_vec, V.element.interpolation_points())
    n_out = dfx.fem.Function(V)
    n_out.interpolate(nh_normalized)

    return n_out


# Case flags
all_facets   = False # Approximate all facet vectors if True, subset if False
tangent_flag = True # Approximate normals if False, approximate tangents if True
interior     = True  # Set to True if the facets are internal (e.g. an interface between two domains)
                     # Set to False if the facets are on the mesh boundary
dim = 2 # Spatial dimension of mesh

if __name__ == '__main__':

    if dim == 2:
        # Create a unit square mesh
        mesh = dfx.mesh.create_unit_square(MPI.COMM_WORLD, nx=8, ny=8, ghost_mode=dfx.mesh.GhostMode.shared_facet)
    elif dim == 3:
        # Create a unit cube mesh
        mesh = dfx.mesh.create_unit_cube(MPI.COMM_WORLD, nx=8, ny=8, nz=8, ghost_mode=dfx.mesh.GhostMode.shared_facet)
    else:
        raise NotImplementedError("Set dim equal to 2 or 3.")
    
    if all_facets:
        # Compute normals/tangents for the whole boundary of the mesh.
        # Initialize None type meshtags and meshtag id.
        facet_tags = None
        ft_id      = None

    else:
        # Compute normals/tangents only on a part of the mesh.
        # Define marker values
        DEFAULT = 2
        SUBSET  = 3

        if interior:
            # Mark the interior facets lying on an interface inside the square/cube.
            # NOTE: this is a pretty ad-hoc way of marking an interface.
            def locator(x):
                """ Marker function that returns True if the x-coordinate is between 0.8 and 0.9. """
                return np.logical_and(x[0] > 0.8, x[0] < 0.9)

        else:
            # Mark one side of the square/cube and approximate only the normals/tangents on that part of the mesh boundary.
            def locator(x):
                """ Marker function that returns True if the x-coordinate is 1. """
                return np.isclose(x[0], 1.0)

        # Create necessary topological entities of the mesh
        mesh.topology.create_entities(mesh.topology.dim-1)
        mesh.topology.create_connectivity(mesh.topology.dim-1, mesh.topology.dim)
        
        # Create an array facet_marker which contains the facet marker value for all facets in the mesh.
        facet_dim = mesh.topology.dim-1 # Topological dimension of facets
        num_facets   = mesh.topology.index_map(facet_dim).size_local + mesh.topology.index_map(facet_dim).num_ghosts # Total number of facets in mesh
        facet_marker = np.full(num_facets, DEFAULT, dtype=np.int32) # Default facet marker value is 2

        # Find the subset of facets to be marked.
        subset_facets = dfx.mesh.locate_entities(mesh, facet_dim, locator) # Get subset of facets to be marked
        facet_marker[subset_facets] = SUBSET # Fill facet marker array with the value SUBSET
        facet_tags = dfx.mesh.meshtags(mesh, facet_dim, np.arange(num_facets, dtype=np.int32), facet_marker) # Create facet meshtags
        ft_id = SUBSET # Set the meshtags id for which we will approximate the facet vectors
    
    # Create a DG1 space for the facet vectors to be approximated.
    DG1 = ufl.VectorElement(family='Discontinuous Lagrange', cell=mesh.ufl_cell(), degree=1)
    space = dfx.fem.FunctionSpace(mesh=mesh, element=DG1)

    # Compute the facet vector approximation.
    nh = facet_vector_approximation(V=space, mt=facet_tags, mt_id=ft_id, interior=interior, tangent=tangent_flag)

    # Write the results to file
    filename = 'tangents.bp' if tangent_flag else 'normals.bp'
    with dfx.io.VTXWriter(mesh.comm, filename, [nh], engine='BP4') as vtx:
        vtx.write(0)