(* :Title: ODE *)

(* :Author: Steven Wilkinson, Northern Kentucky University *)

(* :Summary:  	ODESolve[] generates a table of values for a
						solution function to the given ODE. ODEPlot[]
						plot the solution curve to the given ODE.
*)

(* :Context: Calculus`ODE` *)

(* :Package Version: 1.0 *)

(* :Keywords: ordinary differential equations, Euler method,
	modified Euler method, Runge-Kutta method
*)

BeginPackage[
	"Calculus`ODE`","Utilities`FilterOptions`"
];

ODESolve::usage =
	"ODESolve[eqns,x,{t,tmin,tmax,tStep}] generates a table of
values for the numerical solution function, x[t], to the initial
value problem given by eqns for values of t between tmin and tmax
with stepsize tStep.
	ODESolve[eqns,{x1,x2,...},{t,tmin,tmax,tStep}] generates a
table of values for the numerical solution functions, x1[t],
x2[t],..., to the initial value system given by eqns.";

ODEPlot::usage =
	"ODEPlot[eqns,x,{t,tmin,tmax,tStep}] generates the graph for the
numerical solution function, x[t], to the initial value problem
given by eqn for values of t between tmin and tmax with
stepsize tStep.
	ODEPlot[eqns,{x,y},{t,tmin,tmax,tStep}] generates the
parametric plot of {x[t],y[t]}.";

Method::usage =
	"Method is an option for ODESolve and ODEPlot that is the
algorithm to be used in solving the equations. It can be set
to Euler for the Euler method, ImprovedEuler for the improved
Euler method, or RungeKutta for the Runga-Kutta method.";

Euler::usage =
	"Euler can be used in the option Method so that the Euler
method is used to solve the equations.";

ImprovedEuler::usage =
	"ImprovedEuler can be used in the option Method so that
the improved Euler method is used to solve the equations.";

RungeKutta::usage =
	"RungeKutta can be used in the option Method so that
the Runge-Kutta method is used to solve the equations.";

ODESolve::sol =
	"The equations `1` do not have a unique solution.";

ODESolve::init1 =
	"The equations `1` is not an initial condition problem.
Should be only one value of the independent variable in the
initial conditions.";

ODESolve::init2 =
	"The initial value of the indepent variable, `1`, must lie
in the given interval of values [`2`,`3`].";

ODEPlot::dim =
	"There are too many dependent variables to obtain a 2D plot.";

ODESolve::ind =
	"The third argument of ODESolve must contain the independent
variable name, its minimum value, its maximum value, and the
stepsize to be used.";

Options[ODESolve] =
	{Method->RungeKutta};

Options[ODEPlot] =
	{Method->RungeKutta};

Begin["`Private`"];

odeInfo[eqns_List,func_,{t_,tmin_,tmax_,tStep_}] :=
	Module[{exp,x,initCond,initValues,replacements,t0},
		x = Flatten[{func}];
		replacements =
			Join[
				Map[(#[t]->#)&,x],
				Map[(Derivative[1][#][t]->Derivative[1][#])&,x]
			];
		exp =
			Solve[
				eqns /. replacements,
				Map[Derivative[1],x]
			];
		If[Length[exp] =!= 1,Message[ODESolve::sol,eqns] ];
		exp =
			Map[Derivative[1],x] /. First[exp];
		initCond =
			Select[
				eqns,
				!(MemberQ[Level[#,{-1},Heads->True],Derivative])&
			];
		t0 =
			Apply[Level[#,{2}]&,initCond];
		If[Length[t0]=!=1,
			Message[ODESolve::init1,eqns]
		];
		t0 = First[t0];
		initValues =
			Map[#[t0]&,x] /. Flatten[Solve[initCond,Map[#[t0]&,x]]];
		{a,b} = Sort[{tmin,tmax}];
		If[a>t0 || b<t0,Message[ODESolve::init2,t0,a,b]];
		{exp,t0,initValues,a,b,
		Function[Evaluate[Flatten[{t,x}],exp]]}
	];
		
Euler[eqns_List,x_,tRange:{t_,tmin_,tmax_,tStep_}] :=
	Module[{a,b,exp,f,t0,x0},
		{exp,t0,x0,a,b,f} = odeInfo[eqns,x,tRange];
		Join[
			Reverse[
				Rest[
					NestList[
						# - tStep*{1,Apply[f,Flatten[#]]}&,
						{t0,x0},
						Floor[(t0-a)/tStep]
					]
				]
			],
			NestList[
				# + tStep*{1,Apply[f,Flatten[#]]}&,
				{t0,x0},
				Floor[(b-t0)/tStep]
			]
		]
	];

ImprovedEuler[eqns_List,x_,tRange:{t_,tmin_,tmax_,tStep_}] :=
	Module[{a,b,exp,f,f1,t0,x0},
		{exp,t0,x0,a,b,f} = odeInfo[eqns,x,tRange];
		f1[{s_,{u__}}] := f[s,u];
		Join[
			Reverse[
				Rest[
					NestList[
						N[
							# - tStep*{1,0.5(f1[#]+f1[#-tStep*{1,f1[#]}])}
						]&,
						{t0,x0},
						Floor[(t0-tmin)/tStep]
					]
				]
			],
			NestList[
				N[
					# + tStep*{1,0.5(f1[#]+f1[#+tStep*{1,f1[#]}])}
				]&,
				{t0,x0},
				Floor[(tmax-t0)/tStep]
			]
		]
	];

RungeKutta[eqns_List,x_,tRange:{t_,tmin_,tmax_,tStep_}] :=
	Module[{a,b,exp,f,f1,f2,f3,f4,g,t0,x0},
		{exp,t0,x0,a,b,f} = odeInfo[eqns,x,tRange];
		f1[{s_,{u__}}] := f[s,u];
		f2[h_] := f1[# + .5h*{1,f1[#]}]&;
		f3[h_] := f1[# + .5h*{1,f2[h][#]}]&;
		f4[h_] := f1[# + h*{.5,f3[h][#]}]&;
		g[h_] :=
			N[
				# + h{1,(f1[#] + 2f2[h][#] + 2f3[h][#] + f4[h][#])/6}
			]&;
		Join[
			Reverse[
				Rest[
					NestList[
						g[- tStep],
						{t0,x0},
						Floor[(t0-a)/tStep]
					]
				]
			],
			NestList[
				g[tStep],
				{t0,x0},
				Floor[(b-t0)/tStep]
			]
		]
	];

ODESolve[eqns_List,func_,tRange_,opts___] :=
	Module[{method},
		method =
			Method /. {opts} /. Options[ODESolve];
		method[eqns,Flatten[{func}],tRange]
	] /; Length[tRange]===4;

ODESolve[eqns_List,func_,tRange_,opts___] :=
	Message[ODESolve::ind];

ODEPlot[eqns_List,func_,tRange_,opts___] :=
	ListPlot[
		Map[
			Flatten,
			ODESolve[eqns,func,tRange,opts]
		],
		FilterOptions[ListPlot,opts],
		PlotJoined->True
	] /; Length[func] < 2;

ODEPlot[eqns_List,func_,tRange_,opts___] :=
	ListPlot[
		Map[
			Last,
			ODESolve[eqns,func,tRange,opts]
		],
		FilterOptions[ListPlot,opts],
		PlotJoined->True
	] /; Length[func] == 2;

ODEPlot[eqns_List,func_,tRange_,opts___] :=
	Message[ODEPlot::dim];

End[]; (* "`Private`" *)

EndPackage[];  (* "Graphics`Fields`" *)
