%Written by Jan Christoph Krueger, Hamburg University of Technology
%Short example script for the application of the principal sensitivity
%second-order fourth moment-method
% Here, the error is computed for different step sizes

clc; clear; close all;
rng default; 

%Prepare model
V0=rand(1000,1); x0=rand(10,1);
Cov=rand(10,10); Cov=Cov*Cov';
[CovVec,ew]=eig(Cov); %ev=ev(:,1:10); ew=ew(1:10,1:10);
CovVal=diag(ew);

%% Demonstrate error of PSS and show step size adaption
[mu,si,dmu0,dsig0]=SOFMexact(@(y,x)fun(y,x),V0,x0,ew,CovVec);
stepsize=-9:0.1:0; stepsize=10.^stepsize;
err=zeros(length(stepsize),length(stepsize));
for i=1:length(stepsize)
    for j=1:length(stepsize)
        [~,~,~,dsig1]=SOFM_PS(@(y,x)fun(y,x),V0,x0,CovVec,CovVal,stepsize(i),stepsize(j));
        err(j,i)=norm(dsig1-dsig0)/norm(dsig0);
    end
end

%Use the step size adaption to optimize the step size
[deltax,epshat]=sofmAdaptStep(@(y,x)fun(y,x),V0,x0,CovVec,CovVal,1E-2,1E-2);
[deltax,epshat]=sofmAdaptStep(@(y,x)fun(y,x),V0,x0,CovVec,CovVal,deltax,epshat);
[deltax,epshat]=sofmAdaptStep(@(y,x)fun(y,x),V0,x0,CovVec,CovVal,deltax,epshat);
[~,~,~,dsig1]=SOFM_PS(@(y,x)fun(y,x),V0,x0,CovVec,CovVal,deltax,epshat);
errOpt=norm(dsig1-dsig0)/norm(dsig0);

%Plot error and show optimized step size
figure;
surface(log10(stepsize),log10(stepsize),log10(err));
hold on;
plot3(log10(deltax),log10(epshat),log10(errOpt),'r*');
xlabel('log(deltax)');
ylabel('log(epshat)');
zlabel('error');
figure;
contour(log10(stepsize),log10(stepsize),log10(err),15);
hold on;
xlabel('log(deltax)');
ylabel('log(epshat)');
plot(log10(deltax),log10(epshat),'r*');

%% Optimization of the objective using fmincon
f0=robustObjective(@(y,x) fun(y,x),CovVec,CovVal,deltax,epshat,V0,x0);
options=optimoptions('fmincon','Display','iter','SpecifyObjectiveGradient',true);
[x,f1]=fmincon(@(y) robustObjective(@(y,x) fun(y,x),CovVec,CovVal,deltax,epshat,y,x0),V0,[],[],[],[],[],[],[],options);

%% Functions
function [mu,sig,dmu,dsig]=SOFMexact(fun,y,x,Coveig,Coveigv)
    [f,dfy,dfx,dfxx,dfxy,dfxxy]=fun(y,x);
    dfx=Coveigv'*dfx; dfxx=Coveigv'*dfxx*Coveigv; dfxy=dfxy*Coveigv;
    mu=f+0.5*trace(dfxx*Coveig);
    sig=dfx'*Coveig*dfx+0.5*trace(dfxx*Coveig*dfxx*Coveig);
    dmu=dfy; dsig=2*dfxy*Coveig*dfx;
    for i=1:length(y)
        dfxxyi=Coveigv'*dfxxy(:,:,i)*Coveigv;
        dsig(i)=dsig(i)+trace(dfxxyi*Coveig*dfxx*Coveig);
        dmu(i)=dmu(i)+0.5*trace(dfxxyi*Coveig);
    end
    sig=sqrt(sig); dsig=dsig/(2*sig);
end

function [f,dfy,dfx,dfxx,dfxy,dfxxy]=fun(y,x)
    %Example objective function
    f=(y'*y)^2/(x'*x)^2; dfy=(4*y*(y'*y))/(x'*x)^2; dfx=-(4*(y'*y)^2*x)/(x'*x)^3;
    if nargout>3
        dfxx=24*(y'*y)^2*x*x'/(x'*x)^4-(4*(y'*y)^2*eye(length(x)))/(x'*x)^3; dfxy=-16*(y'*y)*y*x'/(x'*x)^3; 
        dfxxy=zeros(length(x),length(x),length(y));
        for i=1:length(y)
            dfxxy(:,:,i)=96*(y'*y)*y(i)*x*x'/(x'*x)^4-16*(y'*y)*y(i)*eye(length(x))/(x'*x)^3;
        end
    end
end
function [f,df]=robustObjective(fun,CovVec,CovVal,deltax, epshat,y,x)
    [mu,sig,dmu,dsig]=SOFM_PS(@(y,x)fun(y,x),y,x,CovVec,CovVal,deltax,epshat);
    f=mu+sig;
    df=dmu+dsig;
end