! $Id: lmdz_filtre_fft.F90 5159 2024-08-02 19:58:25Z fhourdin $

MODULE lmdz_filtre_fft
  IMPLICIT NONE; PRIVATE
  PUBLIC use_filtre_fft, filtre_u, filtre_v, filtre_inv, init_filtre_fft, filtre_u_fft, &
          filtre_v_fft, filtre_inv_fft

  LOGICAL :: use_filtre_fft
  REAL,ALLOCATABLE :: filtre_u(:,:)
  REAL,ALLOCATABLE :: filtre_v(:,:)
  REAL,ALLOCATABLE :: filtre_inv(:,:)

CONTAINS
  
  SUBROUTINE init_filtre_fft(coeffu,modfrstu,jfiltnu,jfiltsu,coeffv,modfrstv,jfiltnv,jfiltsv)
    USE lmdz_fft
USE lmdz_dimensions, ONLY: iim, jjm, llm, ndm
    IMPLICIT NONE

    REAL,   INTENT(IN) :: coeffu(iim,jjm)
    INTEGER,INTENT(IN) :: modfrstu(jjm)
    INTEGER,INTENT(IN) :: jfiltnu
    INTEGER,INTENT(IN) :: jfiltsu
    REAL,   INTENT(IN) :: coeffv(iim,jjm)
    INTEGER,INTENT(IN) :: modfrstv(jjm)
    INTEGER,INTENT(IN) :: jfiltnv
    INTEGER,INTENT(IN) :: jfiltsv
    
    INTEGER            :: index_vp(iim)
    INTEGER            :: i,j
    INTEGER            :: l,ll_nb

    index_vp(1)=1
    DO i=1,iim/2
      index_vp(i+1)=i*2
    ENDDO
    
    DO i=1,iim/2-1
      index_vp(iim/2+i+1)=iim-2*i+1
    ENDDO
    
    ALLOCATE(filtre_u(iim,jjm))
    ALLOCATE(filtre_v(iim,jjm))
    ALLOCATE(filtre_inv(iim,jjm))
  
    
    DO j=2,jfiltnu
      DO i=1,iim
        IF (index_vp(i) < modfrstu(j)) THEN
          filtre_u(i,j)=0
        ELSE
          filtre_u(i,j)=coeffu(index_vp(i),j)
        ENDIF
      ENDDO
    ENDDO
    
    DO j=jfiltsu,jjm
      DO i=1,iim
        IF (index_vp(i) < modfrstu(j)) THEN
          filtre_u(i,j)=0
        ELSE
          filtre_u(i,j)=coeffu(index_vp(i),j)
        ENDIF
      ENDDO
    ENDDO
 
    DO j=1,jfiltnv
      DO i=1,iim
        IF (index_vp(i) < modfrstv(j)) THEN
          filtre_v(i,j)=0
        ELSE
          filtre_v(i,j)=coeffv(index_vp(i),j)
        ENDIF
      ENDDO
    ENDDO
   
    DO j=jfiltsv,jjm
      DO i=1,iim
        IF (index_vp(i) < modfrstv(j)) THEN
          filtre_v(i,j)=0
        ELSE
          filtre_v(i,j)=coeffv(index_vp(i),j)
        ENDIF
      ENDDO
    ENDDO
         
    DO j=2,jfiltnu
      DO i=1,iim
        IF (index_vp(i) < modfrstu(j)) THEN
          filtre_inv(i,j)=0
        ELSE
          filtre_inv(i,j)=coeffu(index_vp(i),j)/(1.+coeffu(index_vp(i),j))
        ENDIF
      ENDDO
    ENDDO

    DO j=jfiltsu,jjm
      DO i=1,iim
        IF (index_vp(i) < modfrstu(j)) THEN
          filtre_inv(i,j)=0
        ELSE
          filtre_inv(i,j)=coeffu(index_vp(i),j)/(1.+coeffu(index_vp(i),j))
        ENDIF
      ENDDO
    ENDDO
    
#ifdef FFT_FFTW

    WRITE (*,*)"COTH jfiltnu,jfiltsu,jfiltnv,jjm-jfiltsv"
    WRITE (*,*)jfiltnu,jfiltsu,jfiltnv,jjm-jfiltsv
    WRITE (*,*)MAX(jfiltnu-2,jjm-jfiltsu,jfiltnv-2,jjm-jfiltsv)+1
    CALL init_FFT(iim,(llm+1)*(MAX(jfiltnu-2,jjm-jfiltsu,jfiltnv-2,jjm-jfiltsv)+1))
#else    
    CALL init_FFT(iim,(jjm+1)*(llm+1))
#endif        
    
  END SUBROUTINE init_filtre_fft
  
  SUBROUTINE filtre_u_fft(vect_inout,nlat,jj_begin,jj_end,nbniv)
    USE lmdz_fft
#ifdef CPP_PARA
    USE parallel_lmdz,ONLY: OMP_CHUNK
#endif
USE lmdz_dimensions, ONLY: iim, jjm, llm, ndm
    IMPLICIT NONE

    INTEGER,INTENT(IN) :: nlat
    INTEGER,INTENT(IN) :: jj_begin
    INTEGER,INTENT(IN) :: jj_end
    INTEGER,INTENT(IN) :: nbniv
    REAL,INTENT(INOUT) :: vect_inout(iim+1,nlat,nbniv)

    REAL               :: vect(iim+inc,jj_end-jj_begin+1,nbniv)
    COMPLEX         :: TF_vect(iim/2+1,jj_end-jj_begin+1,nbniv)
    INTEGER            :: nb_vect
    INTEGER :: i,j,l
    INTEGER :: ll_nb
    
    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect(i,j,ll_nb)=vect_inout(i,j+jj_begin-1,l)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT

    nb_vect=(jj_end-jj_begin+1)*ll_nb

    CALL FFT_forward(vect,TF_vect,nb_vect)

    DO l=1,ll_nb
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim/2+1
          TF_vect(i,j,l)=TF_vect(i,j,l)*filtre_u(i,jj_begin+j-1)
        ENDDO
      ENDDO
    ENDDO
  
    CALL FFT_backward(TF_vect,vect,nb_vect)
      
      
    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect_inout(i,j+jj_begin-1,l)=vect(i,j,ll_nb)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT

  END SUBROUTINE filtre_u_fft

  SUBROUTINE filtre_v_fft(vect_inout,nlat,jj_begin,jj_end,nbniv)
    USE lmdz_fft
#ifdef CPP_PARA
    USE parallel_lmdz,ONLY: OMP_CHUNK
#endif
USE lmdz_dimensions, ONLY: iim, jjm, llm, ndm
    IMPLICIT NONE

    INTEGER,INTENT(IN) :: nlat
    INTEGER,INTENT(IN) :: jj_begin
    INTEGER,INTENT(IN) :: jj_end
    INTEGER,INTENT(IN) :: nbniv
    REAL,INTENT(INOUT) :: vect_inout(iim+1,nlat,nbniv)

    REAL               :: vect(iim+inc,jj_end-jj_begin+1,nbniv)
    COMPLEX            :: TF_vect(iim/2+1,jj_end-jj_begin+1,nbniv)
    INTEGER            :: nb_vect
    INTEGER :: i,j,l
    INTEGER :: ll_nb
    
    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect(i,j,ll_nb)=vect_inout(i,j+jj_begin-1,l)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT

    
    nb_vect=(jj_end-jj_begin+1)*ll_nb

    CALL FFT_forward(vect,TF_vect,nb_vect)
  
    DO l=1,ll_nb
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim/2+1
          TF_vect(i,j,l)=TF_vect(i,j,l)*filtre_v(i,jj_begin+j-1)
        ENDDO
      ENDDO
    ENDDO
  
    CALL FFT_backward(TF_vect,vect,nb_vect)
    
    
    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect_inout(i,j+jj_begin-1,l)=vect(i,j,ll_nb)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT
  
  END SUBROUTINE filtre_v_fft

  SUBROUTINE filtre_inv_fft(vect_inout,nlat,jj_begin,jj_end,nbniv)
    USE lmdz_fft
#ifdef CPP_PARA
    USE parallel_lmdz,ONLY: OMP_CHUNK
#endif
USE lmdz_dimensions, ONLY: iim, jjm, llm, ndm
    IMPLICIT NONE

    INTEGER,INTENT(IN) :: nlat
    INTEGER,INTENT(IN) :: jj_begin
    INTEGER,INTENT(IN) :: jj_end
    INTEGER,INTENT(IN) :: nbniv
    REAL,INTENT(INOUT) :: vect_inout(iim+1,nlat,nbniv)

     REAL               :: vect(iim+inc,jj_end-jj_begin+1,nbniv)
    COMPLEX            :: TF_vect(iim/2+1,jj_end-jj_begin+1,nbniv)
    INTEGER            :: nb_vect
    INTEGER :: i,j,l
    INTEGER :: ll_nb
    
    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect(i,j,ll_nb)=vect_inout(i,j+jj_begin-1,l)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT

    nb_vect=(jj_end-jj_begin+1)*ll_nb

    CALL FFT_forward(vect,TF_vect,nb_vect)
  
    DO l=1,ll_nb
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim/2+1
          TF_vect(i,j,l)=TF_vect(i,j,l)*filtre_inv(i,jj_begin+j-1)
        ENDDO
      ENDDO
    ENDDO
  
    CALL FFT_backward(TF_vect,vect,nb_vect)

    ll_nb=0
!$OMP DO SCHEDULE(STATIC,OMP_CHUNK)
    DO l=1,nbniv
      ll_nb=ll_nb+1
      DO j=1,jj_end-jj_begin+1
        DO i=1,iim+1
          vect_inout(i,j+jj_begin-1,l)=vect(i,j,ll_nb)
        ENDDO
      ENDDO
    ENDDO
!$OMP END DO NOWAIT

  END SUBROUTINE filtre_inv_fft  
   
END MODULE lmdz_filtre_fft