How to implement the equivalence of `assemble_system` in DOLFINx

Ru_Xiang · April 23, 2022, 11:25pm

Hello, I met an issue when trying to convert my old dolfin code into DOLFINx. Here I was trying to assemble the Jacobian matrix with Dirichlet boundary conditions, but got different results with the two libraries. In the old dolfin code, I used the assemble_system() which would preserve the symmetry of the Jacobian matrix (the A matrix in the code example). While in DOLFINx, I used assemble_matrix() from dolfinx.fem.petsc, where the symmetry is also preserved, but the numerical values are different.

I also noticed that the ordering of function DOFs seems to be different in DOLFINx, but still that would only lead to index changes in the matrix while the values should stay the same.

Here are the minimal examples of my implementations:

DOLFINx code:
version - commit 92202566271a6c80a2ec3a2cade1bbb60c6e5b9e

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

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(0.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.petsc.assemble_matrix(dolfinx.fem.form(dRdu), 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))

where the results are:

 ------ Matrix A by DOLFINx ------- 
[[1. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 4. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 1.]]

old 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(0.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)

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()))

which gives:

 ------ Matrix A by old dolfin - assemble ------- 
[[ 1.  0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  1.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  1.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  1.  0.  0.  0.  0.  0.]
 [ 0. -1. -1.  0.  4.  0. -1. -1.  0.]
 [ 0.  0.  0.  0.  0.  1.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  1.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  1.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.  1.]]
 ------ Matrix A by old dolfin - assemble_system ------- 
 ----------- expected results ---------------- 
[[1. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 3. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 3. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 2. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 4. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 2. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 3. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 3. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 1.]]

The matrix by old dolfin using assemble_system is what I want eventually in DOLFINx. I guess I’m mostly confused by the procedures of assemble_system. My understanding of assembly in FEM would be removing/zero-ing rows and cols in the LHS matrix, and applying lifting to the RHS vector. So, I don’t quite understand how the numbers like 3s and 2s are computed since they were not there in the original stiffness matrix. And I was also looking at this related post, where someone explained the difference between assemble+bc.apply() and assemble_system, but that still doesn’t explain the value changes.

Please leave a comment if you have any thoughts on the assemble methods, or suggestions on how to get the same results in DOLFINx!

Thanks ahead!

dokken · April 24, 2022, 9:07am

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.

Ru_Xiang · April 24, 2022, 11:36pm

Thanks you @dokken ! It does make a lot of sense with the residual vector printed out. And my code is running now in DOLFINx too! (the bug in my code turned out to be something else instead of assemble, so it was very helpful to know that the DOLFINx assemble is actually better than assemble_system in the old DOLFIN!)

arash_goli · July 7, 2022, 12:58am

Hello dr @dokken. I have a question about bilinear form matrix at Diffusion of a Gaussian function — FEniCSx tutorial
After assembling matrix A, I would like to know that if there is any way to extract Mass matrix and Stiffness matrix separately in Dolfinx . I used the above code to print matrix A, I mean: print(convertToDense(A)). but matrix A includes both Mass and Stiffness matrices.

nate · July 7, 2022, 1:45am

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

# mesh and space
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2)
V = dolfinx.fem.FunctionSpace(mesh, ("CG", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)

# Stongly imposed BCs
bdry_facets = dolfinx.mesh.locate_entities_boundary(
    mesh, mesh.topology.dim-1, lambda x: np.full_like(x[0], True, dtype=bool))
bdry_dofs = dolfinx.fem.locate_dofs_topological(
	V, mesh.topology.dim-1, bdry_facets)
bc = dolfinx.fem.dirichletbc(PETSc.ScalarType(0), bdry_dofs, V)

# mass matrix
m = u * v * ufl.dx
M = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(m), bcs=[bc])
M.assemble()

# stiffness matrix
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
K = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(a), bcs=[bc])
K.assemble()

arash_goli · July 7, 2022, 5:12am

@nate , Appreciate it Dr. Sime

Topic		Replies	Views
Reapplying boundary conditions to the matrix Linear Algebra	4	465	October 3, 2023
How to apply boundary conditions to left-hand matrix Errors dolfinx	6	219	June 25, 2024
Porting assemble_system and assemble to dolfinx0.7 dolfinx boundary-conditions , assembly	1	193	January 2, 2024
Assemble with dolfin VS dolfinx dolfinx	9	51	October 31, 2024
Regarding the use of dolfinx.fem.assemble_scalar General	2	773	January 23, 2023

How to implement the equivalence of `assemble_system` in DOLFINx

Related topics