(* :Title: RandomWalkMetropolis *)

(* :Author: Mark Fisher *)

(* :Date: June 2007 *)

(* :Mathematica Version: 6 *)

(* :Package Version: 2.0
	Modified for Version 6.0 June 2007, completely revamped.
*)

(* :Summary:
This package implements the (random-walk) Metropolis version of the Markov
Chain Monte Carlo (MCMC) algorithm. MCMC algorithms are used to draw from
probability distributions. The function BridgeEstimate can be used to compute
the normalizing constant.
*)

(* :Discussion:
The Markov Chain is implemented via NestList. The "state" vector is the
parameter vector with the function value appended. The iterated function
generates a proposal and returns the proposal (and its value) if its value
is greater than a random fraction of the current state value. Otherwise
the current state is returned. Proposals are generated by random-walk
Gaussian or T-distributions.
*)

(* :References:
	RandomWalkMetropolis implements Agorithm A.29 (p. 288) in
	Robert, Christian P. and George Casella (2004) Monte Carlo
		Statistical Methods, Second Edition, Springer.

	For the Gelfand-Dey method of marginal likelihood estimation, see
	Koop, Gary (2003) Bayesian Econometrics, John Wiley & Sons,
		pp. 104-106.
	For additional details, Koop refers the reader to
	Geweke, John (1999) Using Simulation Methods for Bayesian Econometric
		Models: Inference, Development, and Communication (with discussion
		and rejoinder), Econometric Reviews, 18, pp. 1-126.
*)

BeginPackage["RandomWalkMetropolis`",
	{"Utilities`FilterOptions`", "MultivariateStatistics`"}]

RandomWalkMetropolis::usage =
"RandomWalkMetropolis[f, x0, step, nDraws, nThin, nBurn] \
returns nDraws steps of a Markov Chain Monte Carlo (MCMC) simulation from \
the function f computed via the random-walk Metropolis algorithm, \
beginning at the vector x0 of starting values, where step is a scalar or \
vector (with the same dimension as start) that represents standard \
deviations. Alternatively, step can be a covariance matrix. \
nThin specifies the amount of \"thinning\" to do, outputing every \
nThin-th draw. nBurn specifies the number of draws to discard \
from the beginning during the \"burn-in\" period. (Thus the total number \
of function evaluations is nBurn + nThin * nDraws.) The function \
f must have the head Function and take a single vector argument; \
in addition, it should return the log of the kernel value unless the \
option LogLikelihood is set to False. MetropolisMCMC also takes the \
options DiscardValues and BooleanFunction. (See the usage notes.)"

LogLikelihood::usage =
"LogLikelihood is an option for RandomWalkMetropolis that \
specifies whether the function returns values proportional to the \
likelihood or to the loglikelihood. The default setting is \
LogLikelihood -> True."

BooleanFunction::usage =
"BooleanFunction is an option for RandomWalkMetropolis \
that specifies a function that returns True of False to be applied to the \
proposal prior to evaluating the likelihood function. For example, \
BooleanFunction -> (And@@Thread[#>=0]&). The default setting is \
BooleanFunction -> (True&). The option BooleanFunction is provided for \
use with LogLikelihood -> True, in order to avoid taking the log of zero."

RegroupOutput::usage =
"RegroupOutput is an option for RandomWalkMetropolis \
which specifies whether to regroup each row of the output from {x__, v_} \
to {{x}, v}. The default is RegroupOutput -> True."

DiscardValues::usage =
"DiscardValues is an option for RandomWalkMetropolis \
which specifies whether to discard the appended values. The default \
setting is DiscardValues -> True."

CompilePatterns::usage =
"CompilePatterns is an option for \
RandomWalkMetropolis which specifies the list of patterns supplied to \
Compile (as its third argument). These patterns specify the return types \
of external calls. The default setting is CompilePatterns -> {}."

DegreesOfFreedom::usage =
"DegreesOfFreedom is an option for \
RandomWalkMetropolis which specifies the number of degrees of freedom for \
the proposal distribution. The default setting is DegreesOfFreedom -> \
\[Infinity] which specifies the Gaussian distribution. Additional valid \
settings are any positive integer. For example, DegreesOfFreedom -> 1 \
specifies the Cauchy distribtuion (i.e., t distribtion with 1 degree of \
freedom)."

AcceptanceRate::usage =
"AcceptanceRate[mc] computes, the fraction of times \
the proposal is accepted; namely, N[Length[Split[mc]]/Length[mc]]."

MHMFunction::usage =
"MHMFunction[mean, cov, p] returns a function that \
can be used to compute the normalization constant where mean is the \
vector of means of the MCMC parameter matrix from a \
Markov Chain Monte Carlo simulation and cov is the covariance matrix. \
The third argument p specifies \
the truncation for the Gaussian in terms of a p-value. Alternatively, one \
may specify MHMFunction[mcmc, p], where mcmc is the MCMC parameter matrix \
of parameters."

MHMEstimate::usage =
"MHMEstimate[mcmcdata, p] returns the Gelfand-Dey estimate of the \
normalization constant given the Markov Chain Monte Carlo simulation, \
where mcmcdata is produced by MetropolisMCMC with the option \
DiscardValues -> False and where the truncation point is controlled \
by the \"p-value\" p. MHMEsitmate takes the option LogLikelihood, \
which specifies whether the likelihood values are in logs. \
The default value is LogLikelihood -> False. If the option is \
specified the value for p must also be specified."

MarginalLikelihoodMuller::usage = "MarginalLikelihoodMuller[mc, fun, {m, \
V}, n (,opts)] computes the log of the marginal likelihood from the MCMC \
output mc generated from the kernel fun using n draws from a Gaussian \
weighting distribution with mean m and covariance matrix V. (The method is \
due to Ulrich K. Müller.) Since the draws from the Gaussian distribution \
are independent, n can be substantially less than the length of mc. Each \
row of mc output contains the draw and the associated kernel value. (The \
kernel values are assumed to be in logs unless the option LogLikelihood -> \
False is specified.) A function (returning True|False) that determines the \
bounds of the parameter space can be specified via the option \
BooleanFunction."

MullerFunction::usage = "MullerFunction[rf, rg][logC] is a helper function \
called by MarginalLikelihoodMuller. MullerFunction calls the compiled \
ComputeListSum."

MarginalLikelihoodMuller::badbounds = "The interval `1` does not contain
zero. Returning the value closest to zero. Try using a larger number of
draws from the weighting distribution."

MarginalLikelihoodCJ::usage = "MarginalLikelihoodCJ[mc, f, V, xstar, n] \
computes the log of the marginal likelihood of the data given the \
Metropolis output mc according to Chib-Jeliazkov method, where f is the \
kernel of the posterior, V is the covariance matrix of the random-walk \
Gaussian proposal density, xstar is the point at which to estimate the \
density of the poseterior, and n is the number of draws of the Gaussian \
weighting function."

BridgeEstimate::usage = "BridgeEstimate[{f1, f2}, {draws1, draws2}, psi] \
computes the bridge estimate of Log[c1/c2] where ci is the normalizing \
constant for the kernel Exp[fi] given the draws from the two distributions. \
The parameter psi should be set the the ratio of the number of \
\"independent\" draws in draws1 relative to the number of independent draws \
in draws2. (When in doubt, set it to 1.) BridgeEstimate[mc, fun, {m, V}, \
{n, psi}] computes the bridge estimate of the log of the marginal \
likelihood given the posterior kernel fun using the MCMC output mc; n draws \
are made from the Gaussian weighting function which is centered at m with \
covariance matrix V."

BridgeEstimate::badbounds = "The interval `1` does not contain zero.
Returning the value closest to zero. Try using a larger number of draws
from the weighting distribution."

BridgeEstimate::nodraws = "There are no draws from test distribution in
bounds."

BridgeEstimate::ndraws = "There are `1` draws from test distribution in
bounds."

BridgeEstimatePowerFamily::usage = "BridgeEstimatePowerFamily[{z1, z2}, {k, \
A}] returns the bridge estimate of the log of the ratio of normalizing \
constants given the lists of log differences z1 and z2 and the power family \
parameters k and A."

BridgeEstimateFromMetropolis::usage = "BridgeEstimateFromMetropolis[{fun, \
start, stepsize, mcsteps}, {m, V}, {n, psi} (, opts)]."

ForceDraws::usage = "ForceDraws is an option for BridgeEstimate which \
specifies whether to enforce the number of draws from the Gaussian \
to equal the number specified. This is an issue if draws can be \
discarded as out-of-bounds. With the default setting ForceDraws -> \
True, successive draws are made until the number specified are obtained."

Begin["Private`"]

AcceptanceRate[mc_List] := N[Length[Split[mc]]/Length[mc]]

RandomWalkMetropolis::baddf = "DegreesOfFreedom must be either a positive \
integer or \[Infinity]."

Options[RandomWalkMetropolis] = {
	LogLikelihood -> True,
	BooleanFunction -> (True&),
	DiscardValues -> True,
	RegroupOutput -> True,
	DegreesOfFreedom -> \[Infinity],
	CompilePatterns -> {}
	}

RandomWalkMetropolis[
		fun_Function,
		start_?(VectorQ[#, NumericQ]&),
		stepSize_,
		nDraws_Integer?Positive,
		nThin_Integer:1,
		nBurn_Integer:0,
		opts___?OptionQ] :=
	Module[{log, discard, regroup, patt, df, boole,
			k, chol, logfun, sfun, nfun, sv, mc},
	{log, discard, regroup, patt, df, boole} = {LogLikelihood, DiscardValues,
			RegroupOutput, CompilePatterns, DegreesOfFreedom, BooleanFunction} /.
			{opts} /. Options[RandomWalkMetropolis];
	If[Not[(IntegerQ[df] && df > 0) || df === Infinity],
		Message[RandomWalkMetropolis::baddf, df]; Return[$Failed]];
	k = Length[start];
	chol = CholeskyDecomposition[
		Switch[
			TensorRank[stepSize],
			0, stepSize^2 * IdentityMatrix[k],
			1, DiagonalMatrix[stepSize^2],
			2, stepsize (* covariance matrix *)
			]];
	sfun = AssembleStepFunction[chol, k, df];
	logfun = If[TrueQ[log], fun, Log[fun[#]]&];
	nfun = AssembleNestFunction[logfun, boole, sfun];
	sv = Append[start, logfun[start]];
	mc = CompileNestFunction[nfun, sv, nDraws, nThin, nBurn, patt][];
	If[TrueQ[discard], mc[[All, ;;-2]],
		If[TrueQ[regroup], Through[{Most, Last}[#]] & /@ mc, mc]
		]
	]

RandomWalkMetropolis[__] := $Failed

AssembleStepFunction[chol_, k_, df_] :=
	Switch[df,
		Infinity, (* normal proposal *)
			Function[{},
			Table[Sqrt[-2 Log[RandomReal[]]] Cos[2 Pi RandomReal[]], {k}].chol
			],
		_, 				(* t proposal *)
			Function[{},
			(1/Sqrt[(#.#)&[
					Table[Sqrt[-2 Log[RandomReal[]]] Cos[2 Pi RandomReal[]], {df}]]/df] *
			Table[Sqrt[-2 Log[RandomReal[]]] Cos[2 Pi RandomReal[]], {k}]).chol
			]
		]

AssembleNestFunction[vfun_, bfun_, sfun_] :=
		Function[sv,
				Module[{p, pval},
				p = Most[sv] + sfun[];
				If[bfun[p],
					pval = vfun[p];
					If[pval >= Last[sv] + Log[RandomReal[]],
						Append[p, pval],
						sv],
					sv]
			]]

CompileNestFunction[nfun_, start_, nDraws_, nThin_, nBurn_, patt_] :=
	Compile[{}, (* no args *)
	Rest @
		NestList[
			Nest[nfun, #, nThin]&,
			Nest[nfun, start, nBurn],
			nDraws],
		patt]


(***** BridgeEstimate *****)

Options[BridgeEstimate] = {
	BooleanFunction -> (True &),
	LogLikelihood -> True,
	ForceDraws -> True
	}

BridgeEstimate[mc_, fun_, {m_List, V_List}, {n_Integer, psi_?NumericQ},
			opts___?OptionQ] /;
		(Dimensions[mc][[-1]] == 2 &&
		MatrixQ[mc[[All, 1]], NumericQ] &&
		VectorQ[mc[[All, 2]], NumericQ]) :=
	Module[{bfun, loglike, force, k, t, ran, logf, logg,
		fOnMC, gOnMC, fOnRan, gOnRan, z1, z2},
	bfun = BooleanFunction /. {opts} /. Options[BridgeEstimate];
	loglike = TrueQ[LogLikelihood /. {opts} /. Options[BridgeEstimate]];
	force = TrueQ[ForceDraws /. {opts} /. Options[BridgeEstimate]];
	If[force,
		(* then *)
		ran = Table[
			While[
				t = RandomReal[MultinormalDistribution[m, V]];
				!bfun[t]];
			t,
		{n}],
		(* else *)
		ran = Select[RandomReal[MultinormalDistribution[m, V], n], bfun[#]&];
		If[Length[ran] == 0, Message[BridgeEstimate::nodraws]; Return[$Failed]];
		Message[BridgeEstimate::ndraws, Length[ran]];
		];
	logf = If[loglike, fun, Log[fun[#]] &];
	k = Length[m];
	logg = With[{S = Inverse[V], const = -Log[(2Pi)^(k/2)Sqrt[Det[V]]]},
  		Compile[{{z, _Real, 1}}, const - (z - m).S.(z - m)/2]
  		];
	fOnMC = If[loglike, mc[[All, 2]], Log[mc[[All, 2]]]];
	gOnMC = logg /@ mc[[All, 1]];
	fOnRan = If[bfun @ #, logf @ #, -$MaxMachineNumber] & /@ ran;
	gOnRan = logg /@ ran;
	z1 = fOnMC - gOnMC;
	z2 = gOnRan - fOnRan;
	BridgeEstimate[{z1, z2}, psi, opts]
	]

BridgeEstimate[{f1_, f2_}, {draws1_List, draws2_List}, psi_?NumericQ,
		opts___?OptionQ] :=
	With[{
		z1 = (f1[#] - f2[#])& /@ draws1,
		z2 = (f2[#] - f1[#])& /@ draws2
		},
	BridgeEstimate[{z1, z2}, psi, opts]
		]

BridgeEstimate[{z1_List, z2_List}, psi_?NumericQ, opts___?OptionQ] :=
	Block[{x},
	{w1, w2} = {z1, z2};
	(* Print[StringForm["Starting value: `1`",
			NumberForm[LogRatioStartingValue[z1, z2], {6,4}]]]; *)
		x /. FindRoot[
			BridgeFunction[z1, z2, psi][x],
			{x, LogRatioStartingValue[z1, z2]},
			FilterOptions[FindRoot, opts] // Evaluate]
		]

BridgeFunction[z1_, z2_, psi_][x_?NumericQ] :=
	BridgeFunctionCompiled[x, psi, z1, z2]

(* trap for non machine precsion numbers *)
BridgeFunctionCompiled =
	With[{max = Log[$MaxMachineNumber]},
	Compile[{
		{x, _Real, 0},
		{psi, _Real, 0},
		{z1, _Real, 1},
		{z2, _Real, 1}},
	Module[{
		len1 = Length[z1],
		len2 = Length[z2],
		t = 0.
		},
	Sum[
		t = z2[[j]] - x;
		If[t < max,
			1/(psi + Exp[t]),
			0],
		{j, len2}]/len2 -
	Sum[
		t = z1[[j]] + x;
		If[t < max,
			1/(1 + psi * Exp[t]),
			0],
		{j, len1}]/len1
	]]]

(* provide starting value for FindRoot *)
LogRatioStartingValue[z1_, z2_] :=
	BridgeEstimatePowerFamily[{z1, z2}, {1, 1}]

BridgeEstimatePowerFamily[{z1_, z2_}, {k_, A_}] :=
	With[{t = A^(1/k)},
	Log[Mean[(t + Exp[#/k])^(-k)& /@ z1]/
		Mean[(1 + t * Exp[#/k])^(-k)& /@ z2]]
	]

BridgeEstimateFromMetropolis[
		{fun_, start_, stepsize_, mcsteps_},
		{m_, V_}, {n_, psi_}, opts___] :=
	Module[{mc},
	mc = MetropolisMCMC[fun, start, stepsize, mcsteps, opts,
		LogLikelihood -> True, ReturnValues -> True];
	(* Print[StringForm["Acceptance Rate: `1`",
		NumberForm[AcceptanceRate[mc], {3,3}]]]; *)
	BridgeEstimate[mc, fun, {m, V}, {n, psi}, opts, LogLikelihood -> True]
	]

(* Modified Harnomonic Mean (a.k.a. Gelfand-Dey) *)

MHMEstimate[mc_, p_:(.01)] /;
			(Dimensions[mc][[-1]] == 2 &&
			MatrixQ[mc[[All, 1]], NumericQ] &&		(* parameters *)
			VectorQ[mc[[All, 2]], NumericQ]) :=		(* values *)
	1/Mean[
		(MHMFunction[mc[[All, 1]], p] /@ mc[[All, 1]])/mc[[All, 2]]
		]

Options[MHMEstimate] = {LogLikelihood -> False}

MHMEstimate[mc_, p_:(.01), opts___?OptionQ] /;
			(Dimensions[mc][[-1]] == 2 &&
			MatrixQ[mc[[All, 1]], NumericQ] &&		(* parameters *)
			VectorQ[mc[[All, 2]], NumericQ]) :=		(* values *)
	With[{vals = If[
		TrueQ[LogLikelihood /. {opts} /. Options[MHMEstimate]],
		Exp[mc[[All, 2]]],
		mc[[All, 2]]]},
	1/Mean[
		(MHMFunction[mc[[All, 1]], p] /@ mc[[All, 1]])/vals
		]
	]

MHMFunction[mc_?(MatrixQ[#, NumericQ]&), p_:(.01)] :=
	MakeMHMFunction[##, p]& @@ MCMeanAndCovariance[mc]

MHMFunction[___] := $Failed

MHMEstimate[___] := $Failed

MCMeanAndCovariance[mc_?(MatrixQ[#, NumericQ]&)] :=
	{First[#], Rest[#]}& @ MCMeanCov[mc]

MCMeanCov = Compile[{{mc, _Real, 2}},
	Module[{
		len = Length[mc],
		data = Transpose[mc],
		mean = {1.},
		centered = {{1.}},
		cov = {{1.}}
		},
	mean = (Plus @@@ data)/len;
	centered = data - mean;
	cov = (centered.Transpose[centered])/len;
	cov = (cov + Transpose[cov])/2;
	Prepend[cov, mean]
	]];

MakeMHMFunction[mean_, cov_, p_:(.01)] :=
	With[{
		n = Length[mean],
		inv = Inverse[cov]
		},
	With[{
		const = N @ 1/(Sqrt[(2*Pi)^n * Det[cov]] * (1-p)),
		invGamma = N @ 2*InverseGammaRegularized[n/2, 0, 1 - p]
		},
	Compile[{{vec, _Real, 1}},
		With[{
			quad = (mean - vec).inv.(mean - vec)
			},
		If[quad <= invGamma, Exp[-quad/2]const, 0]
	]]]]

(***** MarginalLikelihoodMuller *****)
(* Ulrich K. Müller's method for computing the constant of integration *)

Options[MarginalLikelihoodMuller] =
	{BooleanFunction -> (True &), LogLikelihood -> True}

MarginalLikelihoodMuller[mc_, fun_, {m_, V_}, n_Integer, opts___?OptionQ] /;
		(Dimensions[mc][[-1]] == 2 &&
		MatrixQ[mc[[All, 1]], NumericQ] &&
		VectorQ[mc[[All, 2]], NumericQ]) :=
	Module[{fropts, bfun, loglike, k, ran, logf, logg,
		fOnMC, gOnMC, fOnRan, gOnRan, rf, rg, lo, hi, mlo, mhi, cf},
	fropts = FilterOptions[FindRoot, opts];
	bfun = BooleanFunction /. {opts} /. Options[MarginalLikelihoodMuller];
	loglike = TrueQ[LogLikelihood /. {opts} /. Options[MarginalLikelihoodMuller]];
	k = Length[m];
	ran = RandomArray[MultinormalDistribution[m, V], n];
	logf = If[loglike, fun, Log[fun[#]] &];
	logg = With[{S = Inverse[V], const = -Log[(2Pi)^(k/2)Sqrt[Det[V]]]},
  		Compile[{{z, _Real, 1}}, const - (z - m).S.(z - m)/2]
  		];
	fOnMC = If[loglike, mc[[All, 2]], Log[mc[[All, 2]]]];
	gOnMC = logg /@ mc[[All, 1]];
	fOnRan = If[bfun @ #, logf @ #, -$MaxMachineNumber] & /@ ran;
	gOnRan = logg /@ ran;
	rf = Sort[gOnMC - fOnMC];
	rg = Sort[fOnRan - gOnRan];
	lo = -rf[[-1]];
	hi =  rg[[-1]];
	{mlo, mhi} = MullerFunction[rf, rg] /@ {lo, hi};
	If[mlo <= 0 <= mhi,
		(* then: zero is in bounds *)
		Block[{x}, x /. FindRoot[
					MullerFunction[rf, rg][x],
					{x, .45*lo + .55*hi, .55*lo + .45*hi, lo, hi},
					fropts // Evaluate]],
		(* else: zero is out of bounds *)
		Message[MarginalLikelihoodMuller::badbounds, {mlo, mhi}];
		If[mlo > 0, lo, hi]]
	]

(* rf and rg are sorted logs of ratios *)
MullerFunction[rf_, rg_][x_?NumericQ] :=
	ComputeListSum[-x, rg] - ComputeListSum[x, rf]

(* the list is in logs and has been sorted *)
ComputeListSum =
	Compile[{
		{x, _Real, 0},
		{list, _Real, 1}},
	Module[{
		t = 0.,
		sum = 0.,
		len = Length[list]
		},
	Catch[
		Do[
		t = list[[i]] + x;
		If[t < 0,
			sum += 1 - Exp[t],
			Throw[sum/len]
			],
		{i, len}];
		sum/len]
	]];


(***** Chib-Jeliazkov method *****)

MarginalLikelihoodCJ[mc_, fun_, V_, xstar_, n_] :=
	Module[{ran, valstar, alphavals, qfun, qvals, num, den},
	ran = RandomArray[MultinormalDistribution[xstar, V], n];
	valstar = fun[xstar];
	alphavals = Min[0, valstar - #] & /@ mc[[All, 2]];
	qfun = With[{
			m = xstar,
			S = Inverse[V],
			const = -Log[(2Pi)^(Length[xstar]/2)Sqrt[Det[V]]]},
		Compile[{{z, _Real, 1}}, const - (z - m).S.(z - m)/2]
		];
	qvals = qfun /@ mc[[All, 1]];
	num = Mean[Exp[alphavals + qvals]];
	den = Mean[Exp[Min[0, fun[#] - valstar] & /@ ran]];
	valstar - Log[num/den]
	]


End[]
EndPackage[]