%solves the forward Kalman filter efficiently for a large dendritic tree,
%using a low-rank approximation to the steady-state (zero-SNR) solution.

%this is a basic version of the code used to create the movies
%'low_rank_horiz' and 'low_rank_speckle'

%uses an implicit Euler implementation of the cable equation

%for details, see Paninski 2009, 'Fast Kalman filtering on quasilinear dendritic trees'

clear;

%params/setup
dt=.001; %sec
g=50; %leakiness
a=2e4; %intercompartmental coupling
T=20; %number of time steps
recalc_U_tol=1e-4; %tolerance for updating forward covariance
movie_flag=1; %flag for plotting movie
var_frac=.995; %defines how many singular values to keep
horiz_scan=1; %flag setting which sim to run: speckle or horiz
W=.005; %assume that observation noise covariance is proportional to the identity

disp('loading tree...');
load THINSTAR; %available at http://neuromorpho.org/neuroMorpho/neuron_info.jsp?neuron_name=THINSTAR
neuron=THINSTAR;
parents=neuron(:,7);
parents(1)=0; %soma = root of tree in this case
x=neuron(:,3)';
h=neuron(:,4)';
N=length(parents); %number of compartments

%setup for plotting below
f = find(parents~=0);
pp = parents(f);
X = [x(f); x(pp); repmat(NaN,size(f'))];
H = [h(f); h(pp); repmat(NaN,size(f'))];
X = X(:);
H = H(:);
col=jet(100); %set colormap

%create dynamics matrix (implict Euler implementation of cable equation)
I=speye(N); e = ones(N,1);
K=dt*(-g*I);
for(i=1:N)
	j=parents(i);
	if(j~=0)
		adt=a*dt;
		K(i,j)=adt; K(j,i)=adt; K(i,i)=K(i,i)-adt; K(j,j)=K(j,j)-adt;
	end;
end;
mu=zeros(N,1);
IK2=(I-K)^2;

disp('making data...')
V=zeros(N,1);
Vs=zeros(N,T); %true V matrix
if(~horiz_scan)
	B=zeros(N); %build observation matrix B (note: B is transposed relative to the paper)
	sb2=100; %make this many observations
	for(j=1:sb2)
		B(ceil(N*j/sb2),j)=1;
		if(size(B,1)>N) error('error specifying B'); end;
	end;
	B=B(:,1:j);
else
	B=zeros(N,250); sb2=size(B,2);
	[junk,xx]=hist(x,sb2);
	dx=diff(xx); dx=dx(1)/2;
	ffx=zeros(sb2,1);
	for(i=1:sb2)
		fx=find(abs(x-xx(i))<dx);
		if(~isempty(fx)) B(fx,i)=1; ffx(i)=1; end;
	end;
	B=B(:,find(ffx));
end;
%note that the preceding bit of code is not particularly optimized; it
%would be preferable to represent B as a sparse matrix, for example.
Y=zeros(size(B,2),T); %observations matrix
inj_site=1;
cholIK=chol(I-K); %dividing twice by cholesky factor is faster than dividing by the original matrix. if we do this repeatedly
for(i=1:T)
	if(~mod(i,100)) disp(sprintf('timestep = %i',i)); end;
	V(inj_site)=V(inj_site)+dt*4500*cos(1*2*pi*i*dt); %inject a sinusoid into a highly-connected node
	V=cholIK\(cholIK'\V)+sqrt(dt)*randn(N,1); %noisy dynamics
	y=B'*V+sqrtm(W)*randn(size(B,2),1); %make observations
	Vs(:,i)=V;
	Y(:,i)=y';
end;
%grab some limits for plotting later
if(horiz_scan)
	max_V=max(Vs(:));
	min_V=min(Vs(:));
else
	max_V=max(max(Vs(:)),max(Y(:)));
	min_V=min(min(Vs(:)),min(Y(:)));
end;
Wmat=W*speye(size(B,2));

%do some plotting
figure(3252);
subplot(311);
imagesc(Vs); %caxis([min_V max_V]);
colorbar
set(gca,'XTickLabel',[]);
ylabel('true V');
subplot(312);
imagesc(Y); colorbar
set(gca,'XTickLabel',[]);
ylabel('obs Y');
subplot(313);
cla;

%----------------main loop-----------------
mus=zeros(N,T);
disp('running forward Kalman recursion...')
recalc_U=1; %flag for updating forward covariance matrix
for(i=1:T)
	if(~mod(i,20)) disp(sprintf('timestep = %i',i)); end;
	if(recalc_U)
		%update forward cov matrix
		if(i>1)
			AU_old=cholIK\(cholIK'\U);
			R=AU_old/dt-IK2\AU_old/dt;
			Q=inv(inv(D)+AU_old'*R);
			O=orth([R B]);
		else %handle i=1 case: U,D=0
			AU_old=0; D=0;
			O=orth(B);
			Q=0; R=0;
		end;
		OR=O'*R; OB=O'*B;
		M=-OR*Q*OR'+OB*inv(W)*OB';
		CO=(IK2-I)\(IK2*O)*dt;
		[U_full,D_temp,junk]=svd(CO*(inv(M)+O'*CO)^-.5,'econ');
		d=diag(D_temp).^2;

		if(1) %choose number of singular vals adaptively
			n=min(find(cumsum(d)/sum(d)>=var_frac)); %grab enough singular vals to cover var_frac of variance
		else
			n_svs=10; n=min(size(U,2),n_svs); %choose number of singular vals non-adaptively, with a hard ceiling
		end;

		D_old=D;
		U=U_full(:,1:n);
		D=-spdiags(d(1:n),0,n,n);
		disp(sprintf('number of retained singular values = %i',n));

		if(i>1)
			nd=1:min(n,length(d_old));
			if(norm(-d(nd)-d_old(nd))<recalc_U_tol) %if U (actually, D) has converged, stop updating it
				recalc_U=0;
				disp('forward covariance has converged; now only updating forward mean.');
			end;
		end;
		U_old=U; d_old=-d(1:n);
		%compute reduction in variance at each node
		%(conditional variance minus prior variance = -diag(U*D*U'))
		USD=U*spdiags(sqrt(-diag(D)),0,n,n); var_diff=-sum(USD.^2,2);
		
		%compute some matrices required to update the forward mean vector, just below 
		PB=(IK2-I)\(IK2*B)*dt+AU_old*(D_old*(AU_old'*B));
		IWBPB=inv(Wmat+B'*PB);
	end;

	%update forward mean vector
	%Amu=A_imp*mu;
	Amu=cholIK\(cholIK'\mu);
	y=Y(:,i);
	z=PB*(inv(W)*(B'*mu-y));
	mu=real(Amu+PB*(IWBPB*(B'*(mu-Amu+z)))-z);
	mus(:,i)=mu;
end;
%-----------------end of main loop---------------------

%do some more plotting
figure(3252);
subplot(313);
imagesc(mus); %caxis([min_V max_V]);
colorbar
ylabel('est V');
xlabel('time (ms)')
max_mus=max(mus(:));
min_mus=min(mus(:));


if(movie_flag)
	disp('displaying movie...');
	if(horiz_scan)
		%movie=avifile('low_rank_horiz.avi');
	else
		%movie=avifile('low_rank_speckle.avi');
		fB=find(sum(B'));
	end;
	first_frame=1;
	acc=.99;
	figure(232);
	clf;
	set(gcf,'Position',[260   268   480*2   192*2]);
	
	for(i=1:T)
		if(~mod(i,2))
			figure(232);

			subplot(131);
			if(first_frame) %set up the plot here, then just update colors on subsequent iterations
				hold off;
				%plot (X, H, 'k-','EraseMode','none');
				hold on;
				sub1_data=plot([x;x],[h;h],'.');
				title('true V')
				axis tight;
				set(gca,'YTick',[],'XTick',[]);
			end;
			ylabel(sprintf('timestep = %i',i))
			cc=(Vs(:,i)-min_V)/(max_V-min_V); cc=cc*acc+.5*(1-acc);
			set(sub1_data,{'Color'},mat2cell(col(ceil(cc*100),:),ones(1,length(cc)),3));

			subplot(132);
			if(horiz_scan)
				if(first_frame)
					hold off;
					im_data=imagesc(Y(:,i)'); caxis([min(Y(:)) max(Y(:))]);
					set(gca,'YTick',[]);
					xlabel('slice');
					title('obs Y')
				else
					set(im_data,'CData',Y(:,i)');
				end;
			else
				if(first_frame)
					hold off;
					plot (X, H, 'k-','EraseMode','none');
					hold on;
					sub2_data=plot([x(fB);x(fB)],[h(fB);h(fB)],'.');
					title('obs Y')
					axis tight;
					set(gca,'YTick',[],'XTick',[]);
				end;
				cc=(Y(:,i)-min_V)/(max_V-min_V); cc=cc*acc+.5*(1-acc);
				set(sub2_data,{'Color'},mat2cell(col(ceil(cc*100),:),ones(1,length(cc)),3));
			end;

			subplot(133);
			if(first_frame)
				hold off;
				%plot (X, H, 'k-','EraseMode','none');
				hold on;
				sub3_data=plot([x;x],[h;h],'.');
				title('est V')
				axis tight;
				set(gca,'YTick',[],'XTick',[]);
			end;
			cc=(mus(:,i)-min_mus)/(max_mus-min_mus); cc=cc*acc+.5*(1-acc);
			ceil_cc=ceil(cc*length(col));
			ceil_cc=min(ceil_cc,length(col)); ceil_cc=max(ceil_cc,1);
			set(sub3_data,{'Color'},mat2cell(col(ceil_cc,:),ones(1,length(cc)),3));

			first_frame=0;
			pause(.01);
		end;
	end;
end;

disp('done.');