Custom Newton Solver - How to implement mixed element Dirichlet BCs

Hello,

I’ve implemented the compressible Euler equations for a simple problem of constant flow through a unit square. The function space is specified as a “mixed element” but equal order. Stabilization has not yet been included because this is the simplest test case that should attain the correct solution. The boundary conditions are all Dirichlet conditions:

  • Inlet (left): velocity = (1,0), density = 1

  • Outlet (right): energy ~= 45

  • slip walls (Top/bottom): y-velocity = 0

When I implement a custom Newton solver with Dirichlet BCs, following https://jsdokken.com/dolfinx-tutorial/chapter4/newton-solver.html#newtons-method-with-dirichletbc, the boundary conditions are violated and the solution is not constant and blows up. For example, density is shown here:

I have the MWE here:

import numpy as np
from dolfinx import mesh, fem, nls, plot, cpp
import ufl
from mpi4py import MPI
from petsc4py import PETSc
import pyvista

dt = 0.01

# problem parameters
rho = 1
gamHt = 1.4
Pr = 0.72
Uref=1
Mach=0.2
R=288
cv = R/(gamHt - 1)
Tin = (1/Mach**2)*(1/(gamHt*R))
Pout = (rho/gamHt)*(1/Mach**2)
uin = 1
a = np.sqrt(gamHt*R*Tin)
Econst = rho*Tin*cv + rho/2 * uin**2;

#--------------------------------------------------------------------------------
# Domain
#--------------------------------------------------------------------------------
nx, ny = 5,5
domain = mesh.create_unit_square(MPI.COMM_WORLD, nx, ny, mesh.CellType.triangle)

# Extract Boundaries
#--------------------------------------------------------------------------------
tdim = domain.topology.dim
fdim = tdim - 1
domain.topology.create_connectivity(fdim, tdim)
n = ufl.FacetNormal(domain)

boundaries = [(1, lambda x: np.isclose(x[0], 0)), # inlet
              (2, lambda x: np.isclose(x[0], 1)), # outlet
              (3, lambda x: np.isclose(x[1], 0)), # bottom
              (4, lambda x: np.isclose(x[1], 1))]  # top

facet_indices, facet_markers = [], []
for (marker, locator) in boundaries:
    facets = mesh.locate_entities(domain, fdim, locator)
    facet_indices.append(facets)
    facet_markers.append(np.full_like(facets, marker))
facet_indices = np.hstack(facet_indices).astype(np.int32)
facet_markers = np.hstack(facet_markers).astype(np.int32)
sorted_facets = np.argsort(facet_indices)
facet_tag = mesh.meshtags(domain, fdim, facet_indices[sorted_facets], facet_markers[sorted_facets])

# create boundary measures
ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tag)

#--------------------------------------------------------------------------------
# weak form setup
#--------------------------------------------------------------------------------
# Vector function space
P = ufl.FiniteElement("Lagrange",domain.ufl_cell(),1)
V = ufl.VectorElement("Lagrange",domain.ufl_cell(),1)
E = ufl.FiniteElement("Lagrange",domain.ufl_cell(),1)
ME = ufl.MixedElement([P,V,E])
W = fem.FunctionSpace(domain, ME)

# variables
U_n = fem.Function(W) # conservation vars
U_np1 = fem.Function(W) 
wma, wmo, wen = ufl.TestFunctions(W)

# split functions... 
ma_np1, mo_np1, en_np1 = ufl.split(U_np1)
ma_n, mo_n, en_n = ufl.split(U_n)

# ...and set time integration and midpoint states
ma_h = (ma_np1+ma_n)/2
mo_h = (mo_np1+mo_n)/2
en_h = (en_np1+en_n)/2

dma = (ma_np1-ma_n)/dt
dmo = (mo_np1-mo_n)/dt
den = (en_np1-en_n)/dt

# Residual
#--------------------------------------------------------------------------------
Pval = (R/cv)*( (en_h - ufl.inner(mo_h , mo_h )/(2*ma_h) ))
PRESSURE = Pval*ufl.Identity(2)
f = fem.Constant(domain, PETSc.ScalarType(0))               

# Galerkin - interior
Res = ufl.inner(dma,wma)*ufl.dx # unsteady
Res += - ufl.inner(ufl.grad(wma),mo_h)*ufl.dx  # continuity flux

Res += ufl.inner(dmo,wmo)*ufl.dx # unsteady Momentum
Res += - ufl.inner(ufl.grad(wmo),( 1/ma_h * ufl.outer( mo_h, mo_h ) ) )*ufl.dx \
                   - ufl.inner(PRESSURE,ufl.grad(wmo))*ufl.dx 

Res += ufl.inner(den,wen)*ufl.dx # unsteady Energy
Res += - ufl.inner(ufl.grad(wen), ((1/ma_h)*mo_h*( en_h + Pval )) )*ufl.dx

#--------------------------------------------------------------------------------
# Dirichlet BCs
#--------------------------------------------------------------------------------

# inlet
#--------------------------------------------------------------------------------
#velocity
uD_iu = fem.Function(W)
facets_in = facet_tag.find(1)
dofs_u_in = fem.locate_dofs_topological(W.sub(1), fdim, facets_in)
u_In = fem.Constant(domain, PETSc.ScalarType((1, 0)))
uIn_expr = fem.Expression(u_In, W.sub(1).element.interpolation_points())
uD_iu.sub(1).interpolate(uIn_expr)
bc_u_in = fem.dirichletbc(uD_iu.sub(1), dofs_u_in)

# density
uD_irho = fem.Function(W)
dofs_rho_in = fem.locate_dofs_topological(W.sub(0), fdim, facets_in)
rho = fem.Constant(domain, PETSc.ScalarType(1))
rho_expr = fem.Expression(rho, W.sub(0).element.interpolation_points())
uD_irho.sub(0).interpolate(rho_expr)
bc_rho_in = fem.dirichletbc(uD_irho.sub(0), dofs_rho_in)

# top slip
#--------------------------------------------------------------------------------
uDt = fem.Function(W)
facets_top = facet_tag.find(3)
dofs_u_top = fem.locate_dofs_topological(W.sub(1).sub(1), fdim, facets_top)
uTop = fem.Constant(domain, PETSc.ScalarType(0))
uTop_expr = fem.Expression(uTop, W.sub(1).sub(1).element.interpolation_points())
uDt.sub(1).sub(1).interpolate(uTop_expr)
bc_top = fem.dirichletbc(uDt.sub(1).sub(1), dofs_u_top)

# bottom slip
#--------------------------------------------------------------------------------
uDb = fem.Function(W)
facets_bot = facet_tag.find(4)
dofs_u_bot = fem.locate_dofs_topological(W.sub(1).sub(1), fdim, facets_bot)
uDb.sub(1).sub(1).interpolate(uTop_expr)
bc_bot = fem.dirichletbc(uDb.sub(1).sub(1), dofs_u_bot)

# outlet energy
#-------------------------------------------------------------------------------
uD_eo = fem.Function(W)
facets_out = facet_tag.find(2)
dofs_ene_out = fem.locate_dofs_topological(W.sub(2), fdim, facets_out)
Eval = fem.Constant(domain, PETSc.ScalarType(Econst))
Eval_expr = fem.Expression(uTop, W.sub(1).sub(1).element.interpolation_points())
uD_eo.sub(2).interpolate(Eval_expr)
bc_out = fem.dirichletbc(uD_eo.sub(2), dofs_ene_out)

bcs = [bc_u_in,bc_rho_in,bc_top, bc_bot,bc_out]

#--------------------------------------------------------------------------------
# Nonlinear problem setup
#--------------------------------------------------------------------------------
J = ufl.derivative(Res, U_np1)
residual = fem.form(Res)
jacobian = fem.form(J)

# Iteration dependent structures
A = fem.petsc.create_matrix(jacobian)
L = fem.petsc.create_vector(residual)

solver = PETSc.KSP().create(domain.comm)
solver.setOperators(A)
DeltaA = fem.Function(W)

max_iterations = 10

#  initial conditions
#--------------------------------------------------------------------------------
# used from BC
U_n.sub(0).interpolate(rho_expr)
U_n.sub(1).interpolate(uIn_expr) 
U_n.sub(2).interpolate(Eval_expr) 

#--------------------------------------------------------------------------------
# Time Loop
#--------------------------------------------------------------------------------
for n in range(1):

    # Solve NON-linear problem
    #--------------------------------------------------------------------------------
    A = fem.petsc.create_matrix(jacobian)
    L = fem.petsc.create_vector(residual)

    solver = PETSc.KSP().create(domain.comm)
    solver.setOperators(A)

    # begin newton loop
    i = 0
    while i < max_iterations:
        # Assemble Jacobian and residual
        with L.localForm() as loc_L:
            loc_L.set(0)
        A.zeroEntries()
        fem.petsc.assemble_matrix(A, jacobian, bcs = bcs)
        A.assemble()
        fem.petsc.assemble_vector(L, residual)
        L.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)

        # Scale residual by -1
        L.scale(-1)

        # Compute b - J(u_D-u_(i-1))
        fem.petsc.apply_lifting(L, [jacobian], [bcs], x0=[U_np1.vector], scale=1) 

        # Set dx|_bc = u_{i-1}-u_D
        fem.petsc.set_bc(L, bcs, U_np1.vector, 1.0)
        L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

        # Solve linearized problem
        solver.solve(L, DeltaA.vector)
        DeltaA.x.scatter_forward()

        # Update du_{i+1} = du_i + delta x_i
        U_np1.x.array[:] += DeltaA.x.array
        i += 1

        # Compute norm of update
        correction_norm = DeltaA.vector.norm(0)
        print(f"Iteration {i}: Correction norm {correction_norm}")
        if correction_norm < 1e-10:
            break

    # update for next time step
    U_n.x.array[:] = U_np1.x.array

# plot
#--------------------------------------------------------------------------------
u_topology, u_cell_types, u_geometry = plot.create_vtk_mesh(domain)
u_grid = pyvista.UnstructuredGrid(u_topology, u_cell_types, u_geometry)
u_grid.point_data["density"] = U_n.sub(0).collapse().x.array.real
u_grid.set_active_scalars("density")
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(u_grid, show_edges=True)
u_plotter.view_xy()
u_plotter.show()

I think the problem is how the BCs are applied within the Newton loop. Any ideas on how to fix this? Thanks!