wolffd@0: function [x, V, VV, loglik] = kalman_filter(y, A, C, Q, R, init_x, init_V, varargin)
wolffd@0: % Kalman filter.
wolffd@0: % [x, V, VV, loglik] = kalman_filter(y, A, C, Q, R, init_x, init_V, ...)
wolffd@0: %
wolffd@0: % INPUTS:
wolffd@0: % y(:,t)   - the observation at time t
wolffd@0: % A - the system matrix
wolffd@0: % C - the observation matrix 
wolffd@0: % Q - the system covariance 
wolffd@0: % R - the observation covariance
wolffd@0: % init_x - the initial state (column) vector 
wolffd@0: % init_V - the initial state covariance 
wolffd@0: %
wolffd@0: % OPTIONAL INPUTS (string/value pairs [default in brackets])
wolffd@0: % 'model' - model(t)=m means use params from model m at time t [ones(1,T) ]
wolffd@0: %     In this case, all the above matrices take an additional final dimension,
wolffd@0: %     i.e., A(:,:,m), C(:,:,m), Q(:,:,m), R(:,:,m).
wolffd@0: %     However, init_x and init_V are independent of model(1).
wolffd@0: % 'u'     - u(:,t) the control signal at time t [ [] ]
wolffd@0: % 'B'     - B(:,:,m) the input regression matrix for model m
wolffd@0: %
wolffd@0: % OUTPUTS (where X is the hidden state being estimated)
wolffd@0: % x(:,t) = E[X(:,t) | y(:,1:t)]
wolffd@0: % V(:,:,t) = Cov[X(:,t) | y(:,1:t)]
wolffd@0: % VV(:,:,t) = Cov[X(:,t), X(:,t-1) | y(:,1:t)] t >= 2
wolffd@0: % loglik = sum{t=1}^T log P(y(:,t))
wolffd@0: %
wolffd@0: % If an input signal is specified, we also condition on it:
wolffd@0: % e.g., x(:,t) = E[X(:,t) | y(:,1:t), u(:, 1:t)]
wolffd@0: % If a model sequence is specified, we also condition on it:
wolffd@0: % e.g., x(:,t) = E[X(:,t) | y(:,1:t), u(:, 1:t), m(1:t)]
wolffd@0: 
wolffd@0: [os T] = size(y);
wolffd@0: ss = size(A,1); % size of state space
wolffd@0: 
wolffd@0: % set default params
wolffd@0: model = ones(1,T);
wolffd@0: u = [];
wolffd@0: B = [];
wolffd@0: ndx = [];
wolffd@0: 
wolffd@0: args = varargin;
wolffd@0: nargs = length(args);
wolffd@0: for i=1:2:nargs
wolffd@0:   switch args{i}
wolffd@0:    case 'model', model = args{i+1};
wolffd@0:    case 'u', u = args{i+1};
wolffd@0:    case 'B', B = args{i+1};
wolffd@0:    case 'ndx', ndx = args{i+1};
wolffd@0:    otherwise, error(['unrecognized argument ' args{i}])
wolffd@0:   end
wolffd@0: end
wolffd@0: 
wolffd@0: x = zeros(ss, T);
wolffd@0: V = zeros(ss, ss, T);
wolffd@0: VV = zeros(ss, ss, T);
wolffd@0: 
wolffd@0: loglik = 0;
wolffd@0: for t=1:T
wolffd@0:   m = model(t);
wolffd@0:   if t==1
wolffd@0:     %prevx = init_x(:,m);
wolffd@0:     %prevV = init_V(:,:,m);
wolffd@0:     prevx = init_x;
wolffd@0:     prevV = init_V;
wolffd@0:     initial = 1;
wolffd@0:   else
wolffd@0:     prevx = x(:,t-1);
wolffd@0:     prevV = V(:,:,t-1);
wolffd@0:     initial = 0;
wolffd@0:   end
wolffd@0:   if isempty(u)
wolffd@0:     [x(:,t), V(:,:,t), LL, VV(:,:,t)] = ...
wolffd@0: 	kalman_update(A(:,:,m), C(:,:,m), Q(:,:,m), R(:,:,m), y(:,t), prevx, prevV, 'initial', initial);
wolffd@0:   else
wolffd@0:     if isempty(ndx)
wolffd@0:       [x(:,t), V(:,:,t), LL, VV(:,:,t)] = ...
wolffd@0: 	  kalman_update(A(:,:,m), C(:,:,m), Q(:,:,m), R(:,:,m), y(:,t), prevx, prevV, ... 
wolffd@0: 			'initial', initial, 'u', u(:,t), 'B', B(:,:,m));
wolffd@0:     else
wolffd@0:       i = ndx{t};
wolffd@0:       % copy over all elements; only some will get updated
wolffd@0:       x(:,t) = prevx;
wolffd@0:       prevP = inv(prevV);
wolffd@0:       prevPsmall = prevP(i,i);
wolffd@0:       prevVsmall = inv(prevPsmall);
wolffd@0:       [x(i,t), smallV, LL, VV(i,i,t)] = ...
wolffd@0: 	  kalman_update(A(i,i,m), C(:,i,m), Q(i,i,m), R(:,:,m), y(:,t), prevx(i), prevVsmall, ...
wolffd@0: 			'initial', initial, 'u', u(:,t), 'B', B(i,:,m));
wolffd@0:       smallP = inv(smallV);
wolffd@0:       prevP(i,i) = smallP;
wolffd@0:       V(:,:,t) = inv(prevP);
wolffd@0:     end    
wolffd@0:   end
wolffd@0:   loglik = loglik + LL;
wolffd@0: end
wolffd@0: 
wolffd@0: 
wolffd@0: 
wolffd@0: 
wolffd@0: