(* :Name: NumericalMath`ImplicitRungeKutta` *)

(* :Title: Solver for implicit Runge-Kutta methods *)

(* :Author: Mark Sofroniou *)

(* :Summary:
 This package solves a system of ordinary differential equations
 using any (invertible) implicit Runge-Kutta methods. It also
 derives the coefficients of most implicit Runge-Kutta methods to
 any order and precision.
*)

(* :Context: NumericalMath`ImplicitRungeKutta` *)

(* :Package Version: 1.0 *)

(* :Copyright: Copyright 1997-1999, Wolfram Research, Inc. *)

(* :History:
 Version 1.0 by Mark Sofroniou, August 1997. *)

(* :Keywords:
 Implicit Runge Kutta, Gauss Legendre, Symplectic, Orthogonal and
 Isospectral flows, Numerical methods for Ordinary Differential Equations. *)

(* :Source:
 John C. Butcher: The Numerical Analysis of Ordinary Differential
 Equations: Runge-Kutta and General Linear Methods, John Wiley &
 Sons, New York, 1987.

 E. Hairer, S. P. Norsett and G. Wanner, Solving Ordinary Differential
 Equations I: Nonstiff Problems, Second edition, Springer Verlag,
 New York, 1993.

 E. Hairer, and G. Wanner, Solving Ordinary Differential Equations II:
 Stiff and Differential Algebraic Problems, Second edition, Springer Verlag,
 New York, 1996.

 J. D. Lambert: Numerical Methods for Ordinary Differential Systems,
 The Initial Value Problem, John Wiley & Sons, New York, 1991.

 J. M. Sanz-Serna and M. P. Calvo,
 Numerical Hamiltonian Problems,
 Chapman and Hall, London, 1994.
*)

(* :Mathematica Version: 4.0 *)

(* :Discussion:
 Implicit Runge-Kutta methods require the solution of a nonlinear
 system of equations at each timestep. Many of these methods have
 desirable qualitative properties which are not found in the methods
 currently used in NDSolve. That is, the methods better reflect properties
 of the differential system under consideration. For example the
 Gauss-Legendre methods are symplectic, symmetric (time reversible),
 and preserve orthogonal and isospectral flows. The Gauss methods attain
 maximal order amongst the class of Runge-Kutta methods, namely order 2*s
 for an s-stage method.

 Symplectic Runge-Kutta methods are methods tailored for (but not restricted
 to) the study of Hamiltonian systems. Symplectic methods necessarily require
 a constant stepsize, which is a geometric constraint. Furthermore, symplectic
 systems are conservative therefore we need not be concerned with dissipative
 phenomena such as A-stability. A-stability usually means that we have to
 solve the implicit equations at each step using Newton's method. Instead,
 symplecticity allows us to solve the implicit equations at each stage
 using a simpler iterative scheme.
*)

(* :Limitation:
 The solution method only currently works using a fixed timestep.
 Furthermore, the methods use fixed point iteration which is not
 suitable for stiff differential systems - these require the use
 of Newton's method in order to guarantee convergence.

 The s-stage Gauss-Legendre methods applied to an n-dimensional
 system require the solution of an s*n implicit coupled system
 of equations at each timestep. More efficient methods exist, but
 the implicit solution method here is general and is intended for
 qualitative comparison of methods - not the most efficient for
 each class of method being studied.
*)

BeginPackage["NumericalMath`ImplicitRungeKutta`",
  "LinearAlgebra`Vandermonde`",
  "NumericalMath`OptimizeExpression`" ];

ImplicitRungeKuttaSolve::usage =
"ImplicitRungeKuttaSolve[{dx1,...,dxn}, {x1,...,xn}, {i1,...,in},
{t, ti, tf, dt}]
numerically integrates the dxj as functions of the xj with inital values ij.
The integration proceeds from ti through tf in steps of dt.";

ImplicitRungeKuttaCoefficients::usage =
"ImplicitRungeKuttaCoefficients[p, x, d, s, prec, opts] gives the
coefficients of the implicit Runge Kutta method with s stages defined
by the collocating polonomial p in the variable x to degree d and
precision prec.";

ImplicitRungeKuttaGauss::usage =
"ImplicitRungeKuttaGauss is a method option of ImplicitRungeKuttaSolve.
ImplicitRungeKuttaGauss[p, prec] gives the coefficients of the Gauss Legendre
Runge Kutta method of order >= p to precision prec.";

ImplicitRungeKuttaLobattoIIIA::usage =
"ImplicitRungeKuttaLobattoIIIA is a method option of ImplicitRungeKuttaSolve.
ImplicitRungeKuttaLobattoIIIA[p, prec] gives the coefficients of the
LobattoIIIA Runge Kutta method of order >= p to precision prec.";

ImplicitRungeKuttaLobattoIIIB::usage =
"ImplicitRungeKuttaLobattoIIIB is a method option of ImplicitRungeKuttaSolve.
ImplicitRungeKuttaLobattoIIIB[p, prec] gives the coefficients of the
LobattoIIIB Runge Kutta method of order >= p to precision prec.";

ImplicitRungeKuttaRadauIA::usage =
"ImplicitRungeKuttaRadauIA is a method option of ImplicitRungeKuttaSolve.
ImplicitRungeKuttaRadauIA[p, prec] gives the coefficients of the RadauIA
Runge Kutta method of order >= p to precision prec.";

ImplicitRungeKuttaRadauIIA::usage =
"ImplicitRungeKuttaRadauIIA is a method option of ImplicitRungeKuttaSolve.
ImplicitRungeKuttaRadauIIA[p, prec] gives the coefficients of the RadauIIA
Runge Kutta method of order >= p to precision prec.";

$ImplicitRungeKuttaMethods::usage =
"$ImplicitRungeKuttaMethods gives a list of the built in
methods which can be used in ImplicitRungeKuttaSolve.";

RoundingControl::usage =
"RoundingControl is an option of ImplicitRungeKuttaSolve
that specifies whether rounding errors committed at one timestep
should be added back into the solution at the next timestep.
Valid settings are True, False or Automatic.";

Unprotect[ImplicitRungeKuttaSolve,
ImplicitRungeKuttaCoefficients, ImplicitRungeKuttaGauss,
ImplicitRungeKuttaLobattoIIIA, ImplicitRungeKuttaLobattoIIIB,
ImplicitRungeKuttaRadauIA, ImplicitRungeKuttaRadauIIA,
RoundingControl, $ImplicitRungeKuttaMethods];

Begin["`Private`"];

(* Error and warning messages *)

ImplicitRungeKuttaCoefficients::args =
"The argument `1` in ImplicitRungeKuttaCoefficients did not evaluate to `2`.";

ImplicitRungeKuttaCoefficients::opts =
"The option `1` in ImplicitRungeKuttaCoefficients did not evaluate to `2`";

ImplicitRungeKuttaSolve::arglen =
"The arguments `1`, `2` and `3` in ImplicitRungeKuttaSolve did not evaluate
to vectors of the same length";

ImplicitRungeKuttaSolve::args =
"The argument `1` in ImplicitRungeKuttaSolve did not evaluate
to `2`.";

ImplicitRungeKuttaSolve::comp =
"The `1` in ImplicitRungeKuttaSolve did not evaluate to `2`.";

ImplicitRungeKuttaSolve::ctype =
"Unable to determine the compile type from ``.";

ImplicitRungeKuttaSolve::nsteps =
"Bounding the number of timesteps `1` by MaxSteps";

ImplicitRungeKuttaSolve::opts =
"The option `1` in ImplicitRungeKuttaSolve did not evaluate to `2`";

ImplicitRungeKuttaSolve::sing =
"Singular matrix encountered. Unable to solve for this method.";

(* Packed arrays *)

pack = Developer`ToPackedArray;
packedQ = Developer`PackedArrayQ;

(* Locally defined predicates *)

NumberMatrixQ[x_?packedQ]:= SameQ[TensorRank[ x ], 2];
NumberMatrixQ[x_?MatrixQ]:= Apply[ And, Map[ NumberVectorQ, x ] ];
NumberMatrixQ[_]:= False;

NumberVectorQ[x_?packedQ]:= SameQ[TensorRank[ x ], 1];
NumberVectorQ[{__?NumberQ}] = True;
NumberVectorQ[_]:= False;

NumericVectorQ[{__?NumericQ}]:= True;
NumericVectorQ[_]:= False;

PositiveExtendedRealQ[n_?NumberQ] = Positive[n];
PositiveExtendedRealQ[Infinity] = True;
PositiveExtendedRealQ[_] = False;

RealNumberQ[x_]:= NumberQ[x] && !MatchQ[x, _Complex];
RealNumericQ[x_]:= RealNumberQ[x] || NumericQ[x];

ScalarQ[_List]:= False;
ScalarQ[_]:= True;

UnivariatePolynomialQ[p_, x_]:=
  With[{v = Variables[p]},
    SameQ[Length[v], 1] && SameQ[v, {x}] && PolynomialQ[p, x]
  ];

Numericize[f_, wprec_]:=
  Check[ N[f, wprec], Throw[$Failed], $MaxExtraPrecision::meprec];

(**** Derive implicit Runge Kutta coefficients that are defined by a
 collocating polynomial to arbitrary order and precision. ****)

$ImplicitRungeKuttaMethods =
{ImplicitRungeKuttaGauss,
ImplicitRungeKuttaLobattoIIIA,
ImplicitRungeKuttaLobattoIIIB,
ImplicitRungeKuttaRadauIA,
ImplicitRungeKuttaRadauIIA};

CheckValue[e_, h_]:= If[ MatchQ[e, h], e, Throw[$Failed] ];

(* The localized variable x must not be unique here,
   to enable caching of results *)

ImplicitRungeKuttaGauss[p_Integer?Positive, prec_?PositiveExtendedRealQ]:=
  Block[{res, s, x},
    s = Quotient[p+1, 2];
    res /; (res = Catch[ irkcmain[ x^s (x-1)^s, x, s, s, prec, {Row} ] ];
            res=!=$Failed)
  ];

ImplicitRungeKuttaLobattoIIIA[p_Integer?Positive, prec_?PositiveExtendedRealQ]:=
  Block[{res, s, sm1, x},
    s = Quotient[ p+3, 2];
    sm1 = s-1;
    res /; (res = Catch[ irkcmain[ x^sm1 (x-1)^sm1, x, s-2, s, prec, {Row} ] ];
            res=!=$Failed)
  ];

ImplicitRungeKuttaLobattoIIIB[p_Integer?Positive, prec_?PositiveExtendedRealQ]:=
  Block[{res, s, sm1, x},
    s = Quotient[ p+3, 2];
    sm1 = s-1;
    res /; (res = Catch[ irkcmain[ x^sm1 (x-1)^sm1, x, s-2, s, prec, {Column} ] ];
            res=!=$Failed)
  ];

ImplicitRungeKuttaRadauIA[p_Integer?Positive, prec_?PositiveExtendedRealQ]:=
  Block[{res, s, sm1, x},
    s = Quotient[ p+2, 2];
    sm1 = s-1;
    res /; (res = Catch[ irkcmain[ x^s (x-1)^sm1, x, sm1, s, prec, {Column} ] ];
            res=!=$Failed)
  ];

ImplicitRungeKuttaRadauIIA[p_Integer?Positive, prec_?PositiveExtendedRealQ]:=
  Block[{res, s, sm1, x},
    s = Quotient[ p+2, 2];
    sm1 = s-1;
    res /; (res = Catch[ irkcmain[ x^sm1 (x-1)^s, x, sm1, s, prec, {Row} ] ];
            res=!=$Failed)
  ];

(* General solver *)

IRKCtest =
  If[#1,
    True,
    Message[ImplicitRungeKuttaCoefficients::opts, Apply[Sequence, #2]]; False
  ]&;

IRKCmessages = {
  {Solve, "Column or Row"}
  };

IRKCOptionsTest[p_, x_, d_, s_, prec_, opts___]:=
  Module[{opttests, optvals},

    (* Check arguments *)

    If[ !UnivariatePolynomialQ[p, x],
      Message[ImplicitRungeKuttaCoefficients::args, p,
        SequenceForm["a polynomial in ",x] ];
      Return[ $Failed ]
    ];

    (* Check options *)

    optvals = Options[ImplicitRungeKuttaCoefficients];
    optvals = Map[First, optvals] /. {opts} /. optvals;

    opttests = {
      MatchQ[ Part[optvals,1], Row|Column ] (* Solve *)
    };

    If[ Apply[And, MapThread[IRKCtest, {opttests, IRKCmessages}] ],
      optvals,
      $Failed
    ]
  ];

Options[ImplicitRungeKuttaCoefficients] =
  {Solve -> Row};

ImplicitRungeKuttaCoefficients[p_, x_?ScalarQ, d_Integer?NonNegative,
    s_Integer?Positive, prec_?PositiveExtendedRealQ, opts___?OptionQ]:=
  Module[{res},
    res /; ((res = IRKCOptionsTest[p, x, d, s, prec, opts])=!=$Failed &&
               (res = Catch[ irkcmain[p, x, d, s, prec, res] ];
                res=!=$Failed))
  ];

irkcmain[p_, x_, d_, s_, prec_, rowcol_]:=
  irkcmain[p, x, d, s, prec, rowcol]=
    Module[{a, b, c, r},
      c = IRKCollocation[p, x, d, prec];
      r = N[ Range[ s ], prec ];
      b = IRKWeights[c, r, prec];
      a = IRKMatrix[b, c, r, s, prec, rowcol];
      {a, b, c}
    ];

(* The result can be reused for different precisions *)

IRKEquation[p_, x_, d_]:=
  IRKEquation[p, x, d]=
    D[ Expand[ p ], {x, d}] == 0;

IRKCollocation[p_, x_, d_, Infinity]:=
  Module[{sols},
    sols = x /. {ToRules[ Roots[ IRKEquation[p,x,d], x ] ]};
    sols = RootReduce[sols];
    Sort[ sols, Less ]
  ];

IRKCollocation[p_, x_, d_, prec_]:=
  Module[{sols},
    sols = x /. {ToRules[ NRoots[ IRKEquation[p,x,d], x, prec ] ]};
    Sort[ sols, Less ]
  ];

IRKWeights[c_, r_, p_?(#>$MachinePrecision&)]:=
  CheckValue[
    VandermondeSolve[ c, 1/r, Compiled -> False ],
    _?(Head[#]=!=VandermondeSolve&)
  ];

IRKWeights[c_, r_, p_]:=
  CheckValue[ VandermondeSolve[ c, 1/r ], _?(Head[#]=!=VandermondeSolve&) ];

(* Solve for b^T*A down the columns of A, divide by b and transpose. *)

IRKMatrix[b_, c_, r_, s_, p_?(#>$MachinePrecision&), {Column}]:=
  Transpose[
    Table[
      CheckValue[
        VandermondeSolve[ c, Part[b,i] (1 - Part[c,i]^r)/r, Compiled -> False ],
        _?(Head[#]=!=VandermondeSolve&)
      ]/b,
    {i, s}]
  ];

IRKMatrix[b_, c_, r_, s_, p_, {Column}]:=
  Transpose[
    Table[
      CheckValue[
        VandermondeSolve[ c, Part[b,i] (1 - Part[c,i]^r)/r ],
        _?(Head[#]=!=VandermondeSolve&)
      ]/b,
    {i, s}]
  ];

(* Solve for the rows of A. *)

IRKMatrix[_, c_, r_, s_, p_?(#>$MachinePrecision&), {Row}]:=
  Table[
    CheckValue[
      VandermondeSolve[ c, Part[c,i]^r/r, Compiled -> False ],
      _?(Head[#]=!=VandermondeSolve&)
    ],
  {i, s}];

IRKMatrix[_, c_, r_, s_, p_, {Row}]:=
  Table[
    CheckValue[
      VandermondeSolve[ c, Part[c,i]^r/r ],
      _?(Head[#]=!=VandermondeSolve&)
    ],
  {i, s}];


(**** Implicit Runge-Kutta solver. ****)

(* Argument and option testing *)

IRKStest =
  If[#1,
    True,
    Message[ImplicitRungeKuttaSolve::opts, Apply[Sequence, #2]]; False
  ]&;

IRKSmessages = {
  {Compiled, "True, False or Automatic"},
  {InterpolationOrder, "a positive integer or Automatic"},
  {MaxIterations, "a positive integer"},
  {MaxSteps, "a positive integer"},
  {Method, "a symbol"},
  {Order, "a positive integer"},
  {RoundingControl, "True, False or Automatic"},
  {SameTest, "Automatic, a symbol or a function"},
  {StoppingTest, "None, or a boolean expression"},
  {WorkingPrecision, "a positive number"}
  };

IRKSOptionsTest[dx_, x_, xi_, tdata_, opts___]:=
  Module[{opttests, optvals},

    (* Check arguments *)

    If[ !VectorQ[dx] || !VectorQ[x] || !VectorQ[xi] ||
          !MatchQ[Length[dx], _Integer?Positive] ||
          !SameQ[ Length[dx], Length[x], Length[xi] ],
      Message[ImplicitRungeKuttaSolve::arglen, dx, x, xi];
      Return[ $Failed ]
    ];

    If[ !Apply[ And, Map[ ScalarQ, x] ],
       Message[ImplicitRungeKuttaSolve::args, x, "a list of a scalars"];
       Return[ $Failed ]
    ];

    If[ !NumericVectorQ[xi],
       Message[ImplicitRungeKuttaSolve::args, xi,
         "a list of numeric quantities"];
       Return[ $Failed ]
    ];

    If[ !VectorQ[tdata] || Length[tdata]=!=4 || !ScalarQ[First[tdata]] ||
         !NumericVectorQ[Rest[tdata]],
       Message[ImplicitRungeKuttaSolve::args, tdata,
         "a list of a scalar and three numeric quantities"];
       Return[ $Failed ]
    ];

    (* Check options *)

    optvals = Options[ImplicitRungeKuttaSolve];
    optvals = Map[First, optvals] /. {opts} /. optvals;

    opttests = {
      MatchQ[ Part[optvals, 1], True|False|Automatic ], (* Compiled *)
      MatchQ[ Part[optvals, 2], Automatic|_Integer?Positive ], (* InterpolationOrder *)
      MatchQ[ Part[optvals, 3], _Integer?Positive ], (* MaxIterations *)
      MatchQ[ Part[optvals, 4], _Integer?Positive ], (* MaxSteps *)
      MatchQ[ Part[optvals, 5], _Symbol ], (* Method *)
      MatchQ[ Part[optvals, 6], _Integer?Positive ], (* Order *)
      MatchQ[ Part[optvals, 7], True|False|Automatic ], (* RoundingControl *)
      MatchQ[ Part[optvals, 8], Automatic|_Symbol|_Function ], (* SameTest *)
      MatchQ[ Part[optvals, 9], None|_ ], (* StoppingTest *)
      MatchQ[ Part[optvals, 10], _?(RealNumberQ[#]&&Positive[#]&) ] (* WorkingPrecision *)
      };

    If[ Apply[And, MapThread[IRKStest, {opttests, IRKSmessages}] ],
      optvals,
      $Failed
    ]
  ];

Options[ImplicitRungeKuttaSolve] =
  {Compiled -> Automatic, InterpolationOrder -> Automatic,
   MaxIterations -> 15, MaxSteps -> 1*^6,
   Method -> ImplicitRungeKuttaGauss,
   Order -> 4, RoundingControl -> Automatic, SameTest -> Automatic,
   StoppingTest -> None, WorkingPrecision -> $MachinePrecision};

ImplicitRungeKuttaSolve[dx_, x_, xi_, t_, opts___?OptionQ] :=
  Module[{ans},
    ans /; ((ans = IRKSOptionsTest[dx, x, xi, t, opts])=!=$Failed &&
            (ans = Catch[ IRKSmain[dx, x, xi, t, ans] ])=!=$Failed)
  ];

(* Check consistency of shape the dimensions of the coefficients *)

CheckCoefficients[{amat_, bvec_, cvec_}]:=
  (
   If[ !MatrixQ[amat] || ! VectorQ[bvec] || !VectorQ[cvec] ||
         UnsameQ[ Length[amat], Length[bvec], Length[cvec] ],
     Message[ ImplicitRungeKuttaSolve::comp, "coefficients",
       "a system of appropriate dimension"];
     Throw[ $Failed ]
   ];
   {amat, bvec, cvec}
  );

CheckCoefficients[_]:=
  (
   Message[ ImplicitRungeKuttaSolve::comp, "coefficients",
     "a system of appropriate dimension"];
   Throw[$Failed];
  );

(* Calculate the number of timesteps for the integration. *)

CalculateSteps[nti_, tf_, ndt_, maxsteps_, prec_]:=
  Module[{nsteps},
    Block[{$MaxExtraPrecision},
      nsteps = Round[(tf - nti)/ndt];
      If[ !MatchQ[nsteps, _Integer?Positive],
        Message[ ImplicitRungeKuttaSolve::comp, "number of time steps",
          "a positive integer"];
        Throw[ $Failed ]
      ];
      If[ nsteps > maxsteps,
        Message[ ImplicitRungeKuttaSolve::nsteps, nsteps];
        nsteps = maxsteps;
      ];
      nsteps
    ]
  ];

(* Construct a function for evaluating the right hand side of the
 differential system. Localize variables using a vector consisting
 of part extractions. *)

(* Add time dependency to variables in the differential system.
 Do this even if time dependence was not explicitly specified in the
 variable list. *)

AddTimeDependence[y:(_[t_]), t_]:= y;
AddTimeDependence[x_, t_]:= x[t];

(* Add dependency for interpolating functions in the result *)

AddTimeDependence[x_[t_], y_, t_]:= y[t];
AddTimeDependence[x_, y_, t_]:= y;

MakeLocalFunction[dx_, t_, tsub_, x_, xsub_, wprec_] :=
  Module[{subs, xt},
    xt = Map[ AddTimeDependence[#, t]&, x];
    subs = Union[ Join[ {t->tsub}, Thread[x->xsub], Thread[xt->xsub] ] ];
    subs = Dispatch[ subs ];
    OptimizeExpression[ Numericize[ dx /. subs, wprec ] ]
  ];

(* Make n interpolating functions for the result. *)

InterpolateResult[data_, n_, ord_] :=
  Module[{tdata, xdata},
    tdata = pack[Part[data, All, 1]];
	xdata = pack[Part[data, All, 2]];
    Table[
      Interpolation[
        Transpose[{tdata, Part[ xdata, All, i ]}],
        InterpolationOrder -> ord
      ],
    {i, 1, n}]
  ];

(* Make the list of return rules with time dependencies. *)

MakeReturnRules[x_, t_]:=
  MapThread[ Rule[ #1, AddTimeDependence[#1, #2, t] ]&, x];

(* A single step of the Runge-Kutta method *)

IRKStep[{ti_, xi_}, bvec_, maxit_, rcontrol_]:=
  (
    (*
     Implicit iteration. Could use extrapolation of previous values here,
     but this seems not to be worthwhile. Extrapolation should be scaled by
     the ratio of the timesteps for an adaptive implementation.
    *)

    datadx = FixedPoint[ localdxv[ ndt, nti, xi, # ]&, datadx, maxit ];

    If[ !NumberMatrixQ[datadx], Throw[$Failed] ];

	If[ !rcontrol, Return[ {ti + ndt, xi + Dot[ bvec, datadx ]} ] ];

    (* Add back contribution from rounding errors from the previous step. *)

    temp = rerrx + Dot[ bvec, datadx ];
    datax = xi + temp;
    diffx = xi - datax;
    rerrx = temp + diffx;

	temp = rerrt + ndt;
	datat = ti + temp;
	difft = ti - datat;
	rerrt = temp + difft;

    {datat, datax} (* Update solutions *)
  );

GetType[x_?packedQ]:=
  Block[{$RecursionLimit = Infinity}, ptype[x] ];

GetType[x_?NumberVectorQ]:=
  If[ MatchQ[ x, {__?RealNumberQ} ], _Real, _Complex ];

GetType[x_]:= typeerr[x];

ptype[x_?packedQ]:= ptype[ First[ x ] ];
ptype[x_?RealNumberQ]:= _Real;
ptype[x_?NumberQ]:= _Complex;
ptype[x_]:= typeerr[x];

typeerr[x_]:=
  (Message[ImplicitRungeKuttaSolve::ctype, x]; Throw[$Failed]);

(*
 Form the vector d, where d^T = b^T A^-1.
 This reduces rounding errors at each step and allows the possibility
 of more efficient methods, such as those using simplified iteration
 schemes.
*)

InvertMatrix[amat_, bvec_]:=
  Module[{dvec},
    dvec = LinearSolve[ Transpose[ amat ], bvec ];
    If[ SameQ[ Head[dvec], LinearSolve ],
      Message[ImplicitRungeKuttaSolve::sing];
      Throw[ $Failed ]
    ];
    dvec
  ];

(*
 We make use of the fact that some variables like amat exist in the
 Package`Private` context. This avoids passing them as parameters
 at each integration step.
*)

IRKSmain[dx_, x_, xi_, {t_, ti_, tf_, dt_},
    {comp_, iorder_, maxit_, maxsteps_, method_, order_, rcontrol_, sametest_,
       stoptest_, wprec_}] :=
  Block[{datadx, datat, datax, difft, diffx, localdx, localdxrhs, localdxv,
      localt, localx, localxvars, ndt, nti, rerrt, rerrx, stages, temp},
    Module[{ amat, bvec, cvec, compiledQ, ctype, fpopts, iporder, localtest,
        ndims, nsteps, nxi, partmsgon, prec, res, roundcontrolQ, zero },

      nxi = pack[ CheckValue[ Numericize[ xi, wprec ], _?NumberVectorQ ] ];
      ndims = Length[ nxi ];
	  nti = CheckValue[ Numericize[ ti, wprec ], (_Real|0) ];
      ndt = CheckValue[ Numericize[ dt, wprec ], _Real?(#=!=0&) ];
      nsteps =
        CheckValue[
          CalculateSteps[ nti, tf, ndt, maxsteps, wprec ],
          _Integer?Positive
        ];

      (* Compute the coefficients of the implicit Runge-Kutta method
         through desired order and precision. *)

      prec = wprec + order;
      {amat, bvec, cvec} =
          pack[ Numericize[ CheckCoefficients[ method[order, prec] ], wprec] ];
      bvec = InvertMatrix[ amat, bvec ];
      stages = Length[cvec];

      (* Set any options for FixedPoint *)

      fpopts = Options[FixedPoint, SameTest];
      If[ UnsameQ[ sametest, SameTest /. fpopts ],
        SetOptions[ FixedPoint, SameTest -> sametest ]
      ];

      (* Turn off Part message if it was on and save the state *)

      partmsgon = UnsameQ[ Head[Part::partd], $Off ];
      If[ partmsgon, Off[Part::partd] ];

      (* Make a function of local variables for evaluating the right hand
         side of the differential system *)

      localxvars = Table[ Part[localx, i], {i, 1, ndims}];
      localdxrhs = MakeLocalFunction[ dx, t, localt, x, localxvars, wprec ];
      If[ SameQ[stoptest, None],
        localtest = (False&),
        localtest = MakeLocalFunction[ stoptest, t, localt,
		              x, localxvars, wprec ];
      ];

      compiledQ =
        TrueQ[comp || (SameQ[comp, Automatic] && wprec <= $MachinePrecision)];

      (* Set up the function, vector function and stopping test evaluation *)

	  If[ compiledQ,
        ctype = GetType[nxi];
        localdx = Compile[{{localt, _Real}, {localx, #1, 1}}, #2 ]&[ctype, localdxrhs];
        Function[{ctype, amat, cvec, stages},
		  localdxv =
            Compile[{{ndt, _Real}, {nti, _Real}, {xv, ctype, 1}, {dxv, ctype, 2}},
			  With[{tv = nti + ndt*cvec},
			    ndt*Dot[ amat, Table[ localdx[ Part[tv, i], xv + Part[dxv, i]], {i, stages}]]
	          ]
            ]
        ][ctype, amat, cvec, stages];
        If[ stoptest =!= None,
          localtest =
            Compile[{{localt, _Real}, {localx, #1, 1}}, #2 ]&[ctype, localtest];
        ];,
        (* Else *)
        (localdx[ Pattern[localt, Blank[]], Pattern[localx, Blank[]] ]:= #1)&[localdxrhs];
        Function[{amat, cvec},
		  localdxv[ ndt_, nti_, xv_, dxv_ ]:=
            ndt*Dot[ amat, MapThread[localdx[ #1, xv + #2]&, {nti + ndt*cvec, dxv}] ]
		][amat, cvec];
        If[ stoptest =!= None,
          (localtest[ Pattern[localt, Blank[]], Pattern[localx, Blank[]] ]:= #1)&[localtest];
        ];
      ];

      (* Rounding control *)

      roundcontrolQ =
        Or[
          SameQ[rcontrol, True],
          SameQ[rcontrol, Automatic] &&
             (wprec <= $MachinePrecision || SameQ[$MinPrecision, $MaxPrecision])
        ];

      (* Starting values for implicit iterations and rounding errors *)

      zero = Numericize[0, wprec];
      datadx = pack[ Table[ zero, {stages}, {ndims} ] ];
      diffx = rerrx = pack[ Table[ zero, {ndims} ] ];
	  rerrt = zero;

      (* Carry out the integration by iterating the one step Runge-Kutta map. *)

      res =
        NestWhileList[
          IRKStep[ #, bvec, maxit, roundcontrolQ ]&,
          {nti, nxi}, (* Initial conditions *)
          !Apply[localtest, #]&, (* Stopping test *)
          1,          (* Number of time steps used in stopping test *)
          nsteps,     (* Number of steps *)
          If[ SameQ[stoptest, None], 0, -1 (* Use a pretest *) ]
        ];

      (* Maybe the pretest failed at the outset *)

      If[ SameQ[res, {}], Throw[$Failed]];

	  (* Interpolate the result *)

      iporder = If[ SameQ[ iorder, Automatic ], Min[ nsteps, order ], iorder ];

      res = InterpolateResult[ res, ndims, iporder ];

      (* Reset any options and messages that were localy changed *)

      SetOptions[ FixedPoint, Apply[ Sequence, fpopts] ];
      If[ partmsgon, On[Part::partd] ];

      MakeReturnRules[ {x , res}, t ]
    ]
  ];

End[ ]; (* End `Private` Context. *)

SetAttributes[
{ImplicitRungeKuttaSolve, ImplicitRungeKuttaCoefficients,
ImplicitRungeKuttaGauss, ImplicitRungeKuttaLobattoIIIA,
ImplicitRungeKuttaLobattoIIIB, ImplicitRungeKuttaRadauIA,
ImplicitRungeKuttaRadauIIA},
ReadProtected];

Protect[ImplicitRungeKuttaSolve,
ImplicitRungeKuttaCoefficients, ImplicitRungeKuttaGauss,
ImplicitRungeKuttaLobattoIIIA, ImplicitRungeKuttaLobattoIIIB,
ImplicitRungeKuttaRadauIA, ImplicitRungeKuttaRadauIIA,
RoundingControl, $ImplicitRungeKuttaMethods];

EndPackage[ ]; (* End package Context. *)

(* Examples:
   The idealized (frictionless) nonlinear pendulum on the separatrix:
   ImplicitRungeKuttaSolve[{y, Sin[x]}, {x, y}, {Pi, 2},{t, 0, 100, 1/10}]
   ImplicitRungeKuttaSolve[{y, Sin[x]}, {x, y}, {Pi, 2},{t, 0, 100, 1/10},
     StoppingTest -> x < 2 ]
   ImplicitRungeKuttaSolve[{y, Sin[x]}, {x, y}, {Pi, 2},{t, 0, 100, 1/10},
     MaxIterations -> 5, MaxSteps -> 100 ]
   ImplicitRungeKuttaSolve[{y, Sin[x]}, {x, y}, {Pi, 2},{t, 0, 100, 1/10},
     Compiled -> False ]

   Gauss Legendre methods:
   ImplicitRungeKuttaCoefficients[ x^3 (x-1)^3, x, 3, 3, $MachinePrecision ]
   ImplicitRungeKuttaGauss[6, $MachinePrecision]
   ImplicitRungeKuttaCoefficients[ x^3 (x-1)^3, x, 3, 3, Infinity ]
   ImplicitRungeKuttaGauss[6, Infinity]

   Radau IA methods:
   ImplicitRungeKuttaCoefficients[ x^3 (x-1)^2, x, 2, 3, Infinity ]
   ImplicitRungeKuttaRadauIA[5, Infinity]

   Radau IIA methods:
   ImplicitRungeKuttaCoefficients[ x^2 (x-1)^3, x, 2, 3, Infinity ]
   ImplicitRungeKuttaRadauIIA[5, Infinity]
*)


