!***********************************************************
! Numerical method:                                        *
!     PROGRAM SOLVE                                        *
!     SUBROUTINE STEP( NDEP, METHOD, DX, X, Y )            *
! Equation-dependent routines:                             *
!     SUBROUTINE DERIV( NDEP, X, Y, F )                    *
!     SUBROUTINE OUTPUT( NDEP, X, Y )                      *
!     FUNCTION EXACT( X )                                  *
!***********************************************************


PROGRAM SOLVE
   IMPLICIT NONE
   INTEGER, PARAMETER :: DKIND = KIND( 1.0D0 )   ! use high precision

   ! Main variables
   CHARACTER (LEN=5) METHOD                      ! solution method
   INTEGER NDEP                                  ! number of dependent variables
   INTEGER NSTEP                                 ! number of steps
   INTEGER NPRINT                                ! output frequency
   REAL (KIND=DKIND) DX                          ! step size
   REAL (KIND=DKIND) X                           ! independent variable
   REAL (KIND=DKIND), ALLOCATABLE :: Y(:)        ! vector of dependent variables

   ! Other variables
   INTEGER N                                     ! a counter

   ! Other subprograms
   EXTERNAL STEP                                 ! updates solution by one step
   EXTERNAL OUTPUT                               ! writes any required output


   !******************
   ! Read input data *
   !******************
   OPEN( 10, FILE = 'solve.in' )
   READ( 10, * ) NDEP;   ALLOCATE( Y(NDEP) )
   READ( 10, * ) METHOD
   READ( 10, * ) NSTEP
   READ( 10, * ) DX
   READ( 10, * ) NPRINT
   READ( 10, * ) X
   READ( 10, * ) ( Y(N), N = 1, NDEP )
   CLOSE( 10 )


   !****************************************************
   ! Open file for output and write initial conditions *
   !****************************************************
   OPEN( 20, FILE = 'solve.out' )
   CALL OUTPUT( NDEP, X, Y )


   !****************************************************************
   ! Do NSTEP integration steps, writing output every NPRINT steps *
   !****************************************************************
   DO N = 1, NSTEP
      CALL STEP( NDEP, METHOD, DX, X, Y )
      IF ( MODULO( N, NPRINT ) == 0 ) CALL OUTPUT( NDEP, X, Y )
   END DO


   !*********
   ! Finish *
   !*********
   CLOSE( 20 )
   DEALLOCATE( Y )

END PROGRAM SOLVE


!=======================================================================


SUBROUTINE STEP( NDEP, METHOD, DX, X, Y )
! Updates X and Y() by one step
   IMPLICIT NONE
   INTEGER, PARAMETER :: DKIND = KIND( 1.0D0 )   ! use high precision

   ! Subroutine arguments
   CHARACTER (LEN=5) METHOD                      ! solution method
   INTEGER NDEP                                  ! number of variables
   REAL (KIND=DKIND) DX                          ! step size
   REAL (KIND=DKIND) X                           ! independent variable
   REAL (KIND=DKIND) Y(1:NDEP)                   ! dependent variables

   ! Local variables
   REAL (KIND=DKIND), ALLOCATABLE :: F(:)        ! estimate of dY/dX
   REAL (KIND=DKIND), ALLOCATABLE :: DY1(:), DY2(:), DY3(:), DY4(:)
                                                 ! estimates of increments
   ! Other subprograms
   EXTERNAL DERIV                                ! derivative


   !****************************************************************
   ! Allocate size of vectors according to the number of variables *
   !****************************************************************
   ALLOCATE( F(NDEP) )
   ALLOCATE( DY1(NDEP), DY2(NDEP), DY3(NDEP), DY4(NDEP) )


   !******************************
   ! Update, depending on METHOD *
   !******************************
   SELECT CASE( METHOD )

      CASE( 'EULER' )                  ! Euler method
         CALL DERIV( NDEP, X, Y, F );   DY1 = DX * F
         X = X + DX
         Y = Y + DY1

      CASE( 'MODIF' )                  ! Modified Euler method
         CALL DERIV( NDEP, X     , Y      , F );   DY1 = DX * F
         CALL DERIV( NDEP, X + DX, Y + DY1, F );   DY2 = DX * F
         X = X + DX
         Y = Y + 0.5 * ( DY1 + DY2 )

      CASE( 'RUNGE' )                  ! Runge-Kutta method
         CALL DERIV( NDEP, X           , Y            , F );   DY1 = DX * F
         CALL DERIV( NDEP, X + 0.5 * DX, Y + 0.5 * DY1, F );   DY2 = DX * F
         CALL DERIV( NDEP, X + 0.5 * DX, Y + 0.5 * DY2, F );   DY3 = DX * F
         CALL DERIV( NDEP, X + DX      , Y + DY3      , F );   DY4 = DX * F
         X = X + DX
         Y = Y + ( DY1 + 2.0 * DY2 + 2.0 * DY3 + DY4 ) / 6.0

      CASE DEFAULT                     ! undefined method - PANIC !!!
         PRINT *, 'ERROR: Method ', METHOD, ' is invalid'
         STOP

   END SELECT


   !**************
   ! End of step *
   !**************
   DEALLOCATE( F, DY1, DY2, DY3, DY4 )

END SUBROUTINE STEP


!=======================================================================




!=======================================================================
! EQUATION-DEPENDENT ROUTINES
! Two problems have been set up for testing
!
! Problem A: dy/dx = x + y,    y(0) = 0
!      exact solution: y = exp(x) - x - 1
!
! Problem B: d2y/dx2 + 3dy/dx + 2y = 2x,   y(0) = 1, y'(0) = 1
!      exact solution: y = 5exp(-x) - (5/2)exp(-2x) + x - 3/2
!
! The exact solutions for these test cases can be returned by function
! EXACT(). This is very useful for checking the code.
!=======================================================================


SUBROUTINE DERIV( NDEP, X, Y, F )
! Find derivative at (X,Y); return the result in F()
   IMPLICIT NONE
   INTEGER, PARAMETER :: DKIND = KIND( 1.0D0 )   ! use high precision

   ! Subroutine arguments
   INTEGER           NDEP                        ! number of variables
   REAL (KIND=DKIND) X                           ! value of X
   REAL (KIND=DKIND) Y(NDEP)                     ! value of Y
   REAL (KIND=DKIND) F(NDEP)                     ! derivative at (X,Y)


   !*********************************
   ! Problem A:     dy/dx = x + y   *
   !*********************************
   F(1) = X + Y(1)


   !**********************************************
   ! Problem B:     d2y/dx2 + 3dy/dx + 2y = 2x   *
   !**********************************************
!  F(1) = Y(2)
!  F(2) = 2.0 * X - 2.0 * Y(1) - 3.0 * Y(2)


END SUBROUTINE DERIV


!=======================================================================


SUBROUTINE OUTPUT( NDEP, X, Y )
! Writes output
   IMPLICIT NONE
   INTEGER, PARAMETER :: DKIND = KIND( 1.0D0 )   ! use high precision

   ! Subroutine arguments
   INTEGER           NDEP                        ! number of variables
   REAL (KIND=DKIND) X                           ! value of X
   REAL (KIND=DKIND) Y(NDEP)                     ! value of Y

   ! Local variables
   REAL (KIND=DKIND) YBAR                        ! exact solution
   REAL (KIND=DKIND) ERROR                       ! error

   ! Other subprograms
   REAL (KIND=DKIND), EXTERNAL :: EXACT          ! exact value (if known)

   ! Format for output
   CHARACTER (LEN=*), PARAMETER :: FMTDATA = '(4(1X, 1PE11.4))'


   !***************************************
   ! Output - amend to whatever you want! *
   !***************************************
   YBAR = EXACT( X )
   ERROR = YBAR - Y(1)
   WRITE( * , FMT = FMTDATA ) X, Y(1), YBAR, ERROR
   WRITE( 20, FMT = FMTDATA ) X, Y(1), YBAR, ERROR

END SUBROUTINE OUTPUT


!=======================================================================


FUNCTION EXACT( X )
! Returns exact value of Y(1) (if known)
   IMPLICIT NONE
   INTEGER, PARAMETER :: DKIND = KIND ( 1.0D0 )  ! use high precision

   ! Function type
   REAL (KIND=DKIND) EXACT

   ! Function arguments
   REAL (KIND=DKIND) X                           ! X value


   !******************************
   ! Exact solution of problem A *
   !******************************
   EXACT = EXP( X ) - X - 1.0


   !******************************
   ! Exact solution of problem B *
   !******************************
!  EXACT = 5.0 * EXP( -X ) - 2.5 * EXP( -2.0 * X ) + X - 1.5


END FUNCTION EXACT


!=======================================================================
