"""
This module solves the elliptic measure space control problem
     min 1/2 |y-yd|_L^2 + alpha |u|_M    s.t.   -\Delta y = u
using a semismooth Newton method described in the paper
    "Approximation of elliptic control problems in measure spaces 
     with sparse solutions"
by Eduardo Casas, Christian Clason, and Karl Kunisch, to appear in 
SIAM Journal on Control and Optimization.
Besides NumPy and SciPy, it requires the 'dolfin' module from the
open source FEniCS Project (http://fenicsproject.org/)
"""

__author__ = "Christian Clason <christian.clason@uni-graz.at>"
__date__ = "April 28, 2012"

import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve
from dolfin import *

N     = 128     # number of grid points per dimension
alpha = 1e-3    # penalty parameter
maxit = 20      # max iterations in semismooth Newton method

# construct scipy matrix from bilinear form a, boundary condition bc
parameters['linear_algebra_backend'] = 'uBLAS'
def assemble_csr(a,bc):
    row,col,val = assemble(a,bcs=bc).data()
    return sp.csr_matrix((val,col,row))

# setup grid, target, differential equation
mesh = UnitSquare(N,N)
V = FunctionSpace(mesh, 'CG', 1)
u = TrialFunction(V)
v = TestFunction(V)

# target
yd = Expression('exp(-50*(pow(x[0]-0.5,2) + pow(x[1]-0.5,2)))')
plot(yd,mesh=mesh,title='target')

# Laplace equation with homogeneous Dirichlet conditions
a  = dot(grad(u),grad(v))*dx
bc = DirichletBC(V, Constant(0.0), lambda x,on_boundary: on_boundary)

A  = assemble_csr(a,bc)             # stiffness matrix
AT = assemble_csr(adjoint(a),bc)    # stiffness matrix (adjoint)
M  = assemble_csr(u*v*dx,bc)        # mass matrix
z  = assemble(yd*v*dx,bcs=bc)       # load vector

# initialization
n = V.dim()                 # number of degrees of freedom
Ap = Am = np.zeros(n)       # active sets for upper, lower bound
Ap_old, Am_old = Ap, Am

# continuation strategy
for gamma in 10**np.arange(12):
    print 'Solving for gamma = %1.0e' % (gamma)

    # semismooth Newton iteration
    for it in xrange(maxit):
        # setup Newton step
        As = sp.spdiags(gamma*(Ap+Am),0,n,n)
        H  = sp.bmat([[A,As],[-M,AT]],format='csr')
        g  = np.hstack([alpha*gamma*(Ap-Am), -z])

        # solve for Newton step
        dx = spsolve(H,g)
        y,p = np.split(dx,2)

        # update active sets
        Ap = (p >  alpha).astype(float)
        Am = (p < -alpha).astype(float)

        change = (Ap-Ap_old)+(Am-Am_old)
        update = len(change[change.nonzero()])
        print 'Iteration %d: %d points changed in active set' % (it+1,update)
        if update == 0: break

        Ap_old,Am_old = Ap,Am

    # terminate continuation if Newton iteration converged in one step
    if it == 0: break

# compute starting point for control, active sets
u = -gamma*(np.maximum(0,p-alpha)+np.minimum(0,p+alpha))
Ap = (-u + p >  alpha).astype(float)
Am = (-u + p < -alpha).astype(float)

# compute optimal measure space control
print 'Solving original problem'
I = sp.identity(n,format='csr')

# semismooth Newton iteration
for it in xrange(maxit):
    # setup Newton step
    As = sp.spdiags(Ap+Am,0,n,n)
    C  = sp.bmat([[A,None,-I],[-M,A,None],[None,As,I-As]],format='csr')
    b  = np.hstack([0*z, -z, alpha*(Ap-Am)])

    # solve for Newton step
    dx = spsolve(C,b)
    y,p,u = np.split(dx,3)

    # update active sets
    Ap = (-u + p >  alpha).astype(float)
    Am = (-u + p < -alpha).astype(float)

    change = (Ap-Ap_old)+(Am-Am_old)
    update = len(change[change.nonzero()])
    print 'Iteration %d: %d points changed in active set' % (it+1,update)
    if update == 0: break

    Ap_old,Am_old = Ap,Am

# plot (linear interpolation) of optimal control, state
ufun = Function(V)
ufun.vector()[:] = u
plot(ufun,title='control')
yfun = Function(V)
yfun.vector()[:] = y
plot(yfun,title='state')
interactive()