Calculate sensitivity when the control is a fem function

dolfinx
2

I think there is few things happening here.

When I run your code as is, I get this result:
Convergence rate: [2.0924435576205673, 2.4596218406630284, 4.589509511644071, nan]

The parameter p in the Taylor test should be bigger.
Your code actually works for:
taylor_test(J_cost, grad_J, m_0, p=1e-4)
at which point the result is:
Convergence rate: [2.000010723617184, 2.0000270773462976, 2.0001589607488435, 2.0006124744482032]

You can also set the tolerance of the adjoint problem solver to:
ksp_rtol = 1e-16 and ksp_atol = 1e-16
You do not want any errors coming from there.

The code does not pass the Taylor test when running in parallel though.
There is bunch of stuff that you need to pay attention to. The PETSc is kind of tricky to deal with regarding ghostUpdates, values that are contain in the arrays before you call assemble, and what part of the vectors and matrices exist on each processor. It is not easy to get right.

General advice I would give:

  • Zero the vectors and matrices before you assemble them.
  • In order to test properly, run at least two Taylor tests back to back, so you can be sure that no values from previous test stays in the arrays and mess with the grad calculation between calls
  • If it is not working, try to check if you have done the ghostUpdates correctly.
  • The Taylor test must reflect the gradient and cost globally, not just the local chunks of them.

Having said that here is code that works for me in serial and parallel:

import dolfinx
from ufl import Measure, inner, grad, div, TestFunction, TrialFunction, adjoint, derivative
from dolfinx.mesh import CellType, create_rectangle
from dolfinx import fem, nls, io
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import time

import dolfinx.nls.petsc
import dolfinx.fem.petsc

# mesh and function space
mesh = create_rectangle(MPI.COMM_WORLD, [[0, 0], [1, 1]], [64, 64], CellType.triangle)
V = fem.FunctionSpace(mesh, ("CG", 1))
u = fem.Function(V)
v = TestFunction(V)

# sensitivity will be calculated wrt this
f = fem.Function(V)  # The forcing f is the design variable
f.x.array[:] = 0.0

dx = Measure("dx")

# Residual of the variational form of Poisson problem
R = inner(grad(u), grad(v))*dx - inner(f,v)*dx

# boundary conditions
def left(x):
    return np.isclose(x[0], 0)
def right(x):
    return np.isclose(x[0], 1)
def top(x):
    return np.isclose(x[1], 1)
def bottom(x):
    return np.isclose(x[1], 0)

fdim = mesh.topology.dim - 1

uD = fem.Function(V)

left_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left)
right_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, right)
top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top)
bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom)

with uD.vector.localForm() as loc:
    loc.set(0.0)

bc1 = fem.dirichletbc(uD, dolfinx.fem.locate_dofs_topological(V, fdim, left_facets))
bc2 = fem.dirichletbc(uD, dolfinx.fem.locate_dofs_topological(V, fdim, right_facets))
bc3 = fem.dirichletbc(uD, dolfinx.fem.locate_dofs_topological(V, fdim, top_facets))
bc4 = fem.dirichletbc(uD, dolfinx.fem.locate_dofs_topological(V, fdim, bottom_facets))

bcs = [bc1, bc2, bc3, bc4]

# Solver for the direct problem
problem_dir = fem.petsc.NonlinearProblem(R, u, bcs)
solver = nls.petsc.NewtonSolver(MPI.COMM_WORLD, problem_dir)
solver.rtol = 1e-16
solver.atol = 1e-16

# Compute solution
solver.solve(u)

# Define desired u profile
def profile(x):
    return x[0]*(1-x[0])*x[1]*(1-x[1])
u_obs = fem.Function(V)
u_obs.interpolate(profile)

# Define cost function
J = (1/2)*inner(u-u_obs, u-u_obs)*dx
J_form = fem.form(J)
J_val = fem.assemble_scalar(J_form)

# Bilinear and linear form of the adjoint problem
lhs = adjoint(derivative(R, u))
rhs = -derivative(J, u)

# Create adjoint problem solver
problem = dolfinx.fem.petsc.LinearProblem(lhs, rhs, bcs=bcs, 
                                          petsc_options={"ksp_type": "preonly", 
                                                         "pc_type": "lu"
                                                         })

# solve adjoint vector
lmbda = problem.solve()

# Compute sensitivity: dJ/dm
# Partial derivative of J w.r.t. m
dJdm_form = fem.form(derivative(J, f))
dJdm = fem.petsc.create_vector(dJdm_form)
dJdm.assemble()

# partial derivative of R w.r.t. m
dRdm_form = fem.form(adjoint(derivative(R, f)))
dRdm = fem.petsc.create_matrix(dRdm_form)

# Calculate the gradient of J with respect to m (the forcing f)
#dJdm += dRdm*lmbda.vector
#dJdm.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

# Function to compute the cost
def J_cost(m_n):
    f.x.array[:] = m_n
    solver.solve(u)
    J_val = fem.assemble_scalar(J_form)

    return J_val

# Function to compute the gradient of the cost w.r.t. f
def grad_J(m_n,dJdm=dJdm):
    f.x.array[:] = m_n

    # Direct solution
    solver.solve(u)

    # Adjoint solution
    lmbda = problem.solve()

    # Reassemble dRdm
    dRdm.zeroEntries()
    fem.petsc.assemble_matrix_mat(dRdm, dRdm_form)
    dRdm.assemble()

    # Compute gradient
    # zero gradient first
    with dJdm.localForm() as dJdm_loc:
        dJdm_loc.set(0)

    fem.petsc.assemble_vector(dJdm, dJdm_form)
    dJdm.assemble()

    dJdm += dRdm*lmbda.vector
    dJdm.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)

    return dJdm

# Taylor test
def taylor_test(cost, grad, m_0, p=1e-6, n=5):
        g0 = grad(m_0)
        l0 = cost(m_0)
        l0 = MPI.COMM_WORLD.allreduce(l0, op=MPI.SUM)
        reminder = []
        perturbance = []
        for i in range(0, n):
            l1 = cost(m_0+p)
            l1 = MPI.COMM_WORLD.allreduce(l1, op=MPI.SUM)
            g = np.sum(g0.array)
            g = MPI.COMM_WORLD.allreduce(g, op=MPI.SUM)
            reminder.append(l1 - l0 - g*p)
            perturbance.append(p)
            p /= 2
        conv_rate = convergence_rates(reminder, perturbance)
        return reminder, perturbance, conv_rate

# Compute convergence rate (they should be 2.0 for small p)
def convergence_rates(r, p):
    cr = [] # convergence rates
    for i in range(1, len(p)):
        cr.append(np.log(r[i] / r[i - 1])
                 / np.log(p[i] / p[i - 1]))
    return cr

m_0 = np.zeros(len(u.x.array))
r = taylor_test(J_cost, grad_J, m_0, p=1e-4)
r = taylor_test(J_cost, grad_J, m_0, p=1e-4)
print("Convergence rate:", r[2])

which for mpirun -n 1 gives:
Convergence rate: [2.0000107236531175, 2.000027077418167, 2.000158960892612, 2.0006124747359526]

and for mpirun -n 10 gives:

Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]
Convergence rate: [1.9999988379660412, 1.999996020138476, 1.9999705520429036, 1.9999611215295605]

Hope this helps.

3 Likes