Adjoint for function evaluation at point

Hi, I’m trying to implement the interpolation of a coefficient as part of an adjoint optimisation problem. Say, for example a coefficient z that is a function of two variables, x_1 and x_2

z=f(x_1,x_2) ,

where these two variables are control inputs and the coefficient is to be used in the bigger problem. There is no analytical form of the coefficient calculation.

Evaluation of the function at a point does not seems to be annotated for the adjoint derivation. So this does not work:

J = z(x1,x2)

I tried to write a custom function and overload it to work with the adjoint, as in the MWE below. However, the Taylor test returns a convergence around 1 instead of 2, so this seems to be incorrect.

How would I be able to fix the code below to yield the correct behaviour? Or is there perhaps a more effective method to include the point evaluation of a function in the adjoint?

MWE:

from fenics import *
from fenics_adjoint import *
import numpy as np
from pyadjoint import Block
from pyadjoint.overloaded_function import overload_function

def get_coefficient(func, coord1, coord2):
    return func(coord1, coord2)

backend_get_coefficient = get_coefficient

class CoefficientBlock(Block):
    def __init__(self, func, coord1, coord2, **kwargs):
        super(CoefficientBlock, self).__init__()
        self.kwargs = kwargs
        self.func = func
        self.add_dependency(coord1)
        self.add_dependency(coord2)

    def __str__(self):
        return "CoefficientBlock"

    def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
        grad_idx = project(self.func.dx(idx), self.func.function_space())
        return grad_idx(inputs[0], inputs[1]) * adj_inputs[0]

    def recompute_component(self, inputs, block_variable, idx, prepared):
        return backend_get_coefficient(self.func, inputs[0], inputs[1])

get_coefficient = overload_function(get_coefficient, CoefficientBlock)

mesh = UnitSquareMesh(10, 10)
V1 = FunctionSpace(mesh, "Lagrange", 1)

u = Function(V1)
x = SpatialCoordinate(u)
z = project(x[0]*x[1], V1)

x1 = Constant(np.random.rand())
x2 = Constant(np.random.rand())

J = get_coefficient(z, x1, x2)
controls = [x1, x2]
m = [Control(c) for c in controls]
h = [Constant(0.01*np.random.rand()) for c in controls]

Jhat = ReducedFunctional(J, m)
Jhat.derivative()

taylor_test(Jhat, controls, h)

Results from Taylor test:

Computed residuals: [7.3138620463933775e-06, 3.6569310231966887e-06, 1.8284655115428332e-06, 9.142327556603943e-07]
Computed convergence rates: [1.0, 1.0000000000437994, 1.0000000001751974]

Any help is much appreciated!

Thanks,
Maarten

The problem is that you are projecting the gradient into the wrong space, as it is a first order Lagrange space, the gradient is a zeroth order DG function. See the below code for a correction of this, as well as making your functional quadratic:

from fenics import *
from fenics_adjoint import *
import numpy as np
from pyadjoint import Block
from pyadjoint.overloaded_function import overload_function

def get_coefficient(func, coord1, coord2):
    return func(coord1, coord2)

backend_get_coefficient = get_coefficient

class CoefficientBlock(Block):
    def __init__(self, func, coord1, coord2, **kwargs):
        super(CoefficientBlock, self).__init__()
        self.kwargs = kwargs
        self.func = func
        self.add_dependency(coord1)
        self.add_dependency(coord2)
        degree = func.function_space().ufl_element().degree()
        family = func.function_space().ufl_element().family()
        if np.isin(family, ["CG", "Lagrange"]):
            self.V = FunctionSpace(mesh, "DG", degree-1)
        else:
            raise NotImplementedError(
                "Not implemented for other elements than Lagrange")
    def __str__(self):
        return "CoefficientBlock"

    def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
        grad_idx = project(self.func.dx(idx), self.V)
        return grad_idx(inputs[0], inputs[1]) * adj_inputs[0]

    def recompute_component(self, inputs, block_variable, idx, prepared):
        return backend_get_coefficient(self.func, inputs[0], inputs[1])

get_coefficient = overload_function(get_coefficient, CoefficientBlock)

mesh = UnitSquareMesh(10, 10)
V1 = FunctionSpace(mesh, "CG", 1)

u = Function(V1)
x = SpatialCoordinate(u)
z = project(x[0]*x[1], V1)

x1 = Constant(np.random.rand())
x2 = Constant(np.random.rand())

J = get_coefficient(z, x1, x2)**2
controls = [x1, x2]
m = [Control(c) for c in controls]
h = [Constant(0.01*np.random.rand()) for c in controls]

Jhat = ReducedFunctional(J, m)
Jhat.derivative()

taylor_test(Jhat, controls, h)
3 Likes

Thanks! That solves it.