function KRG
%
rosenbrock=@(x1,x2) 100*(x2-x1.^2).^2+(1-x1).^2;
 %
 % (a)
 % Generate training samples
 ny=40;
 np=101;
 xtr=lhsdesign(ny,2,'iteration',2000,'criterion','maximin')';
 xtr(1,:) = (xtr(1,:)-0.5)*6;
 xtr(2,:) = (xtr(2,:)-1/3)*6;
 y=rosenbrock(xtr(1,:),xtr(2,:));
 %
 x1 = linspace(-3, 3, np);
 x2 = linspace(-2, 4, np);
 [X1,X2] = meshgrid(x1,x2);
 F = rosenbrock(X1,X2);

% Sample data
% xs = [np, ndim] np training samples
% y = [np, 1]     np output 
%
xs=xtr'
y=y'
% Scaling variables
[m,n]=size(xs);
mx=mean(xs); sx=std(xs);
xs=(xs-repmat(mx,m,1))./repmat(sx,m,1);
my=mean(y); sy=std(y);
y=(y-my)/sy;

nparx=[mx;sx];
npary=[my;sy];
%
% Estimating hyper-parameters
nij=m*(m-1)/2;
diff=zeros(nij,n);
ij=zeros(nij,2);
l=0;
for i=1:m-1
   l=l(end)+(1:m-i);
   ij(l,:)=[repmat(i,m-i,1),(i+1:m)'];
   diff(l,:)=repmat(xs(i,:),m-i,1)-xs(i+1:m,:);
end
lb=[0.01; 0.01]; ub=[10; 10]; % up and lower bounds of roughness parameter
par = struct('xs',xs,'y',y,'d',diff,'ij',ij,'nparx',nparx,'npary',npary,...
   'trendf_poly','Hmat_poly0','corr_model','corr_Gauss','estim','GLS',...
   'post','posteri');
hyp0=[5 5];
ndvar = length(lb); % # of hyperparameters to be estimated
options = optimset('Algorithm','active-set','Display','off');
hypmle=fmincon(@(hyp) posteri(hyp,par),hyp0,[],[],[],[],lb,ub,[],options);
fprintf(1,'Optimal hyper-parameters: %10.3e %10.3e\n',hypmle);
%
% Fitting Kriging surrogate
[f,fit]=posteri(hypmle,par);
%
 %
 % Evaluation at grid
  zKRG=zeros(np);
  for i=1:np
   for j=1:np
    zKRG(j,i)=GPeval(fit,[x1(i) x2(j)]);
   end
  end
  figure;
contour(X1,X2,zKRG,[0, 1, 5, 10, 50, 100, 200, 400, 800, 1600, 3200, 6400]);
  rmsKRG=sqrt(sum(sum((F-zKRG).^2))/np^2)

end


function [fl,fit]=posteri(w,par)
% Objective function for given roughness parameter vector 
%
% w             : Roughness parameter vector
% par           : Auxiliary information structure
% 
r=corrGauss(par.d,w);
ny=size(par.y,1);
idx=find(r>0);
ii=(1:ny)';
mu=(10+ny)*eps;
S=sparse([par.ij(idx,1);ii],[par.ij(idx,2);ii],[r(idx);ones(ny,1)+mu]);
[C,rd]=chol(S); 

C=C';
H = feval(par.trendf_poly,par.xs);
Ft=C\H;
[Q,G]=qr(Ft,0);
Yt=C\par.y;
b=G\(Q'*Yt);
rho=(Yt-Ft*b);
Winv=G'*G;
gamma=rho'/C;
detR = prod( full(diag(C)) .^ (2/ny) );

if ( strcmp( par.estim , 'BAYES' ) )
%    Based on the Bayesian method
%    O'Hagan, 1992, "Some Bayesian Numerical Analysis"
   s0=0; a0=0; s=s0+sum(rho.^2); a=a0+ny-size(H,2);
   sigma=sqrt(s*1/(a+2));
   fl = sigma^2 * prod( full(diag(C)) .^ ( 2/(ny-size(H,2)) ) ) .* det(Winv) .^ ( 1/(ny-size(H,2)) ); % BAYES_GA3
elseif ( strcmp( par.estim , 'GLS' ) )
%    Based on the generalized least square method
%    Sacks et. al, 1989, "Design and Analysis of Computer Experiments"
   sigma2=sum(rho.^2)/ny;
   sigma=sqrt(sigma2);
   fl = sum(sigma2) * detR;
end

if  nargout > 1
   fit = struct('sigma',par.npary(2,:).*sigma,'w',w./par.nparx(2,:).^2,...
                  'beta',b,'gamma',gamma,'C',C,'Ft',Ft,...
                  'G',G','xs',par.xs,...
                  'nparx',par.nparx,'npary',par.npary,'rho',[],...
                  'reciplik',fl,...
                  'detR',detR,...
                  'trendf_poly',par.trendf_poly,...
                  'corr_model',par.corr_model);
end
end

function [ybar,vary]=GPeval(fit,xgrid)
% Make predictions at xgrid
% 
% fit           : Surrofate information structure
% xgrid         : Surrogate evaluation point vector
% 
 [mg,ng]=size(xgrid);
 mx=fit.nparx(1,:); sx=fit.nparx(2,:); 
 my=fit.npary(1,:); sy=fit.npary(2,:);
 xgrid=(xgrid-repmat(mx,mg,1))./repmat(sx,mg,1);
%
 beta1bar=fit.beta;
 [mx,nx]=size(xgrid);
 [m,n]=size(fit.xs);
 if (nx~=n)
   error('dimension of data is not equal to that of function evaluation points')
 end
 H1 = feval(fit.trendf_poly,xgrid);
 xrep=repmat(xgrid',m,1);
 xrep=transpose(reshape(xrep,n,mx*m));
 d=xrep-repmat(fit.xs,mx,1);
 r=corrGauss(d,fit.w.*fit.nparx(2,:).^2);
 r = reshape(r,m,mx);
 ybar=H1*beta1bar+(fit.gamma*r)';
 ybar=ybar*sy+my;
%
 rt=fit.C\r;
 q=fit.G\(H1'-fit.Ft'*rt);
 yvar=repmat(fit.sigma.^2,mx,1).*(1-sum(rt.^2,1)+sum(q.^2,1))';
 vary=sqrt(yvar);
end

function r=corrGauss(d,theta)
% Build a correlation matrix using the Gaussian kernel
%
% d             : Distances between sampling points
% theta         : Roughness parameter vector
% 
 [m, n] = size(d);
 if  length(theta) == 1
    theta = repmat(theta,1,n);
 elseif  length(theta) ~= n
    error(sprintf('Length of theta must be 1 or %d',n))
 end
 td = d.^2 .* repmat(-theta(:).',m,1);
 r = exp(sum(td, 2));
end

function H=Hmat_poly0(x,arg1)
% Location matrix based on a constant trend function
% 
 switch nargin
    case {1, 2}
        H=ones(size(x,1),1);
%     otherwise
%         c = 0;
 end
end