sr_fit.m

function [out,in] = sr_fit(in,opt)
% Super-resolution/denoising of multichannel images.
%
%   . This model assumes that data are generated by projecting noise-free
%     high-res images to a low-resolution space and adding Gaussian noise.
%   . A multi-channel total-variation (MTV) prior is set over the unknown
%     high-res images and a maximum a posteriori (MAP) or an approximate
%     Gaussian posterior (mean + variance) is estimated.
%   . The non-differentiable L1 problem is solved using a reweighted
%     least-squares (RLS) approach.
%   . Additionally, the unknown images can evenutally be encoded in
%     log-space, such that the forward model is 'x = A*exp(y) + eps' 
%     instead of 'x = A*y + eps'. This ensures that the reconstructed 
%     images are positive and that the prior is scale-invariant. The
%     problem becomes non-convex but still possesses a unique optimum.
%
% FORMAT out = sr_fit(in,opt)
%
% in   - Input volumes.
% opt  - Structure of parameters.
% out  - Output data structure.
%
% Input volumes can be:
% - (cell|array) of (cell|array) of (string|char|nifti):
%       . Each filename corresponds to one 3D volume.
%       . Outer loop: contasts | Inner loop: repeats
%       . 4D volumes will be treated as series of 3D volumes.
%       . Specific volumes from a 4D series can be specified using SPM's
%         notation (e.g., 'path/to/file.nii,2' for the 2nd volume)
% - (numeric|file) array:
%       . For denoising only.
%       . All inputs must have the same lattice and are assumed to have the
%         same orientation matrix. This orientation matrix can be specified
%         in the option structure.
%__________________________________________________________________________
%
% The use of Multi-Channel Total Variation in the context of MRI
% super-resolution is described in:
%
% . "MRI Super-Resolution using Multi-Channel Total Variation"
%   Brudfors, M., Balbastre, Y., Nachev, P., and Ashburner, J.
%   MIUA 2018
%   https://arxiv.org/abs/1810.03422
%
% . "A Tool for Super-Resolving Multimodal Clinical MRI"
%   Brudfors, M., Balbastre, Y., Nachev, P., and Ashburner, J.
%   Preprint
%   https://arxiv.org/abs/1909.01140
%__________________________________________________________________________

% -------------------------------------------------------------------------
%
%                           Initialisation
% 
% -------------------------------------------------------------------------


% -------------------------------------------------------------------------
% Add subdirectory to path
addpath(fullfile(fileparts(which('sr_fit')), 'sub'));
try
    spm_dir = spm('dir');
catch
    error('SPM12 is not on your path.');
end
addpath(fullfile(spm_dir, 'toolbox', 'Longitudinal'));

% -------------------------------------------------------------------------
% Options
if nargin < 2, opt = struct; end
opt = sr_opt_defaults(opt);

% -------------------------------------------------------------------------
% GUI
if opt.verbose > 1
    spm_figure('GetWin','Interactive');
end
if nargin == 0 || isempty(in)
    row = 1;
    Nc   = spm_input('Number of contrasts',row,'n','',1); row = row+1;
    opt.mode = find(strcmpi(["denoise";"superres"], opt.mode));
    opt.mode = char(spm_input('Mode',row,'b','denoising|super-resolution',{'denoise';'superres'},opt.mode)); row = row+1;
    if strcmpi(opt.mode, 'superres')
        opt.vs = opt.vs(1);
        opt.vs = spm_input('Voxel size',row,'r',opt.vs,1); row = row+1;
        opt.vs = opt.vs * ones(1,3);
        opt.slice.thickest = find(strcmpi([0 2], opt.slice.thickest));
        opt.slice.thickest = spm_input('Profile (slice)',row,'b','rect|gauss',[0 2],opt.slice.thickest); row = row+1;
        opt.slice.other    = find(strcmpi([0 2], opt.slice.other));
        opt.slice.other    = spm_input('Profile (in plane)',row,'b','rect|gauss',[0 2],opt.slice.other); row = row+1;
        opt.slice.gap      = spm_input('Slice gap',row,'r',opt.slice.gap,1); row = row+1;
        opt.slice.accumulate = spm_input('Model',row,'b','integrate|average',[true;false],1+(~opt.slice.accumulate)); row = row+1;
    end
    opt.coreg.do  = spm_input('Co-register',row,'b','yes|no',[true;false],1+(~opt.coreg.do)); row = row+1;
    opt.reg.value = spm_input('Regularisation',row,'r',opt.reg.value,1); row = row+1;
    opt.itermax   = spm_input('Iterations',row,'r',opt.itermax,1); row = row+1;
    in = cell(1,Nc);
    for c=1:Nc
        msg   = sprintf('Select files for contrast %d...',c);
        in{c} = num2cell(spm_select(Inf, 'image', msg), 2)';
    end
end
if nargout == 0 && isempty(opt.out.folder)
    opt.out.folder = spm_select(1, 'dir', 'Select output directory...');
end
    
% -------------------------------------------------------------------------
% Multithread SPM
[threads0.spm,threads0.matlab] = sr_threads();
sr_threads(opt.threads, 'both');

% -------------------------------------------------------------------------
% Read and format input volumes
if opt.verbose > 0, fprintf('Read input volumes\n'); end
in = sr_in_format(in,opt);

% > Here, dat is a cell of cell of struct with fields:
% . dat - observed data (numeric or file array)
% . mat - orientation matrix
% . dim - lattice dimensions
% . var - observation uncertainty (numeric or file array, usually a scalar)
% . lam - noise precision
% . mu  - mean tissue intensity

% -------------------------------------------------------------------------
% Co-registration
if opt.coreg.do
    if opt.verbose > 0, fprintf('Coregister volumes\n'); end
    in = sr_in_coregister(in, opt.coreg.fwhm);
end

% -------------------------------------------------------------------------
% Create output structure
[dim,mat,vs] = compute_mean_space(in,opt);
if opt.verbose > 0, fprintf('Prepare output structure\n'); end
out = sr_out_format(dim, mat, in, opt);

% Here, out is a structure with fields:
% . dat   - denoised/upsampled data, as a 4D (numeric|file) array
% . mat   - orientation matrix
% . dim   - lattice dimensions
% . lam   - regularisation value
% . rls   - image of weights (RLS)

% -------------------------------------------------------------------------
% Precompute AtA(1). Faster, but uses more memory
Ha = sr_hessian_precomp(in, out.dim, out.mat, opt);

% -------------------------------------------------------------------------
%
%                         Iterative optimisation
% 
% -------------------------------------------------------------------------

if opt.verbose > 1, sr_plot_progress(in, out,[],opt); end
ll  = [];
K   = size(out.dat, 4);
vol = prod(vs);
for it=1:opt.itermax

    if opt.verbose > 0, fprintf('Iteration %i\n', it); end

    if it <= numel(opt.armijo), armijo = 1/opt.armijo(it); end
    
    % ---------------------------------------------------------------------
    % Update maps
    % ---------------------------------------------------------------------
    g    = zeros([dim K], 'single');           % Gradient
    H    = zeros([dim K], 'single');           % Hessian
    llx  = 0;                                  % Log-likelihood: data term
    lly  = 0;                                  % Log-likelihood: prior term
    lltv = 0;                                  % True MTV term
    
    % ---------------------------------------------------------------------
    % Gradient: data term
    % ---------------------------------------------------------------------
    if opt.verbose > 0, fprintf('Gradient: data '); end
    for c=1:numel(in) % Loop over contrasts
        
        % - Compute gradient
        [llx1,g1,H1] = sr_gradient(c, in{c}, out, opt, Ha{c});
        llx          = llx + llx1;

        % - Add to full gradient
        g(:,:,:,c) = g(:,:,:,c) + g1; clear g1
        H(:,:,:,c) = H(:,:,:,c) + H1; clear H1

    end
    if opt.verbose > 0, fprintf('\n'); end
    
    % ---------------------------------------------------------------------
    % Gradient: Membrane
    % ---------------------------------------------------------------------
    if opt.reg.mode > 0
        if opt.verbose > 0, fprintf('Gradient: membrane '); end
        switch opt.reg.mode
            % -------------------------------------------------------------
            % L1 regularisation
            case 1
                w    = single(out.rls());
                lly  = 0.5 * vol * sum(w(:), 'double');
                w    = 1./w;
                wnew = 0;
                for k=1:K
                    if opt.verbose > 0, fprintf('.'); end
                    y          = single(out.dat(:,:,:,k));
                    [Ly,Dy]    = sr_vel2mom_l1(y, out.lam(k), vs, w);
                    Dy         = out.lam(k) * sum(sum(Dy.^2,5),4);
                    wnew       = wnew + Dy;
                    lly        = lly + 0.5 * vol * sum(y(:).*Ly(:), 'double');
                    g(:,:,:,k) = g(:,:,:,k) + vol * Ly;
                    clear Ly
                end
                lltv = vol * sum(sqrt(wnew(:)), 'double');
                wnew = sqrt(wnew + opt.reg.smo);
                out.rls(:,:,:) = wnew; clear wnew
            
            % -------------------------------------------------------------
            % L2 regularisation
            case 2
                for k=1:K
                    if opt.verbose > 0, fprintf('.'); end
                    y          = single(out.dat(:,:,:,k));
                    Ly         = sr_vel2mom_l2(y, out.lam(k), vs);
                    lly        = lly + 0.5 * vol * sum(y(:).*Ly(:), 'double');
                    g(:,:,:,k) = g(:,:,:,k) + vol * Ly;
                    clear y Ly
                end
        end
        if opt.verbose > 0, fprintf('\n'); end
    end
    
    % ---------------------------------------------------------------------
    % Gauss-Newton
    % ---------------------------------------------------------------------
    switch opt.reg.mode
        % -----------------------------------------------------------------
        % No regularisation
        case 0
            dy = sr_solve_l0(H, g);
            
        % -----------------------------------------------------------------
        % L1 regularisation
        case 1
            if sum(opt.solver.fmg) > 0
                % - Initialise using a majoriser of the true Hessian and  
                %   the full multigrid solver.
                %   It captures low frequencies and provides a robust 
                %   starting estimate for the conjugate gradient solver.
                optsolver = struct;
                optsolver.nbcycle = opt.solver.fmg(1);
                optsolver.nbiter  = opt.solver.fmg(2);
                dy = sr_solve_l1_fmg(H, g, w, vol * out.lam, vs, optsolver);
            end
            if any(w(:)~=1) || sum(opt.solver.fmg) > 0
                if opt.solver.cg > 0
                    % - Continue with conjugate gradient.
                    %   It is fast and generic.
                    optsolver           = struct;
                    optsolver.verbose   = opt.verbose;
                    optsolver.precond   = false;
                    optsolver.tolerance = 0;
                    optsolver.nbiter    = opt.solver.cg;
                    dy = sr_solve_l1_cg(H, g, w, vol * out.lam, vs, optsolver, dy);
                end
                if opt.solver.relax > 0
                    % - Refine with relax.
                    %   It is slow to get the smooth part correct but
                    %   efficient with crisp edges.
                    optsolver           = struct;
                    optsolver.verbose   = opt.verbose;
                    optsolver.tolerance = 0.01;
                    optsolver.nbiter    = opt.solver.relax;
                    dy = sr_solve_l1_relax(H, g, w, vol * out.lam, vs, optsolver, dy);
                end
            end
        % -----------------------------------------------------------------
        % L2 regularisation
        case 2
            dy = sr_solve_l2(H, g, vol * out.lam, vs);
    end
    clear H g w
    
    % ---------------------------------------------------------------------
    % Update
    if opt.verbose > 0, fprintf('Update\n'); end
    for k=1:K
        out.dat(:,:,:,k) = out.dat(:,:,:,k) - armijo * dy(:,:,:,k);
    end
    clear dy
    
    % ---------------------------------------------------------------------
    % Log-likelihood + Gain
    % ---------------------------------------------------------------------
    ll = [ll, [llx+lly;llx;lly;lltv]];
    if size(ll,2) > 1
        gain = abs((ll(1,end-1) - ll(1,end)) / (max(ll(1,:)) - ll(1,end-1)));
    else
        gain = Inf;
    end
    
    % ---------------------------------------------------------------------
    % Plot
    % ---------------------------------------------------------------------
    if opt.verbose > 0, fprintf('%s\n', repmat('-',[1 80])); end
    if opt.verbose > 0, fprintf('ll = %7.3g | llx = %7.3g | lly = %7.3g | gain = %7.3g\n', llx+lly, llx, lly, gain); end
    if opt.verbose > 0, fprintf('%s\n', repmat('-',[1 80])); end

    % ---------------------------------------------------------------------
    % Out?
    % ---------------------------------------------------------------------
    if abs(gain) < opt.tolerance
        if opt.verbose > 1, sr_plot_progress(in, out, ll, opt); end
        break
    end
    if opt.verbose > 1, sr_plot_progress(in, out, ll, opt); end

end

out = sr_out_write(out, opt);

sr_threads(threads0.spm,    'spm');
sr_threads(threads0.matlab, 'matlab');

% =========================================================================
function [alldim,allmat] = get_orientation_matrices(in)
num = zeros(1,numel(in));
for c=1:numel(in), num(c) = numel(in{c}); end
allmat = zeros(4,4,sum(num));
alldim = zeros(3,sum(num));
i = 0;
for c=1:numel(in)
    for n=1:numel(in{c})
        i = i + 1;
        allmat(:,:,i) = in{c}{n}.mat;
        alldim(:,i)   = in{c}{n}.dim(:);
    end
end
% =========================================================================
function [dim,mat,vs] = compute_mean_space(in,opt)

[alldim,allmat] = get_orientation_matrices(in);

if ~isfinite(opt.fov)
    allsame = true;
    for i=1:3
        allsame = allsame && (numel(unique(alldim(i,:))) == 1);
        if ~allsame, break; end
    end
    for i=1:4
        for j=1:4
            allsame = allsame && (numel(unique(allmat(i,j,:))) == 1);
        if ~allsame, break; end
        end
    end
    if allsame
        fov = 1;
    else
        fov = 0;
    end
else
    fov = opt.fov;
end

if fov == 0
    % Estimate a mean orientation matrix
    if opt.verbose > 0, fprintf('Estimate mean space\n'); end
    [dim,mat,vs] = sr_mean_space(allmat, alldim, opt.vs);
else
    % Use orientation matrix of the n-th volume
    mat = in{opt.fov}{1}.mat;
    dim = in{opt.fov}{1}.dim;
    vs0 = opt.vs;
    vs  = sqrt(sum(mat(1:3,1:3).^2));
    vs0(~isfinite(vs0)) = vs(~isfinite(vs0));
    scl  = vs0(:)./vs(:);
    mat = mat * diag([scl(:); 1]);
    dim = ceil(dim(:)./scl(:))';
    vs  = sqrt(sum(mat(1:3,1:3).^2));
end