function multibangcontrol
% This function solves the multibang control problem
%  min 1/2 \|y-yd\|^2 + alpha/2 \int |u|^2 + beta \prod_{i=1}^d |u-u_i|_0
%      s.t. -\Delta y = u,  u_1 <= u(x) <= u_d
% using the approach described in the paper
%    "Multibang control of elliptic systems"
% by Christian Clason and Karl Kunisch, see
% http://math.uni-graz.at/mobis/publications/SFB-Report-2013-006.pdf.
%
% March 27, 2013            Christian Clason <christian.clason@uni-graz.at>
%                                            http://www.uni-graz.at/~clason

%% setup
% problem parameters
N     = 128;                           % number of nodes per dimension
maxit = 10;                            % max number of Newton steps
alpha = 5e-3;                          % control cost parameter (L^2)
beta  = 1e-3;                          % control cost parameter (N^0)
ub = linspace(-2,2,5)';                % vector of control states
d  = length(ub);                       % number of control states

% setup grid, assemble stiffness and mass matrix
[A,M,xx,yy] = assembleFEM(N);
tplot = @(n,f,s) tplot_(n,f,s,N,xx,yy);

% setup target
z = peaks(N)/10;
tplot(1,z,'target');

% precompute some terms
Mz = M*z(:);    AT = A';    N2 = N*N;    sqab = sqrt(2*alpha*beta);

%% compute control
% initialize iterates
y  = zeros(N2,1);                      % state variable
p  = zeros(N2,1);                      % dual variable
as = zeros(4*d-2,N2);                  % active sets

% continuation: start with sufficiently eps^0
eps  = min(sqrt(2*alpha*beta),min(alpha/4*(ub(2:end)-ub(1:d-1))))*0.9;
while eps > 1e-12
    fprintf('\nCompute solution for epsilon = %1.3e:\n',eps);
    for it = 1:maxit
        % update active sets
        as_old = as;
        % P_i^eps
        as(1,:) = (p <= alpha*ub(1) + sqab - eps) & ...
                  (p <= (alpha/2*(ub(1)+ub(2)) - eps));
        for i = 2:d-1
            as(i,:) = (abs(p-alpha*ub(i)) <= (sqab - eps)) & ...
                      (p >= (alpha/2*(ub(i-1)+ub(i)) + eps)) & ...
                      (p <= alpha/2*(ub(i)+ub(i+1)) - eps);
        end
        as(d,:) =  (p >= alpha*ub(d) - sqab + eps) & ...
                   (p >= (alpha/2*(ub(d-1)+ub(d)) + eps));
        He = as(1:d,:)'*ub;
        % P_i,i+1^eps
        for i = 1:d-1
            ind = d+1+i;
            if alpha/2*(ub(i+1)-ub(i)) <= sqab
                as(ind,:) = (abs(p-alpha/2*(ub(i)+ub(i+1))) < eps);
                He  = He + (ub(i)*(alpha/2*(ub(i)+ub(i+1))+eps-p) + ...
                    ub(i+1)*(p-alpha/2*(ub(i)+ub(i+1))+eps)).*as(ind,:)'/(2*eps);
            end
        end
        DHe = as(d+2:2*d,:)'*diff(ub)/(2*eps);
        % P_i,0-^eps, P_i,0+^eps (P_1,0-^eps, P_d,0+^eps always empty!)
        for i = 2:d-1
            ind = 2*d+2*(i-1)-1;
            if  alpha/2*(ub(i)-ub(i-1)) > sqab
                as(ind,:) = (abs(p-alpha*ub(i)+sqab) < eps);
                He  = He + ((alpha*ub(i)-sqab-eps).*(alpha*ub(i)-sqab+eps-p)/alpha +...
                    ub(i)*(p-alpha*ub(i)+sqab+eps)).*as(ind,:)'/(2*eps);
            end
            if alpha/2*(ub(i+1)-ub(i)) > sqab
                as(ind+1,:) = (abs(p-alpha*ub(i)-sqab) < eps);
                He  = He + (ub(i).*(alpha*ub(i)+sqab+eps-p) + ...
                    (alpha*ub(i)+sqab+eps)*(p-alpha*ub(i)-sqab+eps)/alpha).*as(ind+1,:)'/(2*eps);
            end
        end
        DHe = DHe + (sqab+eps)/(2*eps*alpha)*sum(as(2*d+1:end,:))';
        % P_0^eps: all remaining nodes
        as(d+1,:) = 1-sum(as(1:d,:))-sum(as(d+2:end,:));
        He  = He + p.*as(d+1,:)'/alpha;
        DHe = DHe + as(d+1,:)'/alpha;

        % system matrix, right hand side
        C   = [M AT; A -M*spdiags(DHe,0,N2,N2)];
        rhs = [Mz-M*y-AT*p; -A*y + M*He];
        nr  = norm(rhs(:));

        update = nnz((as-as_old));
        fprintf('It# %i: update = %i,\t residual = %1.3e\n',it,update,nr);
        if update == 0; 
            break; 
        end

        % semismooth Newton step
        dx = C\rhs;
        y = y+dx(1:N2);
        p = p+dx(1+N2:end);        
    end
       
    % check convergence
    if it < maxit                      % converged: accept iterate
        u = He;                  
        regnodes = nnz(as(d+2:end,:));
        fprintf('Solution has %i node(s) in regularized active sets\n',regnodes);
        if regnodes == 0               % solution optimal: terminate
            break;
        else                           % reduce eps, continue
            eps = eps/10;
        end
    else                               % not converged: reject, terminate
        fprintf('Iterate rejected, returning u_eps for eps = %1.3e\n',eps*10);
        break;
    end 
    
end 

%% plot control
tplot(2,u,'control')
end % main function


function [K,M,xx,yy] = assembleFEM(n)
a   = 0;    b = 1;       % computational domain [a,b]^2
nel = 2*(n-1)^2;         % number of nodes
h2  = ((b-a)/(n-1))^2;   % Jacobi determinant of transformation (2*area(T))

% nodes
[xx,yy] = meshgrid(linspace(0,1,n));

% triangulation
tri = zeros(nel,3);
ind = 1;
for i = 1:n-1
    for j = 1:n-1
        node         = (i-1)*n+j+1;              % two triangles at node
        tri(ind,:)   = [node node-1 node+n];     % triangle 1 (lower left)
        tri(ind+1,:) = [node+n-1 node+n node-1]; % triangle 2 (upper right)
        ind = ind+2;
    end
end

% Mass and stiffness matrices
Ke = 1/2 * [2 -1 -1 -1 1 0 -1 0 1]';   % elemental stiffness matrix
Me = h2/24 * [2 1 1 1 2 1 1 1 2]';     % elemental mass matrix

ent = 9*nel;
row = zeros(ent,1);
col = zeros(ent,1);
valk = zeros(ent,1);
valm = zeros(ent,1);

ind = 1;
for el=1:nel
    ll       = ind:(ind+8);            % local node indices
    gl       = tri(el,:);              % global node indices
    row(ll)  = gl([1;1;1],:); rg = gl';
    col(ll)  = rg(:,[1 1 1]);
    valk(ll) = Ke;
    valm(ll) = Me;
    ind      = ind+9;
end
M = sparse(row,col,valm);
K = sparse(row,col,valk);

% modify matrices for homogenenous Dirichlet conditions
bdnod = [find(abs(xx-a) < eps); find(abs(yy-a) < eps); ...
         find(abs(xx-b) < eps); find(abs(yy-b) < eps)];
M(bdnod,:) = 0;
K(bdnod,:) = 0;  K(:,bdnod) = 0;
for j = bdnod'
    K(j,j) = 1; %#ok<SPRIX>
end
end


function tplot_(n,f,s,N,x,y)
figure(n); surf(x,y,reshape(f,N,N));
shading interp; camlight headlight; lighting gouraud; alpha(0.8);
title(s); xlabel('x_1'); ylabel('x_2');
set(gca,'FontSize',12); drawnow;
end % tplot_ function
