/**

	RKF.java
	Peter N. Saeta
	2/13/01
	2/19/01		Commented more thoroughly.
	2/23/01		Added option to run fourth-order fixed step RK4.
	3/4/01		Fixed some documentation flaws.
	
	Class RKF implements the Runge-Kutta-Fehlberg 4-5 order adaptive step
	size algorithm for integrating a (set of) differential equation(s).
 
	See DeVries Chapter 5 for a discussion of the algorithm.


	Example application: integrate the equation y''(x) = -y(x) between
	x = 0 to x = 6, starting with y(0) = 1 and y'(0) = 0.
	
	
	public class SHO extends ODE {
		public SHO() {
			super( 2 ); 			// we have 2 ODEs to solve simultaneously
		}
		
		void derivatives( double xx, double[] yy, double[] ff ) {
			ff[0] = yy[1];
			ff[1] = -yy[0];
		 }
		void intermediateResult( double[] xyy ) {	// write results to the screen
			System.out.print( xyy[0] );
			for ( int i = 1; i < xyy.length; i++ )
				System.out.print( "\t" + xyy[i] );
			System.out.println();
		}
						
	}
	
	SHO s = new SHO();				// create a new SHO object
	RKF r = new RKF( s );			// pass the ODE into the new RKF
	
	Vector res = new Vector();		// create a vector to hold the results
	double [] IC = new double[2];	// create an array for the initial condition
	
	IC[0] = 1;						// y(0) = 1
	IC[1] = 0;						// y'(0) = 0
	
	try {
		r.Integrate( 0, 6, IC, res, false );
									// we don't need the results at the end,
									// since they have been echoed to the screen, so
									// last parameter is false. Something might go wrong
									// during integration, so we need a surrounding try.
	} catch ( Exception e ) {
		System.out.println( "I caught exception " + e );
	}

	// If we had called Integrate() with the final value set to true,
	// at this point, the Vector res would have all the intermediate points. We
	// could access them with a loop such as the following:
	
	Enumeration e = res.elements();
	while ( e.hasMoreElements() ) {
		double t[] = ( double[] ) e.nextElement();
		// now do something with the data in t[]
	}
	
	See the routine StartIntegration() below for instructions on integrating
	a system of equations over an indefinite interval. 
	
*/

import java.awt.*;
import java.util.*;

/** This interface defines a function that will be called with the results
 * 	of each step of the integration, in the form of an array of doubles,
 *	x, y0, y1, y2...
 */
public interface ODEStepResult {
	public void intermediateResult( double[] xyy );
}

/** This class defines the ODE function. You subclass this function, as shown
 *	in the example above, to solve the ODE (or set of ODEs) that interests you.
 *	Notes: the protected variable dim stores the number of first-order equations
 *	you are solving. Its value is set when you instantiate your ODE class. Therefore,
 *	your class's constructor must call the ODE constructor with the number of
 *	dimensions (e.g., super(2);).
 */
public class ODE {
	
	protected int dim;
	
	/** public void derivatives( double xx, double[] yy, double[] ff ) calculates
		the vector of derivatives. Override this function in your ODE, and be
		sure to call super(n) from your function's constructor, where n is the
		number of equations.
	*/
	
	void derivatives( double xx, double[] yy, double[] ff ) { }
	public int Dimension() { return dim; }
	ODE( int dims ) { dim = dims; }
}

/**
 *  Class to integrate a set of ODEs using the Runge-Kutta-Fehlberg 4th-5th order
 *  algorithm. Results are returned in a Vector of double[]. Each double array holds
 *  the value of the independent variable (x) and each of the dependent variables
 *  y[0], y[1], ... 
 */
public class RKF {
	
	ODE				theODE;						// subclass to integrate
	ODEStepResult	stepRoutine;				// routine to call with intermediate results
	
	double	precision;							// precision of solution
	double	h;									// step size
	double	hMax, hMin;							// maximum and minimum step sizes
	double	cumulativeError;					// sum of |error| for each step
	boolean	useRK4;								// false by default
	
	protected double	x;
	protected double []	y, y0, yHat;
	protected double []	f0, f1, f2, f3, f4, f5;
	
	protected boolean backward;					// are we going towards negative x?
	protected boolean fixedStepSize;			// set this by calling setStepSize()
	
	
	/** StartIntegration prepares to perform an integration by successive
	 *  calls to Step(). Use this routine to prepare for an open-ended integration,
	 *  such as would be appropriate for integrating ODEs forward in time with
	 *  no definite stopping point envisioned.
	 */
	public void StartIntegration( double fromX, double [] initialCondition ) {
		x = fromX;
		System.arraycopy( initialCondition, 0, y0, 0, initialCondition.length );
		h = hMax;
		cumulativeError = 0;
	}
	

	/** Workhorse routine to perform a Runge-Kutta-Fehlberg integration
	 *	over the range ( fromX - toX ).
	 *	
	 *	After each step, the routine calls the following routine
	 *	
	 *	intermediateResult( double[] xyy )
	 *	
	 *	with an array of values { x, y0, y1, y2... }. By default, the routine
	 *	does nothing; override this routine in your subclass if you wish to do
	 *	something more interesting with the intermediate values (such as plotting them).
	 *	
	 *	If the value of the boolean logResults is true, then the Vector results is
	 *	stuffed with arrays for each intermediate point just as they are passed to
	 *	intermediateResult(). If logResults is false, then results is stuffed with a
	 *	single array, the final values of x and y[i].
	 */
	
	public void Integrate( double fromX, double toX, double [] initialCondition,
		Vector results, boolean logResults ) throws Exception
	{
		x = fromX;
		System.arraycopy( initialCondition, 0, y0, 0, initialCondition.length );
		
		if ( results == null )
			results = new Vector();
		results.removeAllElements();
		cumulativeError = 0;
		
		if ( !fixedStepSize )
			h = hMax;
		backward = ( toX < fromX );				// are we propagating to negative x?
		if ( backward ) h = -h;					// if so, reverse the sign of the step h
		int rows = 0, i;
		
		// Initialize the output matrix with the starting condition
		
		double tmp[] = new double[ theODE.Dimension() + 1 ];
		tmp[0] = x;
		for ( i = 0; i < theODE.Dimension(); i++ )
			tmp[i + 1] = y0[i];
//		echo( y0 );
//		echo( tmp );
		results.addElement( tmp );				// add the initial condition
			
		while( greater( toX, x ) ) {
			Step( toX );
			tmp = new double[ theODE.Dimension() + 1 ];
			tmp[0] = x;
			for ( i = 0; i < theODE.Dimension(); i++ )
				tmp[i+1] = y0[i];
			
			if ( stepRoutine != null )
				stepRoutine.intermediateResult( tmp );			// report the step

			if ( logResults ) {
				results.addElement ( tmp );
			}
		}
		if ( !logResults )						// store the ending result if we
		{										// aren't accumulating the matrix
			tmp = new double[ theODE.Dimension() + 1 ];
			tmp[0] = x;
			for ( i = 0; i < y0.length; i++ )
				tmp[i + 1] = y0[i];
			results.setElementAt( tmp, 0 );
		}
	}

	void Step( double target ) throws Exception {
		if ( useRK4 ) StepRK4( target );
		else StepRKF( target );
	}
	
	/**
	 *  Take the next step in the integration, stopping at or before 'target'.
	 *  This routine is called by Integrate(). For looping through the integration,
	 *  first call StartIntegration(), then make successive calls to Step() with a
	 *  large value for the target.
	 */

	void StepRKF( double target ) throws Exception
	{
		int i, tries = 0;
		double hNew, maxError, bigError, error = 0;
		double absh;
		
		while ( true )
		{
			absh = h;
			if ( backward ) {
				absh = -h;
				if ( x + h < target ) {
					h = target - x;
					absh = -h;
				}
			}
			else if ( (x + h) > target ) {
				h = target - x;
			}
			
			theODE.derivatives( x, y0, f0 );
			
			for ( i = 0; i < y.length; i++ )
				y[i] = y0[i] + h * b10 * f0[i];
			theODE.derivatives( x + a1 * h, y, f1 );
			
			for ( i = 0; i < y.length; i++ )
				y[i] = y0[i] + h * (b20 * f0[i] + b21 * f1[i]);
			theODE.derivatives( x + a2 * h, y, f2 );
			
			for ( i = 0; i < y.length; i++ )
				y[i] = y0[i] + h * (b30 * f0[i] + b31 * f1[i] + b32 * f2[i]);
			theODE.derivatives( x + a3 * h, y, f3 );
			
			for ( i = 0; i < y.length; i++ )
				y[i] = y0[i] + h * (b40 * f0[i] + b41 * f1[i]
					+ b42 * f2[i] + b43 * f3[i]);
			theODE.derivatives( x + a4 * h, y, f4 );
			
			for ( i = 0; i < y.length; i++ )
				y[i] = y0[i] + h * ( b50 * f0[i] + b51 * f1[i]
					+ b52 * f2[i] + b53 * f3[i] + b54 * f4[i] );
			theODE.derivatives( x + a5 * h, y, f5  );
			
			bigError = 0;								// maximum error for any dimensions
			for ( i = 0; i < yHat.length; i++ )
			{
				yHat[i] = y0[i] + h * ( c0 * f0[i] + c2 * f2[i] +
						c3 * f3[i] + c4 * f4[i] + c5 * f5[i] );
				error = Math.abs( h * ( d0 * f0[i] + d2 * f2[i] + d3 * f3[i]
						+ d4 * f4[i] + d5 * f5[i] ) );
				if ( error > bigError )
					bigError = error;
			}
			if ( fixedStepSize ) {
				hNew = h;
				break;
			}
			
			maxError = absh * precision;
			hNew = 0.9 * Math.sqrt( Math.sqrt( maxError / error ) );
				
			// Bound the increases or decreases in the step size
			
			if ( hNew > 4.0 )	hNew = 4.0;
			if ( hNew < 0.1 )	hNew = 0.1;
			
			hNew *= absh;
			if ( hNew > hMax )		hNew = hMax;
			if ( hNew < hMin )
				throw new Exception( "Step size too small in RKF45::Step()" );
			if ( backward )	hNew = -hNew;
			
			if ( error < maxError )
				break;
			if ( ++tries > 15 )
				throw new Exception( "Failure to converge in RKF45::Step() at x = " + x +
				 " h = " + h + " error = " + error  );
			h = hNew;
		}
		
		// The error was small enough so we can accept the step
		
		cumulativeError += error;
		x += h;
		System.arraycopy( yHat, 0, y0, 0, yHat.length );
		h = hNew;
	}

		/**
	 *  Take the next step in the integration, stopping at or before 'target'.
	 *  This routine is called by Integrate(). For looping through the integration,
	 *  first call StartIntegration(), then make successive calls to Step() with a
	 *  large value for the target.
	 */

	void StepRK4( double target ) throws Exception
	{
		int i;
		double absh;
		
		absh = h;
		if ( backward ) {
			absh = -h;
			if ( x + h < target ) {
				h = target - x;
				absh = -h;
			}
		}
		else if ( (x + h) > target ) {
			h = target - x;
		}
		
		theODE.derivatives( x, y0, f0 );
		
		for ( i = 0; i < y.length; i++ )
			y[i] = y0[i] + h * 0.5 * f0[i];
		theODE.derivatives( x + 0.5 * h, y, f1 );
		
		for ( i = 0; i < y.length; i++ )
			y[i] = y0[i] + h * 0.5 * f1[i];
		theODE.derivatives( x + 0.5 * h, y, f2 );
		
		for ( i = 0; i < y.length; i++ )
			y[i] = y0[i] + h * f2[i];
		theODE.derivatives( x + h, y, f3 );
		
		for ( i = 0; i < y.length; i++ )
			y0[i] += h / 6.0 * ( f0[i] + 2*f1[i] + 2*f2[i] + f3[i] );

		x += h;
	}

	protected void init(  )
	{
		// set the dimensions of the f vectors
		
		int dim = theODE.Dimension();
		useRK4 = false;					// use RKF by default
		cumulativeError = 0;
		y = new double[ dim ];
		y0 = new double[ dim ];
		yHat = new double[ dim ];
		f0 = new double[ dim ];
		f1 = new double[ dim ];
		f2 = new double[ dim ];
		f3 = new double[ dim ];
		f4 = new double[ dim ];
		f5 = new double[ dim ];	
	}
	
	public RKF( ODE ode )
	{
		fixedStepSize = false;
		precision = 1e-8;
		hMax = 1;
		hMin = 1e-5;
		backward = false;
		theODE = ode;
		init( );
	}
	
	public void setStepSize( double dh ) {
		h = dh;
		fixedStepSize = true;
	}
	
	public void setStepRoutine( ODEStepResult osr ) {
		stepRoutine = osr;
	}
	
	public double [] getState() {
		double [] d = new double[ y0.length + 1];
		System.arraycopy( y0, 0, d, 1, y0.length );
		d[0] = x;
		return d;
	}
	
	private boolean greater( double x1, double x2 ) {
		if ( backward )
			return x2 > x1;
		else
			return x1 > x2;
	}
	
	/** public void ech( double[] x ) prints a representation of the array
		to System.out
	*/
	public void echo( double [] x ) {
		System.out.print( x[0] );
		for ( int i = 1; i < x.length; i++ )
			System.out.print( "\t" + x[i] );
		System.out.println();
	}
	
	/** constants used in the RKF45 routines */
	 
	static final double a1 = 0.25;
	static final double a2 = 0.375;
	static final double a3 = 12.0 / 13.0;
	static final double a4 = 1.0;
	static final double a5 = 0.5;
		
	static final double b10 = 0.25;
	static final double b20 = 3 / 32.0;
	static final double b21 = 9 / 32.0;
	static final double b30 = 1932.0 / 2197.0;
	static final double b31 = -7200.0 / 2197.0;
	static final double b32 = 7296.0 / 2197.0;
	static final double b40 = 439.0 / 216.0;
	static final double b41 = -8;
	static final double b42 = 3680.0 / 513.0;
	static final double b43 = -845.0 / 4104.0;
	static final double b50 = -8.0 / 27.0;
	static final double b51 = 2;
	static final double b52 = -3544.0 / 2565.0;
	static final double b53 = 1859.0 / 4104.0;
	static final double b54 = -11.0 / 40.0;
		
	static final double c0 = 16.0 / 135.0;
	static final double c2 = 6656.0 / 12825.0;
	static final double c3 = 28561.0 / 56430.0;
	static final double c4 = -0.18;
	static final double c5 = 2.0 / 55.0;
		
	static final double d0 = 1.0 / 360.0;
	static final double d2 = -128.0 / 4275.0;
	static final double d3 = -2197.0 / 75240.0;
	static final double d4 = 0.02;
	static final double d5 = 2.0 / 55.0;
}