% Nonlinear REML - with smoothing
% -------------------------------
% The objective is to estimate the hyper-parameters used to
% simulate the data:
w = [1 5e-6];
m = 256;
n = ceil(m/1);

B    = dctmtx(m)'; B = B(:,1:n); % Basis functions
XX2  = diag((1:n).^3.0); % Form of priors
x    = (1:m)';

% Simulated images
% ----------------
f         = zeros(m,1);
f(60:200) = exp(randn(length(60:200),1))*10;
f         = conv2(f,exp(-0.5 * (-40:40)'.^2/50)/sqrt(2*pi*50),'same');

p = (XX2*w(2))^(-0.5)*randn(n,1);
t = x+B*p;
g = interp1(x,f,t,'*linear'); g(~finite(g))=0;

% Add noise and smooth
% --------------------
g  = g +randn(size(t))*w(1)^(-0.5);
f  = f +randn(size(t))*w(1)^(-0.5);
vr = 3;
kk = exp(-0.5/vr*[0:(m-1)].^2)/sqrt(2*pi*vr);
K  = toeplitz(kk,kk);   
g  = K*g;
f  = K*f;

n = 64;
B    = dctmtx(m)'; B = B(:,1:n); % Basis functions
XX2  = diag((1:n).^3.0); % Form of priors
x    = (1:m)';

clf;
% Starting estimates
% ------------------
beta = zeros(n,1);
w    = [0.01 0.0001];
rml  = [];
for it=1:20,
	fprintf('*** Iteration %d (w = [%g %g]) ***\n', it, w);
	% E-Step: obtain MAP solution using current hyper-parameters
	% ----------------------------------------------------------
	f1    = interp1(x,f,x+B*beta,'*linear'); f1(~finite(f1))=0;
	df1   = interp1(x,gradient(f),x+B*beta,'*linear'); df1(~finite(df1))=0;

	Y1    = f1 - g;
	r1    = Y1'*Y1;
	r2    = beta'*XX2*beta;
	r     = r1*w(1)+r2*w(2);

	for i=1:20,
		obeta = beta; oY1 = Y1;
		or1   = r1;   or2 = r2;or = r;

		X1    = repmat(df1,[1 size(B,2)]).*B;
		XX1   = X1'*X1;
		XY1   = X1'*Y1;
		T     = inv(w(1)*XX1 + w(2)*XX2);
		beta  = T* (w(1)*XX1*beta - w(1)*XY1);
		f1    = interp1(x,f,x+B*beta,'*linear'); f1(~finite(f1))=0;
		df1   = interp1(x,gradient(f),x+B*beta,'*linear'); df1(~finite(df1))=0;
		Y1    = f1 - g;
		r1    = Y1'*Y1;
		r2    = beta'*XX2*beta;
		r     = r1*w(1)+r2*w(2);
		fprintf('%g\t%g\t| %g\n', r1*w(1),r2*w(2),r1*w(1)+r2*w(2));
		subplot(1,2,1);plot(x,g,'r',x,f1,'b',x,f,'b:');
		title('"Images"');
		subplot(1,2,2);plot(x,x+B*beta,'b',x,t,'r');axis([1 m 1 m]);
		title('"Deformations"');
		drawnow;

		if r>or,
			beta = obeta; Y1 = oY1;
			r1   = or1;   r2 = or2; r = or;
		end;
		if (or-r)/(r+or)<0.00001, break; end;
	end;

	% M-Step: Re-estimate hyper-parameters
	% ----------------------------------------------------------
	% Smoothness estimate
	vr   = sum(Y1.^2)./sum(diff(Y1).^2)/2;
	kk   = exp(-0.5/vr*[0:(m-1)].^2)/sqrt(2*pi*vr);
	K    = toeplitz(kk,kk);     
	V    = K*K';
	nu1  = trace(V) - sum(sum(T.*(X1'*V*X1)))*w(1);
	w(1) = nu1/(sum((K*Y1).^2)+eps);

	w(2) = (n-trace(T*XX2)*w(2))/(r2+eps);
	drawnow;
end;

