! ===================================================================================
PROGRAM Homogeneous

! ===================================================================================
! Modules to demonstrate a model of homogeneous background error covariances using
! the Fourier transform method.
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90
!
! On Unix in the Meteorology Dept
! ------------------
! setup studio12
! To compile on Unix
! f90 homogeneous.f90 -o homogeneous.out
! To run on Unix
! homogeneous.out
! ------------------
! Input files : None
! Output files: input_state (the function that the modelled B-matrix acts upon)
!               output_state (the result of B acting on the input state)
!               variance_spectrum (the imposed variance spectr
!
! Note about some compilers
! ------------------
! The complex exponential function CEXP is used here.
! Some compilers will compile with this but will not produce correct output.
! Found that replacing CEXP with EXP solves this problem.


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

  ! Declare variables

  IMPLICIT NONE

  ! Constants used in the code
  INTEGER,          PARAMETER :: Nx = 100             ! No. of x-points
  DOUBLE PRECISION, PARAMETER :: dx = 1000.0          ! Grid length (m)
  DOUBLE PRECISION, PARAMETER :: width = 0.0001       ! Width of variance spectrum (/m)
  DOUBLE PRECISION, PARAMETER :: sigma = 10.0         ! Parameter used to define variance spec
  INTEGER,          PARAMETER :: ControlPt = 50       ! Index of control point

  ! Variables used in the code
  INTEGER                     :: i, k
  DOUBLE PRECISION            :: InputState(1:Nx)     ! Exists in position space
  DOUBLE PRECISION            :: OutputState(1:Nx)    ! Exists in position space
  ! Note: OutputState = B InputState
  DOUBLE PRECISION            :: Variancespec(1:Nx)   ! Exists in spectral space
  DOUBLE COMPLEX              :: Intermediate(1:Nx)   ! Exists in spectral space
  DOUBLE PRECISION            :: GridPosns(1:Nx)      ! Values of grid posns (pos space)
  DOUBLE PRECISION            :: Wavenumbers(1:Nx)    ! Values of wavenumbers (spec space)
  DOUBLE PRECISION            :: mag



  ! **********************************************************************************
  ! Add code here to call routines to set-up the above arrays, and to call CVT CVT_hat
  ! **********************************************************************************



  ! Output data to files
  ! --------------------
  PRINT *, 'Outputting data in files input_state, output_state, variance_spectrum, intermediate'

  ! Put the input state to file as a function of grid position
  OPEN (12, file='input_state')
  DO i = 1, Nx
    WRITE (12,*) GridPosns(i), InputState(i)
  ENDDO
  CLOSE (12)

  ! Put the output state to file as a function of grid position
  OPEN (12, file='output_state')
  DO i = 1, Nx
    WRITE (12,*) GridPosns(i), OutputState(i)
  ENDDO
  CLOSE (12)

  ! Put the variance spectrum to file
  OPEN (12, file='variance_spectrum')
  DO k = 1, Nx
    WRITE (12,*) Wavenumbers(k), Variancespec(k)
  ENDDO
  CLOSE (12)

  ! Put the intermediate state to file
  OPEN (12, file='intermediate')
  DO k = 1, Nx
    mag = SQRT(REAL(Intermediate(k))**2+CMPLX(Intermediate(k))**2)
    WRITE (12,*) Wavenumbers(k), mag
  ENDDO
  CLOSE (12)


CONTAINS

! ===================================================================================
SUBROUTINE SetUp ( Nx,           & ! IN    No. of gridpoints and wavenumbers
                   dx,           & ! IN    Grid length
                   width,        & ! IN    Width of variance spectrum
                   sigma,        & ! IN    Parameter used to define variance spectrum
                   ControlPt,    & ! IN    Control point (in intput array)
                   GridPosns,    & ! INOUT Values of grid positions
                   Wavenumbers,  & ! INOUT Values of wavenumbers
                   Spec,         & ! INOUT Variance spectrum
                   InputState )    ! INOUT Prepare input state as delta fn

! Subroutine to set-up values of grid positions, values of wavenumbers and variance spec
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90

  IMPLICIT NONE

  ! Subroutine arguments
  INTEGER,          INTENT(IN)    :: Nx
  DOUBLE PRECISION, INTENT(IN)    :: dx
  DOUBLE PRECISION, INTENT(IN)    :: width
  DOUBLE PRECISION, INTENT(IN)    :: sigma
  INTEGER,          INTENT(IN)    :: ControlPt
  DOUBLE PRECISION, INTENT(INOUT) :: GridPosns(1:Nx)
  DOUBLE PRECISION, INTENT(INOUT) :: Wavenumbers(1:Nx)
  DOUBLE PRECISION, INTENT(INOUT) :: Spec(1:Nx)
  DOUBLE PRECISION, INTENT(INOUT) :: InputState(1:Nx)

  ! Local variables
  INTEGER                         :: i, id2, k
  DOUBLE PRECISION                :: pi, fac, ratio, total

  pi  = 4. * ATAN(1.)
  id2 = Nx / 2

  !Set-up the grid positions and the wavenumbers
  fac = 2. * pi / (REAL(Nx)*dx)
  DO i = 1, Nx
    GridPosns(i)   = REAL(i) * dx
    Wavenumbers(i) = fac * REAL(i - id2)
  ENDDO


  ! **********************************************************************************
  ! Add code here to set-up the variance spectrum, Spec
  ! **********************************************************************************


  ! **********************************************************************************
  ! Add code here to set-up the input state, InputState
  ! (to be acted upon by B in the main code)
  ! **********************************************************************************



END SUBROUTINE SetUp


! ===================================================================================
SUBROUTINE FT    ( Nx,           & ! IN    No. of gridpoints and wavenumbers
                   Input,        & ! IN    Input (position space)
                   Output,       & ! INOUT Output (spectral space)
                   GridPosns,    & ! IN    Values of grid positions
                   Wavenumbers )   ! IN    Values of wavenumbers

! Subroutine to do Fourier transform (inverse or adjoint of iFT)
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90

  IMPLICIT NONE

  ! Subroutine arguments
  INTEGER,            INTENT(IN)    :: Nx
  DOUBLE PRECISION,   INTENT(IN)    :: Input(1:Nx)
  DOUBLE COMPLEX,     INTENT(INOUT) :: Output(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: GridPosns(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Wavenumbers(1:Nx)

  ! Local variables
  INTEGER                           :: i, k
  DOUBLE PRECISION                  :: fac
  DOUBLE COMPLEX                    :: total, arg

  fac = 1. / SQRT(REAL(Nx))

  DO k = 1, Nx
    total  = CMPLX(0.0, 0.0)
    DO i = 1, Nx
      arg   = CMPLX(0.0, -1.0 * GridPosns(i) * Wavenumbers(k))
      total = total + Input(i) * CEXP(arg)
    ENDDO
    Output(k) = total * fac
  END DO
    
END SUBROUTINE FT




! ===================================================================================
SUBROUTINE iFT   ( Nx,           & ! IN    No. of gridpoints and wavenumbers
                   Input,        & ! IN    Input (spectral space)
                   Output,       & ! INOUT Output (position space)
                   GridPosns,    & ! IN    Values of grid positions
                   Wavenumbers )   ! IN    Values of wavenumbers

! Subroutine to do inverse Fourier transform
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90

  IMPLICIT NONE

  ! Subroutine arguments
  INTEGER,            INTENT(IN)    :: Nx
  DOUBLE COMPLEX,     INTENT(IN)    :: Input(1:Nx)
  DOUBLE PRECISION,   INTENT(INOUT) :: Output(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: GridPosns(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Wavenumbers(1:Nx)

  ! Local variables
  INTEGER                           :: i, k
  DOUBLE PRECISION                  :: fac
  DOUBLE COMPLEX                    :: total, arg

  fac = 1. / SQRT(REAL(Nx))

  DO i = 1, Nx
    total  = CMPLX(0.0, 0.0)
    DO k = 1, Nx
      arg   = CMPLX(0.0, GridPosns(i) * Wavenumbers(k))
      total = total + Input(k) * CEXP(arg)
    ENDDO
    Output(i) = REAL(total) * fac
  END DO
    
END SUBROUTINE iFT



! ===================================================================================
SUBROUTINE CVT_hat ( Nx,           & ! IN    No. of gridpoints and wavenumbers
                     Input,        & ! IN    Input (position space)
                     Output,       & ! INOUT Output (spectral space)
                     GridPosns,    & ! IN    Values of grid positions
                     Wavenumbers,  & ! IN    Values of wavenumbers
                     Spec )          ! IN    Variance spectrum

! Subroutine to do adjoint control variable transform
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90

  IMPLICIT NONE

  ! Subroutine arguments
  INTEGER,            INTENT(IN)    :: Nx
  DOUBLE PRECISION,   INTENT(IN)    :: Input(1:Nx)
  DOUBLE COMPLEX,     INTENT(INOUT) :: Output(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: GridPosns(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Wavenumbers(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Spec(1:Nx)

  ! Local variables
  INTEGER                           :: k

  ! Stage 1: Adjoint of inverse Fourier transform (this is the Fourier transform)
  CALL FT ( Nx,                 &
            Input(1:Nx),        &
            Output(1:Nx),       &
            GridPosns(1:Nx),    &
            Wavenumbers(1:Nx) )

  ! Stage 2: Multiply by the square-root of the variance spectrum
  DO k = 1, Nx
    Output(k) = Output(k) * SQRT(Spec(k))
  END DO
    
END SUBROUTINE CVT_hat



! ===================================================================================
SUBROUTINE CVT ( Nx,           & ! IN    No. of gridpoints and wavenumbers
                 Input,        & ! IN    Input (spectral space)
                 Output,       & ! INOUT Output (position space)
                 GridPosns,    & ! IN    Values of grid positions
                 Wavenumbers,  & ! IN    Values of wavenumbers
                 Spec )          ! IN    Variance spectrum

! Subroutine to do control variable transform
! Ross Bannister
! National Centre for Earth Observation
! January 2012
! FORTRAN 90

  IMPLICIT NONE

  ! Subroutine arguments
  INTEGER,            INTENT(IN)    :: Nx
  DOUBLE COMPLEX,     INTENT(IN)    :: Input(1:Nx)
  DOUBLE PRECISION,   INTENT(INOUT) :: Output(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: GridPosns(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Wavenumbers(1:Nx)
  DOUBLE PRECISION,   INTENT(IN)    :: Spec(1:Nx)

  ! Local variables
  INTEGER                           :: k
  DOUBLE COMPLEX                    :: Input1(1:Nx)

  ! Stage 1: Multiply by the square-root of the variance spectrum
  DO k = 1, Nx
    Input1(k) = Input(k) * SQRT(Spec(k))
  END DO

  ! Stage 2: Inverse Fourier transform
  CALL iFT ( Nx,                 &
             Input1(1:Nx),       &
             Output(1:Nx),       &
             GridPosns(1:Nx),    &
             Wavenumbers(1:Nx) )
    
END SUBROUTINE CVT

END PROGRAM Homogeneous
! ===================================================================================
