function [varargout] = prt_rvr(varargin)
% Optimisation for Relevance Vector Regression
%
% [w,alpha,beta,ll] = prt_rvr(Phi,t)
% Phi   - MxM matrix derived from kernel function of vector pairs
% t     - the values to be matched
% w     - weights
% alpha - 1/variance for the prior part of the model
% beta  - 1/variance for the likelihood part of the model
% ll    - the negative log-likelihood.
%
% [w,alpha,beta,nu,ll]=spm_rvr(K,t,opt)
% K     - a cell-array of MxM dot-product matrices.
% t     - the values to be matched
% opt   - either 'Linear' or 'Gaussian RBF'
%         'Linear'       is for linear regression models, where
%                        the optimal kernel is generated by
%                        [nu(1)*K{1} + nu(1)*K{2}... ones(size(K{1},1),1)]
%         'Gaussian RBF' is for regression using Gaussian radial basis
%                        functions.  The kernel is generated from
%                        P1  = nu(1)*K{1} + nu(1)*K{2} ... ;
%                        P2  = repmat(diag(P1) ,1,size(P1,2)) +...
%                              repmat(diag(P1)',size(P1,1),1) - 2*P1;
%                        Phi = exp([-0.5*P2 ones(size(P1,1),1)]);
% w     - weights
% alpha - 1/variance for the prior part of the model
% beta  - 1/variance for the likelihood part of the model
% nu    - parameters that convert the dot-product matrices into
%         a kernel matrix (Phi).
% ll    - the negative log-likelihood.
%
% The first way of calling the routine simply optimises the
% weights.  This involves estimating a restricted maximum
% likelihood (REML) solution, which maximises P(alpha,beta|t,Phi).
% Note that REML is also known as Type II Maximum Likelihood
% (ML-II). The ML-II solution tends towards infinite weights for
% some the regularisation terms (i.e. 1/alpha(i) approaches 0).
% The appropriate columns are removed from the model when
% this happens.
%
% The second way of calling the routine also estimates additional
% input scale parameters as described in Appendix C of Tipping (2001).
% This method is much slower, as a full optimisation for the scale
% parameters is done after each update of the alphas and beta.
%
% see: http://research.microsoft.com/mlp/RVM/relevance.htm
%
% Refs:
% The Relevance Vector Machine.
% In S. A. Solla, T. K. Leen, and K.-R. Mller (Eds.),
% Advances in Neural Information Processing Systems 12,
% pp.  652-658. Cambridge, Mass: MIT Press.
%
% Michael E. Tipping
% Sparse Bayesian Learning and the Relevance Vector Machine
% Journal of Machine Learning Research 1 (2001) 211-244
%________________________________________________________________________
% Copyright (C) 2011 Wellcome Department of Imaging Neuroscience & 
% Machine Learning & Neuroimaging Laboratory

% Written by John Ashburner
% $Id: prt_rvr.m 741 2013-07-22 14:07:36Z mjrosa $

if isnumeric(varargin{1}),
    [varargout{1:nargout}]=regression0(varargin{:});
elseif iscell(varargin{1}),
    [varargout{1:nargout}]=regression1(varargin{:});
else
    error('Incorrect usage');
end;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,ll]=regression0(Phi,t)
[N,M]  = size(Phi);
if N==M,
    Phi = [Phi ones(N,1)];
elseif M~=N+1,
    error('Phi must be N x (N+1)');
end;
scale             = sqrt(sum(sum(Phi(1:N,1:N).^2))/N^2);
scale             = [ones(N,1)*scale ; 1];
Phi               = Phi/spdiags(scale,0,numel(scale),numel(scale));
alpha             = ones(size(Phi,2),1)/N;
%beta             = N/sum((t-mean(t)).^2);
beta              = 1e6;
[w,alpha,beta,ll] = rvr1a(Phi,t,alpha,beta);
alpha             = [alpha(1)*ones(N,1) ; alpha(2)];
[w,alpha,beta,ll] = rvr2a(Phi,t,alpha,beta);
w                 = w./scale;
alpha             = alpha.*scale.^2;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,ll] = rvr1a(Phi,t,alpha,beta)
% This function is not actually used
%spm_chi2_plot('Init','ML-II (non-sparse)','-Log-likelihood','Iteration');
[N,M]   = size(Phi);
ll      = Inf;
PP      = Phi'*Phi;
Pt      = Phi'*t;
for subit=1:10,
    alpha_old = alpha;
    beta_old  = beta;

    % E-step
    S         = inv(PP*beta + spdiags([ones(N,1)*alpha(1) ; alpha(2)],0,N+1,N+1));
    w         = S*(Pt*beta);

   % figure(3); plot(t,Phi*w,'.'); drawnow;

    tmp       = t-Phi*w;
    ll        = ...
         -0.5*log(alpha(1))*N-0.5*log(alpha(2))-0.5*N*log(beta)-0.5*logdet(S)...
         +0.5*tmp'*tmp*beta + 0.5*sum(w.^2.*[repmat(alpha(1),N,1) ; alpha(2)])...
         +0.5*(M-N)*log(2*pi);
  %  if subit>1, spm_chi2_plot('Set',ll); end;
    %fprintf('%g\n',ll);

    % M-step
    ds        = diag(S);
    dfa1      = sum(ds(1:N))*alpha(1);
    dfa2      = sum(ds(N+1))*alpha(2);
    alpha(1)  = max(N-dfa1,eps)/(sum(w(1:N).^2)   +eps);
    alpha(2)  = max(1-dfa2,eps)/(sum(w(N+1).^2)   +eps);
    beta      = max(dfa1+dfa2-1,eps)/(sum((Phi*w-t).^2)+eps);

    % Convergence
    if max(max(abs(log((alpha+eps)./(alpha_old+eps)))),log(beta/beta_old)) < 1e-9,
        break;
    end;
end;
%spm_chi2_plot('Clear');
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,ll]=rvr2a(Phi,t,alpha,beta)
%spm_chi2_plot('Init','ML-II (sparse)','-Log-likelihood','Iteration');
[N,M] = size(Phi);
nz    = true(M,1);

PP    = Phi'*Phi;
Pt    = Phi'*t;

for subit=1:200,
    th         = min(alpha)*1e9;
    nz         = alpha<th;
    alpha(~nz) = th*1e9;
    alpha_old  = alpha;
    beta_old   = beta;

    % E-step
    S         = inv(PP(nz,nz)*beta + diag(alpha(nz)));
    w         = S*Pt(nz)*beta;

  %  figure(3); plot(t,Phi(:,nz)*w,'.'); drawnow;

    tmp = t-Phi(:,nz)*w;
    ll  = ...
        -0.5*sum(log(alpha(nz)+1e-32))-0.5*N*log(beta+1e-32)-0.5*logdet(S)...
        +0.5*tmp'*tmp*beta + 0.5*sum(w.^2.*alpha(nz))...
        +0.5*(sum(nz)-N)*log(2*pi);
%    if subit>0, spm_chi2_plot('Set',ll); end;
    %fprintf('%g\t%g\n',ll,exp(mean(log(alpha)))/beta);

    % M-step
    gam       = 1 - alpha(nz).*diag(S);
    alpha(nz) = max(gam,eps)./(w.^2+1e-32);
    beta      = max(N-sum(gam),eps)./(sum((Phi(:,nz)*w-t).^2)+1e-32);

    % Convergence
    if max(max(abs(log((alpha(nz)+eps)./(alpha_old(nz)+eps)))),log(beta/beta_old)) < 1e-6*N,
        break;
    end;
end;
w(nz)  = w;
w(~nz) = 0;
w      = w(:);
%spm_chi2_plot('Clear');
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,nu,ll]=regression1(K,t,opt)
% Relevance vector regression
if nargin<3, opt = 'Linear'; end;
switch opt,
case {'Linear','linear','lin'},
    dkrn_f = @make_dphi;
    krn_f  = @make_phi;
case {'Gaussian RBF','nonlinear','nonlin'},
    dkrn_f = @make_dphi_rbf;
    krn_f  = @make_phi_rbf;
otherwise
    error('Unknown option');
end;
[N,M]  = size(K{1});
nu     = ones(numel(K),1);
rescal = ones(numel(K),1);
for i=1:numel(K),
    if strcmpi(opt,'Gaussian RBF') || strcmpi(opt,'nonlinear') || strcmpi(opt,'nonlin'),
        d         = 0.5*diag(K{i});
        K{i}      = repmat(d,[1 size(K{i},1)]) + repmat(d',[size(K{i},1),1]) - K{i};
        K{i}      = max(K{i},0);
        K{i}      = -K{i};
        nu(i)     = 1/sqrt(sum(K{i}(:).^2)/(size(K{i},1).^2-size(K{i},1)));
    else
        rescal(i) = sqrt(size(K{i},1)/sum(K{i}(:).^2));
        K{i}      = K{i}.*rescal(i);
    end;
end;

alpha  = [1 1]';
%beta  = 1/sum((t-mean(t)).^2);
beta   = 1e6;
[w,alpha,beta,nu,ll]=rvr1(K,t,alpha,beta,nu,krn_f,dkrn_f);
alpha  = [alpha(1)*ones(N,1) ; alpha(2)];
[w,alpha,beta,nu,ll]=rvr2(K,t,alpha,beta,nu,krn_f,dkrn_f);
nu    = nu.*rescal;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,nu,ll] = rvr1(K,t,alpha,beta,nu,krn_f,dkrn_f)
% This function is not actually used
spm_chi2_plot('Init','ML-II (non-sparse)','-Log-likelihood','Iteration');
[N,M]   = size(K{1});
ll      = Inf;
for iter=1:50,
    Phi     = feval(krn_f,nu,K);
    for subit=1:1,
        alpha_old = alpha;
        beta_old  = beta;

        % E-step
        S         = inv(Phi'*Phi*beta + spdiags([ones(N,1)*alpha(1) ; alpha(2)],0,N+1,N+1));
        w         = S*(Phi'*t*beta);

        % M-step
        ds        = diag(S);
        dfa1      = sum(ds(1:N))*alpha(1);
        dfa2      = sum(ds(N+1))*alpha(2);
        alpha(1)  = max(N-dfa1,eps)/(sum(w(1:N).^2)   +eps);
        alpha(2)  = max(1-dfa2,eps)/(sum(w(N+1).^2)   +eps);
        beta      = max(dfa1+dfa2-1,eps)/(sum((Phi*w-t).^2)+eps);

        % Convergence
        if max(max(abs(log((alpha+eps)./(alpha_old+eps)))),log(beta/beta_old)) < 1e-9,
            break;
        end;
    end;


    % Update nu
    oll       = ll;
    al1       = [ones(N,1)*alpha(1) ; alpha(2)];
    [nu,ll]   = re_estimate_nu(K,t,nu,al1,beta,krn_f,dkrn_f);

%    scale = sqrt(sum(nu.^2));
%    nu    = nu/scale;
%    alpha = alpha/scale^2;

    spm_chi2_plot('Set',ll);
    if abs(oll-ll) < 1e-6*N, break; end;
end;
spm_chi2_plot('Clear');
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [w,alpha,beta,nu,ll]=rvr2(K,t,alpha,beta,nu,krn_f,dkrn_f)
spm_chi2_plot('Init','ML-II (sparse)','-Log-likelihood','Iteration');
[N,M] = size(K{1});
w     = zeros(N+1,1);
ll    = Inf;
for iter=1:100,
    for subits=1:1,
        % Suboptimal estimates of nu if the alphas and weights are pruned
        % too quickly.

        th         = min(alpha)*1e9;
        nz         = alpha<th;
        alpha(~nz) = th*1e9;

        alpha_old  = alpha;
        beta_old   = beta;
        Phi        = feval(krn_f,nu,K,nz);

        % E-step
        S         = inv(beta*Phi'*Phi + diag(alpha(nz)));
        w(nz)     = S*Phi'*t*beta;
        w(~nz)    = 0;

        % figure(3); plot(t,Phi*w(nz),'.');drawnow;

        % M-step
        gam       = 1 - alpha(nz).*diag(S);
        alpha(nz) = max(gam,eps)./(w(nz).^2+1e-32);
        beta      = max(N-sum(gam),eps)./(sum((Phi*w(nz)-t).^2)+1e-32);

       % Convergence
        if max(max(abs(log((alpha+eps)./(alpha_old+eps)))),log(beta/beta_old)) < 1e-6,
            break;
        end;
    end;

    oll       = ll;
    [nu,ll]   = re_estimate_nu(K,t,nu,alpha,beta,krn_f,dkrn_f,nz);

    % scale     = sqrt(sum(nu.^2));
    % nu        = nu/scale;
    % alpha     = alpha/scale^2;

    spm_chi2_plot('Set',ll);

    % Convergence
    if abs(oll-ll) < 1e-9*N,
        break;
    end;
end;
spm_chi2_plot('Clear');
return;
%__________________________________________________________________________

%__________________________________________________________________________
function Phi = make_phi(nu,K,nz)
% Dot product matrix, generated from linear combination of dot-product
% matrices.
if nargin>2 && ~isempty(nz),
    nz1 = nz(1:size(K{1},1));
    nz2 = nz(size(K{1},1)+1);
    Phi = K{1}(:,nz1)*nu(1);
    for i=2:numel(K),
        Phi=Phi+K{i}(:,nz1)*nu(i);
    end;
    Phi = [Phi ones(size(Phi,1),sum(nz2))];
else
    Phi = K{1}*nu(1);
    for i=2:numel(K),
        Phi=Phi+K{i}*nu(i);
    end;
    Phi = [Phi ones(size(Phi,1),1)];
end;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [dPhi,d2Phi] = make_dphi(nu,K,nz)
% First and second derivatives of Phi with respect to nu, where
% Phi is a dot-product matrix.
dPhi  = cell(size(K));
d2Phi = cell(numel(K));
if nargin>2 && ~isempty(nz),
    nz1 = nz(1:size(K{1},1));
    nz2 = nz(size(K{1},1)+1);
    for i=1:numel(K),
        dPhi{i} = [K{i}(:,nz1),zeros(size(K{i},1),sum(nz2))];
        dPhi{i} = dPhi{i}*nu(i);
    end;
else
    for i=1:numel(K),
        dPhi{i} = [K{i},zeros(size(K{i},1),1)];
        dPhi{i} = dPhi{i}*nu(i);
    end;
end;
%z = sparse([],[],[],size(dPhi{1},1),size(dPhi{1},2));
z = zeros(size(dPhi{1}));
for i=1:numel(K),
    d2Phi{i,i} = nu(i)*dPhi{i};
    for j=(i+1):numel(K),
        d2Phi{i,j} = z;
        d2Phi{j,i} = z;
    end;
    dPhi{i} = nu(i)*dPhi{i};
end;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function Phi = make_phi_rbf(nu,K,nz)
% Radial basis function kernel.
if nargin>2 && ~isempty(nz),
    nz1 = nz(1:size(K{1},1));
    nz2 = nz(size(K{1},1)+1);
    Phi       = K{1}(:,nz1)*nu(1);
    for i=2:numel(K),
        Phi=Phi+K{i}(:,nz1)*nu(i);
    end;
    Phi = [exp(Phi) ones(size(Phi,1),sum(nz2))];
else
    Phi       = K{1}*nu(1);
    for i=2:numel(K),
        Phi=Phi+K{i}*nu(i);
    end;
    Phi = [exp(Phi) ones(size(Phi,1),1)];
end;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [dPhi,d2Phi] = make_dphi_rbf(nu,K,nz)
% First and second derivatives of Phi with respect to nu, where
% Phi is defined by radial basis functions.
Phi   = make_phi_rbf(nu,K,nz);
dPhi  = cell(size(K));
d2Phi = cell(numel(K));
if nargin>2 && ~isempty(nz),
    nz1 = nz(1:size(K{1},1));
    nz2 = nz(size(K{1},1)+1);
    for i=1:numel(K),
        dPhi{i} = [K{i}(:,nz1),zeros(size(K{i},1),sum(nz2))];
    end;
else
    for i=1:numel(K),
        dPhi{i} = [K{i},zeros(size(K{i},1),1)];
    end;
end;

for i=1:numel(K),
    d2Phi{i,i} = nu(i)*dPhi{i}.*Phi.*(1+nu(i)*dPhi{i});
    for j=(i+1):numel(K),
        d2Phi{i,j} = (nu(i)*nu(j))*dPhi{i}.*dPhi{j}.*Phi;
        d2Phi{j,i} = d2Phi{i,j};
    end;
    dPhi{i} = nu(i)*dPhi{i}.*Phi;
end;
return;
%__________________________________________________________________________

%__________________________________________________________________________
function [nu,ll] = re_estimate_nu(K,t,nu,alpha,beta,krn_f,dkrn_f,nz)
% See Appendix C of Tipping (2001).  Note that a Levenberg-Marquardt
% approach is used for the optimisation.
if nargin<8, nz = true(size(K{1},2)+1,1); end;

ll   = Inf;
lam  = 1e-6;
Phi  = feval(krn_f,nu,K,nz);
S    = inv(Phi'*Phi*beta+diag(alpha(nz)));
w    = beta*S*Phi'*t;
N    = size(Phi,1);
ll   = ...
     -0.5*sum(log(alpha(nz)))-0.5*N*log(beta)-0.5*logdet(S)...
     +0.5*(t-Phi*w)'*(t-Phi*w)*beta + 0.5*sum(w.^2.*alpha(nz))...
     +0.5*(sum(nz)-N)*log(2*pi);

for subit=1:30,

    % Generate 1st and second derivatives of the objective function (ll).
    % These are derived from the first and second partial derivatives
    % of Phi with respect to nu.
    g    = zeros(numel(K),1);
    H    = zeros(numel(K));
    [dPhi,d2Phi] = feval(dkrn_f,nu,K,nz);
    for i=1:numel(K),
        tmp1  = Phi'*dPhi{i};
        tmp1  = tmp1+tmp1';
        g(i)  = 0.5*beta*(sum(sum(S.*tmp1)) + w'*tmp1*w - 2*t'*dPhi{i}*w);
        for j=i:numel(K),
            tmp    = dPhi{j}'*dPhi{i} + Phi'*d2Phi{i,j};
            tmp    = tmp+tmp';
            tmp2   = Phi'*dPhi{j};
            tmp2   = tmp2+tmp2';
            H(i,j) = sum(sum(S.*(tmp - tmp1*S*tmp2))) + w'*tmp*w - 2*w'*d2Phi{i,j}'*t;
            H(i,j) = 0.5*beta*H(i,j);
            H(j,i) = H(i,j);
        end;
    end;

    oll  = ll;
    onu  = nu;

    % Negative eigenvalues indicate that the solution is tending
    % towards a saddle point or a maximum rather than a minimum
    % of the negative log-likelihood
    lam  = max(lam,-real(min(eig(H)))*1.5);

    for subsubit=1:30,
        drawnow;

        % Levenberg-Marquardt update of log(nu)
        warning off
        nu        = exp(log(nu) - (H+lam*speye(size(H)))\g);
        warning on

        % Make sure the values are within a semi-reasonable range
        nu        = max(max(nu,1e-12),max(nu)*1e-9);
        nu        = min(min(nu,1e12) ,min(nu)*1e9);

        % Re-compute the log-likelihood
        Phi1      = feval(krn_f,nu,K,nz);
        warning off
        S1        = inv(Phi1'*Phi1*beta+diag(alpha(nz))); % Sigma
        warning on
        w1        = beta*S1*Phi1'*t; % weights
        ll        = ...
             -0.5*sum(log(alpha(nz)+eps))-0.5*N*log(beta+eps)-0.5*logdet(S1)...
             +0.5*(t-Phi1*w1)'*(t-Phi1*w1)*beta + 0.5*sum(w1.^2.*alpha(nz))...
             +0.5*(sum(nz)-N)*log(2*pi);

        % Terminate if no difference
        if abs(ll-oll)<1e-9,
            nu  = onu;
            ll  = oll;
            break;
        end;

        if ll>oll, % Solution gets worse
            lam = lam*10; % More regularisation required
            nu  = onu; % Discard new estimates of nu and ll
            ll  = oll;
        else % Solution improves
            lam = lam/10; % Try even less regularisation
            lam = max(lam,1e-12);

            % Use the new Phi, S and w for recomputing the
            % derivatives in the next iteration
            Phi = Phi1;
            S   = S1;
            w   = w1;
            break;
        end;
    end;
    if abs(ll-oll)<1e-9*N, break; end;
end;
%__________________________________________________________________________

%__________________________________________________________________________
function [ld,C] = logdet(A)
A  = (A+A')/2;
C  = chol(A);
d  = max(diag(C),eps);
ld = sum(2*log(d));

