"""

PDE:

-laplace(y) + y^3 = beta*u
y=0 on the boundary

0. Solve the problem
1. Evaluate a given objective function
2. Compute an approximation to directional derivative using finite differences
3. Evaluate the directional derivative using the "direct" method
4. Evaluate the directional derivative using the adjoint method
"""

from dolfin import *

N = 64  # number of elements in the mesh
mesh = UnitSquareMesh(N,N)

# state space
Y = FunctionSpace(mesh,"Lagrange",1)
y_ = TrialFunction(Y)
v_ = TestFunction(Y)
# control space
U = FunctionSpace(mesh,"Discontinuous Lagrange",0)
u  = Function(U)

# let beta be the indicator function for a circle in the middle
beta = Expression("(x[0]-0.5)*(x[0]-0.5)+(x[1]-0.5)*(x[1]-0.5)<=0.5*0.5", degree=0)


# initialize u
project(Constant(1.0),U,function=u)

y = Function(Y)

# define the residual
F = inner(grad(y),grad(v_))*dx + inner(y*y*y,v_)*dx - inner(beta*u,v_)*dx
# compute the Jacobian
jac = derivative(F,y,y_)

bc = DirichletBC(Y, Constant(0), "on_boundary")
# Solve the non-linear system
solve(F==0,y,bc,J=jac,\
      solver_parameters={"newton_solver": \
                         {"relative_tolerance": 1e-6}})


# define the objective functional
lam = 1.0
y_Omega = Expression("x[0]*x[0] + x[1]*x[1]", degree=2)
J = 0.5*inner(y-y_Omega,y-y_Omega)*dx + lam/2*inner(u,u)*dx


# evaluate the objective functional at a present control point
J0 = assemble(J)

# perturb the solution a little bit in some random direction h
from numpy import random, arange
h = Function(U)
h.vector().set_local(0.5-random.random(h.vector().size()))


"""
We will now compute the directional derivative of J using a different method.
J = 0.5 ||y-y_Omega||^2 + lam/2 ||u||^2,
and y = y0 + Su, where S is the solution operator for the state problem with homogeneous boundary conditions.

Thus we have: 
J'h = J'_y Sh + J'_u h = 
    = (y-y_Omega, Sh) + lam(u,h)
"""

# first, we homogenize the boundary conditions
bc0 = DirichletBC(bc)
bc0.homogenize()
# second, we change the right-hand side of the problem to "beta*h"
yh = Function(Y)
solve(inner(grad(y_),grad(v_))*dx + 3*y*y*y_*v_*dx == beta*h*v_*dx, yh,bc0)
# evaluate the derivative now
dJdh = assemble(inner(y-y_Omega,yh)*dx + lam*inner(u,h)*dx)
print("dJdh(direct) = %e" % (dJdh))

"""
We use adjoint method now.

First, adjoint problem is

-div(grad(p)) + 3*y^2 p = J'_y = (y-y_Omega,.)
p_dOmega = 0

and then
J'h = S*(J'_y)h + J'_u h = 
    = (beta*p,h) + lam(u,h) 
"""

# Adjoint problem
p = Function(Y)
solve(inner(grad(y_),grad(v_))*dx + 3*y*y*y_*v_*dx == (y-y_Omega)*v_*dx, p,bc0)
# evaluate the derivative now
dJdh = assemble(inner(beta*p+lam*u,h)*dx)
print("dJdh(adjoint) = %e" % (dJdh))


# save the old u
u_old = Function(U)
u_old.vector().set_local(u.vector().array())

# run through a sequence of perturbations
for delta in pow(10.,-arange(10)):
    u.vector().set_local(u_old.vector().array() + delta*h.vector().array())
    # resolve the non-linear PDE
    solve(F==0,y,bc,J=jac,\
      solver_parameters={"newton_solver": \
                         {"relative_tolerance": 1e-6}})
    
    # re-evaluate the functional
    J1 = assemble(J)
    # compute the FD approximation
    dJdh = (J1-J0)/delta
    print("delta = %e, dJdh = %e" % (delta,dJdh))