% Non-uniformly spaced 2d/3d FFT matrix: matrix object computations
% 
% Code was inspired by Jeff Fessler's library
% http://www.eecs.umich.edu/~fessler/irt/irt/nufft/
%
% (c) by Hannes Nickisch & Alexander Loktyushin, 
%                               MPI for Biological Cybernetics, 2011 February 06
%
%
%
% Initialize structure for NUFFT using KB interpolator,
% particularly the interpolation matrix in sparse format.
% caution: this routine can require a lot of memory!
% INPUT
%	k	      k-space coordinates [-1/2,1/2]; k-space is assumed to be cyclic i.e. 
%                             0.6 is equivalent to -0.4
%	imsz 		image size
% fftsz   Fourier grid size
% kbwsz   Kaiser-Bessel window size
% deriv   indicates whether we are interested in derivatives (applies to P, s)
%
% OUTPUT
%	f	[prod(imsz),1]  filter corresponding to the Fourier transforms of the
%	                  Kaiser-Bessel window function i.e. the spectrum is smoothed
%                   by a Kaiser-Bessel filter before linear interpolation is
%                   applied
%	P	[M,prod(fftsz)] sparse interpolation matrix
% s [M,1]           scaling vector to circumvent the circularity assumption
%
% Like fft(), the NUFFT expects the signals to be x(0,0), ...
%
% (c) by Jeff Fessler, The University of Michigan

function [f,P,s] = nufft_init(k, imsz, fftsz, kbwsz, deriv)
  ndims = numel(imsz); if nargin<5, deriv = 0; end
  
  %% 1) (real) Kaiser-Bessel filter f: outer product of vectors
  f = 1;
  for i=1:ndims
    [ff, alpha{i}] = nufft_scale(imsz(i), kbwsz(i), fftsz(i)); f = f(:)*ff';
  end
  f = reshape(f,imsz); clear ff               % f and alpha are independent of k
  
  %% 2) interpolation coefficient matrix P: huge sparse matrix
  for id=1:ndims
    N=imsz(id); J=kbwsz(id); K=fftsz(id); al=alpha{id}; L=length(al)-1;   % dims
    cssc = zeros(J,J);  % matrix T = [C' S S' C]\inv for NUFFT, independent of k
    [j1 j2] = ndgrid(1:J, 1:J);
    for l1 = -L:L
      alf1 = al(abs(l1)+1); if l1 < 0, alf1 = conj(alf1); end
      for l2 = -L:L
        alf2 = al(abs(l2)+1); if l2 < 0, alf2 = conj(alf2); end
        cssc = cssc + alf1 * conj(alf2) * nufft_sinc((j2-j1+l1-l2)/(K/N));
      end
    end
    T = inv(cssc); clear tmp cssc alf1 alf2 j1 j2 l1 l2, T = (T+T')/2; % tidy up

    % compute NUFFT r vector
    M = length(k(:,id));
    % offset for NUFFT
    if mod(J,2)	% odd J
      koff = round(K*k(:,id)) - (J+1)/2;                          % integer part
    else		    % even J
      koff = floor(K*k(:,id)) - J/2;                              % integer part
    end
    dk = K*k(:,id) - koff;                                     % fractional part
    arg = outer_sum(-(1:J)', dk');
    r = zeros(J,M); dr = zeros(J,M);
    for l1 = -L:L
       r1 =     nufft_sinc(  (arg+l1)*(N/K) );
      dr1 = N * nufft_dsinc( (arg+l1)*(N/K) );
      alf = al(abs(l1)+1);
      if l1<0, alf = conj(alf); end
       r =  r + alf *  r1;
      dr = dr + alf * dr1;
    end   
    const = 2i*pi*(N-1)/(2*K); ud{id} = exp(const*arg);
    if deriv==id
      ud{id} = K*const*ud{id}.*(T*r) + ud{id}.*(T*dr); % linear phase derivative
    else
      ud{id} = ud{id}.*(T*r);                                     % linear phase
    end
    
    % indices into oversampled FFT components; pre-convert indices into offsets
    kd{id} = mod(outer_sum((1:J)', koff'), K) + 1;
    if id>1, kd{id} = (kd{id}-1) * prod(fftsz(1:(id-1))); end
  end

  % build sparse matrix with entries per frequency point
  kk = kd{1};	uu = ud{1};
  for id = 2:ndims
    kk = block_outer_sum(kk, kd{id});	                    % outer sum of indices
    kk = reshape(kk, prod(kbwsz(1:id)), M);
    uu = block_outer_prod(uu, ud{id});	         % outer product of coefficients
    uu = reshape(uu, prod(kbwsz(1:id)), M);
  end
  mm = (1:M); mm = mm(ones(prod(kbwsz),1),:);                 % pre-do transpose
  P = sparse(mm(:), kk(:), conj(uu(:)), M, prod(fftsz));	       % sparse matrix

  %% 3) scaling matrix to circumvent circularity
  marg = 1/max(imsz);    % 1 pix margin 2 pix: k     in          [-1,1]/2 => s=1
  id1 = abs(k)<=1/2; id0 = abs(k)>1/2+marg;  % k not in (1+marg)*[-1,1]/2 => s=0
  sgn = @(x) 2*(x>0)-1;
  s = (-1/marg)*sgn(k).*(k-sgn(k).*(1/2+marg)); s(id0) = 0; s(id1) = 1;  % [0,1]
  if deriv>0
    s(:,deriv) = (-1/marg)*sign(k(:,deriv)).*(0<s(:,deriv) & s(:,deriv)<1);
  end
  s = prod(s,2);

% scaling factors for 1D NUFFT (from Fourier series coefficients)
function [d,al] = nufft_scale(N, J, K)
  % Compute scaling factors using the "do no harm" strategy.
  ssn = 1 ./ kaiser_bessel_ft( J, (1/2-N/2:N/2-1/2)'/K );
  % use regression to match NUFFT with best Kaiser scalings
  nn = -(N-1)/2:(N-1)/2;  L = min(ceil(N/3),13);
  X = cos(2*pi*nn'*(0:L)/K);
  al = (pinv(X)*ssn)';	al(2:end) = al(2:end)/2; % LS regress
  L = length(al)-1;
  d = exp( 1i*2*pi/K * nn'*(-L:L) ) * [fliplr(al(2:end)),al]'; d = real(d);

function y = kaiser_bessel_ft(J, x)        % Fourier transform; real in real out
  o2 = [1 1];                        % hardwired, because it was nearly the best
  zn = [2.5 2.27 2.31 2.34 2.32*o2 2.35 2.34*o2 2.35 2.34 2.35*[1 1 1] 2.33]';  
  alpha = J*zn(J==(2:16)');        % real positive scalar value independent of x
  z = sqrt( (2*pi*(J/2)*x).^2 - alpha^2 );             % z can be complex valued 
  y = real( sqrt(2*pi)*J/2 * besselj(1/2,z) ./ besseli(0,alpha) ./ sqrt(z) );

function y = nufft_sinc(x)                                          % sinc(pi*x)
  y = pi*x; y = sin(y)./y;  y(abs(x)<1e-9) = 1;
  
function y = nufft_dsinc(x)                                  % d sinc(pi*x) / dx
  y = pi*x; y = pi*(y.*cos(y)-sin(y)) ./ (y.*y);  y(abs(x)<1e-9) = 0;

function y = block_outer_sum(x1, x2)
  x1 = reshape(x1, size(x1,1), 1, []);
  x2 = reshape(x2, 1, size(x2,1), []);
  y = x1(:,ones(size(x2,2),1),:) + x2(ones(size(x1,1),1),:,:);

function y = block_outer_prod(x1, x2)
  x1 = reshape(x1, size(x1,1), 1, []);
  x2 = reshape(x2, 1, size(x2,1), []);
  y = x1(:,ones(size(x2,2),1),:) .* x2(ones(size(x1,1),1),:,:);

function s = outer_sum(x,y)
	x = x(:); y = y(:)'; s = x(:, ones(numel(y),1)) + y(ones(numel(x),1),:);
