Skip to content
Snippets Groups Projects
dynspg_ts.F90 74.2 KiB
Newer Older
Guillaume Samson's avatar
Guillaume Samson committed
MODULE dynspg_ts

   !! Includes ROMS wd scheme with diagnostic outputs ; puu(:,:,:,Kmm) and puu(:,:,:,Krhs) updates are commented out ! 

   !!======================================================================
   !!                   ***  MODULE  dynspg_ts  ***
   !! Ocean dynamics:  surface pressure gradient trend, split-explicit scheme
   !!======================================================================
   !! History :   1.0  ! 2004-12  (L. Bessieres, G. Madec)  Original code
   !!              -   ! 2005-11  (V. Garnier, G. Madec)  optimization
   !!              -   ! 2006-08  (S. Masson)  distributed restart using iom
   !!             2.0  ! 2007-07  (D. Storkey) calls to BDY routines
   !!              -   ! 2008-01  (R. Benshila)  change averaging method
   !!             3.2  ! 2009-07  (R. Benshila, G. Madec) Complete revisit associated to vvl reactivation
   !!             3.3  ! 2010-09  (D. Storkey, E. O'Dea) update for BDY for Shelf configurations
   !!             3.3  ! 2011-03  (R. Benshila, R. Hordoir, P. Oddo) update calculation of ub_b
   !!             3.5  ! 2013-07  (J. Chanut) Switch to Forward-backward time stepping
   !!             3.6  ! 2013-11  (A. Coward) Update for z-tilde compatibility
   !!             3.7  ! 2015-11  (J. Chanut) free surface simplification
   !!              -   ! 2016-12  (G. Madec, E. Clementi) update for Stoke-Drift divergence
   !!             4.0  ! 2017-05  (G. Madec)  drag coef. defined at t-point (zdfdrg.F90)
   !!---------------------------------------------------------------------

   !!----------------------------------------------------------------------
   !!   dyn_spg_ts     : compute surface pressure gradient trend using a time-splitting scheme 
   !!   dyn_spg_ts_init: initialisation of the time-splitting scheme
   !!   ts_wgt         : set time-splitting weights for temporal averaging (or not)
   !!   ts_rst         : read/write time-splitting fields in restart file
   !!----------------------------------------------------------------------
   USE oce             ! ocean dynamics and tracers
   USE dom_oce         ! ocean space and time domain
   USE sbc_oce         ! surface boundary condition: ocean
   USE isf_oce         ! ice shelf variable (fwfisf)
   USE zdf_oce         ! vertical physics: variables
   USE zdfdrg          ! vertical physics: top/bottom drag coef.
   USE sbcapr          ! surface boundary condition: atmospheric pressure
   USE dynadv    , ONLY: ln_dynadv_vec
   USE dynvor          ! vortivity scheme indicators
   USE phycst          ! physical constants
   USE dynvor          ! vorticity term
   USE wet_dry         ! wetting/drying flux limter
   USE bdy_oce         ! open boundary
   USE bdyvol          ! open boundary volume conservation
   USE bdytides        ! open boundary condition data
   USE bdydyn2d        ! open boundary conditions on barotropic variables
   USE tide_mod        !
   USE sbcwave         ! surface wave
#if defined key_agrif
   USE agrif_oce_interp ! agrif
   USE agrif_oce
#endif
#if defined key_asminc   
   USE asminc          ! Assimilation increment
#endif
   !
   USE in_out_manager  ! I/O manager
   USE lib_mpp         ! distributed memory computing library
   USE lbclnk          ! ocean lateral boundary conditions (or mpp link)
   USE prtctl          ! Print control
   USE iom             ! IOM library
   USE restart         ! only for lrst_oce

   USE iom   ! to remove

   IMPLICIT NONE
   PRIVATE

   PUBLIC dyn_spg_ts        ! called by dyn_spg 
   PUBLIC dyn_spg_ts_init   !    -    - dyn_spg_init

   !! Time filtered arrays at baroclinic time step:
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   un_adv , vn_adv   !: Advection vel. at "now" barocl. step
   !
   INTEGER, SAVE :: icycle      ! Number of barotropic sub-steps for each internal step nn_e <= 2.5 nn_e
   REAL(dp),SAVE :: rDt_e       ! Barotropic time step
Guillaume Samson's avatar
Guillaume Samson committed
   !
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:)    :: wgtbtp2! 1st & 2nd weights used in time filtering of barotropic fields
   REAL(dp), ALLOCATABLE, SAVE, DIMENSION(:)    :: wgtbtp1! 1st & 2nd weights used in time filtering of barotropic fields
Guillaume Samson's avatar
Guillaume Samson committed
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zwz                ! ff_f/h at F points
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   ftnw, ftne         ! triad of coriolis parameter
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   ftsw, ftse         ! (only used with een vorticity scheme)

   REAL(wp) ::   r1_12 = 1._wp / 12._wp   ! local ratios
   REAL(wp) ::   r1_8  = 0.125_wp         !
   REAL(wp) ::   r1_4  = 0.25_wp          !
   REAL(wp) ::   r1_2  = 0.5_wp           !

   !! * Substitutions
#  include "do_loop_substitute.h90"
#  include "single_precision_substitute.h90"
Guillaume Samson's avatar
Guillaume Samson committed
#  include "domzgr_substitute.h90"
   !!----------------------------------------------------------------------
   !! NEMO/OCE 4.0 , NEMO Consortium (2018)
   !! $Id: dynspg_ts.F90 15489 2021-11-10 09:18:39Z jchanut $
   !! Software governed by the CeCILL license (see ./LICENSE)
   !!----------------------------------------------------------------------
CONTAINS

   INTEGER FUNCTION dyn_spg_ts_alloc()
      !!----------------------------------------------------------------------
      !!                  ***  routine dyn_spg_ts_alloc  ***
      !!----------------------------------------------------------------------
      INTEGER :: ierr(3)
      !!----------------------------------------------------------------------
      ierr(:) = 0
      !
      ALLOCATE( wgtbtp1(3*nn_e), wgtbtp2(3*nn_e), zwz(jpi,jpj), STAT=ierr(1) )
      IF( ln_dynvor_een .OR. ln_dynvor_eeT )   &
         &     ALLOCATE( ftnw(jpi,jpj) , ftne(jpi,jpj) , ftsw(jpi,jpj) , ftse(jpi,jpj), STAT=ierr(2)   )
         !
      ALLOCATE( un_adv(jpi,jpj), vn_adv(jpi,jpj)                    , STAT=ierr(3) )
      !
      dyn_spg_ts_alloc = MAXVAL( ierr(:) )
      !
      CALL mpp_sum( 'dynspg_ts', dyn_spg_ts_alloc )
      IF( dyn_spg_ts_alloc /= 0 )   CALL ctl_stop( 'STOP', 'dyn_spg_ts_alloc: failed to allocate arrays' )
      !
   END FUNCTION dyn_spg_ts_alloc


   SUBROUTINE dyn_spg_ts( kt, Kbb, Kmm, Krhs, puu, pvv, pssh, puu_b, pvv_b, Kaa, k_only_ADV )
      !!----------------------------------------------------------------------
      !!
      !! ** Purpose : - Compute the now trend due to the explicit time stepping
      !!              of the quasi-linear barotropic system, and add it to the
      !!              general momentum trend. 
      !!
      !! ** Method  : - split-explicit schem (time splitting) :
      !!      Barotropic variables are advanced from internal time steps
      !!      "n"   to "n+1" if ln_bt_fw=T
      !!      or from 
      !!      "n-1" to "n+1" if ln_bt_fw=F
      !!      thanks to a generalized forward-backward time stepping (see ref. below).
      !!
      !! ** Action :
      !!      -Update the filtered free surface at step "n+1"      : pssh(:,:,Kaa)
      !!      -Update filtered barotropic velocities at step "n+1" : puu_b(:,:,:,Kaa), vv_b(:,:,:,Kaa)
      !!      -Compute barotropic advective fluxes at step "n"     : un_adv, vn_adv
      !!      These are used to advect tracers and are compliant with discrete
      !!      continuity equation taken at the baroclinic time steps. This 
      !!      ensures tracers conservation.
      !!      - (puu(:,:,:,Krhs), pvv(:,:,:,Krhs)) momentum trend updated with barotropic component.
      !!
      !! References : Shchepetkin and McWilliams, Ocean Modelling, 2005. 
      !!---------------------------------------------------------------------
      INTEGER                             , INTENT( in )  ::  kt                  ! ocean time-step index
      INTEGER                             , INTENT( in )  ::  Kbb, Kmm, Krhs, Kaa ! ocean time level indices
      REAL(dp), DIMENSION(jpi,jpj,jpk,jpt), INTENT(inout) ::  puu, pvv            ! ocean velocities and RHS of momentum equation
      REAL(wp), DIMENSION(jpi,jpj,jpt)    , INTENT(inout)  :: puu_b, pvv_b! SSH and barotropic velocities at main time levels
      REAL(dp), DIMENSION(jpi,jpj,jpt)    , INTENT(inout)  :: pssh! SSH and barotropic velocities at main time levels
Guillaume Samson's avatar
Guillaume Samson committed
      INTEGER , OPTIONAL                  , INTENT( in )  ::  k_only_ADV          ! only Advection in the RHS
      !
      INTEGER  ::   ji, jj, jk, jn        ! dummy loop indices
      LOGICAL  ::   ll_fw_start           ! =T : forward integration 
      LOGICAL  ::   ll_init               ! =T : special startup of 2d equations
      INTEGER  ::   noffset               ! local integers  : time offset for bdy update
      REAL(wp) ::   r1_Dt_b, z1_hu, z1_hv          ! local scalars
      REAL(wp)  :: za0, za2, za3!   -      -
      REAL(dp)  :: za1!   -      -
Guillaume Samson's avatar
Guillaume Samson committed
      REAL(wp) ::   zztmp, zldg               !   -      -
      REAL(wp)  :: zhu_bck, zhv_bck!   -      -
      REAL(dp)  :: zhdiv!   -      -
Guillaume Samson's avatar
Guillaume Samson committed
      REAL(wp) ::   zun_save, zvn_save              !   -      -
      REAL(wp), DIMENSION(jpi,jpj) :: zu_trd, zu_frc, zu_spg
Guillaume Samson's avatar
Guillaume Samson committed
      REAL(wp), DIMENSION(jpi,jpj) :: zv_trd, zv_frc, zv_spg
      REAL(wp), DIMENSION(jpi,jpj) :: zsshu_a, zhup2_e, zhtp2_e
      REAL(wp), DIMENSION(jpi,jpj) :: zsshv_a, zhvp2_e, zsshp2_e
      REAL(wp), DIMENSION(jpi,jpj) :: zCdU_u, zCdU_v   ! top/bottom stress at u- & v-points
      REAL(wp), DIMENSION(jpi,jpj)  :: zhV! fluxes
      REAL(wp), DIMENSION(jpi,jpj)  :: zhU! fluxes
Guillaume Samson's avatar
Guillaume Samson committed
!!st#if defined key_qco 
!!st      REAL(wp), DIMENSION(jpi, jpj, jpk) :: ze3u, ze3v
!!st#endif
      !
      REAL(wp) ::   zwdramp                     ! local scalar - only used if ln_wd_dl = .True. 

      INTEGER  :: iwdg, jwdg, kwdg   ! short-hand values for the indices of the output point

      REAL(wp) ::   zepsilon, zgamma            !   -      -
      REAL(wp), ALLOCATABLE, DIMENSION(:,:) :: ztwdmask, zuwdmask, zvwdmask ! ROMS wetting and drying masks at t,u,v points
      REAL(wp), ALLOCATABLE, DIMENSION(:,:) :: zuwdav2, zvwdav2    ! averages over the sub-steps of zuwdmask and zvwdmask
      REAL(wp) ::   zt0substep !   Time of day at the beginning of the time substep
      !!----------------------------------------------------------------------
      !
      !                                         !* Allocate temporary arrays
      IF( ln_wd_dl ) ALLOCATE( ztwdmask(jpi,jpj), zuwdmask(jpi,jpj), zvwdmask(jpi,jpj), zuwdav2(jpi,jpj), zvwdav2(jpi,jpj))
      !
      zwdramp = r_rn_wdmin1               ! simplest ramp 
!     zwdramp = 1._wp / (rn_wdmin2 - rn_wdmin1) ! more general ramp
      !                                         ! inverse of baroclinic time step 
      r1_Dt_b = 1._wp / rDt 
      !
      ll_init     = ln_bt_av                    ! if no time averaging, then no specific restart 
      ll_fw_start = .FALSE.
      !                                         ! time offset in steps for bdy data update
      IF( .NOT.ln_bt_fw ) THEN   ;   noffset = - nn_e
      ELSE                       ;   noffset =   0 
      ENDIF
      !
      IF( kt == nit000 ) THEN                   !* initialisation
         !
         IF(lwp) WRITE(numout,*)
         IF(lwp) WRITE(numout,*) 'dyn_spg_ts : surface pressure gradient trend'
         IF(lwp) WRITE(numout,*) '~~~~~~~~~~   free surface with time splitting'
         IF(lwp) WRITE(numout,*)
         !
         IF( l_1st_euler )   ll_init=.TRUE.
         !
         IF( ln_bt_fw .OR. l_1st_euler ) THEN
            ll_fw_start =.TRUE.
            noffset     = 0
         ELSE
            ll_fw_start =.FALSE.
         ENDIF
         !                    ! Set averaging weights and cycle length:
         CALL ts_wgt( ln_bt_av, ll_fw_start, icycle, wgtbtp1, wgtbtp2 )
         !
      ELSEIF( kt == nit000 + 1 ) THEN           !* initialisation 2nd time-step
         !
         IF( .NOT.ln_bt_fw ) THEN
            ! If we did an Euler timestep on the first timestep we need to reset ll_fw_start
            ! and the averaging weights. We don't have an easy way of telling whether we did
            ! an Euler timestep on the first timestep (because l_1st_euler is reset to .false.
            ! at the end of the first timestep) so just do this in all cases. 
            ll_fw_start = .FALSE.
            CALL ts_wgt( ln_bt_av, ll_fw_start, icycle, wgtbtp1, wgtbtp2 )
         ENDIF
         !
      ENDIF
      !
      ! -----------------------------------------------------------------------------
      !  Phase 1 : Coupling between general trend and barotropic estimates (1st step)
      ! -----------------------------------------------------------------------------
      !      
      !
      !                                   !=  zu_frc =  1/H e3*d/dt(Ua)  =!  (Vertical mean of Ua, the 3D trends)
      !                                   !  ---------------------------  !
#if defined key_qco 
      zu_frc(:,:) = SUM( e3u_0(:,:,:  ) * puu(:,:,:,Krhs) * umask(:,:,:), DIM=3 ) * r1_hu_0(:,:)
      zv_frc(:,:) = SUM( e3v_0(:,:,:  ) * pvv(:,:,:,Krhs) * vmask(:,:,:), DIM=3 ) * r1_hv_0(:,:)
#else
      zu_frc(:,:) = SUM( e3u(:,:,:,Kmm) * puu(:,:,:,Krhs) * umask(:,:,:), DIM=3 ) * r1_hu(:,:,Kmm)
      zv_frc(:,:) = SUM( e3v(:,:,:,Kmm) * pvv(:,:,:,Krhs) * vmask(:,:,:), DIM=3 ) * r1_hv(:,:,Kmm)
#endif 
      !
      !
      !                                   !=  U(Krhs) => baroclinic trend  =!   (remove its vertical mean)
      DO jk = 1, jpkm1                    !  -----------------------------  !
         puu(:,:,jk,Krhs) = ( puu(:,:,jk,Krhs) - zu_frc(:,:) ) * umask(:,:,jk)
         pvv(:,:,jk,Krhs) = ( pvv(:,:,jk,Krhs) - zv_frc(:,:) ) * vmask(:,:,jk)
      END DO
      
!!gm  Question here when removing the Vertically integrated trends, we remove the vertically integrated NL trends on momentum....
!!gm  Is it correct to do so ?   I think so...
      
      !                                   !=  remove 2D Coriolis trend  =!
      !                                   !  --------------------------  !
      !
      IF( kt == nit000 .OR. .NOT. ln_linssh )   CALL dyn_cor_2D_init( Kmm )   ! Set zwz, the barotropic Coriolis force coefficient
      !                      ! recompute zwz = f/depth  at every time step for (.NOT.ln_linssh) as the water colomn height changes
      !
      IF( .NOT. PRESENT(k_only_ADV) ) THEN   !* remove the 2D Coriolis trend  
         zhU(:,:) = puu_b(:,:,Kmm) * hu(:,:,Kmm) * e2u(:,:)        ! now fluxes 
         zhV(:,:) = pvv_b(:,:,Kmm) * hv(:,:,Kmm) * e1v(:,:)        ! NB: FULL domain : put a value in last row and column
         !
         CALL dyn_cor_2d( CASTSP(ht(:,:)), hu(:,:,Kmm), hv(:,:,Kmm), puu_b(:,:,Kmm), pvv_b(:,:,Kmm), zhU, zhV,  &   ! <<== in
Guillaume Samson's avatar
Guillaume Samson committed
            &                                                                          zu_trd, zv_trd   )   ! ==>> out
         !
         DO_2D( 0, 0, 0, 0 )                          ! Remove coriolis term (and possibly spg) from barotropic trend
            zu_frc(ji,jj) = zu_frc(ji,jj) - zu_trd(ji,jj) * ssumask(ji,jj)
            zv_frc(ji,jj) = zv_frc(ji,jj) - zv_trd(ji,jj) * ssvmask(ji,jj)
         END_2D
      ELSE
         ! Ensure zhU and zhV are initialised to SOMETHING at all points to avoid referencing
         ! uninitialsed values in halos later on!
         zhU(:,:) = 0._wp
         zhV(:,:) = 0._wp
Guillaume Samson's avatar
Guillaume Samson committed
      ENDIF
      !
      !                                   !=  Add bottom stress contribution from baroclinic velocities  =!
      !                                   !  -----------------------------------------------------------  !
      IF( PRESENT(k_only_ADV) ) THEN         !* only Advection in the RHS : provide the barotropic bottom drag coefficients
         DO_2D( 0, 0, 0, 0 )
            zCdU_u(ji,jj) = r1_2*( rCdU_bot(ji+1,jj)+rCdU_bot(ji,jj) )
            zCdU_v(ji,jj) = r1_2*( rCdU_bot(ji,jj+1)+rCdU_bot(ji,jj) )
         END_2D
      ELSE				     !* remove baroclinic drag AND provide the barotropic drag coefficients
         CALL dyn_drg_init( Kbb, Kmm, puu, pvv, puu_b, pvv_b, zu_frc, zv_frc, zCdU_u, zCdU_v )
      ENDIF
      !
      !                                   !=  Add atmospheric pressure forcing  =!
      !                                   !  ----------------------------------  !
      IF( ln_apr_dyn ) THEN
         IF( ln_bt_fw ) THEN                          ! FORWARD integration: use kt+1/2 pressure (NOW+1/2)
            DO_2D( 0, 0, 0, 0 )
               zu_frc(ji,jj) = zu_frc(ji,jj) + grav * (  ssh_ib (ji+1,jj  ) - ssh_ib (ji,jj) ) * r1_e1u(ji,jj)
               zv_frc(ji,jj) = zv_frc(ji,jj) + grav * (  ssh_ib (ji  ,jj+1) - ssh_ib (ji,jj) ) * r1_e2v(ji,jj)
            END_2D
         ELSE                                         ! CENTRED integration: use kt-1/2 + kt+1/2 pressure (NOW)
            zztmp = grav * r1_2
            DO_2D( 0, 0, 0, 0 )
               zu_frc(ji,jj) = zu_frc(ji,jj) + zztmp * (  ssh_ib (ji+1,jj  ) - ssh_ib (ji,jj)  &
                    &                                   + ssh_ibb(ji+1,jj  ) - ssh_ibb(ji,jj)  ) * r1_e1u(ji,jj)
               zv_frc(ji,jj) = zv_frc(ji,jj) + zztmp * (  ssh_ib (ji  ,jj+1) - ssh_ib (ji,jj)  &
                    &                                   + ssh_ibb(ji  ,jj+1) - ssh_ibb(ji,jj)  ) * r1_e2v(ji,jj)
            END_2D
         ENDIF
      ENDIF
      !
      !                                   !=  Add wind forcing  =!
      !                                   !  ------------------  !
      IF( ln_bt_fw ) THEN
         DO_2D( 0, 0, 0, 0 )
            zu_frc(ji,jj) =  zu_frc(ji,jj) + r1_rho0 * utau(ji,jj) * r1_hu(ji,jj,Kmm)
            zv_frc(ji,jj) =  zv_frc(ji,jj) + r1_rho0 * vtau(ji,jj) * r1_hv(ji,jj,Kmm)
         END_2D
      ELSE
         zztmp = r1_rho0 * r1_2
         DO_2D( 0, 0, 0, 0 )
            zu_frc(ji,jj) =  zu_frc(ji,jj) + zztmp * ( utau_b(ji,jj) + utau(ji,jj) ) * r1_hu(ji,jj,Kmm)
            zv_frc(ji,jj) =  zv_frc(ji,jj) + zztmp * ( vtau_b(ji,jj) + vtau(ji,jj) ) * r1_hv(ji,jj,Kmm)
         END_2D
      ENDIF  
      !
      !              !----------------!
      !              !==  sssh_frc  ==!   Right-Hand-Side of the barotropic ssh equation   (over the FULL domain)
      !              !----------------!
      !                                   !=  Net water flux forcing applied to a water column  =!
      !                                   ! ---------------------------------------------------  !
      IF (ln_bt_fw) THEN                          ! FORWARD integration: use kt+1/2 fluxes (NOW+1/2)
         ssh_frc(:,:) = r1_rho0 * ( emp(:,:) - rnf(:,:) - fwfisf_cav(:,:) - fwfisf_par(:,:) )
Guillaume Samson's avatar
Guillaume Samson committed
      ELSE                                        ! CENTRED integration: use kt-1/2 + kt+1/2 fluxes (NOW)
         zztmp = r1_rho0 * r1_2
         ssh_frc(:,:) = zztmp * (   emp(:,:)        + emp_b(:,:)          &
            &                     - rnf(:,:)        - rnf_b(:,:)          &
            &                     - fwfisf_cav(:,:) - fwfisf_cav_b(:,:)   &
            &                     - fwfisf_par(:,:) - fwfisf_par_b(:,:)   )
Guillaume Samson's avatar
Guillaume Samson committed
      ENDIF
      !                                   !=  Add Stokes drift divergence  =!   (if exist)
      IF( ln_sdw ) THEN                   !  -----------------------------  !
         ssh_frc(:,:) = ssh_frc(:,:) + div_sd(:,:)
Guillaume Samson's avatar
Guillaume Samson committed
      ENDIF
      !
      !                                         ! ice sheet coupling
      IF ( ln_isf .AND. ln_isfcpl ) THEN
         !
         ! ice sheet coupling
         IF( ln_rstart .AND. kt == nit000 ) THEN
            ssh_frc(:,:) = ssh_frc(:,:) + risfcpl_ssh(:,:)
Guillaume Samson's avatar
Guillaume Samson committed
         END IF
         !
         ! conservation option
         IF( ln_isfcpl_cons ) THEN
            ssh_frc(:,:) = ssh_frc(:,:) + risfcpl_cons_ssh(:,:)
Guillaume Samson's avatar
Guillaume Samson committed
         END IF
         !
      END IF
      !
#if defined key_asminc
      !                                   !=  Add the IAU weighted SSH increment  =!
      !                                   !  ------------------------------------  !
      IF( lk_asminc .AND. ln_sshinc .AND. ln_asmiau ) THEN
         ssh_frc(:,:) = ssh_frc(:,:) - ssh_iau(:,:)
Guillaume Samson's avatar
Guillaume Samson committed
      ENDIF
#endif
      !                                   != Fill boundary data arrays for AGRIF
      !                                   ! ------------------------------------
#if defined key_agrif
         IF( .NOT.Agrif_Root() ) CALL agrif_dta_ts( kt )
#endif
      !
      ! -----------------------------------------------------------------------
      !  Phase 2 : Integration of the barotropic equations 
      ! -----------------------------------------------------------------------
      !
      !                                             ! ==================== !
      !                                             !    Initialisations   !
      !                                             ! ==================== !  
      ! Initialize barotropic variables:      
      IF( ll_init )THEN
         sshbb_e(:,:) = 0._wp
         ubb_e  (:,:) = 0._wp
         vbb_e  (:,:) = 0._wp
         sshb_e (:,:) = 0._wp
         ub_e   (:,:) = 0._wp
         vb_e   (:,:) = 0._wp
      ENDIF
      !
      IF( ln_linssh ) THEN    ! mid-step ocean depth is fixed (hup2_e=hu_n=hu_0)
         zhup2_e(:,:) = hu_0(:,:)
         zhvp2_e(:,:) = hv_0(:,:)
         zhtp2_e(:,:) = ht_0(:,:)
      ENDIF
      !
      IF( ln_bt_fw ) THEN                 ! FORWARD integration: start from NOW fields                    
         sshn_e(:,:) =    pssh (:,:,Kmm)            
         un_e  (:,:) =    puu_b(:,:,Kmm)            
         vn_e  (:,:) =    pvv_b(:,:,Kmm)
         !
         hu_e  (:,:) =    hu(:,:,Kmm)       
         hv_e  (:,:) =    hv(:,:,Kmm) 
         hur_e (:,:) = r1_hu(:,:,Kmm)    
         hvr_e (:,:) = r1_hv(:,:,Kmm)
      ELSE                                ! CENTRED integration: start from BEFORE fields
         sshn_e(:,:) =    pssh (:,:,Kbb)
         un_e  (:,:) =    puu_b(:,:,Kbb)         
         vn_e  (:,:) =    pvv_b(:,:,Kbb)
         !
         hu_e  (:,:) =    hu(:,:,Kbb)       
         hv_e  (:,:) =    hv(:,:,Kbb) 
         hur_e (:,:) = r1_hu(:,:,Kbb)    
         hvr_e (:,:) = r1_hv(:,:,Kbb)
      ENDIF
      !
      ! Initialize sums:
      puu_b (:,:,Kaa) = 0._wp       ! After barotropic velocities (or transport if flux form)          
      pvv_b (:,:,Kaa) = 0._wp
      pssh  (:,:,Kaa) = 0._wp       ! Sum for after averaged sea level
      un_adv(:,:)     = 0._wp       ! Sum for now transport issued from ts loop
      vn_adv(:,:)     = 0._wp
      !
      IF( ln_wd_dl ) THEN
         zuwdmask(:,:) = 0._wp  ! set to zero for definiteness (not sure this is necessary) 
         zvwdmask(:,:) = 0._wp  ! 
         zuwdav2 (:,:) = 0._wp 
         zvwdav2 (:,:) = 0._wp   
      END IF 

      !                                             ! ==================== !
      DO jn = 1, icycle                             !  sub-time-step loop  !
         !                                          ! ==================== !
         !
         l_full_nf_update = jn == icycle   ! false: disable full North fold update (performances) for jn = 1 to icycle-1
         !
         !                    !==  Update the forcing ==! (BDY and tides)
         !
         IF( ln_bdy      .AND. ln_tide )   CALL bdy_dta_tides( kt, kit=jn, pt_offset= REAL(noffset+1,wp) )
         ! Update tide potential at the beginning of current time substep
         IF( ln_tide_pot .AND. ln_tide ) THEN
            zt0substep = REAL(nsec_day, wp) - 0.5_wp*rn_Dt + (jn + noffset - 1) * rn_Dt / REAL(nn_e, wp)
            CALL upd_tide(zt0substep, Kmm)
         END IF
         !
         !                    !==  extrapolation at mid-step  ==!   (jn+1/2)
         !
         !                       !* Set extrapolation coefficients for predictor step:
         IF ((jn<3).AND.ll_init) THEN      ! Forward           
           za1 = 1._wp                                          
           za2 = 0._wp                        
           za3 = 0._wp                        
         ELSE                              ! AB3-AM4 Coefficients: bet=0.281105 
           za1 =  1.781105_wp              ! za1 =   3/2 +   bet
           za2 = -1.06221_wp               ! za2 = -(1/2 + 2*bet)
           za3 =  0.281105_wp              ! za3 = bet
         ENDIF
         !
         !                       !* Extrapolate barotropic velocities at mid-step (jn+1/2)
         !--        m+1/2               m                m-1           m-2       --!
         !--       u      = (3/2+beta) u   -(1/2+2beta) u      + beta u          --!
         !-------------------------------------------------------------------------!
         ua_e(:,:) = za1 * un_e(:,:) + za2 * ub_e(:,:) + za3 * ubb_e(:,:)
         va_e(:,:) = za1 * vn_e(:,:) + za2 * vb_e(:,:) + za3 * vbb_e(:,:)

         IF( .NOT.ln_linssh ) THEN                        !* Update ocean depth (variable volume case only)
            !                                             !  ------------------
            ! Extrapolate Sea Level at step jit+0.5:
            !--         m+1/2                 m                  m-1             m-2       --!
            !--      ssh      = (3/2+beta) ssh   -(1/2+2beta) ssh      + beta ssh          --!
            !--------------------------------------------------------------------------------!
            zsshp2_e(:,:) = za1 * sshn_e(:,:)  + za2 * sshb_e(:,:) + za3 * sshbb_e(:,:)
            
            ! set wetting & drying mask at tracer points for this barotropic mid-step
            IF( ln_wd_dl )   CALL wad_tmsk( zsshp2_e, ztwdmask )
            !
            !                          ! ocean t-depth at mid-step
            zhtp2_e(:,:) = ht_0(:,:) + zsshp2_e(:,:)
            !
            !                          ! ocean u- and v-depth at mid-step   (separate DO-loops remove the need of a lbc_lnk)
#if defined key_qcoTest_FluxForm
            !                                ! 'key_qcoTest_FluxForm' : simple ssh average
            DO_2D( 1, 0, 1, 1 )   ! not jpi-column
               zhup2_e(ji,jj) = hu_0(ji,jj) + r1_2 * (  zsshp2_e(ji,jj) + zsshp2_e(ji+1,jj  )  ) * ssumask(ji,jj)
            END_2D
            DO_2D( 1, 1, 1, 0 )
               zhvp2_e(ji,jj) = hv_0(ji,jj) + r1_2 * (  zsshp2_e(ji,jj) + zsshp2_e(ji  ,jj+1)  ) * ssvmask(ji,jj)
            END_2D
#else
            !                                ! no 'key_qcoTest_FluxForm' : surface weighted ssh average
            DO_2D( 1, 0, 1, 1 )   ! not jpi-column
               zhup2_e(ji,jj) = hu_0(ji,jj) + r1_2 * r1_e1e2u(ji,jj)                        &
                    &                              * (  e1e2t(ji  ,jj) * zsshp2_e(ji  ,jj)  &
                    &                                 + e1e2t(ji+1,jj) * zsshp2_e(ji+1,jj)  ) * ssumask(ji,jj)
            END_2D
            DO_2D( 1, 1, 1, 0 )   ! not jpj-row
               zhvp2_e(ji,jj) = hv_0(ji,jj) + r1_2 * r1_e1e2v(ji,jj)                        &
                    &                              * (  e1e2t(ji,jj  ) * zsshp2_e(ji,jj  )  &
                    &                                 + e1e2t(ji,jj+1) * zsshp2_e(ji,jj+1)  ) * ssvmask(ji,jj)
            END_2D
#endif               
            !
         ENDIF
         !
         !                    !==  after SSH  ==!   (jn+1)
         !
         !                             ! update (ua_e,va_e) to enforce volume conservation at open boundaries
         !                             ! values of zhup2_e and zhvp2_e on the halo are not needed in bdy_vol2d
         IF( ln_bdy .AND. ln_vol ) CALL bdy_vol2d( kt, jn, ua_e, va_e, zhup2_e, zhvp2_e )
         !      
         !                             ! resulting flux at mid-step (not over the full domain)
Guillaume Samson's avatar
Guillaume Samson committed
         DO_2D( 1, 0, 1, 1 )   ! not jpi-column
            zhU(ji,jj) = e2u(ji,jj) * ua_e(ji,jj) * zhup2_e(ji,jj)
         END_2D
         DO_2D( 1, 1, 1, 0 )   ! not jpj-row
            zhV(ji,jj) = e1v(ji,jj) * va_e(ji,jj) * zhvp2_e(ji,jj)
         END_2D
#if defined key_agrif
         ! Set fluxes during predictor step to ensure volume conservation
         IF( ln_bt_fw )   CALL agrif_dyn_ts_flux( jn, zhU, zhV )
Guillaume Samson's avatar
Guillaume Samson committed
#endif

         IF( ln_wd_dl ) THEN           ! un_e and vn_e are set to zero at faces where 
            !                          ! the direction of the flow is from dry cells
            CALL wad_Umsk( ztwdmask, zhU, zhV, un_e, vn_e, zuwdmask, zvwdmask )   ! not jpi colomn for U, not jpj row for V
            !
         ENDIF    

         ! It seems safest to do this here since zhU and zhV are not initially calculated in halos
         ! by this code or by wad_Umsk, but halo values (ji-1 and jj-1) ARE required in the zhdiv 
         ! sea level calculation. The current trunk (Feb 2024) has resolved all these issues by rewriting.
         CALL lbc_lnk( 'dynspg_ts',  zhU, 'U', -1._wp)
         CALL lbc_lnk( 'dynspg_ts',  zhV, 'V', -1._wp)
Guillaume Samson's avatar
Guillaume Samson committed
         !
         !
         !     Compute Sea Level at step jit+1
         !--           m+1        m                               m+1/2          --!
         !--        ssh    =  ssh   - delta_t' * [ frc + div( flux      ) ]      --!
         !-------------------------------------------------------------------------!
         DO_2D( 0, 0, 0, 0 )
            zhdiv = (   zhU(ji,jj) - zhU(ji-1,jj) + zhV(ji,jj) - zhV(ji,jj-1)   ) * r1_e1e2t(ji,jj)
            ssha_e(ji,jj) = (  sshn_e(ji,jj) - rDt_e * ( ssh_frc(ji,jj) + zhdiv )  ) * ssmask(ji,jj)
Guillaume Samson's avatar
Guillaume Samson committed
         END_2D

         CALL lbc_lnk( 'dynspg_ts', ssha_e, 'T', 1._dp)

Guillaume Samson's avatar
Guillaume Samson committed
         !
         ! Duplicate sea level across open boundaries (this is only cosmetic if linssh=T)
         IF( ln_bdy )   CALL bdy_ssh( ssha_e )
#if defined key_agrif
Guillaume Samson's avatar
Guillaume Samson committed
#endif
Guillaume Samson's avatar
Guillaume Samson committed
         !
         !                             ! Sum over sub-time-steps to compute advective velocities
         za2 = wgtbtp2(jn)             ! zhU, zhV hold fluxes extrapolated at jn+0.5
         un_adv(:,:) = un_adv(:,:) + za2 * zhU(:,:) * r1_e2u(:,:)
         vn_adv(:,:) = vn_adv(:,:) + za2 * zhV(:,:) * r1_e1v(:,:)
         ! sum over sub-time-steps to decide which baroclinic velocities to set to zero (zuwdav2 is only used when ln_wd_dl_bc=True) 
         IF ( ln_wd_dl_bc ) THEN
            DO_2D( 1, 0, 1, 1 )   ! not jpi-column
               zuwdav2(ji,jj) = zuwdav2(ji,jj) + za2 * zuwdmask(ji,jj)
            END_2D
            DO_2D( 1, 1, 1, 0 )   ! not jpj-row
               zvwdav2(ji,jj) = zvwdav2(ji,jj) + za2 * zvwdmask(ji,jj)
            END_2D
         END IF
         !
         !  
         ! Sea Surface Height at u-,v-points (vvl case only)
         IF( .NOT.ln_linssh ) THEN
Guillaume Samson's avatar
Guillaume Samson committed
#if defined key_qcoTest_FluxForm
            !                                ! 'key_qcoTest_FluxForm' : simple ssh average
            DO_2D( 1, 0, 1, 1 )
               zsshu_a(ji,jj) = r1_2 * (  ssha_e(ji,jj) + ssha_e(ji+1,jj  )  ) * ssumask(ji,jj)
            END_2D
            DO_2D( 1, 1, 1, 0 )
               zsshv_a(ji,jj) = r1_2 * (  ssha_e(ji,jj) + ssha_e(ji  ,jj+1)  ) * ssvmask(ji,jj)
            END_2D
#else
            DO_2D( 0, 0, 0, 0 )
               zsshu_a(ji,jj) = r1_2 * r1_e1e2u(ji,jj) * ( e1e2t(ji  ,jj  ) * ssha_e(ji  ,jj  )   &
                  &                                      + e1e2t(ji+1,jj  ) * ssha_e(ji+1,jj  ) ) * ssumask(ji,jj)
               zsshv_a(ji,jj) = r1_2 * r1_e1e2v(ji,jj) * ( e1e2t(ji  ,jj  ) * ssha_e(ji  ,jj  )   &
                  &                                      + e1e2t(ji  ,jj+1) * ssha_e(ji  ,jj+1) ) * ssvmask(ji,jj)
            END_2D
#endif
         ENDIF
         !         
         ! Half-step back interpolation of SSH for surface pressure computation at step jit+1/2
         !--            m+1/2           m+1              m               m-1              m-2     --!
         !--        ssh'    =  za0 * ssh     +  za1 * ssh   +  za2 * ssh      +  za3 * ssh        --!
         !------------------------------------------------------------------------------------------!
         CALL ts_bck_interp( jn, ll_init, za0, za1, za2, za3 )   ! coeficients of the interpolation
         zsshp2_e(:,:) = za0 *  ssha_e(:,:) + za1 *  sshn_e (:,:)   &
            &          + za2 *  sshb_e(:,:) + za3 *  sshbb_e(:,:)
         !
         !                             ! Surface pressure gradient
         zldg = ( 1._wp - rn_scal_load ) * grav    ! local factor
         DO_2D( 0, 0, 0, 0 )
            zu_spg(ji,jj) = - zldg * ( zsshp2_e(ji+1,jj) - zsshp2_e(ji,jj) ) * r1_e1u(ji,jj)
            zv_spg(ji,jj) = - zldg * ( zsshp2_e(ji,jj+1) - zsshp2_e(ji,jj) ) * r1_e2v(ji,jj)
         END_2D
         !
         ! Add Coriolis trend:
         ! zwz array below or triads normally depend on sea level with ln_linssh=F and should be updated
         ! at each time step. We however keep them constant here for optimization.
         ! Recall that zhU and zhV hold fluxes at jn+0.5 (extrapolated not backward interpolated)
         CALL dyn_cor_2d( zhtp2_e, zhup2_e, zhvp2_e, ua_e, va_e, zhU, zhV,    zu_trd, zv_trd   )
         !
         ! Add tidal astronomical forcing if defined
         IF ( ln_tide .AND. ln_tide_pot ) THEN
            DO_2D( 0, 0, 0, 0 )
               zu_trd(ji,jj) = zu_trd(ji,jj) + grav * ( pot_astro(ji+1,jj) - pot_astro(ji,jj) ) * r1_e1u(ji,jj)
               zv_trd(ji,jj) = zv_trd(ji,jj) + grav * ( pot_astro(ji,jj+1) - pot_astro(ji,jj) ) * r1_e2v(ji,jj)
            END_2D
         ENDIF
         !
         ! Add bottom stresses:
!jth do implicitly instead
         IF ( .NOT. ll_wd ) THEN ! Revert to explicit for bit comparison tests in non wad runs
            DO_2D( 0, 0, 0, 0 )
               zu_trd(ji,jj) = zu_trd(ji,jj) + zCdU_u(ji,jj) * un_e(ji,jj) * hur_e(ji,jj)
               zv_trd(ji,jj) = zv_trd(ji,jj) + zCdU_v(ji,jj) * vn_e(ji,jj) * hvr_e(ji,jj)
            END_2D
         ENDIF
         !
         ! Set next velocities:
         !     Compute barotropic speeds at step jit+1    (h : total height of the water colomn)
         !--                              VECTOR FORM
         !--   m+1                 m               /                                                       m+1/2           \    --!
         !--  u     =             u   + delta_t' * \         (1-r)*g * grad_x( ssh') -         f * k vect u      +     frc /    --!
         !--                                                                                                                    --!
         !--                             FLUX FORM                                                                              --!
         !--  m+1   __1__  /  m    m               /  m+1/2                             m+1/2              m+1/2    n      \ \  --!
         !-- u    =   m+1 |  h  * u   + delta_t' * \ h     * (1-r)*g * grad_x( ssh') - h     * f * k vect u      + h * frc /  | --!
         !--         h     \                                                                                                 /  --!
         !------------------------------------------------------------------------------------------------------------------------!
Guillaume Samson's avatar
Guillaume Samson committed
         IF( ln_dynadv_vec .OR. ln_linssh ) THEN      !* Vector form
            DO_2D( 0, 0, 0, 0 )
               ua_e(ji,jj) = (                                 un_e(ji,jj)   & 
                         &     + rDt_e * (                   zu_spg(ji,jj)   &
                         &                                 + zu_trd(ji,jj)   &
                         &                                 + zu_frc(ji,jj) ) & 
                         &   ) * ssumask(ji,jj)

               va_e(ji,jj) = (                                 vn_e(ji,jj)   &
                         &     + rDt_e * (                   zv_spg(ji,jj)   &
                         &                                 + zv_trd(ji,jj)   &
                         &                                 + zv_frc(ji,jj) ) &
                         &   ) * ssvmask(ji,jj)
            END_2D
            !
         ELSE                           !* Flux form
            DO_2D( 0, 0, 0, 0 )
               !                    ! hu_e, hv_e hold depth at jn,  zhup2_e, zhvp2_e hold extrapolated depth at jn+1/2
               !                    ! backward interpolated depth used in spg terms at jn+1/2
#if defined key_qcoTest_FluxForm
            !                                ! 'key_qcoTest_FluxForm' : simple ssh average
               zhu_bck = hu_0(ji,jj) + r1_2 * (  zsshp2_e(ji,jj) + zsshp2_e(ji+1,jj  )  ) * ssumask(ji,jj)
               zhv_bck = hv_0(ji,jj) + r1_2 * (  zsshp2_e(ji,jj) + zsshp2_e(ji  ,jj+1)  ) * ssvmask(ji,jj)
#else
               zhu_bck = hu_0(ji,jj) + r1_2*r1_e1e2u(ji,jj) * (  e1e2t(ji  ,jj) * zsshp2_e(ji  ,jj)    &
                    &                                          + e1e2t(ji+1,jj) * zsshp2_e(ji+1,jj)  ) * ssumask(ji,jj)
               zhv_bck = hv_0(ji,jj) + r1_2*r1_e1e2v(ji,jj) * (  e1e2t(ji,jj  ) * zsshp2_e(ji,jj  )    &
                    &                                          + e1e2t(ji,jj+1) * zsshp2_e(ji,jj+1)  ) * ssvmask(ji,jj)
#endif
               !                    ! inverse depth at jn+1
               z1_hu = ssumask(ji,jj) / ( hu_0(ji,jj) + zsshu_a(ji,jj) + 1._wp - ssumask(ji,jj) )
               z1_hv = ssvmask(ji,jj) / ( hv_0(ji,jj) + zsshv_a(ji,jj) + 1._wp - ssvmask(ji,jj) )
               !
               ua_e(ji,jj) = (               hu_e  (ji,jj) *   un_e (ji,jj)      & 
                    &            + rDt_e * (  zhu_bck        * zu_spg (ji,jj)  &   !
                    &                       + zhup2_e(ji,jj) * zu_trd (ji,jj)  &   !
                    &                       +  hu(ji,jj,Kmm) * zu_frc (ji,jj)  )   ) * z1_hu
               !
               va_e(ji,jj) = (               hv_e  (ji,jj) *   vn_e (ji,jj)      &
                    &            + rDt_e * (  zhv_bck        * zv_spg (ji,jj)  &   !
                    &                       + zhvp2_e(ji,jj) * zv_trd (ji,jj)  &   !
                    &                       +  hv(ji,jj,Kmm) * zv_frc (ji,jj)  )   ) * z1_hv
            END_2D
         ENDIF
!jth implicit bottom friction:
         IF ( ll_wd ) THEN ! revert to explicit for bit comparison tests in non wad runs
            DO_2D( 0, 0, 0, 0 )
               ua_e(ji,jj) =  ua_e(ji,jj) / ( 1._wp - rDt_e * zCdU_u(ji,jj) * hur_e(ji,jj) )
               va_e(ji,jj) =  va_e(ji,jj) / ( 1._wp - rDt_e * zCdU_v(ji,jj) * hvr_e(ji,jj) )
            END_2D
         ENDIF
       
         IF( .NOT.ln_linssh ) THEN !* Update ocean depth (variable volume case only)
Guillaume Samson's avatar
Guillaume Samson committed
               hu_e (ji,jj) =    hu_0(ji,jj) + zsshu_a(ji,jj)
               hur_e(ji,jj) = ssumask(ji,jj) / (  hu_e(ji,jj) + 1._wp - ssumask(ji,jj)  )
               hv_e (ji,jj) =    hv_0(ji,jj) + zsshv_a(ji,jj)
               hvr_e(ji,jj) = ssvmask(ji,jj) / (  hv_e(ji,jj) + 1._wp - ssvmask(ji,jj)  )
            END_2D
         ENDIF
         !
         IF( .NOT.ln_linssh ) THEN   !* Update ocean depth (variable volume case only)
            CALL lbc_lnk( 'dynspg_ts', ua_e , 'U', -1._wp, va_e , 'V', -1._wp  &
                                     , hu_e , 'U',  1._wp, hv_e , 'V',  1._wp  &
                                     , hur_e, 'U',  1._wp, hvr_e, 'V',  1._wp  )
Guillaume Samson's avatar
Guillaume Samson committed
         ELSE
            CALL lbc_lnk( 'dynspg_ts', ua_e , 'U', -1._wp, va_e , 'V', -1._wp )
Guillaume Samson's avatar
Guillaume Samson committed
         ENDIF
         !                                                 ! open boundaries
         IF( ln_bdy )   CALL bdy_dyn2d( jn, ua_e, va_e, un_e, vn_e, hur_e, hvr_e, ssha_e )
#if defined key_agrif                                                           
Guillaume Samson's avatar
Guillaume Samson committed
#endif
         !                                             !* Swap
         !                                             !  ----
         ubb_e  (:,:) = ub_e  (:,:)
         ub_e   (:,:) = un_e  (:,:)
         un_e   (:,:) = ua_e  (:,:)
         !
         vbb_e  (:,:) = vb_e  (:,:)
         vb_e   (:,:) = vn_e  (:,:)
         vn_e   (:,:) = va_e  (:,:)
         !
         sshbb_e(:,:) = sshb_e(:,:)
         sshb_e (:,:) = sshn_e(:,:)
         sshn_e (:,:) = ssha_e(:,:)

         !                                             !* Sum over whole bt loop
         !                                             !  ----------------------
         za1 = wgtbtp1(jn)                                    
         IF( ln_dynadv_vec .OR. ln_linssh ) THEN    ! Sum velocities
            puu_b  (:,:,Kaa) = puu_b  (:,:,Kaa) + za1 * ua_e  (:,:) 
            pvv_b  (:,:,Kaa) = pvv_b  (:,:,Kaa) + za1 * va_e  (:,:) 
         ELSE                                       ! Sum transports
            IF ( .NOT.ln_wd_dl ) THEN  
               puu_b  (:,:,Kaa) = puu_b  (:,:,Kaa) + za1 * ua_e  (:,:) * hu_e (:,:)
               pvv_b  (:,:,Kaa) = pvv_b  (:,:,Kaa) + za1 * va_e  (:,:) * hv_e (:,:)
            ELSE 
               puu_b  (:,:,Kaa) = puu_b  (:,:,Kaa) + za1 * ua_e  (:,:) * hu_e (:,:) * zuwdmask(:,:)
               pvv_b  (:,:,Kaa) = pvv_b  (:,:,Kaa) + za1 * va_e  (:,:) * hv_e (:,:) * zvwdmask(:,:)
            END IF 
         ENDIF
         !                                          ! Sum sea level
         pssh(:,:,Kaa) = pssh(:,:,Kaa) + za1 * ssha_e(:,:)

      END DO                                               !        end loop      !
      !                                                    ! ==================== !
      ! -----------------------------------------------------------------------------
      ! Phase 3. update the general trend with the barotropic trend
      ! -----------------------------------------------------------------------------
      !
      ! Set advection velocity correction:
      IF (ln_bt_fw) THEN
         IF( .NOT.( kt == nit000 .AND. l_1st_euler ) ) THEN
            DO_2D( nn_hls, nn_hls, nn_hls, nn_hls )
               zun_save = un_adv(ji,jj)
               zvn_save = vn_adv(ji,jj)
               !                          ! apply the previously computed correction 
               un_adv(ji,jj) = r1_2 * ( ub2_b(ji,jj) + zun_save - rn_atfp * un_bf(ji,jj) )
               vn_adv(ji,jj) = r1_2 * ( vb2_b(ji,jj) + zvn_save - rn_atfp * vn_bf(ji,jj) )
               !                          ! Update corrective fluxes for next time step
               un_bf(ji,jj)  = rn_atfp * un_bf(ji,jj) + ( zun_save - ub2_b(ji,jj) )
               vn_bf(ji,jj)  = rn_atfp * vn_bf(ji,jj) + ( zvn_save - vb2_b(ji,jj) )
               !                          ! Save integrated transport for next computation
               ub2_b(ji,jj) = zun_save
               vb2_b(ji,jj) = zvn_save
            END_2D
         ELSE
            un_bf(:,:) = 0._wp            ! corrective fluxes for next time step set to zero
            vn_bf(:,:) = 0._wp
            ub2_b(:,:) = un_adv(:,:)      ! Save integrated transport for next computation
            vb2_b(:,:) = vn_adv(:,:)
         END IF
      ENDIF


      !
      ! Update barotropic trend:
      IF( ln_dynadv_vec .OR. ln_linssh ) THEN
         DO jk=1,jpkm1
            puu(:,:,jk,Krhs) = puu(:,:,jk,Krhs) + ( puu_b(:,:,Kaa) - puu_b(:,:,Kbb) ) * r1_Dt_b
            pvv(:,:,jk,Krhs) = pvv(:,:,jk,Krhs) + ( pvv_b(:,:,Kaa) - pvv_b(:,:,Kbb) ) * r1_Dt_b
         END DO
      ELSE
         ! At this stage, pssh(:,:,:,Krhs) has been corrected: compute new depths at velocity points
#if defined key_qcoTest_FluxForm
         !                                ! 'key_qcoTest_FluxForm' : simple ssh average
         DO_2D( 1, 0, 1, 0 )
            zsshu_a(ji,jj) = r1_2 * ( pssh(ji,jj,Kaa) + pssh(ji+1,jj  ,Kaa) ) * ssumask(ji,jj)
            zsshv_a(ji,jj) = r1_2 * ( pssh(ji,jj,Kaa) + pssh(ji  ,jj+1,Kaa) ) * ssvmask(ji,jj)
         END_2D
#else
         DO_2D( 1, 0, 1, 0 )
            zsshu_a(ji,jj) = r1_2 * r1_e1e2u(ji,jj) * ( e1e2t(ji  ,jj) * pssh(ji  ,jj,Kaa)   &
               &                                      + e1e2t(ji+1,jj) * pssh(ji+1,jj,Kaa) ) * ssumask(ji,jj)
            zsshv_a(ji,jj) = r1_2 * r1_e1e2v(ji,jj) * ( e1e2t(ji,jj  ) * pssh(ji,jj  ,Kaa)   &
               &                                      + e1e2t(ji,jj+1) * pssh(ji,jj+1,Kaa) ) * ssvmask(ji,jj)
         END_2D
#endif   
         CALL lbc_lnk( 'dynspg_ts', zsshu_a, 'U', 1._wp, zsshv_a, 'V', 1._wp ) ! Boundary conditions
         !
         DO jk=1,jpkm1
            puu(:,:,jk,Krhs) = puu(:,:,jk,Krhs) + r1_hu(:,:,Kmm)   &
               &             * ( puu_b(:,:,Kaa) - puu_b(:,:,Kbb) * hu(:,:,Kbb) ) * r1_Dt_b
            pvv(:,:,jk,Krhs) = pvv(:,:,jk,Krhs) + r1_hv(:,:,Kmm)   &
               &             * ( pvv_b(:,:,Kaa) - pvv_b(:,:,Kbb) * hv(:,:,Kbb) ) * r1_Dt_b
         END DO
         ! Save barotropic velocities not transport:
         puu_b(:,:,Kaa) =  puu_b(:,:,Kaa) / ( hu_0(:,:) + zsshu_a(:,:) + 1._wp - ssumask(:,:) )
         pvv_b(:,:,Kaa) =  pvv_b(:,:,Kaa) / ( hv_0(:,:) + zsshv_a(:,:) + 1._wp - ssvmask(:,:) )
      ENDIF


      ! Correct velocities so that the barotropic velocity equals (un_adv, vn_adv) (in all cases)  
      DO jk = 1, jpkm1
         puu(:,:,jk,Kmm) = ( puu(:,:,jk,Kmm) + un_adv(:,:)*r1_hu(:,:,Kmm) - puu_b(:,:,Kmm) ) * umask(:,:,jk)
         pvv(:,:,jk,Kmm) = ( pvv(:,:,jk,Kmm) + vn_adv(:,:)*r1_hv(:,:,Kmm) - pvv_b(:,:,Kmm) ) * vmask(:,:,jk)
      END DO

      IF ( ln_wd_dl .and. ln_wd_dl_bc) THEN 
         DO jk = 1, jpkm1
            puu(:,:,jk,Kmm) = ( un_adv(:,:)*r1_hu(:,:,Kmm) &
                       & + zuwdav2(:,:)*(puu(:,:,jk,Kmm) - un_adv(:,:)*r1_hu(:,:,Kmm)) ) * umask(:,:,jk) 
            pvv(:,:,jk,Kmm) = ( vn_adv(:,:)*r1_hv(:,:,Kmm) & 
                       & + zvwdav2(:,:)*(pvv(:,:,jk,Kmm) - vn_adv(:,:)*r1_hv(:,:,Kmm)) ) * vmask(:,:,jk)  
         END DO
      END IF 

      
      CALL iom_put(  "ubar", un_adv(:,:)*r1_hu(:,:,Kmm) )    ! barotropic i-current
      CALL iom_put(  "vbar", vn_adv(:,:)*r1_hv(:,:,Kmm) )    ! barotropic i-current
      !
#if defined key_agrif
      ! Save time integrated fluxes during child grid integration
      ! (used to update coarse grid transports at next time step)
      !
      IF( .NOT.Agrif_Root() .AND. ln_bt_fw .AND. ln_agrif_2way ) THEN
         IF( Agrif_NbStepint() == 0 ) THEN
            ub2_i_b(:,:) = 0._wp
            vb2_i_b(:,:) = 0._wp
         END IF
         !
         za1 = 1._wp / REAL(Agrif_rhot(), wp)
         ub2_i_b(:,:) = ub2_i_b(:,:) + za1 * ub2_b(:,:)
         vb2_i_b(:,:) = vb2_i_b(:,:) + za1 * vb2_b(:,:)
      ENDIF
#endif      
      !                                   !* write time-spliting arrays in the restart
      IF( lrst_oce .AND.ln_bt_fw )   CALL ts_rst( kt, 'WRITE' )
      !
      IF( ln_wd_dl )   DEALLOCATE( ztwdmask, zuwdmask, zvwdmask, zuwdav2, zvwdav2 )
      !
      CALL iom_put( "baro_u" , puu_b(:,:,Kmm) )  ! Barotropic  U Velocity
      CALL iom_put( "baro_v" , pvv_b(:,:,Kmm) )  ! Barotropic  V Velocity
      !
   END SUBROUTINE dyn_spg_ts


   SUBROUTINE ts_wgt( ll_av, ll_fw, jpit, zwgt1, zwgt2)
      !!---------------------------------------------------------------------
      !!                   ***  ROUTINE ts_wgt  ***
      !!
      !! ** Purpose : Set time-splitting weights for temporal averaging (or not)
      !!----------------------------------------------------------------------
      LOGICAL, INTENT(in) ::   ll_av      ! temporal averaging=.true.
      LOGICAL, INTENT(in) ::   ll_fw      ! forward time splitting =.true.
      INTEGER, INTENT(inout) :: jpit      ! cycle length    
      REAL(wp), DIMENSION(3*nn_e), INTENT(inout)  :: zwgt2
      REAL(dp), DIMENSION(3*nn_e), INTENT(inout)  :: zwgt1

Guillaume Samson's avatar
Guillaume Samson committed
      
      INTEGER ::  jic, jn, ji                      ! temporary integers
Guillaume Samson's avatar
Guillaume Samson committed
      !!----------------------------------------------------------------------

      zwgt1(:) = 0._wp
      zwgt2(:) = 0._wp

      ! Set time index when averaged value is requested
      IF (ll_fw) THEN 
         jic = nn_e
      ELSE
         jic = 2 * nn_e
      ENDIF

      ! Set primary weights:
      IF (ll_av) THEN
           ! Define simple boxcar window for primary weights 
           ! (width = nn_e, centered around jic)     
         SELECT CASE ( nn_bt_flt )
              CASE( 0 )  ! No averaging
                 zwgt1(jic) = 1._wp
                 jpit = jic

              CASE( 1 )  ! Boxcar, width = nn_e
                 DO jn = 1, 3*nn_e
                    za1 = ABS(float(jn-jic))/float(nn_e) 
                    IF (za1 < 0.5_wp) THEN
                      zwgt1(jn) = 1._wp
                      jpit = jn
                    ENDIF
                 ENDDO

              CASE( 2 )  ! Boxcar, width = 2 * nn_e
                 DO jn = 1, 3*nn_e
                    za1 = ABS(float(jn-jic))/float(nn_e) 
                    IF (za1 < 1._wp) THEN
                      zwgt1(jn) = 1._wp
                      jpit = jn
                    ENDIF
                 ENDDO
              CASE DEFAULT   ;   CALL ctl_stop( 'unrecognised value for nn_bt_flt' )
         END SELECT

      ELSE ! No time averaging
         zwgt1(jic) = 1._wp
         jpit = jic
      ENDIF
    
      ! Set secondary weights
      DO jn = 1, jpit
        DO ji = jn, jpit
             zwgt2(jn) = zwgt2(jn) + zwgt1(ji)
        END DO
      END DO

      ! Normalize weigths:
      za1 = 1._wp / SUM(zwgt1(1:jpit))
      za2 = 1._wp / SUM(zwgt2(1:jpit))
      DO jn = 1, jpit
        zwgt1(jn) = zwgt1(jn) * za1
        zwgt2(jn) = zwgt2(jn) * za2
      END DO
      !
   END SUBROUTINE ts_wgt


   SUBROUTINE ts_rst( kt, cdrw )
      !!---------------------------------------------------------------------
      !!                   ***  ROUTINE ts_rst  ***
      !!
      !! ** Purpose : Read or write time-splitting arrays in restart file
      !!----------------------------------------------------------------------
      INTEGER         , INTENT(in) ::   kt     ! ocean time-step
      CHARACTER(len=*), INTENT(in) ::   cdrw   ! "READ"/"WRITE" flag
      !!----------------------------------------------------------------------
      !
      IF( TRIM(cdrw) == 'READ' ) THEN        ! Read/initialise 
         !                                   ! ---------------
         IF( ln_rstart .AND. ln_bt_fw .AND. .NOT.l_1st_euler ) THEN    !* Read the restart file
            CALL iom_get( numror, jpdom_auto, 'ub2_b'  , ub2_b  (:,:), cd_type = 'U', psgn = -1._wp )   
            CALL iom_get( numror, jpdom_auto, 'vb2_b'  , vb2_b  (:,:), cd_type = 'V', psgn = -1._wp ) 
            CALL iom_get( numror, jpdom_auto, 'un_bf'  , un_bf  (:,:), cd_type = 'U', psgn = -1._wp )   
            CALL iom_get( numror, jpdom_auto, 'vn_bf'  , vn_bf  (:,:), cd_type = 'V', psgn = -1._wp ) 
            IF( .NOT.ln_bt_av ) THEN
               CALL iom_get( numror, jpdom_auto, 'sshbb_e'  , sshbb_e(:,:), cd_type = 'T', psgn =  1._wp )   
               CALL iom_get( numror, jpdom_auto, 'ubb_e'    ,   ubb_e(:,:), cd_type = 'U', psgn = -1._wp )   
               CALL iom_get( numror, jpdom_auto, 'vbb_e'    ,   vbb_e(:,:), cd_type = 'V', psgn = -1._wp )
               CALL iom_get( numror, jpdom_auto, 'sshb_e'   ,  sshb_e(:,:), cd_type = 'T', psgn =  1._wp ) 
               CALL iom_get( numror, jpdom_auto, 'ub_e'     ,    ub_e(:,:), cd_type = 'U', psgn = -1._wp )   
               CALL iom_get( numror, jpdom_auto, 'vb_e'     ,    vb_e(:,:), cd_type = 'V', psgn = -1._wp )
            ENDIF
#if defined key_agrif
            ! Read time integrated fluxes
            IF ( .NOT.Agrif_Root() ) THEN
               CALL iom_get( numror, jpdom_auto, 'ub2_i_b'  , ub2_i_b(:,:), cd_type = 'U', psgn = -1._wp )   
               CALL iom_get( numror, jpdom_auto, 'vb2_i_b'  , vb2_i_b(:,:), cd_type = 'V', psgn = -1._wp )
            ELSE
               ub2_i_b(:,:) = 0._wp   ;   vb2_i_b(:,:) = 0._wp   ! used in the 1st update of agrif
            ENDIF
#endif
         ELSE                                   !* Start from rest
            IF(lwp) WRITE(numout,*)
            IF(lwp) WRITE(numout,*) '   ==>>>   start from rest: set barotropic values to 0'
            ub2_b  (:,:) = 0._wp   ;   vb2_b  (:,:) = 0._wp   ! used in the 1st interpol of agrif
            un_adv (:,:) = 0._wp   ;   vn_adv (:,:) = 0._wp   ! used in the 1st interpol of agrif
            un_bf  (:,:) = 0._wp   ;   vn_bf  (:,:) = 0._wp   ! used in the 1st update   of agrif
#if defined key_agrif
            ub2_i_b(:,:) = 0._wp   ;   vb2_i_b(:,:) = 0._wp   ! used in the 1st update of agrif
#endif
         ENDIF
         !
      ELSEIF( TRIM(cdrw) == 'WRITE' ) THEN   ! Create restart file
         !                                   ! -------------------
         IF(lwp) WRITE(numout,*) '---- ts_rst ----'
         CALL iom_rstput( kt, nitrst, numrow, 'ub2_b'   , ub2_b  (:,:) )
         CALL iom_rstput( kt, nitrst, numrow, 'vb2_b'   , vb2_b  (:,:) )
         CALL iom_rstput( kt, nitrst, numrow, 'un_bf'   , un_bf  (:,:) )
         CALL iom_rstput( kt, nitrst, numrow, 'vn_bf'   , vn_bf  (:,:) )
         !
         IF (.NOT.ln_bt_av) THEN
            CALL iom_rstput( kt, nitrst, numrow, 'sshbb_e'  , sshbb_e(:,:) ) 
            CALL iom_rstput( kt, nitrst, numrow, 'ubb_e'    ,   ubb_e(:,:) )
            CALL iom_rstput( kt, nitrst, numrow, 'vbb_e'    ,   vbb_e(:,:) )
            CALL iom_rstput( kt, nitrst, numrow, 'sshb_e'   ,  sshb_e(:,:) )
            CALL iom_rstput( kt, nitrst, numrow, 'ub_e'     ,    ub_e(:,:) )
            CALL iom_rstput( kt, nitrst, numrow, 'vb_e'     ,    vb_e(:,:) )
         ENDIF
#if defined key_agrif
         ! Save time integrated fluxes