!
!__________________________________________________________________________________________
!  Aleksandar Donev
!  PHY201--Runge-Kutta integrators for physics applications
!  last modified: 11/5/2000
!__________________________________________________________________________________________
!
!__________________________________________________________________________________________
!  This module provides named constants for single and double precision kinds
!__________________________________________________________________________________________
MODULE reals
      INTEGER, PARAMETER, PUBLIC :: dp = KIND (0.0D0), sp = KIND (0.0E0)
      INTEGER, PARAMETER, PUBLIC :: pr = dp
END MODULE reals
!__________________________________________________________________________________________
!
!__________________________________________________________________________________________
! This module provides Runge-Kutta procedures for integrating a system of first order
! ODE's of the form y_'=f_(x,y_), where the underscores denote an array of values (vectors).
! The main generic routine runge_kutta can be called with y_ being either a 1D, 2D or a 3D
! array. This changes nothing in the routine but the declarations, but 3D(2D) arrays are easier
! to use in Molecular Dynamics codes where we have two indices--the particle number and
! the component, and a possibly a third one to choose between position and velocity.
!
! NOTE: This is modeled after the code available on the web from the
! Alberquerque High Performance Computing Center
!
! The calling sequence for runge_kutta is:
!   runge_kutta(f,y_0,x_out,y_out,n_steps,n_points,h,method)
! where f is the above function of x and y, y_0 has the initial conditions y(0), x_out
! and y_out contain the output values for every n_points steps with a total of n_steps
! integration steps. The step size is h, and method can be either 2,3 or 4 to specify
! the order of the Runge-Kutta method used.
! NOTE: No run-time argument validation is done!
!__________________________________________________________________________________________
MODULE RK4
      USE reals
      IMPLICIT NONE
!
      INTERFACE runge_kutta
         MODULE PROCEDURE runge_kutta1D
         MODULE PROCEDURE runge_kutta2D
         MODULE PROCEDURE runge_kutta3D
      END INTERFACE
!
CONTAINS
  !___________________________________________________________
  ! This is the 1D runge_kutta integration routine.
  ! It is written using array-syntax so as to be very
  ! easy to follow.
  !___________________________________________________________
      SUBROUTINE runge_kutta1D (f, y_0, x_out, y_out, n_steps, n_points, h, method)
!
    ! Declaration of the dummy arguments.
    ! Note: We are doing compile-time size validation of the arguments:
         INTEGER, INTENT (IN) :: n_steps, n_points, method
         REAL (KIND=pr), DIMENSION (:), INTENT (IN) :: y_0
         REAL (KIND=pr), DIMENSION (Int(n_steps/n_points)), INTENT (OUT) :: x_out
         REAL (KIND=pr), DIMENSION (size(y_0), Int(n_steps/n_points)), INTENT (OUT) :: y_out
         REAL (KIND=pr), INTENT (IN) :: h
    ! This is the interface for the function f(x,y):
         INTERFACE
            FUNCTION f (x, y)
               USE reals
               REAL (KIND=pr), INTENT (IN) :: x
               REAL (KIND=pr), DIMENSION (:), INTENT (IN) :: y
               REAL (KIND=pr), DIMENSION (size(y)) :: f
            END FUNCTION f
         END INTERFACE
!
    ! Local variables:
         INTEGER :: i, j, k
         REAL (KIND=pr) :: x
         REAL (KIND=pr), DIMENSION (size(y_0)) :: y, k1, k2, k3, k4
    ! Initialization of y:
         y = y_0
         k = 0
!
    ! The solution is computed inside a do loop:
         DO i = 1, n_steps
            x = (i-1) * h
!
      ! Every n_points we save the results:
            IF (Mod(i, n_points) == 0) THEN
               k = k + 1
               x_out (k) = x
               y_out (:, k) = y
            END IF
!
      ! Depending on method, we use an appropriate Runge-Kutta:
            SELECT CASE (method)
            CASE (2)
               k1 = y + h * f (x, y) / 2
               y = y + h * f (x+h/2, k1)
            CASE (3)
               k1 = f (x, y)
               k2 = f (x+h/3, y+h*k1/3)
               k3 = f (x+2*h/3, y+2*h*k2/3)
               y = y + h * (k1+3*k3) / 4
            CASE (4)
               k1 = f (x, y)
               k2 = f (x+h/2, y+h*k1/2)
               k3 = f (x+h/2, y+h*k2/2)
               k4 = f (x+h, y+h*k3)
               y = y + h * (k1+2*k2+2*k3+k4) / 6
            END SELECT
         END DO
!
      END SUBROUTINE runge_kutta1D
  !___________________________________________________________
!
  !___________________________________________________________
  ! This is the 2D runge_kutta integration routine.
  ! It is written using array-syntax so as to be very
  ! easy to follow. Comments are given above in runge_kutta1D.
  !___________________________________________________________
      SUBROUTINE runge_kutta2D (f, y_0, x_out, y_out, n_steps, n_points, h, method)
!
         INTEGER, INTENT (IN) :: n_steps, n_points, method
         REAL (KIND=pr), DIMENSION (:, :), INTENT (IN) :: y_0
         REAL (KIND=pr), DIMENSION (Int(n_steps/n_points)), INTENT (OUT) :: x_out
         REAL (KIND=pr), DIMENSION (size(y_0, 1), size(y_0, 2), Int(n_steps/n_points)), INTENT (OUT) :: y_out
         REAL (KIND=pr), INTENT (IN) :: h
         INTERFACE
            FUNCTION f (x, y)
               USE reals
               REAL (KIND=pr), INTENT (IN) :: x
               REAL (KIND=pr), DIMENSION (:, :), INTENT (IN) :: y
               REAL (KIND=pr), DIMENSION (size(y, 1), size(y, 2)) :: f
            END FUNCTION f
         END INTERFACE
!
         INTEGER :: i, j, k
         REAL (KIND=pr) :: x
         REAL (KIND=pr), DIMENSION (size(y_0, 1), size(y_0, 2)) :: y, k1, k2, k3, k4
!
         y = y_0
         k = 0
!
         DO i = 1, n_steps
            x = (i-1) * h
!
            IF (Mod(i, n_points) == 0) THEN
               k = k + 1
               x_out (k) = x
               y_out (:, :, k) = y
            END IF
!
            SELECT CASE (method)
            CASE (2)
               k1 = y + h * f (x, y) / 2
               y = y + h * f (x+h/2, k1)
            CASE (3)
               k1 = f (x, y)
               k2 = f (x+h/3, y+h*k1/3)
               k3 = f (x+2*h/3, y+2*h*k2/3)
               y = y + h * (k1+3*k3) / 4
            CASE (4)
               k1 = f (x, y)
               k2 = f (x+h/2, y+h*k1/2)
               k3 = f (x+h/2, y+h*k2/2)
               k4 = f (x+h, y+h*k3)
               y = y + h * (k1+2*k2+2*k3+k4) / 6
            END SELECT
         END DO
!
      END SUBROUTINE runge_kutta2D
  !___________________________________________________________
!
  !___________________________________________________________
  ! This is the 3D runge_kutta integration routine.
  ! It is written using array-syntax so as to be very
  ! easy to follow. Comments are given above in runge_kutta1D.
  !___________________________________________________________
      SUBROUTINE runge_kutta3D (f, y_0, x_out, y_out, n_steps, n_points, h, method)
!
         INTEGER, INTENT (IN) :: n_steps, n_points, method
         REAL (KIND=pr), DIMENSION (:, :, :), INTENT (IN) :: y_0
         REAL (KIND=pr), DIMENSION (Int(n_steps/n_points)), INTENT (OUT) :: x_out
         REAL (KIND=pr), DIMENSION (size(y_0, 1), size(y_0, 2), size(y_0, 3), Int(n_steps/n_points)), INTENT (OUT) :: y_out
         REAL (KIND=pr), INTENT (IN) :: h
         INTERFACE
            FUNCTION f (x, y)
               USE reals
               REAL (KIND=pr), INTENT (IN) :: x
               REAL (KIND=pr), DIMENSION (:, :, :), INTENT (IN) :: y
               REAL (KIND=pr), DIMENSION (size(y, 1), size(y, 2), size(y, 3)) :: f
            END FUNCTION f
         END INTERFACE
!
         INTEGER :: i, j, k
         REAL (KIND=pr) :: x
         REAL (KIND=pr), DIMENSION (size(y_0, 1), size(y_0, 2), size(y_0, 3)) :: y, k1, k2, k3, k4
!
         y = y_0
         k = 0
!
         DO i = 1, n_steps
            x = (i-1) * h
!
            IF (Mod(i, n_points) == 0) THEN
               k = k + 1
               x_out (k) = x
               y_out (:, :, :, k) = y
            END IF
!
            SELECT CASE (method)
            CASE (2)
               k1 = y + h * f (x, y) / 2
               y = y + h * f (x+h/2, k1)
            CASE (3)
               k1 = f (x, y)
               k2 = f (x+h/3, y+h*k1/3)
               k3 = f (x+2*h/3, y+2*h*k2/3)
               y = y + h * (k1+3*k3) / 4
            CASE (4)
               k1 = f (x, y)
               k2 = f (x+h/2, y+h*k1/2)
               k3 = f (x+h/2, y+h*k2/2)
               k4 = f (x+h, y+h*k3)
               y = y + h * (k1+2*k2+2*k3+k4) / 6
            END SELECT
         END DO
!
      END SUBROUTINE runge_kutta3D
  !___________________________________________________________
!
END MODULE RK4
!__________________________________________________________________________________________
!
!