How to implement the equivalence of `assemble_system` in DOLFINx

dolfinx
2

To illustrate the difference between the two Jacobians, I have modified your problem slightly, to have a non-zero DirichletBC. If you then consider the residual vector, you will understand what is different between the two codes.
Modified DOLFINx code

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

n = 2
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, n, n)
V = dolfinx.fem.FunctionSpace(mesh, ('CG', 1))
VF = dolfinx.fem.FunctionSpace(mesh, ('DG', 0))

u = dolfinx.fem.Function(V)
v = ufl.TestFunction(V)

# Define the source term


class Expression_f:
    def __init__(self):
        self.alpha = 1e-6

    def eval(self, x):
        return (1/(1+self.alpha*4*np.power(np.pi, 4)) *
                np.sin(np.pi*x[0])*np.sin(np.pi*x[1]))


f_analytic = Expression_f()
f = dolfinx.fem.Function(VF)
f.interpolate(f_analytic.eval)

# Apply zero boundary condition on the outer boundary
tdim = mesh.topology.dim
fdim = tdim - 1
mesh.topology.create_connectivity(fdim, tdim)
boundary_facets = np.flatnonzero(
    dolfinx.mesh.compute_boundary_facets(
        mesh.topology))

boundary_dofs = dolfinx.fem.locate_dofs_topological(V, fdim, boundary_facets)
ubc = dolfinx.fem.Function(V)
ubc.vector.set(1.0)
bc = [dolfinx.fem.dirichletbc(ubc, boundary_dofs)]

# Variational form of Poisson's equation
res = (ufl.inner(ufl.grad(u), ufl.grad(v))-f*v)*ufl.dx

# Quantity of interest: will be used as the Jacobian in adjoint method
dRdu = ufl.derivative(res, u)

# Option 1: assemble A and b seperately
a = dolfinx.fem.form(dRdu)
A = dolfinx.fem.petsc.assemble_matrix(a, bcs=bc)


def convertToDense(A_petsc):
    """
    Convert the PETSc matrix to a dense numpy array
    for debugging purposes
    """
    A_petsc.assemble()
    A_dense = A_petsc.convert("dense")
    return A_dense.getDenseArray()


print(" ------ Matrix A by DOLFINx ------- ")
print(convertToDense(A))
L = dolfinx.fem.form(res)
b = dolfinx.fem.petsc.assemble_vector(L)
dolfinx.fem.petsc.apply_lifting(b, [a], [bc])
b.ghostUpdate(PETSc.InsertMode.ADD_VALUES, PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(b, bc)
print(b.array)

Modified DOLFIN code

from dolfin import *

n = 2
mesh = UnitSquareMesh(n, n)
V = FunctionSpace(mesh, 'CG', 1)
VF = FunctionSpace(mesh, 'DG', 0)

u = Function(V)
v = TestFunction(V)

# Define the source term
w = Expression("sin(pi*x[0])*sin(pi*x[1])", degree=3)
alpha = Constant(1e-6)
f_analytic = Expression("1/(1+alpha*4*pow(pi,4))*w",
                        w=w, alpha=alpha, degree=3)
f = interpolate(f_analytic, VF)

# Apply zero boundary condition on the outer boundary
bc = DirichletBC(V, Constant(1.0), "on_boundary")

# Variational form of Poisson's equation
res = (inner(grad(u), grad(v))-f*v)*dx

# Quantity of interest: will be used as the Jacobian in adjoint method
dRdu = derivative(res, u)

# Option 1: assemble A and b seperately
A = assemble(dRdu)
b = assemble(res)
bc.apply(A, b)

# Option 2: assemble the system
A_, b_ = assemble_system(dRdu, res, bcs=bc)

print(b_.get_local())


def convertToDense(A_petsc):
    """
    Convert the PETSc matrix to a dense numpy array
    for debugging purposes
    """
    A_petsc.assemble()
    A_dense = A_petsc.convert("dense")
    return A_dense.getDenseArray()


print(" ------ Matrix A by old dolfin - assemble ------- ")
print(convertToDense(as_backend_type(A).mat()))
print(" ------ Matrix A by old dolfin - assemble_system ------- ")
print(" ----------- expected results ---------------- ")
print(convertToDense(as_backend_type(A_).mat()))

Looking at the output vector you obtain:
DOLFINx

[1.         1.         1.         3.86538367 1.         1.
 1.         1.         1.        ]

DOLFIN

[ 1.          3.          3.          2.          3.86538367  2.          3.
  3.          1.        ]

Here you observe that in DOLFIN, the DirichletBC values are accumulated in the matrix when apply lifting inside assemble_system, while in DOLFINx, the values are not accumulated,and the values are only set once, to the actual Dirichlet BC value.

The key here is that the way DOLFINx does this is better than DOLFIN, as it each diagonal value of the matrix is not dependent on the number of neighbouring cells (which dolfin does).
Therefore I would not try to obtain the matrices from DOLFIN, and rather go for the improved version in DOLFINx.

3 Likes