Index: NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_top_sponge.F90
===================================================================
--- NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_top_sponge.F90	(revision 13961)
+++ NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_top_sponge.F90	(revision 13987)
@@ -67,14 +67,14 @@
       !
       INTEGER  ::   ji, jj, jk, jn   ! dummy loop indices
-      REAL(wp) ::   zabe1, zabe2, ztrelax
-      REAL(wp), DIMENSION(i1:i2,j1:j2)               ::   ztu, ztv
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2,1:jptra) ::   trbdiff
+      INTEGER  ::   iku, ikv
+      REAL(wp) ::   ztra, zabe1, zabe2, zbtr
+      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk)  ::   ztu, ztv
+      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2,n1:n2) ::   trbdiff
       ! vertical interpolation:
-      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk,1:jptra) ::tabres_child
+      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk,n1:n2) ::tabres_child
       REAL(wp), DIMENSION(k1:k2,1:jptra) :: tabin
       REAL(wp), DIMENSION(k1:k2) :: h_in
       REAL(wp), DIMENSION(1:jpk) :: h_out
       INTEGER :: N_in, N_out
-      REAL(wp) :: h_diff
       !!----------------------------------------------------------------------
       !
@@ -94,5 +94,5 @@
             DO jj=j1,j2
                DO ji=i1,i2
-                  tabres(ji,jj,jk,jpts+1) = tmask(ji,jj,jk) * e3t(ji,jj,jk,Kbb_a) 
+                  tabres(ji,jj,jk,jpts+1) = tmask(ji,jj,jk) * e3t(ji,jj,jk,Kbb_a)
                END DO
             END DO
@@ -128,7 +128,7 @@
                DO jk=1,jpkm1
 # if defined key_vertical
-                  trbdiff(ji,jj,jk,1:jptra) = tr(ji,jj,jk,1:jptra,Kbb_a) - tabres_child(ji,jj,jk,1:jptra)
+                  trbdiff(ji,jj,jk,1:jptra) = ( tr(ji,jj,jk,1:jptra,Kbb_a) - tabres_child(ji,jj,jk,1:jptra) ) * tmask(ji,jj,jk)
 # else
-                  trbdiff(ji,jj,jk,1:jptra) = tr(ji,jj,jk,1:jptra,Kbb_a) - tabres(ji,jj,jk,1:jptra)
+                  trbdiff(ji,jj,jk,1:jptra) = ( tr(ji,jj,jk,1:jptra,Kbb_a) - tabres(ji,jj,jk,1:jptra)) * tmask(ji,jj,jk)
 # endif
                ENDDO
@@ -136,32 +136,48 @@
          ENDDO
 
-         !* set relaxation time scale
-         IF( l_1st_euler .AND. lk_agrif_fstep ) THEN   ;   ztrelax =   rn_trelax_tra  / (        rn_Dt )
-         ELSE                                          ;   ztrelax =   rn_trelax_tra  / (2._wp * rn_Dt )
-         ENDIF
-
          DO jn = 1, jptra
             DO jk = 1, jpkm1
-               DO jj = j1,j2-1
+               ztu(i1:i2,j1:j2,jk) = 0._wp
+               DO jj = j1,j2
                   DO ji = i1,i2-1
-                     zabe1 = rn_sponge_tra * fspu(ji,jj) * e2_e1u(ji,jj) * e3u(ji,jj,jk,Kmm_a) * umask(ji,jj,jk)
-                     zabe2 = rn_sponge_tra * fspv(ji,jj) * e1_e2v(ji,jj) * e3v(ji,jj,jk,Kmm_a) * vmask(ji,jj,jk)
-                     ztu(ji,jj) = zabe1 * ( trbdiff(ji+1,jj  ,jk,jn) - trbdiff(ji,jj,jk,jn) )
-                     ztv(ji,jj) = zabe2 * ( trbdiff(ji  ,jj+1,jk,jn) - trbdiff(ji,jj,jk,jn) )
+                     zabe1 = rn_sponge_tra * r1_Dt * fspu(ji,jj) * umask(ji,jj,jk) * e2_e1u(ji,jj) * e3u(ji,jj,jk,Kmm_a)
+                     ztu(ji,jj,jk) = zabe1 * ( trbdiff(ji+1,jj,jk,jn) - trbdiff(ji,jj,jk,jn) )
+                  END DO
+               END DO
+               ztv(i1:i2,j1:j2,jk) = 0._wp
+               DO ji = i1,i2
+                  DO jj = j1,j2-1
+                     zabe2 = rn_sponge_tra * r1_Dt * fspv(ji,jj) * vmask(ji,jj,jk) * e1_e2v(ji,jj) * e3v(ji,jj,jk,Kmm_a)
+                     ztv(ji,jj,jk) = zabe2 * ( trbdiff(ji,jj+1,jk,jn) - trbdiff(ji,jj,jk,jn) )
                   END DO
                END DO
                !
+               IF( ln_zps ) THEN      ! set gradient at partial step level
+                  DO jj = j1,j2
+                     DO ji = i1,i2
+                        ! last level
+                        iku = mbku(ji,jj)
+                        ikv = mbkv(ji,jj)
+                        IF( iku == jk )   ztu(ji,jj,jk) = 0._wp
+                        IF( ikv == jk )   ztv(ji,jj,jk) = 0._wp
+                     END DO
+                  END DO
+               ENDIF
+            END DO
+            !
+            DO jk = 1, jpkm1
                DO jj = j1+1,j2-1
                   DO ji = i1+1,i2-1
-                     IF( .NOT. tabspongedone_trn(ji,jj) ) THEN 
-                        tr(ji,jj,jk,jn,Krhs_a) = tr(ji,jj,jk,jn,Krhs_a) + (  ztu(ji,jj) - ztu(ji-1,jj  )     &
-                           &                                   + ztv(ji,jj) - ztv(ji  ,jj-1)  )  &
-                           &                                * r1_e1e2t(ji,jj) / e3t(ji,jj,jk,Kmm_a)  &
-                           &                                - ztrelax * fspt(ji,jj) * trbdiff(ji,jj,jk,jn)
+                     IF( .NOT. tabspongedone_trn(ji,jj) ) THEN
+                        zbtr = r1_e1e2t(ji,jj) / e3t(ji,jj,jk,Kmm_a)
+                                                ! horizontal diffusive trends
+                        ztra = zbtr * (  ztu(ji,jj,jk) - ztu(ji-1,jj,jk) + ztv(ji,jj,jk) - ztv(ji,jj-1,jk)  ) &
+                             &  - rn_trelax_tra * r1_Dt * fspt(ji,jj) * trbdiff(ji,jj,jk,jn)
+                        ! add it to the general tracer trends
+                        tr(ji,jj,jk,jn,Krhs_a) = tr(ji,jj,jk,jn,Krhs_a) + ztra
                      ENDIF
                   END DO
                END DO
             END DO
-            !
          END DO
          !
Index: NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_user.F90
===================================================================
--- NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_user.F90	(revision 13961)
+++ NEMO/branches/2020/dev_r13312_AGRIF-03-04_jchanut_vinterp_tstep/src/NST/agrif_user.F90	(revision 13987)
@@ -367,6 +367,6 @@
          CALL Agrif_Bc_variable( unb_interp_id,calledweight=1.,procname=interpunb )
          CALL Agrif_Bc_variable( vnb_interp_id,calledweight=1.,procname=interpvnb )
-         CALL Agrif_Bc_variable(ub2b_interp_id,calledweight=1.,procname=interpub2b)
-         CALL Agrif_Bc_variable(vb2b_interp_id,calledweight=1.,procname=interpvb2b)
+         CALL Agrif_Bc_variable(ub2b_interp_id,calledweight=1.,procname=interpub2b)   ! must be called before unb_id to define ubdy
+         CALL Agrif_Bc_variable(vb2b_interp_id,calledweight=1.,procname=interpvb2b)   ! must be called before vnb_id to define vbdy
          use_sign_north = .FALSE.
          ubdy(:,:) = 0._wp
@@ -726,5 +726,5 @@
       CALL Agrif_Bc_variable(trn_sponge_id,calledweight=1.,procname=interptrn_sponge)
       ! reset tsa to zero
-      tra(:,:,:,:) = 0._wp
+      tr(:,:,:,:,Krhs_a) = 0._wp
 
       ! 3. Some controls
