(* :Title: AnnuitySolve.m *)

(* :Context: AnnuitySolve` *)

(* :Author: Mark Fisher *)

(* :Summary:
In an economy where the representative agent has homothetic recursive
preferences (homothetic stochastic differential utility), the optimal
wealth-consumption ratio is the value of an annuity that satisfies a quasi-
linear 2nd-order PDE. This package provides tools to compute a series
solution to the PDE.
*)

(* :Package Version: 2.0 (June 2000) *)

(* :Package History: 
Version 1.0 (August 1998) Initial version.
Version 1.1 (December 1998)	Fix multivariate related problems.
Version 1.2 (June 1999) Deal with DSolve bug in Version 4.
Version 2.0 (June 2000) Major rewriting. 
	This package version requires Version 4 of Mathematica.
*)

(* :Mathematica Version: 4.0 *)

(* :Sources:
Fisher, M. and C. Gilles (1999) "Consumption and asset prices with
homothetic recursive preferences." Federal Reserve Bank of Atlanta.
*)

(* :Discussion:
NAnnuitySolve is designed to compute the series solution to a quasi-linear
2nd-order PDE in terms of an unknown function A[x,t], where the data are
real analytic and where the boundary condition is of the form A[x,0] == C,
for C constant. For an annuity, C >= 0, while for a zero-coupon bond, C ==
1. The solution provided by NAnnuitySolve is of the form Sum[a[i][t](x-x0)
^i, {i, 0, order}] where the a[i][t] are functions of t. The boundary
condition is specified indirectly via the option InitialCondition -> C,
which imposes the condition a[0][0] == C (along with a[i][0] == 0 for i >=
1). For exponential-polynomial bond prices, the option FunctionalForm ->
Exp allows one to solve for Log[A[x,t]] subject to a[0][0] == 0.
*)

(* :Requirements: Utilities`FilterOptions` *)

(* :Examples:
For examples, see the notebook AnnuitySolveExamples.nb, which is
available from the author upon request.
*)

BeginPackage["AnnuitySolve`", {"Utilities`FilterOptions`"}]

AnnuitySolve::usage = "AnnuitySolve[pde, A, t, {x, x0}, order] calculates a
symbolic series solution for the pde to the order specified, where A is an
unknown function of x and t. It expands a polynomial of factor loadings in
t around the point x0. Multiple state variables can be specified as in
AnnuitySolve[pde, A, t, {x, x0}, {y, y0}, order]. As specified,
AnnuitySolve returns a pure function; if A is replaced by A[outargs],
AnnuitySolve returns the function evaluated at outargs. AnnuitySolve has
the options FunctionalForm and InitialCondition. In addition, options can
be passed to DSolve."

NAnnuitySolve::usage = "NAnnuitySolve[pde, A, {t, min, max}, {x, x0},
order] calculates a numerical series solution for the pde to the order
specified, where A is an unknown function of x and t. It expands a
polynomial of factor loadings in t around the point x0. Multiple state
variables can be specified as in AnnuitySolve[pde, A, {t, min, max}, {x,
x0}, {y, y0}, order]. As specified, NAnnuitySolve returns a pure function;
if A is replaced by A[outargs], NAnnuitySolve returns the function
evaluated at outargs. NAnnuitySolve has the options FunctionalForm,
InitialCondition, FactorLoadingSymbol, PolynomialOrderDifferential, and
DifferentialLoadings. In addition, options can be passed to NDSolve."

FunctionalForm::usage = "FunctionalForm is an option for AnnuitySolve and
NAnnuitySolve. The default setting is FunctionalForm -> Identity. The
setting FunctionalForm -> Exp can be used to solve for exponential-
polynomial bond prices."

InitialCondition::usage = "InitialCondition is an option for AnnuitySolve
and NAnnuitySolve. It specifies the value of the zero-order factor loading
function at t = 0. The default setting for NAnnuitySolve is
InitialCondition -> 10^-100, which avoids division by zero when the PDE is
quasi-linear."

PolynomialOrderDifferential::usage = "PolynomialOrderDifferential is an
option for NAnnuitySolve. The default setting is
PolynomialOrderDifferential -> 0. This option is used in conjunction with
the option DifferentialLoadings. Typically for this purpose the setting
would be PolynomialOrderDifferential -> 2."

DifferentialLoadings::usage = "DifferentialLoadings is an option for
NAnnuitySolve. The default setting is DifferentialLoadings -> {}. This
option is used in conjunction with the option PolynomialOrderDifferential.
Typical use would involve computing exponential-polynomial bond prices
(using NAnnuitySolve with FunctionalForm -> Exp), and then calling
NBondToAnnuitySeries, the output of which would be used as the
DifferentialLoadings."

FactorLoadingSymbol::usage = "FactorLoadingSymbol is an option for
NAnnuitySolve. It is used to coordinate passing the output of
NBondToAnnuitySeries to NAnnuitySolve via the option DifferentialLoadings.
The default setting is FactorLoadingSymbol -> $a, which need not be changed
unless there is a symbol conflict."

BondToAnnuitySeries::usage = "BondToAnnuitySeries[bond, t, {x, x0}, order]
takes an expression for zero-coupon bond prices as a function of x and t
and returns series coefficients (as functions of t that have been
symbolically integrated) up to the order specified for the associated
annuity. Multiple state variables can be specified as in
BondToAnnuitySeries[bond, t, {x, x0}, {y, y0}, order]. Options can be
passed to Integrate."

NBondToAnnuitySeries::usage = "NBondToAnnuitySeries[bond, {t, min, max},
{x, x0}, order] takes an expression for zero-coupon bond prices as a
function of x and t and returns series coefficients (as functions of t that
have been numerically integrated) up to the order specified for the
associated annuity. Multiple state variables can be specified as in
NBondToAnnuitySeries[bond, {t, min, max}, {x, x0}, {y, y0}, order]. The
output is designed to be used with the option DifferentialLoadings (in
conjunction with PolynomialOrderDifferential) for NAnnuitySolve. Options
can be passed to NDSolve."

ResidualFromPDE::usage = "ResidualFromPDE[pde, soln] returns the
\"residual\" of the pde (as a pure function), where the PDE has been
evaluated according to the given solution function."

MakeIndices::usage = "MakeIndices[n, order] returns a list of indices that
represent a complete polynomial basis."

MakePolynomial::usage = "MakePolynomial[xvars, t, order, sym] returns a
complete polynomial in xvars (a list) and t of the given order using the
given symbol. If the last argument is omitted, it defaults to $a."

$a::usage = "$a is the symbol for the factor loadings."

Begin["`Private`"]

AnnuitySolve::badargs = "The arguments to the function in the PDE are not
properly specified."

Options[AnnuitySolve] = {FunctionalForm -> Identity,
	InitialCondition -> 0}

AnnuitySolve[pde_, A_Symbol, t_Symbol, xargs:{_Symbol, _}..,
		order_Integer?Positive, opts___?OptionQ] :=
	Module[{ff, ic0, sym, x, x0, args, indices, coeffs,
		initconds, poly, odes},
	{ff, ic0} = {FunctionalForm, InitialCondition} /. {opts} /.
		Options[AnnuitySolve];
	Catch[{x, x0, args, indices, coeffs, initconds} =	
		AnnuitySolveSetup[pde, A, t, {xargs}, order, sym, ic0]];
	poly = ff @ MakePolynomial[x - x0, t, order, sym];
	odes = MakeODEs[pde, A, args, poly, x, x0, order, indices];
	Function @@ {args, poly /. 
		( (DSolve[Join[odes, initconds], coeffs, t] /. {{x__}} -> {x}) )}
	]

(* auxillary functions for AnnuitySolve and NAnnuitySolve *)
AnnuitySolveSetup[pde_, A_, t_, xargs_, order_, sym_, ic0_] :=
	Module[{x, x0, arglist, args, indices, coeffs, coeffs0, initconds},
	{x, x0} = Transpose[xargs];
	arglist = Union @ Cases[pde, f_[A][x__] | A[x__] :> {x}, Infinity];
	If[arglist == {} || 
		(args = First @ arglist;
		Length[arglist] != 1 || Union[args] =!= Union[Join[x, {t}]]),
			Message[AnnuitySolve::badargs]; Throw[Return[$Failed]]];
	indices = MakeIndices[Length[x], order];
	coeffs = sym @@@ indices;
	coeffs0 = Through[coeffs[0]];
	initconds = Join[{First[coeffs0] == ic0}, Thread[Rest[coeffs0] == 0]];
	{x, x0, args, indices, coeffs, initconds}
	]

MakeODEs[pde_, A_, args_, poly_, x_, x0_, order_, indices_] :=
	Module[{polyeqn, ser},
	polyeqn = pde /. Equal -> Subtract /. A -> Function @@ {args, poly};
	ser = Series[polyeqn, Sequence @@ Thread[{x, x0, order}]];
	Thread[(SeriesCoefficient[ser, #] & /@ indices) == 0]
	]

Options[NAnnuitySolve] = {FunctionalForm -> Identity,
	InitialCondition -> 10^-100, PolynomialOrderDifferential -> 0,
	DifferentialLoadings -> {}, FactorLoadingSymbol -> $a}

NAnnuitySolve[pde_, A_Symbol,
		range:{t_Symbol, tmin_?NumericQ, tmax_?NumericQ},
		xargs:{_Symbol, _?NumericQ}.., order_Integer?Positive,
		opts___?OptionQ] :=
	Module[{ff, ic0, diff, diffloads, sym, ndopts, x, x0,
		args, indices, coeffs, initconds, poly, odes},
	{ff, ic0, diff, diffloads, sym} = {FunctionalForm,
		InitialCondition, PolynomialOrderDifferential,
		DifferentialLoadings, FactorLoadingSymbol} /. {opts} /.
		Options[NAnnuitySolve];
	ndopts = FilterOptions[NDSolve, opts];
	Catch[{x, x0, args, indices, coeffs, initconds} =	
		AnnuitySolveSetup[pde, A, t, {xargs}, order, sym, ic0]];
	poly = ff @ MakePolynomial[x - x0, t, order + diff, sym] /. 
		Cases[diffloads, f_[g_[d__ /; Plus[d] > order], h_] :> f[sym[d], h]];
	odes = MakeODEs[pde, A, args, poly, x, x0, order, indices];
	Function @@ {args, poly /.
		( (NDSolve[Join[odes, initconds], coeffs, range, Evaluate[ndopts],
			StartingStepSize -> 10^-6, MaxSteps -> 10^6] /. {{x__}} -> {x}) ) }
	]

NAnnuitySolve[pde_, A_Symbol[outargs__],
		range:{t_Symbol, tmin_?NumericQ, tmax_?NumericQ},
		xargs:{_Symbol, _?NumericQ}.., order_Integer?Positive,
		opts___?OptionQ] :=
	NAnnuitySolve[pde, A, range, xargs, order, opts][outargs]

AnnuitySolve[pde_Equal, A_Symbol[outargs__], t_Symbol,
		xargs:{_Symbol, _}.., order_Integer?Positive, opts___?OptionQ] :=
	AnnuitySolve[pde, A, t, xargs, order, opts][outargs]

BondToAnnuitySeries[bond_, t_Symbol, xargs:{_Symbol, _}..,
		order_Integer?Positive, opts___?OptionQ] :=
	Module[{sym, iopts, indices, sercoeffs, loadings, s},
	sym = FactorLoadingSymbol /. {opts} /. Options[NAnnuitySolve];
	iopts = FilterOptions[Integrate, opts];
	{sercoeffs, indices} = BondToAnnuitySeriesSetup[bond, {xargs}, order];
	loadings = Function @@ {t, 
		Integrate[# /. t -> s, {s, 0, t}, Evaluate[iopts]]}& /@ sercoeffs;
	Thread[(sym @@@ indices) -> loadings]
	]

BondToAnnuitySeriesSetup[bond_, xargs_, order_] :=
	Module[{x, x0, ser, indices, sercoeffs},
	{x, x0} = Transpose[xargs];
	ser = Series[bond, Sequence @@ Thread[{x, x0, order}]];
	indices = MakeIndices[Length[x], order];
	sercoeffs = SeriesCoefficient[ser, #] & /@ indices;	
	{sercoeffs, indices}
	]

NBondToAnnuitySeries[bond_,
		range:{t_Symbol, tmin_?NumericQ, tmax_?NumericQ},
		xargs:{_Symbol, _?NumericQ}..,
		order_Integer?Positive, opts___?OptionQ] :=
	Module[{sym, ndopts, x, x0, ser, indices, sercoeffs, loadings, g},
	sym = FactorLoadingSymbol /. {opts} /. Options[NAnnuitySolve];
	ndopts = FilterOptions[NDSolve, opts];
	{sercoeffs, indices} = BondToAnnuitySeriesSetup[bond, {xargs}, order];
	loadings = (g /. First @ NDSolve[{g'[t] == #, g[0] == 0},
			g, range, Evaluate[ndopts], StartingStepSize -> 10^-6,
			MaxSteps -> 10^6])& /@ sercoeffs;
	Thread[(sym @@@ indices) -> loadings]
	]

ResidualFromPDE[pde_, soln_Function] :=
	(Function @@ {#1, pde /. Equal -> Subtract /. #2 -> soln})& @@
		Cases[pde, _[f_][x__] :> {{x}, f}, Infinity, 1][[1]]

MakeIndices[n_, r_] := 
	With[{dummy = Unique[]& /@ Range[n]},
	Flatten[
		Table[dummy, ##]& @@ 
		Thread[{dummy, 0, r - FoldList[Plus, 0, Drop[dummy, -1]]}], 
		n-1]
	]

MakePolynomial[xvars_List, t_, r_Integer, sym_:$a] :=
	With[{indices = MakeIndices[Length[xvars], r]},
	(Times @@@ (xvars^# & /@ indices)).Through[(sym @@@ indices)[t]]
	]

End[]
EndPackage[]
