%% Bayesian Optimization (BO): full examples
% © 2024, Antonio Sala. Universitat Politecnica de Valencia, Spain. All rights reserved.
%
% Presentations in video:
% https://personales.upv.es/asala/YT/V/boloop1EN.html
% https://personales.upv.es/asala/YT/V/boloop2EN.html
% https://personales.upv.es/asala/YT/V/boloop3EN.html
%
% Objectives:Program and understand the operation of the complete Bayesian optimization methodology: acquire, posterior update, decide sample (acquisition function) and repeat.
% Table of Contents
% Motivation, preliminary discussion
% Example Gaussian process setup
%     Prior: mean, kernel, confidence interval
%     Loading the GP with a set of "past" observations
%     Bayesian Optimization main loop
%% Motivation, preliminary discussion
% Goal of BO: obtaining the global minimum of f(x) via a "statistical model" of it (a Gaussian Process, actually).
% Methodology:
% Set a "prior" GP model (+ maybe some initial preexisting samples)
% Decide which is the best "next x" to try, x_k.
% Acquire samples (usually one by one, costly process), pairs (x_k,y_k=f(x_k)+\epsilon_k).
% Update the GP model (obtain the "posterior")
% Evaluate termination criteria and go to step 2 if not met.
%
% Step 2, detail: we associate to each x an "acquisition function" a(x) related to how "good" the candidate sampling point is assumed to be, based on some statistical quantities.
% The x achieving "best" value of a(x) is the one proposed to actually try in step 2.
% Key idea: evaluating a(x) must be much faster than evaluating f(x), otherwise we would resort to standard deterministic optimization.
% Popular options for a(x) are:
% EV: Expected value, to minimize
% PI: Probability of improvement, to maximize
% EI: Expected improvement, to maximize
% LCB: Lower confidence bound, to minimize
% These are the ones to be discussed in this material, via an example.

yn=@(mean,var,y) (y-mean)/sqrt(var); %Normalize to standard normal (utility)

% \bullet Plainly, EV expected value (posterior mean) may be the easiest thing to think of... but we have other options
% \bullet PI, probability of improvement:

PI=@(m,v,level) normcdf(yn(m,v,level)); %prob of. standard normal distribution below "level"

% \bullet EI, expected improvement (it is conceptually better than "expected value" if there is no "cost" in being unlucky, i.e., sampling a point which does not work well)

EIN=@(c) -normpdf(c)-c*normcdf(c); %minimizing: the more negative the better.
EI=@(m,v,ybest) EIN(yn(m,v,ybest))*sqrt(v);

% \bullet Lower confidence bound
%
%% Example Gaussian process setup
% Let us invent a "ground truth" function to be optimized (several, for different tests):

example=1; %
rng(100)
switch example
    case 1
        TrueF=@(x) 1.05*sin(1.8*x+1.35)+.52*x.^2-.15+.45.*x; %perfect
        ybestopt="mean"; MaxSamples=8;
        Data.priormean=@(x) 0.15*x.^2+.1;
    case 2
         TrueF=@(x) -sin(1.6*x+0.5)-.2*x.^2+.2*x+0.5*log(x+3.2)-.2; %perfect
        ybestopt="mean"; MaxSamples=18;
        Data.priormean=@(x) zeros(size(x)); %LCB ok, but PI local minimum
    case 3
        TrueF=@(x) 1.05*sin(1.71*x+1.35)+.52*x.^2+.025+.45.*x; 
        ybestopt="measurement"; MaxSamples=26; %cambiar a PI/EI...
        Data.priormean=@(x) 0.15*x.^2+.1;
    case 4
        TrueF=@(x) 0.75*sin(3.1*x+.5)-.2*x.^2+.3-.2.*x; %PSD ok, amplitude NOT
        ybestopt="mean"; maxSamples=10;
        Data.priormean=@(x) 0.45*x.^2+0.45*x-.1;
    case 5
        TrueF=@(x) sin(15.2*x+.5)-.52*x.^2+.7-.6.*x; %PSD wrong, amplitude wrong
        Data.priormean=@(x) 0.15*x.^2+.1;
               ybestopt="mean"; maxSamples=12; %try PI as acquisition
 
end

%% Prior: mean, kernel, confidence interval

Data.X=[]; Data.Y=[]; %To test "prior" model, empty data
Data.K=@K;
STDMeasureNoise=4e-2; %measurement noise standard deviation
VarMeasureNoise=STDMeasureNoise^2;
Xtest=-2:(1/50):2; %uniform grid to test stuff
TrueData=TrueF(Xtest);
[TrueOptimum,TrueOptIdx]=min(TrueData);
TrueXOpt=Xtest(TrueOptIdx);
LL=length(Xtest);
[MeanPrior,CovMatrixPrior]=predictGP(Xtest,Data,VarMeasureNoise);
sg=sqrt(diag(CovMatrixPrior));
PredPrior=[MeanPrior-1.96*sg MeanPrior MeanPrior+1.96*sg];
NTrials=4;
FunctionRealization=mvnrnd(MeanPrior,CovMatrixPrior,NTrials);%clean
for i=1:NTrials
    col=[.5 .6 .65]+rand()*[.3 0 0]+rand()*[0 .3 0]+rand()*[0 0 0.2];
    plot(Xtest,FunctionRealization(i,:),Color=col,LineWidth=1.5), hold on
end
plot(Xtest,PredPrior(:,2),'k',LineWidth=1.5), grid on
plot(Xtest,PredPrior(:,[1 3]),'r:',LineWidth=2)
plot(Xtest,TrueData,'-.',LineWidth=3.5,Color=[.4 .5 .2]), hold off
xline(TrueXOpt,':g',LineWidth=2,label="True optimal x (unknown)")
title("prior, some realizations, and true F")
xlabel("input variable"), ylabel("Function value")

%% Loading the GP with a set of "past" observations

Data.X=[-.1 .7];
Data.Y=TrueF(Data.X)+randn(size(Data.X))*STDMeasureNoise;Data.Y
plot(Xtest,PredPrior(:,2),'k',LineWidth=1.5), grid on, hold on
plot(Xtest,PredPrior(:,[1 3]),'r:',LineWidth=2)
plot(Data.X,Data.Y,'*k',LineWidth=3), 
plot(Xtest,TrueData,'-.',LineWidth=3,Color=[.4 .5 .2])
hold off
xlabel("input variable"), ylabel("Function value")
title("Prior confidence bounds + observed values")
legend("mean","conf. int.","","samples","true F",Location="best")

%
%% Bayesian Optimization main loop

ybestDataPLOT=[];ybestFilteredPLOT=[]; %cosmético
while length(Data.X)<MaxSamples 
    ybestData=min(Data.Y);
    ybestDataPLOT=[ybestDataPLOT ybestData];
    yFiltered=predictGP(Data.X,Data,VarMeasureNoise); %si hay ruido de medida, posiblemente nos creamos un "promedio filtrado".
    ybestFiltered=min(yFiltered);
    ybestFilteredPLOT=[ybestFilteredPLOT ybestFiltered];
    if ybestopt=="mean"
      ybest=ybestFiltered;
      ytarget=ybest-STDMeasureNoise; %PI o EI respecto a "media-1sigmanoise"
    else
        ybest=ybestData
        ytarget=ybest;
    end

%
%
% For all test points, we'll do:

    [post_mean,post_var]=predictGP(Xtest,Data,VarMeasureNoise);
    SPIplt=zeros(1,LL); EIplt=zeros(1,LL); %for brute-force optimization
    for i=1:LL
        PIplt(i)=PI(post_mean(i),post_var(i,i),ytarget);
        EIplt(i)=EI(post_mean(i),post_var(i,i),ytarget);
    end

% We will select the best abscissa of each acquisition function proposal... so we try several (really only ONE is needed).
% Note: minimizing by looking the optimal acquisition function value in a "very fine mesh" is obviously NOT the best way to do it, because it does not scale to more dimensions of the input variable and we depend on the "granularity" of said mesh. Optimization software (fmincon or similar) should be used, but we do not do so here for clarity/brevity in the presentation.

    [min_ev,ev_idx]=min(post_mean);
    [max_pi,pi_idx]=max(PIplt); %Maximum probability of improvement
    [min_ei,ei_idx]=min(EIplt);
    sg=sqrt(diag(post_var));
    Pred=[post_mean-1.96*sg post_mean post_mean+1.96*sg];
    [min_lcb,lcb_idx]=min(Pred(:,1));
    fprintf("Best PI: %g, EV: %g, EI: %g",max_pi,min_ev-ybest,min_ei)
    
    AFname='PI';%chosen acquisition function, to label plots
    NewSample=Xtest(pi_idx) %chosen acquisition function sample
    NewMeasure=TrueF(NewSample)+randn()*STDMeasureNoise
  

% Let us plot something...

    figure()
    plot(Xtest,Pred(:,2),LineWidth=1.5)
    hold on
    plot(Xtest,Pred(:,[1 3],:),'r:',LineWidth=1.5)
    plot(Data.X,Data.Y,'*k',LineWidth=3,MarkerSize=9);
    plot(Xtest,TrueData,'-.',LineWidth=2,Color=[.4 .5 0.2])
    plot(Xtest,EIplt+ybest)
   % plot(Xtest,EIplt/max(abs(EIplt))) %scaling to see better, if so wished
    plot(Xtest,PIplt);
    plot(NewSample,NewMeasure,'*',Color=[.4 .8 .2],LineWidth=3.5,MarkerSize=10);
    hold off
    grid on, xlabel("input variable"), ylabel("Function value")
    title(['Posterior ' mat2str(length(Data.X)-1)])
    yline(ybest,label="ybest",LabelHorizontalAlignment="left")
    xline(NewSample,LineWidth=1.5,label=['New proposed sample' AFname])
    xline(TrueXOpt,':g',LineWidth=2,label="Optimal x (unknown)")
    legend("mean (EV)","lcb/ucb","","samples","true F (unknown)","EI+ybest","PI","New")
    Data.X=[Data.X NewSample]; %for next iter
    Data.Y=[Data.Y NewMeasure]; %for next iter

end %end of Bayesian Optimization loop

%
% Let us represent optimization progress in a new plot:

plot(Data.Y(3:end)), hold on
plot(0:(MaxSamples-3), ybestDataPLOT,'*'),
plot(0:(MaxSamples-3), ybestFilteredPLOT,'*'), hold off
yline(TrueOptimum)
yline(TrueOptimum+1.96*STDMeasureNoise,'g')
yline(TrueOptimum-1.96*STDMeasureNoise,'g'),
xline(1,':',label="decisions BO")
title("Search evolution"), grid on
legend("samples BO","best sample","best filtered estim.","True optimum","95% confidence (noise)")

%
% Note about "overexploitation": when we are stuck at hopefully the global minimum withing measurement noise precision, there is not much else to do... it's all over bar the shouting. However, since the "prior" is really an invention, modifying it to explore elsewhere can be beneficial. See the "ExplorationRatio" parameter in https://es.mathworks.com/help/stats/bayesian-optimization-algorithm.html .
%

function [mu,Var]=predictGP(x1,Data,VarMeasureNoise)
%set up the PRIOR mean and covariance at prediction points x1
priormean=Data.priormean;
K=@Data.K;
mu=priormean(x1)';
Var=K(x1,x1);
N=length(Data.X);
if(N>0) %set up the POSTERIOR with some data
    K1X=K(x1,Data.X);
    Gain=K1X/(K(Data.X,Data.X)+VarMeasureNoise*eye(N));
    mu=mu+Gain*(Data.Y-priormean(Data.X))';
    Var=Var-Gain*K1X';
    Var=0.5*(Var+Var'); %roundoff-errors, correct them
end
end




function km=K(x1,x2)
M=0.6; sg=0.65;%hyperparameters stddev x and y %PI was 0.65 in 1st example
kernel=@(x1,x2)   M^2*exp(-norm(x2-x1)^2/2/sg^2); %covariance kernel function
N1=length(x1); N2=length(x2); %1D case
km=zeros(N1,N2);
for i=1:N1
    for j=1:N2
        km(i,j)=kernel(x1(:,i),x2(:,j)); %do it for all pairs (x1_i,x2_j)
    end
end
end



