(* :Title: Particle Filter *)

(* :Author: Mark Fisher *)

(* :Context: ParticleFilter` *)

(* :Mathematica Version: 5.0 *)

(* :History:
	Version 0.0 April 2005
*)

(* :Summary:
	SIR (Sampling Importance Resampling) particle filter and
	backward simulation particle smoother
*)

(* :Discussion:
	This package implements ParticleFilter (an SIR particle filter)
	and ParticleSmoother (a backward simulation particle smoother that
	takes the output from ParticleFilter as input). In particular, given
		{loglikelihood, states} =
			ParticleFilter[data, {x0, w0}, f, gp, IncludeStates -> True];
	the marginal loglikelihood of the parameters (implicit in f and gp) is
	given by loglikelihood. To compute n draws of paths from the smoothed
	distribution, use
		collected = CollectStates[states];
		smoothed = ParticleSmoother[collected, fp, n];

*)

(* :References:
	Doucet, de Freitas, and Gordon (editors),
		"Sequenctial Monte Carlo Methods in Practice",
		Springer, 2001.

	Arulampalam, Maskell, Gordon, and Clapp,
		"A tutorial on particle filters for online nonlinear/non-Gaussian
		Bayesian tracking",
		IEEE Transactions on Signal Processing, vol. 50, no. 2, February 2002.

	Godsill, Doucet, and West,
		"Monte Carlo smoothing for nonlinear time series",
		Journal of the American Statistical Association, vol. 99, no. 465,
		pp. 156-168, March 2004.
*)

(* :Example:
(*
This is the now classic nonlinear example.
The functions are compiled for speed.
The parameters are accessible via the functions.
*)

Clear[signal, data, f, fCompiled, gp, gpCompiled, fp, fpCompiled,
	swarmsize, prior, loglikelihood, states, collected, smoothed]

signal = Rest @ FoldList[
	{.5#1[[1]] + 25#1[[1]]/(1 + #1[[1]]^2) + 8 Cos[1.2#1[[2]]] + #2, #1[[2]] + 1} &,
	{Random[NormalDistribution[0, Sqrt[10]]], 0},
	RandomArray[NormalDistribution[0, Sqrt[10]], 100]];

data = Partition[signal[[All, 1]]^2/20, 1] +
	RandomArray[NormalDistribution[0, 1], 100];

f[x_, {a_, b_, c_, s2v_}] :=
	fCompiled[
		x[[All, 1]],
		x[[1, 2]],
        RandomArray[NormalDistribution[0, Sqrt[s2v]], Length[x]],
        a, b, c
        ]

fCompiled = Compile[
	{{x, _Real, 1}, {t, _Real, 0}, {v, _Real, 1},
		{a, _Real, 0}, {b, _Real, 0}, {c, _Real, 0}},
    MapThread[{a #1 + b #1/(1 + #1^2) + c*Cos[1.2t] + #2, t + 1} &,
    {x, v}]];

gp[x_, y_, {d_, s2w_}] := gpCompiled[x[[All, 1]], y[[1]], d, s2w]

gpCompiled = Compile[
	{{x, _Real, 1}, {y, _Real, 0}, {d, _Real, 0}, {s2w, _Real, 0}},
	E^-((y - d*#1^2)^2/(2*s2w^2))/(Sqrt[2*Pi]*s2w) & /@ x
	];

fp[x1_, x0_, {a_, b_, c_, s2v_}] :=
	fpCompiled[x1[[1]], x0[[All, 1]], x0[[1, 2]], a, b, c, s2v]

fpCompiled = Compile[
	{{x1, _Real, 0}, {x0, _Real, 1}, {t, _Real, 0},
		{a, _Real, 0}, {b, _Real, 0}, {c, _Real, 0}, {s2v, _Real, 0}},
	1/(E^((x1 - c*Cos[1.2*t] - a*#1 - (b*#1)/(1 + #1^2))^2/(2*s2v^2))*s2v) & /@ x0
	];

swarmsize = 10^4;
prior = Thread[{RandomArray[NormalDistribution[0, Sqrt[10]], swarmsize], 0}];

{loglikelihood, states} =
	ParticleFilter[data, prior, f[#, {.5, 25, 8, 10}]&, gp[##, {.05, 1}]&,
		IncludeStates -> True];

collected = CollectStates[states];

smoothed = ParticleSmoother[collected, fp[##, {.5, 25, 8, 10}] &, 10];

*)

BeginPackage["ParticleFilter`"]

ParticleFilter::usage = "ParticleFilter[data, x0, f, gp] takes a list of
data vectors, a list of state vectors x0 that represent the prior
information (the particle swarm), a transition function f that takes the
list state vectors as an argument and returns the list of predictions for
the new state vectors, and a function gp that takes two arguments, the list
of state vectors and the current a single data point, that returns
likelihood of the data given the state from the measurement equation.
ParticleFilter returns the log of the marginal likelihood of the (implicit)
parameters. By setting the option IncludeStates -> True, ParticleFilter
also returns the list of states, suitable as input to ParticleSmoother.
ParticleFilter also takes the option MCMCSmooth."

IncludeStates::usage = "IncludeStates is an option for ParticleFilter,
which specifies whether to include the filtered particle swarms in
the output. The default setting is IncludeStates -> False."

ParticleSmoother::usage = "ParticleSmoother[states, fp, n] takes the list
of states produced by ParticleFilter and a function fp that takes two
arguments (the first of which is a single current state vector and the
second of which is a list of previous state vectors) and returns a list of
likelihoods of the previous states given the value of the current state.
ParticleSmoother returns n paths from the joint smoothed distribution of
the state vectors."

FilterStep::usage = "FilterStep[{x0, loglike0}, y, f, gp] is called by
ParticleFilter and returns {x1, loglike1}."

SmootherStep::usage = "SmootherStep[x1, {x0, w0}, fp] is called by
ParticleSmoother. SmootherStep returns a list of previous states sampled
from x0 corresponding to the list of current states x1."

Resample::usage = "Resample[{x, w}, n] returns a list of n samples drawn
with replacement from x with probabilities proportional to the nonnegative
weights w (made with SystematicResampling). Resample[{x, w}] returns a
single draw (made with BinarySerach). Resample is called by ParticleFilter
and ParticleSmoother."

SystematicResampling::usage = "SystematicResampling[weights, n] returns a
list of n positions computed from the nonnegative weights according to
systematic resampling. SystematicResampling is called by FilterStep (within
ParticleFilter)."

BinarySearch::usage = "BinarySearch[weights] returns a single position
computed from the nonnegative weights using binary search. BinarySearch is
called by SmootherStep (within ParticleSmoother)."

CollectStates::usage = "CollectStates is an option for ParticleSmoother
that specifies whether to run the function CollectStates on the input.
CollectStates[states] take a list of states (as repreented by particle
swarms such as can be produced by ParticleFilter) and returns a list of the
union of the particles paired with their weights. The output is suitable as
input for ParticleSmoother."

MCMCSmooth::usage = "MCMCSmooth is an option for ParticleFilter that
specifies whether to run the function MCMCSmooth at the end of FilterStep
after resampling. The default setting is MCMCSmooth -> False.
MCMCSmooth[{x1, w}, y, x0, f, gp] compares the likelihood each element of
x1 with a draw from f[x0] and replaces it with the draw according to the
ususal Metropolis criterion. The purpose is to obtain more distinct values
while maintaining the same distribution."

PriorToPosterior::usage = "PriorToPosterior[prior, likelihood] takes a list
of points that represent a prior distribution and a likelihood function and
returns a list of points that represents the posterior distribution."

Begin["`Private`"]

PriorToPosterior[prior_, Lfun_] :=
	prior[[ SystematicResampling[Lfun /@ prior, Length[prior]] ]]

(***** filtering *****)

Options[ParticleFilter] = {IncludeStates -> False, MCMCSmooth -> False}

ParticleFilter[data_, x0_, f_, gp_, opts___?OptionQ] :=
	If[TrueQ[IncludeStates /. {opts} /. Options[ParticleFilter]],
		(* then *)
		{#[[-1, -1]], #[[All, 1]]}& @
			Rest[FoldList[FilterStep[##, f, gp, opts]&, {x0, 0}, data]],
		(* else *)
			Fold[FilterStep[##, f, gp, opts]&, {x0, 0}, data][[2]]
		]

FilterStep[{x0_, loglike_}, y_, f_, gp_, opts___] :=
	Module[{x1, w, pos},
	x1 = f[x0];
	w = gp[x1, y];
	pos = SystematicResampling[w, Length[w]];
	{If[TrueQ[MCMCSmooth /. {opts} /. Options[ParticleFilter]],
		(* then *)
(*		MCMCSmooth[{x1[[ pos ]], w[[ pos ]]}, y, x0, f, gp], *)
		MCMCChoose[x1[[ pos ]], w[[ pos ]], x1, w],
		(* else *)
		x1[[ pos ]]
		],
		Log[Mean[w]] + loglike}
	]

(*  c[[i]] is the number of children of the i-th parent
	p[[j]] is the position of the parent of the j-th child
	w is the list of nonnegative weights (need not be normalized)
	M is the desired number of children
	*)
SystematicResampling = Compile[{{w, _Real, 1}, {M, _Integer, 0}},
	Module[{
	c = (Rest[#] - Most[#])& @
		Floor[FoldList[Plus, 0, M * w/(Plus@@w)] - Random[]],
	p = Table[0, {M}],
	j = 1},
	Do[
		Do[p[[j]] = i; j++, {c[[i]]}],
	{i, Length[w]}];
	p
	]]

(* use {x1[[ pos ]], w[[ pos ]]} as first argument,
	where pos = SystematicResampling[w, Length[w]] *)
MCMCSmooth[{x1_, w_}, y_, x0_, f_, gp_] :=
	With[{proposal = f[x0]},
	MCMCChoose[x1, w, proposal, gp[proposal, y]]
	]

MCMCChoose = Compile[
	{{x1, _Real, 2}, {w, _Real, 1}, {xp, _Real, 2}, {wp, _Real, 1}},
	With[{alpha =
		If[# >= 0, 1, 0] & /@ (wp - w * Table[Random[], {Length[w]}])},
	alpha * xp + (1 - alpha) * x1
	]]

(***** smoothing *****)

Options[ParticleSmoother] = {CollectStates -> False}

ParticleSmoother[states_, fp_, n_Integer:1, opts___?OptionQ] :=
	Module[{
	collected = Reverse @
		If[TrueQ[CollectStates /. {opts} /. Options[ParticleSmoother]],
		(* then *)
			CollectStates[states],
		(* else *)
			states],
	result},
	result = Transpose @ Reverse @
		FoldList[
			SmootherStep[##, fp]&,
			Resample[First[collected], n],
			Rest[collected]
		];
	If[n == 1, result[[1]], result]
	]

SmootherStep[x1_, {x0_, w0_}, fp_] :=
	x0[[ BinarySearch[w0 * fp[#, x0]] ]]& /@ x1


CollectStates[states_] :=
	Transpose[{First[#], Length[#]}& /@ Split[Sort[#]]]& /@ states

Resample[{x_, w_}, n_] := x[[ SystematicResampling[w, n] ]]

Resample[{x_, w_}] := x[[ BinarySearch[w] ]]

BinarySearch = Compile[{{w, _Real, 1}},
	Module[{
		lo = 0,
		hi = Length[w],
		mid = 0,
		r = Random[],
		c = Most@FoldList[Plus, 0, w/Plus @@ w]},
	While[hi - lo > 1,
		mid = lo + Floor[(hi - lo)/2];
		If[r > c[[mid + 1]],
			lo = mid,
			hi = mid
			]
		];
	hi
	]]

End[]
EndPackage[]