function data = Sorad(data)
%SORAD  SORAD: A systems biology approach to predict and modulate
%       dynamic signaling pathway response from phosphoproteome
%       time-course measurements
%   
%   usage: data = Sorad(data)
%
%   input (the fields of the data structure):
%       data            a cell matrix (number of proteins by number of
%                       conditions), where each element has a row vector
%                       (number of time points) containing the protein and
%                       condition specific training data
%       protNames       a cell vector containing the names of the proteins
%       conditionNames  a cell vector containing the names of the training
%                       conditions
%       viz             a boolean variable controlling the visualization
%       infer           a boolean variable controlling whether to calculate
%                       the marginal likelihoods (true) or predict the 
%                       dynamics of the system given the topology (false)
%
%       IF infer == true
%       maxRegulators   a variable controlling the maximum number of 
%                       regulatory proteins (<= number of proteins)
%
%       IF infer == false
%       networkMatrix   a binary matrix containing the connectivity matrix
%                       of the signaling pathway, prot_j -> prot_i := (i,j)
%       testData        a cell matrix (number of proteins by number of
%                       test conditions), where each element has a row
%                       vector (number of time points) containing the
%                       protein and condition specific test data
%       conditionNamesTest
%                       a cell vector containing names of the test
%                       conditions
%
%   output (the fields of the data structure):
%       X               a cell vector (number of proteins) containing the
%                       fitted x functions (number of time points by
%                       number of conditions)
%       f               a cell vector (number of proteins) containing the
%                       fitted f functions (number of time points by
%                       number of conditions)
%       param           a cell vector (number of proteins) containing the
%                       protein specific alpha and lambda parameters
%
%       IF infer == true
%       subPowerset     a cell vector containing the model definitions, in
%                       other words, the regulatory proteins of the 
%                       different models
%       margLikelihoods a matrix (number of models by number of proteins)
%                       containing the calculated marginal likelihoods of
%                       the models (the order is the same as in 
%                       subPowerset) 
%
%       IF infer == false
%       hyperparameters a cell vector (number of proteins) containing the
%                       logarithms of the protein and model specific sigma,
%                       l and sigma_n^2 hyperparameters
%       trajectories    a cell matrix (number of test conditions) in which
%                       each element has a matrix (number of time points by
%                       number of proteins) containing the predictions
%
%   notes:              - please see example.m for an example of the
%                         use
%                       - it might be a good idea to scale the data in
%                         such a way that it is between 0 and 1  
%                       - current implementation requires that all the
%                         measurement sets share the same time points
%                       - current implementation uses the time points in
%                         the training data also in the prediction
%                       - implementation of the intervention estimation
%                         procedure for a general setting is not
%                         straightforward so it is not included in this
%                         routine, however, upon a request the author will
%                         provide an example, which can be modified based 
%                         on the given requirements
%
%   The GPML Matlab code is written by Carl Edward Rasmussen. The simulated
%   annealing implementation is written by Joachim Vandekerckhove.
%
%   Author: Tarmo Aijo <tarmo.aijo@aalto.fi>
%
%   Updated: 07/21/2013

% estimate f functions with respect to time
disp('estimating f functions ..')
data = solvefWrtTime(data);
disp('done')

if data.infer
    % estimate the model specific functions with respect to the activity levels
    disp('estimating the model specific g functions and their marginal likelihoods ..')
    data = solvefWrtConcentrationsInfer(data);
    disp('done')
else
    % estimate g functions with respect to the activity levels
    disp('estimating g functions ..')
    data.hyperparameters = solvefWrtConcentrations(data);
    disp('done')
    % predict the trajectories on the unseen conditions
    disp('predicting the behaviour of the system under the test conditions ..')
    data = solveTestDataTrajectories(data);
    disp('done')
end
end

function data = solveTestDataTrajectories(data)
% use the same time points as in the training data set
tSolve = data.t;
% calculate the time steps
tStep = diff(tSolve);

% find the conditions that are used as the test sets
testCondIndices = 1:size(data.testData,2);

trajectories = cell(length(testCondIndices),1); 
timeVector = tSolve;

for testCondIdx = testCondIndices
    % loop over these proteins
    protIndices = 1:size(data.testData,1);
    
    % get the initial values
    trajectories{find(testCondIndices == testCondIdx)} = zeros(length(timeVector),length(protIndices));
    tmp = [data.testData{:,testCondIdx}];
    trajectories{find(testCondIndices == testCondIdx)}(1,:) = tmp(1,:);
    
    % loop over the set of proteins
    for tIdx=1:length(timeVector)-1
        for rowIdx = protIndices
            % response data (training data)
            y3 = cat(1,data.f{rowIdx}(:));

            % explanatory data (training data)
            regIndices = find(data.networkMatrix(rowIdx,:));
            X = gatherExplanatoryData(data,y3,regIndices);

            % explanatory data (test data)
            Xs = gatherExplanatoryDataForSolving(regIndices,tIdx,trajectories{find(testCondIndices == testCondIdx)});
            % estimate the activity level of the next time point
            trajectories{find(testCondIndices == testCondIdx)}(tIdx+1,rowIdx) = ...
                solveIncreaseEuler(tStep(tIdx),data.param{rowIdx},trajectories{find(testCondIndices == testCondIdx)}(tIdx,rowIdx),X,Xs,data.hyperparameters{rowIdx},y3);

            % make sure that the activity levels are >= 0
            if trajectories{find(testCondIndices == testCondIdx)}(tIdx+1,rowIdx) < 0
                trajectories{find(testCondIndices == testCondIdx)}(tIdx+1,rowIdx) = 0;
            end
        end
    end
end
data.trajectories = trajectories;

if data.viz
    for rowIdx = protIndices
        figure
        suptitle(data.protNames{rowIdx})
        for idx=1:length(testCondIndices)
            subplot(length(testCondIndices),1,idx)
            plot(tSolve,data.trajectories{idx}(:,1),'r-')
            title(data.conditionNamesTest{idx})
            legend({'predicted'})
        end
    end
end

end

function xIncrease = solveIncreaseEuler(tStep,param,xPresent,X,Xs,hyperparameters,y)
xIncrease = xPresent + tStep * (param(1) - param(2) * xPresent + gpr(hyperparameters,{'covSum',{'covNNone','covNoise'}},X,y,Xs));
end

function X = gatherExplanatoryData(data,y,regIndices)
X = zeros(length(y),length(regIndices));
% loop over the set of regulatory proteins
for regIdx = regIndices    
    tmp = [data.data{regIdx,:}];
    X(:,regIndices == regIdx) = tmp(:);    
end
end

function X = gatherExplanatoryDataForSolving(regIndices,tIdx,trajectories)
X = zeros(1,length(regIndices));
% loop over the set of regulatory proteins
for regIdx = regIndices
    tmp = trajectories(:,regIdx);
    X(1,regIndices == regIdx) = tmp(tIdx);    
end
end

function data = solvefWrtTime(data)
tstar = data.t;
muX = cell(length(data.protNames),1);
muF = cell(length(data.protNames),1);
param = cell(length(data.protNames),1);
for rowIdx = 1:length(data.protNames)
    measIndices = 1:size(data.data,2);
    % response data
    y = cat(1,data.data{rowIdx,measIndices});
    % construct the mean vectors and the residual vector
    meanY2 = reshape(y,[length(data.t) length(measIndices)]);
    meanY = repmat(meanY2(1,:),length(data.t),1);
    meanYs = repmat(meanY2(1,:),length(tstar),1);
    meanY = meanY(:);
    meanYs = meanYs(:);

    % optimize the values of the hyperparameters with simulated annealing
    options.Verbosity = 0;
    tmp = data.t;
    loss = @(bar)foo([bar(1) 0.001 bar(2) bar(3)],tmp,tmp,measIndices,meanY,y);
    hyperparameters = abs(anneal(loss,[1 0.001 0.001],options));
    l = abs(hyperparameters(1));
    nVariance = 1e-3;
    alpha = real(hyperparameters(2));
    lambda = real(hyperparameters(3));
    r = y - m(data.t,lambda,alpha,meanY);

    meanYs = m(tstar,lambda,alpha,meanYs(:));

    % construct the (block) covariance matrix 
    Kxx = [];
    Kxsx = [];
    Kxsxs = [];
    Kxf = [];
    Kff = [];
    for measIdx = measIndices
        Kxx = blkdiag(Kxx,kxx(data.t,data.t,l,lambda));
        Kxsx = blkdiag(Kxsx,kxx(tstar,data.t,l,lambda));
        Kxsxs = blkdiag(Kxsxs,kxx(tstar,tstar,l,lambda));
        Kxf = blkdiag(Kxf,kxf(data.t,tstar,l,lambda));
        Kff = blkdiag(Kff,kff(tstar,tstar,l));
    end

    % calculate the mean of x (for a given protein)
    mux = meanYs(:) + Kxsx*((Kxx+nVariance*eye(size(Kxx)))\r);
    mux = reshape(mux,[length(tstar) length(measIndices)]);

    % calculate the mean of f (for a given protein)
    muf = Kxf'*((Kxx+nVariance*eye(size(Kxx)))\r);
    muf = reshape(muf,[length(tstar) length(measIndices)]);

    % take the results 
    muX{rowIdx} = mux; % interpolated x
    muF{rowIdx} = muf; % % interpolated f
    param{rowIdx} = [alpha lambda];
end
data.X = muX;
data.f = muF;
data.param = param;

if data.viz
    for idx = 1:length(data.protNames)
        figure,
        subplot(length(measIndices),1,1)
        suptitle(data.protNames{idx})
        for jdx = 1:length(measIndices)
            subplot(length(measIndices),1,jdx)
            plot(data.t,data.data{idx,jdx},'x')
            hold on
            plot(tstar,data.X{idx}(:,jdx),'r-')
            title(data.conditionNames{jdx})
            legend({'measurements','fit'})
        end
    end
end
end

function data = solvefWrtConcentrationsInfer(data)

data.subPowerset = powerset(1:length(data.protNames),data.maxRegulators);

data.margLikelihoods = zeros(length(data.subPowerset),length(data.protNames));

for rowIdx = 1:length(data.protNames)
    % response data
    y = cat(1,data.f{rowIdx}(:));

    % explanatory data
    for modelIdx = 1:length(data.subPowerset)
        regIndices = data.subPowerset{modelIdx};
        X = gatherExplanatoryData(data,y,regIndices);

        covFunc = {'covSum',{'covNNone','covNoise'}};
        foo = true;
        logHyperStart = [0; 0; 0];
        while foo
            foo = false;
            % optimize the set of hyperparameters (l, sigma_n^2)
            try logHyper= minimize(logHyperStart, 'gpr', -1000, covFunc, X, y);
            catch
                foo = true;
                logHyperStart = [-1; -1; -1] + ([0; 0; 0]-[-1; -1; -1]).*rand(3,1);
            end
        end
        data.margLikelihoods(modelIdx,rowIdx) = exp(-gpr(logHyper, covFunc, X, y));
    end
end
end

function hyperparameters = solvefWrtConcentrations(data)
tstar = data.t;

hyperparameters = cell(length(data.protNames),1);

measIndices = 1:size(data.data,2);

for rowIdx = 1:length(data.protNames)
    % response data
    y = cat(1,data.f{rowIdx}(:));

    % explanatory data
    regIndices = find(data.networkMatrix(rowIdx,:));
    X = gatherExplanatoryData(data,y,regIndices);

    covFunc = {'covSum',{'covNNone','covNoise'}};
    foo = true;
    logHyperStart = [0; 0; 0];
    while foo
        foo = false;
        % optimize the set of hyperparameters (l, sigma_n^2)
        try logHyper= minimize(logHyperStart, 'gpr', -1000, covFunc, X, y);
        catch
            foo = true;
            logHyperStart = [-1; -1; -1] + ([0; 0; 0]-[-1; -1; -1]).*rand(3,1);
        end
    end
    hyperparameters{rowIdx} = logHyper;
    
    foo = gpr(logHyper, covFunc, X, y, X);
    foo = reshape(foo,[length(tstar) length(measIndices)]);
    
    if data.viz
        figure
        subplot(length(measIndices),1,1)
        suptitle(data.protNames{rowIdx})
        for jdx = 1:length(measIndices)
            subplot(length(measIndices),1,jdx)
            plot(tstar,data.f{rowIdx}(:,jdx),'r-')
            hold on
            plot(tstar,foo(:,jdx),'g-')
            title(data.conditionNames{jdx})
            legend({'estimated f','estimated g'})
        end
    end
end
end

function mlld = foo(param,t,tp,measIndices,meanY,y)
% a wrapper for the minus log likelihood
param(1:2) = abs(param(1:2));
param(3:4) = real(param(3:4));
mlld = real(mLldX(t,tp,param(1),param(2),param(3),param(4),...
    measIndices,meanY,y));
end

function y = mLldX(t,tp,l,nVariance,alpha,lambda,measIndices,meanY,y)
% returns the minus log likelihood
Kxx = [];
for measIdx = measIndices
        Kxx = blkdiag(Kxx,kxx(t,tp,l,lambda));
end
r = y - m(t,lambda,alpha,meanY);
Kxx = Kxx + nVariance*eye(size(Kxx));
y = 1/2*r'*(Kxx\r)+1/2*log(det(Kxx))+length(r)/2*log(2*pi);
end

function y = m(t,lambda,alpha,k)
% mean function m
lambda = real(lambda);
alpha = real(alpha);
k = k - alpha/lambda;
y = alpha/lambda+k.*exp(-lambda*repmat(t,[length(k)/length(t) 1]));
end

function y = kxx(t,tp,l,lambda)
% covariance function k_{x,x}
lambda = real(lambda);
l = real(l);
t = repmat(t,[1 size(tp,1)]);
tp = repmat(tp',[size(t,1) 1]);
y = (sqrt(pi).*l.*exp((1/4).*lambda.*(l^2.*lambda-4.*(t+tp))).*...
    (exp(2.*t.*lambda).*(erf(t/l+(l.*lambda)/2)-erf((t-tp)/l+...
    (l.*lambda)/2))+exp(2.*tp.*lambda).*(erf(tp/l+(l.*lambda)/2)-...
    erf((tp-t)/l+(l.*lambda)/2))+erf((l.*lambda)/2-t/l)+erf((l.*lambda)...
    /2-tp/l)-2.*erf((l.*lambda)/2)))/(4.*lambda);
end

function y = kxf(t,tp,l,lambda)
% cross-covariance function k_{x,f}
lambda = real(lambda);
l = real(l);
t = repmat(t,[1 size(tp,1)]);
tp = repmat(tp',[size(t,1) 1]);
y = 1/2.*exp(1/4.*lambda.*(-4.*t+4.*tp+l.^2.*lambda)).*l.*sqrt(pi).*...
    (erf(tp./l+(l.*lambda)/2)-erf((-t+tp)/l+(l.*lambda)./2));
end

function y = kff(t,tp,l)
% covariance function k_{f,f}
if size(t,2) > 1
    y = exp(-1/l^2*(repmat(sum(t'.*t')',[1,size(tp,1)])+...
        repmat(sum(tp'.*tp'),[size(t,1),1])-2*t*tp'));
else
    t = repmat(t,[1 size(tp,1)]);
    tp = repmat(tp',[size(t,1) 1]);
    y = exp(-1/l^2*(t-tp).^2);
end
end

function subsets = powerset(set, lub)
%POWERSET  A Matlab function to generate the powerset of a set
%
%   usage: subsets = powerset(set)
%
%   where:
%       set         is a vector which holds the members of a set
%       lub         maximum cardinality of a subset
%       subsets     is a cell vector which holds all subsets of a set
%       
%   Author: Tarmo ?ij? <tarmo.aijo@tut.fi>

n_elements = 0;
for i=1:lub
    n_elements = n_elements + nchoosek(length(set), i);
end
subsets = cell(n_elements,1); % initialize the power set
card_set = length(set);
k = 1;
for i=1:lub
    subset_indices = nchoosek(1:card_set, i);
    l = 1;
    for j=k:k+size(subset_indices, 1) - 1
        subsets{j} = set(subset_indices(l,:));
        l = l + 1;
    end
    k = j + 1;
end
end

function [out1, out2] = gpr(logtheta, covfunc, x, y, xstar)
% gpr - Gaussian process regression, with a named covariance function. Two
% modes are possible: training and prediction: if no test data are given, the
% function returns minus the log likelihood and its partial derivatives with
% respect to the hyperparameters; this mode is used to fit the hyperparameters.
% If test data are given, then (marginal) Gaussian predictions are computed,
% whose mean and variance are returned. Note that in cases where the covariance
% function has noise contributions, the variance returned in S2 is for noisy
% test targets; if you want the variance of the noise-free latent function, you
% must substract the noise variance.
%
% usage: [nlml dnlml] = gpr(logtheta, covfunc, x, y)
%    or: [mu S2]  = gpr(logtheta, covfunc, x, y, xstar)
%
% where:
%
%   logtheta is a (column) vector of log hyperparameters
%   covfunc  is the covariance function
%   x        is a n by D matrix of training inputs
%   y        is a (column) vector (of size n) of targets
%   xstar    is a nn by D matrix of test inputs
%   nlml     is the returned value of the negative log marginal likelihood
%   dnlml    is a (column) vector of partial derivatives of the negative
%                 log marginal likelihood wrt each log hyperparameter
%   mu       is a (column) vector (of size nn) of prediced means
%   S2       is a (column) vector (of size nn) of predicted variances
%
% For more help on covariance functions, see "help covFunctions".
%
% (C) copyright 2006 by Carl Edward Rasmussen (2006-03-20).

if ischar(covfunc), covfunc = cellstr(covfunc); end % convert to cell if needed
[n, D] = size(x);
if eval(feval(covfunc{:})) ~= size(logtheta, 1)
  error('Error: Number of parameters do not agree with covariance function')
end

K = feval(covfunc{:}, logtheta, x);    % compute training set covariance matrix
% K = K + 1e-0*eye(size(K));

L = chol(K)';                        % cholesky factorization of the covariance
alpha = solve_chol(L',y);

if nargin == 4 % if no test cases, compute the negative log marginal likelihood

  out1 = 0.5*y'*alpha + sum(log(diag(L))) + 0.5*n*log(2*pi);

  if nargout == 2               % ... and if requested, its partial derivatives
    out2 = zeros(size(logtheta));       % set the size of the derivative vector
    W = L'\(L\eye(n))-alpha*alpha';                % precompute for convenience
    for i = 1:length(out2)
      out2(i) = sum(sum(W.*feval(covfunc{:}, logtheta, x, i)))/2;
    end
  end

else                    % ... otherwise compute (marginal) test predictions ...

  [Kss, Kstar] = feval(covfunc{:}, logtheta, x, xstar);     %  test covariances

  out1 = Kstar' * alpha;                                      % predicted means

  if nargout == 2
    v = L\Kstar;
    out2 = Kss - sum(v.*v)';
  end  

end
end

function [A, B] = covNNone(loghyper, x, z)

% Neural network covariance function with a single parameter for the distance
% measure. The covariance function is parameterized as:
%
% k(x^p,x^q) = sf2 * asin(x^p'*P*x^q / sqrt[(1+x^p'*P*x^p)*(1+x^q'*P*x^q)])
%
% where the x^p and x^q vectors on the right hand side have an added extra bias
% entry with unit value. P is ell^-2 times the unit matrix and sf2 controls the
% signal variance. The hyperparameters are:
%
% loghyper = [ log(ell)
%              log(sqrt(sf2) ]
%
% For more help on design of covariance functions, try "help covFunctions".
%
% (C) Copyright 2006 by Carl Edward Rasmussen (2006-03-24)

if nargin == 0, A = '2'; return; end              % report number of parameters

persistent Q K;                 
[n D] = size(x);
ell = exp(loghyper(1)); em2 = ell^(-2);
sf2 = exp(2*loghyper(2));
x = x/ell;

if nargin == 2                                             % compute covariance
  Q = x*x';
  K = (em2+Q)./(sqrt(1+em2+diag(Q))*sqrt(1+em2+diag(Q)'));
  A = sf2*asin(K);                 
elseif nargout == 2                              % compute test set covariances
  z = z/ell; 
  A = sf2*asin((em2+sum(z.*z,2))./(1+em2+sum(z.*z,2)));
  B = sf2*asin((em2+x*z')./sqrt((1+em2+sum(x.*x,2))*(1+em2+sum(z.*z,2)')));
else                                                % compute derivative matrix
  % check for correct dimension of the previously calculated kernel matrix
  if any(size(Q)~=n)  
    Q = x*x';
  end
  % check for correct dimension of the previously calculated kernel matrix
  if any(size(K)~=n)  
    K = (em2+Q)./(sqrt(1+em2+diag(Q))*sqrt(1+em2+diag(Q)'));
  end
  if z == 1                                                   % first parameter
    v = (em2+sum(x.*x,2))./(1+em2+diag(Q));
    A = -2*sf2*((em2+Q)./(sqrt(1+em2+diag(Q))*sqrt(1+em2+diag(Q)'))- ...
                            K.*(repmat(v,1,n)+repmat(v',n,1))/2)./sqrt(1-K.^2);
    clear Q;
  else                                                       % second parameter
    A = 2*sf2*asin(K);
    clear K;
  end
end
end

function [A, B] = covSum(covfunc, logtheta, x, z)
% covSum - compose a covariance function as the sum of other covariance
% functions. This function doesn't actually compute very much on its own, it
% merely does some bookkeeping, and calls other covariance functions to do the
% actual work.
%
% For more help on design of covariance functions, try "help covFunctions".
%
% (C) Copyright 2006 by Carl Edward Rasmussen, 2006-03-20.

for i = 1:length(covfunc)                   % iterate over covariance functions
  f = covfunc(i);
  if iscell(f{:}), f = f{:}; end          % dereference cell array if necessary
  j(i) = cellstr(feval(f{:}));
end

if nargin == 1,                                   % report number of parameters
  A = char(j(1)); for i=2:length(covfunc), A = [A, '+', char(j(i))]; end
  return
end

[n, D] = size(x);

v = [];              % v vector indicates to which covariance parameters belong
for i = 1:length(covfunc), v = [v repmat(i, 1, eval(char(j(i))))]; end

switch nargin
case 3                                              % compute covariance matrix
  A = zeros(n, n);                       % allocate space for covariance matrix
  for i = 1:length(covfunc)                  % iteration over summand functions
    f = covfunc(i);
    if iscell(f{:}), f = f{:}; end        % dereference cell array if necessary
    A = A + feval(f{:}, logtheta(v==i), x);            % accumulate covariances
  end

case 4                      % compute derivative matrix or test set covariances
  if nargout == 2                                % compute test set cavariances
    A = zeros(size(z,1),1); B = zeros(size(x,1),size(z,1));    % allocate space
    for i = 1:length(covfunc)
      f = covfunc(i);
      if iscell(f{:}), f = f{:}; end      % dereference cell array if necessary
      [AA BB] = feval(f{:}, logtheta(v==i), x, z);   % compute test covariances
      A = A + AA; B = B + BB;                                  % and accumulate
    end
  else                                            % compute derivative matrices
    i = v(z);                                       % which covariance function
    j = sum(v(1:z)==i);                    % which parameter in that covariance
    f = covfunc(i);
    if iscell(f{:}), f = f{:}; end        % dereference cell array if necessary
    A = feval(f{:}, logtheta(v==i), x, j);                 % compute derivative
  end

end
end

function [A, B] = covNoise(logtheta, x, z)
% Independent covariance function, ie "white noise", with specified variance.
% The covariance function is specified as:
%
% k(x^p,x^q) = s2 * \delta(p,q)
%
% where s2 is the noise variance and \delta(p,q) is a Kronecker delta function
% which is 1 iff p=q and zero otherwise. The hyperparameter is
%
% logtheta = [ log(sqrt(s2)) ]
%
% For more help on design of covariance functions, try "help covFunctions".
%
% (C) Copyright 2006 by Carl Edward Rasmussen, 2006-03-24.

if nargin == 0, A = '1'; return; end              % report number of parameters

s2 = exp(2*logtheta);                                          % noise variance

if nargin == 2                                      % compute covariance matrix
  A = s2*eye(size(x,1));
elseif nargout == 2                              % compute test set covariances
  A = s2;
  B = 0;                               % zeros cross covariance by independence
else                                                % compute derivative matrix
  A = 2*s2*eye(size(x,1));
end
end

function [X, fX, i] = minimize(X, f, length, varargin)
% Minimize a differentiable multivariate function. 
%
% Usage: [X, fX, i] = minimize(X, f, length, P1, P2, P3, ... )
%
% where the starting point is given by "X" (D by 1), and the function named in
% the string "f", must return a function value and a vector of partial
% derivatives of f wrt X, the "length" gives the length of the run: if it is
% positive, it gives the maximum number of line searches, if negative its
% absolute gives the maximum allowed number of function evaluations. You can
% (optionally) give "length" a second component, which will indicate the
% reduction in function value to be expected in the first line-search (defaults
% to 1.0). The parameters P1, P2, P3, ... are passed on to the function f.
%
% The function returns when either its length is up, or if no further progress
% can be made (ie, we are at a (local) minimum, or so close that due to
% numerical problems, we cannot get any closer). NOTE: If the function
% terminates within a few iterations, it could be an indication that the
% function values and derivatives are not consistent (ie, there may be a bug in
% the implementation of your "f" function). The function returns the found
% solution "X", a vector of function values "fX" indicating the progress made
% and "i" the number of iterations (line searches or function evaluations,
% depending on the sign of "length") used.
%
% The Polack-Ribiere flavour of conjugate gradients is used to compute search
% directions, and a line search using quadratic and cubic polynomial
% approximations and the Wolfe-Powell stopping criteria is used together with
% the slope ratio method for guessing initial step sizes. Additionally a bunch
% of checks are made to make sure that exploration is taking place and that
% extrapolation will not be unboundedly large.
%
% See also: checkgrad 
%
% Copyright (C) 2001 - 2006 by Carl Edward Rasmussen (2006-09-08).

INT = 0.1;    % don't reevaluate within 0.1 of the limit of the current bracket
EXT = 3.0;                  % extrapolate maximum 3 times the current step-size
MAX = 20;                         % max 20 function evaluations per line search
RATIO = 10;                                       % maximum allowed slope ratio
SIG = 0.1; RHO = SIG/2; % SIG and RHO are the constants controlling the Wolfe-
% Powell conditions. SIG is the maximum allowed absolute ratio between
% previous and new slopes (derivatives in the search direction), thus setting
% SIG to low (positive) values forces higher precision in the line-searches.
% RHO is the minimum allowed fraction of the expected (from the slope at the
% initial point in the linesearch). Constants must satisfy 0 < RHO < SIG < 1.
% Tuning of SIG (depending on the nature of the function to be optimized) may
% speed up the minimization; it is probably not worth playing much with RHO.

% The code falls naturally into 3 parts, after the initial line search is
% started in the direction of steepest descent. 1) we first enter a while loop
% which uses point 1 (p1) and (p2) to compute an extrapolation (p3), until we
% have extrapolated far enough (Wolfe-Powell conditions). 2) if necessary, we
% enter the second loop which takes p2, p3 and p4 chooses the subinterval
% containing a (local) minimum, and interpolates it, unil an acceptable point
% is found (Wolfe-Powell conditions). Note, that points are always maintained
% in order p0 <= p1 <= p2 < p3 < p4. 3) compute a new search direction using
% conjugate gradients (Polack-Ribiere flavour), or revert to steepest if there
% was a problem in the previous line-search. Return the best value so far, if
% two consecutive line-searches fail, or whenever we run out of function
% evaluations or line-searches. During extrapolation, the "f" function may fail
% either with an error or returning Nan or Inf, and minimize should handle this
% gracefully.

if max(size(length)) == 2, red=length(2); length=length(1); else red=1; end
if length>0, S='Linesearch'; else S='Function evaluation'; end 

i = 0;                                            % zero the run length counter
ls_failed = 0;                             % no previous line search has failed
[f0 df0] = feval(f, X, varargin{:});          % get function value and gradient
fX = f0;
i = i + (length<0);                                            % count epochs?!
s = -df0; d0 = -s'*s;           % initial search direction (steepest) and slope
x3 = red/(1-d0);                                  % initial step is red/(|s|+1)

while i < abs(length)                                      % while not finished
  i = i + (length>0);                                      % count iterations?!

  X0 = X; F0 = f0; dF0 = df0;                   % make a copy of current values
  if length>0, M = MAX; else M = min(MAX, -length-i); end

  while 1                             % keep extrapolating as long as necessary
    x2 = 0; f2 = f0; d2 = d0; f3 = f0; df3 = df0;
    success = 0;
    while ~success && M > 0
      try
        M = M - 1; i = i + (length<0);                         % count epochs?!
        [f3 df3] = feval(f, X+x3*s, varargin{:});
        if isnan(f3) || isinf(f3) || any(isnan(df3)+isinf(df3)), error(''), end
        success = 1;
      catch                                % catch any error which occured in f
        x3 = (x2+x3)/2;                                  % bisect and try again
      end
    end
    if f3 < F0, X0 = X+x3*s; F0 = f3; dF0 = df3; end         % keep best values
    d3 = df3'*s;                                                    % new slope
    if d3 > SIG*d0 || f3 > f0+x3*RHO*d0 || M == 0  % are we done extrapolating?
      break
    end
    x1 = x2; f1 = f2; d1 = d2;                        % move point 2 to point 1
    x2 = x3; f2 = f3; d2 = d3;                        % move point 3 to point 2
    A = 6*(f1-f2)+3*(d2+d1)*(x2-x1);                 % make cubic extrapolation
    B = 3*(f2-f1)-(2*d1+d2)*(x2-x1);
    x3 = x1-d1*(x2-x1)^2/(B+sqrt(B*B-A*d1*(x2-x1))); % num. error possible, ok!
    if ~isreal(x3) || isnan(x3) || isinf(x3) || x3 < 0 % num prob | wrong sign?
      x3 = x2*EXT;                                 % extrapolate maximum amount
    elseif x3 > x2*EXT                  % new point beyond extrapolation limit?
      x3 = x2*EXT;                                 % extrapolate maximum amount
    elseif x3 < x2+INT*(x2-x1)         % new point too close to previous point?
      x3 = x2+INT*(x2-x1);
    end
  end                                                       % end extrapolation

  while (abs(d3) > -SIG*d0 || f3 > f0+x3*RHO*d0) && M > 0  % keep interpolating
    if d3 > 0 || f3 > f0+x3*RHO*d0                         % choose subinterval
      x4 = x3; f4 = f3; d4 = d3;                      % move point 3 to point 4
    else
      x2 = x3; f2 = f3; d2 = d3;                      % move point 3 to point 2
    end
    if f4 > f0           
      x3 = x2-(0.5*d2*(x4-x2)^2)/(f4-f2-d2*(x4-x2));  % quadratic interpolation
    else
      A = 6*(f2-f4)/(x4-x2)+3*(d4+d2);                    % cubic interpolation
      B = 3*(f4-f2)-(2*d2+d4)*(x4-x2);
      x3 = x2+(sqrt(B*B-A*d2*(x4-x2)^2)-B)/A;        % num. error possible, ok!
    end
    if isnan(x3) || isinf(x3)
      x3 = (x2+x4)/2;               % if we had a numerical problem then bisect
    end
    x3 = max(min(x3, x4-INT*(x4-x2)),x2+INT*(x4-x2));  % don't accept too close
    [f3 df3] = feval(f, X+x3*s, varargin{:});
    if f3 < F0, X0 = X+x3*s; F0 = f3; dF0 = df3; end         % keep best values
    M = M - 1; i = i + (length<0);                             % count epochs?!
    d3 = df3'*s;                                                    % new slope
  end                                                       % end interpolation

  if abs(d3) < -SIG*d0 && f3 < f0+x3*RHO*d0          % if line search succeeded
    X = X+x3*s; f0 = f3; fX = [fX' f0]';                     % update variables
    %fprintf('%s %6i;  Value %4.6e\r', S, i, f0);
    s = (df3'*df3-df0'*df3)/(df0'*df0)*s - df3;   % Polack-Ribiere CG direction
    df0 = df3;                                               % swap derivatives
    d3 = d0; d0 = df0'*s;
    if d0 > 0                                      % new slope must be negative
      s = -df0; d0 = -s'*s;                  % otherwise use steepest direction
    end
    x3 = x3 * min(RATIO, d3/(d0-realmin));          % slope ratio but max RATIO
    ls_failed = 0;                              % this line search did not fail
  else
    X = X0; f0 = F0; df0 = dF0;                     % restore best point so far
    if ls_failed || i > abs(length)         % line search failed twice in a row
      break;                             % or we ran out of time, so we give up
    end
    s = -df0; d0 = -s'*s;                                        % try steepest
    x3 = 1/(1-d0);                     
    ls_failed = 1;                                    % this line search failed
  end
end
%fprintf('\n');
end


function x = solve_chol(A, B)
% solve_chol - solve linear equations from the Cholesky factorization.
% Solve A*X = B for X, where A is square, symmetric, positive definite. The
% input to the function is R the Cholesky decomposition of A and the matrix B.
% Example: X = solve_chol(chol(A),B);
%
% NOTE: The program code is written in the C language for efficiency and is
% contained in the file solve_chol.c, and should be compiled using matlabs mex
% facility. However, this file also contains a (less efficient) matlab
% implementation, supplied only as a help to people unfamiliar with mex. If
% the C code has been properly compiled and is avaiable, it automatically
% takes precendence over the matlab code in this file.
%
% Copyright (c) 2004, 2005, 2006 by Carl Edward Rasmussen. 2006-02-08.

if nargin ~= 2 | nargout > 1
  error('Wrong number of arguments.');
end

if size(A,1) ~= size(A,2) | size(A,1) ~= size(B,1)
  error('Wrong sizes of matrix arguments.');
end

x = A\(A'\B);
end

function [minimum,fval] = anneal(loss, parent, options)
% ANNEAL  Minimizes a function with the method of simulated annealing
% (Kirkpatrick et al., 1983)
%
%  ANNEAL takes three input parameters, in this order:
%
%  LOSS is a function handle (anonymous function or inline) with a loss
%  function, which may be of any type, and needn't be continuous. It does,
%  however, need to return a single value.
%
%  PARENT is a vector with initial guess parameters. You must input an
%  initial guess.
%
%  OPTIONS is a structure with settings for the simulated annealing. If no
%  OPTIONS structure is provided, ANNEAL uses a default structure. OPTIONS
%  can contain any or all of the following fields (missing fields are
%  filled with default values):
%
%       Verbosity: Controls output to the screen.
%                  0 suppresses all output
%                  1 gives final report only [default]
%                  2 gives temperature changes and final report 
%       Generator: Generates a new solution from an old one.
%                  Any function handle that takes a solution as input and
%                  gives a valid solution (i.e. some point in the solution
%                  space) as output.
%                  The default function generates a row vector which
%                  slightly differs from the input vector in one element:
%                  @(x) (x+(randperm(length(x))==length(x))*randn/100)
%                  Other examples of possible solution generators:
%                  @(x) (rand(3,1)): Picks a random point in the unit cube
%                  @(x) (ceil([9 5].*rand(2,1))): Picks a point in a 9-by-5
%                                                 discrete grid
%                  Note that if you use the default generator, ANNEAL only
%                  works on row vectors. For loss functions that operate on
%                  column vectors, use this generator instead of the
%                  default:
%                  @(x) (x(:)'+(randperm(length(x))==length(x))*randn/100)'
%        InitTemp: The initial temperature, can be any positive number.
%                  Default is 1.
%        StopTemp: Temperature at which to stop, can be any positive number
%                  smaller than InitTemp. 
%                  Default is 1e-8.
%         StopVal: Value at which to stop immediately, can be any output of
%                  LOSS that is sufficiently low for you.
%                  Default is -Inf.
%       CoolSched: Generates a new temperature from the previous one.
%                  Any function handle that takes a scalar as input and
%                  returns a smaller but positive scalar as output. 
%                  Default is @(T) (.8*T)
%      MaxConsRej: Maximum number of consecutive rejections, can be any
%                  positive number.
%                  Default is 1000.
%        MaxTries: Maximum number of tries within one temperature, can be
%                  any positive number.
%                  Default is 300.
%      MaxSuccess: Maximum number of successes within one temperature, can
%                  be any positive number.
%                  Default is 20.
%
%
%  Usage:
%     [MINIMUM,FVAL] = ANNEAL(LOSS,NEWSOL,[OPTIONS]);
%          MINIMUM is the solution which generated the smallest encountered
%          value when input into LOSS.
%          FVAL is the value of the LOSS function evaluated at MINIMUM.
%     OPTIONS = ANNEAL();
%          OPTIONS is the default options structure.
%
%
%  Example:
%     The so-called "six-hump camelback" function has several local minima
%     in the range -3<=x<=3 and -2<=y<=2. It has two global minima, namely
%     f(-0.0898,0.7126) = f(0.0898,-0.7126) = -1.0316. We can define and
%     minimise it as follows:
%          camel = @(x,y)(4-2.1*x.^2+x.^4/3).*x.^2+x.*y+4*(y.^2-1).*y.^2;
%          loss = @(p)camel(p(1),p(2));
%          [x f] = ANNEAL(loss,[0 0])
%     We get output:
%               Initial temperature:     	1
%               Final temperature:       	3.21388e-007
%               Consecutive rejections:  	1027
%               Number of function calls:	6220
%               Total final loss:        	-1.03163
%               x =
%                  -0.0899    0.7127
%               f =
%                  -1.0316
%     Which reasonably approximates the analytical global minimum (note
%     that due to randomness, your results will likely not be exactly the
%     same).

%  Reference:
%    Kirkpatrick, S., Gelatt, C.D., & Vecchi, M.P. (1983). Optimization by
%    Simulated Annealing. _Science, 220_, 671-680.

%   joachim.vandekerckhove@psy.kuleuven.be
%   $Revision: v5 $  $Date: 2006/04/26 12:54:04 $

def = struct(...
        'CoolSched',@(T) (.8*T),...
        'Generator',@(x) (x+(randperm(length(x))==length(x))*randn/100),...
        'InitTemp',1,...
        'MaxConsRej',1000,...
        'MaxSuccess',20,...
        'MaxTries',300,...
        'StopTemp',1e-8,...
        'StopVal',-Inf,...
        'Verbosity',1);

% Check input
if ~nargin %user wants default options, give it and stop
    minimum = def;
    return
elseif nargin<2, %user gave only objective function, throw error
    error('MATLAB:anneal:noParent','You need to input a first guess.');
elseif nargin<3, %user gave no options structure, use default
    options=def;
else %user gave all input, check if options structure is complete
    if ~isstruct(options)
        error('MATLAB:anneal:badOptions',...
            'Input argument ''options'' is not a structure')
    end
    fs = {'CoolSched','Generator','InitTemp','MaxConsRej',...
        'MaxSuccess','MaxTries','StopTemp','StopVal','Verbosity'};
    for nm=1:length(fs)
        if ~isfield(options,fs{nm}), options.(fs{nm}) = def.(fs{nm}); end
    end
end

% main settings
newsol = options.Generator;      % neighborhood space function
Tinit = options.InitTemp;        % initial temp
minT = options.StopTemp;         % stopping temp
cool = options.CoolSched;        % annealing schedule
minF = options.StopVal;
max_consec_rejections = options.MaxConsRej;
max_try = options.MaxTries;
max_success = options.MaxSuccess;
report = options.Verbosity;
k = 1;                           % boltzmann constant

% counters etc
itry = 0;
success = 0;
finished = 0;
consec = 0;
T = Tinit;
initenergy = loss(parent);
oldenergy = initenergy;
total = 0;
if report==2, fprintf(1,'\n  T = %7.5f, loss = %10.5f\n',T,oldenergy); end

while ~finished;
    itry = itry+1; % just an iteration counter
    current = parent; 
    
    % % Stop / decrement T criteria
    if itry >= max_try || success >= max_success;
        if T < minT || consec >= max_consec_rejections;
            finished = 1;
            total = total + itry;
            break;
        else
            T = cool(T);  % decrease T according to cooling schedule
            if report==2, % output
                fprintf(1,'  T = %7.5f, loss = %10.5f\n',T,oldenergy);
            end
            total = total + itry;
            itry = 1;
            success = 1;
        end
    end
    
    newparam = newsol(current);
    newenergy = loss(newparam);
    
    if (newenergy < minF),
        parent = newparam; 
        oldenergy = newenergy;
        break
    end
    
    if (oldenergy-newenergy > 1e-6)
        parent = newparam;
        oldenergy = newenergy;
        success = success+1;
        consec = 0;
    else
        if (rand < exp( (oldenergy-newenergy)/(k*T) ));
            parent = newparam;
            oldenergy = newenergy;
            success = success+1;
        else
            consec = consec+1;
        end
    end
end

minimum = parent;
fval = oldenergy;

if report;
    fprintf(1, '\n  Initial temperature:     \t%g\n', Tinit);
    fprintf(1, '  Final temperature:       \t%g\n', T);
    fprintf(1, '  Consecutive rejections:  \t%i\n', consec);
    fprintf(1, '  Number of function calls:\t%i\n', total);
    fprintf(1, '  Total final loss:        \t%g\n', fval);
end
end
