Index: branches/UKMO/restart_datestamp/NEMOGCM/NEMO/OPA_SRC/ZDF/zdftke.F90
===================================================================
--- branches/UKMO/restart_datestamp/NEMOGCM/NEMO/OPA_SRC/ZDF/zdftke.F90	(revision 5420)
+++ branches/UKMO/restart_datestamp/NEMOGCM/NEMO/OPA_SRC/ZDF/zdftke.F90	(revision 5462)
@@ -26,4 +26,5 @@
    !!                 !                                + cleaning of the parameters + bugs correction
    !!            3.3  !  2010-10  (C. Ethe, G. Madec) reorganisation of initialisation phase
+   !!            3.6  !  2014-11  (P. Mathiot) add ice shelf capability
    !!----------------------------------------------------------------------
 #if defined key_zdftke   ||   defined key_esopa
@@ -236,14 +237,18 @@
       zfact3 = 0.5_wp       * rn_ediss
       !
+      !
       !                     !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
       !                     !  Surface boundary condition on tke
       !                     !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+      IF ( ln_isfcav ) THEN
+         DO jj = 2, jpjm1            ! en(mikt(ji,jj))   = rn_emin
+            DO ji = fs_2, fs_jpim1   ! vector opt.
+               en(ji,jj,mikt(ji,jj))=rn_emin * tmask(ji,jj,1)
+            END DO
+         END DO
+      END IF
       DO jj = 2, jpjm1            ! en(1)   = rn_ebb taum / rau0  (min value rn_emin0)
          DO ji = fs_2, fs_jpim1   ! vector opt.
-            IF (mikt(ji,jj) .GT. 1) THEN
-               en(ji,jj,mikt(ji,jj))=rn_emin
-            ELSE
-               en(ji,jj,1) = MAX( rn_emin0, zbbrau * taum(ji,jj) ) * tmask(ji,jj,1)
-            END IF
+            en(ji,jj,1) = MAX( rn_emin0, zbbrau * taum(ji,jj) ) * tmask(ji,jj,1)
          END DO
       END DO
@@ -301,6 +306,9 @@
          END DO
          zcof = 0.016 / SQRT( zrhoa * zcdrag )
+!CDIR NOVERRCHK
          DO jk = 2, jpkm1         !* TKE Langmuir circulation source term added to en
-            DO jj = 2, jpjm1
+!CDIR NOVERRCHK
+            DO jj = 2, jpjm1
+!CDIR NOVERRCHK
                DO ji = fs_2, fs_jpim1   ! vector opt.
                   zus  = zcof * SQRT( taum(ji,jj) )           ! Stokes drift
@@ -309,5 +317,5 @@
                   zwlc = zind * rn_lc * zus * SIN( rpi * fsdepw(ji,jj,jk) / zhlc(ji,jj) )
                   !                                           ! TKE Langmuir circulation source term
-                  en(ji,jj,jk) = en(ji,jj,jk) + rdt * ( zwlc * zwlc * zwlc ) / zhlc(ji,jj) * tmask(ji,jj,jk)
+                  en(ji,jj,jk) = en(ji,jj,jk) + rdt * ( zwlc * zwlc * zwlc ) / zhlc(ji,jj) * wmask(ji,jj,jk) * tmask(ji,jj,1)
                END DO
             END DO
@@ -328,6 +336,6 @@
                avmu(ji,jj,jk) = avmu(ji,jj,jk) * ( un(ji,jj,jk-1) - un(ji,jj,jk) )   &
                   &                            * ( ub(ji,jj,jk-1) - ub(ji,jj,jk) )   & 
-                  &           / (  fse3uw_n(ji,jj,jk)         &
-                  &              * fse3uw_b(ji,jj,jk) )
+                  &                            / (  fse3uw_n(ji,jj,jk)               &
+                  &                              *  fse3uw_b(ji,jj,jk)  )
                avmv(ji,jj,jk) = avmv(ji,jj,jk) * ( vn(ji,jj,jk-1) - vn(ji,jj,jk) )   &
                   &                            * ( vb(ji,jj,jk-1) - vb(ji,jj,jk) )   &
@@ -338,7 +346,7 @@
       END DO
       !
-      DO jj = 2, jpjm1
-         DO ji = fs_2, fs_jpim1   ! vector opt.
-            DO jk = mikt(ji,jj)+1, jpkm1           !* Matrix and right hand side in en
+      DO jk = 2, jpkm1           !* Matrix and right hand side in en
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1   ! vector opt.
                zcof   = zfact1 * tmask(ji,jj,jk)
                zzd_up = zcof * ( avm  (ji,jj,jk+1) + avm  (ji,jj,jk  ) )   &  ! upper diagonal
@@ -357,26 +365,48 @@
                en(ji,jj,jk) = en(ji,jj,jk) + rdt * (  zesh2  -   avt(ji,jj,jk) * rn2(ji,jj,jk)    &
                   &                                 + zfact3 * dissl(ji,jj,jk) * en (ji,jj,jk)  ) &
-                  &                                 * tmask(ji,jj,jk) * tmask(ji,jj,jk-1)
-            END DO
-            !                          !* Matrix inversion from level 2 (tke prescribed at level 1)
-            DO jk = mikt(ji,jj)+2, jpkm1                             ! First recurrence : Dk = Dk - Lk * Uk-1 / Dk-1
+                  &                                 * wmask(ji,jj,jk)
+            END DO
+         END DO
+      END DO
+      !                          !* Matrix inversion from level 2 (tke prescribed at level 1)
+      DO jk = 3, jpkm1                             ! First recurrence : Dk = Dk - Lk * Uk-1 / Dk-1
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1    ! vector opt.
                zdiag(ji,jj,jk) = zdiag(ji,jj,jk) - zd_lw(ji,jj,jk) * zd_up(ji,jj,jk-1) / zdiag(ji,jj,jk-1)
             END DO
-            ! Second recurrence : Lk = RHSk - Lk / Dk-1 * Lk-1
-            zd_lw(ji,jj,mikt(ji,jj)+1) = en(ji,jj,mikt(ji,jj)+1) - zd_lw(ji,jj,mikt(ji,jj)+1) * en(ji,jj,mikt(ji,jj))    ! Surface boudary conditions on tke
-            !
-            DO jk = mikt(ji,jj)+2, jpkm1
+         END DO
+      END DO
+      !
+      ! Second recurrence : Lk = RHSk - Lk / Dk-1 * Lk-1
+      DO jj = 2, jpjm1
+         DO ji = fs_2, fs_jpim1   ! vector opt.
+            zd_lw(ji,jj,2) = en(ji,jj,2) - zd_lw(ji,jj,2) * en(ji,jj,1)    ! Surface boudary conditions on tke
+         END DO
+      END DO
+      DO jk = 3, jpkm1
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1    ! vector opt.
                zd_lw(ji,jj,jk) = en(ji,jj,jk) - zd_lw(ji,jj,jk) / zdiag(ji,jj,jk-1) *zd_lw(ji,jj,jk-1)
             END DO
-            !
-            ! thrid recurrence : Ek = ( Lk - Uk * Ek+1 ) / Dk
+         END DO
+      END DO
+      !
+      ! thrid recurrence : Ek = ( Lk - Uk * Ek+1 ) / Dk
+      DO jj = 2, jpjm1
+         DO ji = fs_2, fs_jpim1   ! vector opt.
             en(ji,jj,jpkm1) = zd_lw(ji,jj,jpkm1) / zdiag(ji,jj,jpkm1)
-            !
-            DO jk = jpk-2, mikt(ji,jj)+1, -1
+         END DO
+      END DO
+      DO jk = jpk-2, 2, -1
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1    ! vector opt.
                en(ji,jj,jk) = ( zd_lw(ji,jj,jk) - zd_up(ji,jj,jk) * en(ji,jj,jk+1) ) / zdiag(ji,jj,jk)
             END DO
-            !
-            DO jk = mikt(ji,jj), jpkm1                             ! set the minimum value of tke
-               en(ji,jj,jk) = MAX( en(ji,jj,jk), rn_emin ) * tmask(ji,jj,jk)
+         END DO
+      END DO
+      DO jk = 2, jpkm1                             ! set the minimum value of tke
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1   ! vector opt.
+               en(ji,jj,jk) = MAX( en(ji,jj,jk), rn_emin ) * wmask(ji,jj,jk)
             END DO
          END DO
@@ -391,5 +421,5 @@
                DO ji = fs_2, fs_jpim1   ! vector opt.
                   en(ji,jj,jk) = en(ji,jj,jk) + rn_efr * en(ji,jj,1) * EXP( -fsdepw(ji,jj,jk) / htau(ji,jj) )   &
-                     &                                 * ( 1._wp - fr_i(ji,jj) )  * tmask(ji,jj,jk) * tmask(ji,jj,1)
+                     &                                 * ( 1._wp - fr_i(ji,jj) )  * wmask(ji,jj,jk) * tmask(ji,jj,1)
                END DO
             END DO
@@ -400,5 +430,5 @@
                jk = nmln(ji,jj)
                en(ji,jj,jk) = en(ji,jj,jk) + rn_efr * en(ji,jj,1) * EXP( -fsdepw(ji,jj,jk) / htau(ji,jj) )   &
-                  &                                 * ( 1._wp - fr_i(ji,jj) )  * tmask(ji,jj,jk) * tmask(ji,jj,1)
+                  &                                 * ( 1._wp - fr_i(ji,jj) )  * wmask(ji,jj,jk) * tmask(ji,jj,1)
             END DO
          END DO
@@ -416,5 +446,5 @@
                   zdif = rhftau_scl * MAX( 0._wp, zdif + rhftau_add )  ! apply some modifications...
                   en(ji,jj,jk) = en(ji,jj,jk) + zbbrau * zdif * EXP( -fsdepw(ji,jj,jk) / htau(ji,jj) )   &
-                     &                        * ( 1._wp - fr_i(ji,jj) ) * tmask(ji,jj,jk) * tmask(ji,jj,jk-1) * tmask(ji,jj,1)
+                     &                        * ( 1._wp - fr_i(ji,jj) ) * wmask(ji,jj,jk) * tmask(ji,jj,1)
                END DO
             END DO
@@ -484,34 +514,28 @@
       !                     !* Buoyancy length scale: l=sqrt(2*e/n**2)
       !
+      ! initialisation of interior minimum value (avoid a 2d loop with mikt)
+      zmxlm(:,:,:)  = rmxl_min    
+      zmxld(:,:,:)  = rmxl_min
+      !
       IF( ln_mxl0 ) THEN            ! surface mixing length = F(stress) : l=vkarmn*2.e5*taum/(rau0*g)
          DO jj = 2, jpjm1
             DO ji = fs_2, fs_jpim1
-               IF (mikt(ji,jj) .GT. 1) THEN
-                  zmxlm(ji,jj,mikt(ji,jj)) = rmxl_min
-               ELSE
-                  zraug = vkarmn * 2.e5_wp / ( rau0 * grav )
-                  zmxlm(ji,jj,mikt(ji,jj)) = MAX( rn_mxl0, zraug * taum(ji,jj) )
-               END IF
+               zraug = vkarmn * 2.e5_wp / ( rau0 * grav )
+               zmxlm(ji,jj,1) = MAX( rn_mxl0, zraug * taum(ji,jj) * tmask(ji,jj,1) )
             END DO
          END DO
       ELSE 
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1                         ! surface set to the minimum value
-               zmxlm(ji,jj,mikt(ji,jj)) = MAX( tmask(ji,jj,1) * rn_mxl0, rmxl_min)
-            END DO
-         END DO
+         zmxlm(:,:,1) = rn_mxl0
       ENDIF
-      zmxlm(:,:,jpk)  = rmxl_min     ! last level set to the interior minium value
-      !
-!CDIR NOVERRCHK
-      DO jj = 2, jpjm1
-!CDIR NOVERRCHK
-         DO ji = fs_2, fs_jpim1   ! vector opt.
-            !CDIR NOVERRCHK
-            DO jk = mikt(ji,jj)+1, jpkm1              ! interior value : l=sqrt(2*e/n^2)
+      !
+!CDIR NOVERRCHK
+      DO jk = 2, jpkm1              ! interior value : l=sqrt(2*e/n^2)
+!CDIR NOVERRCHK
+         DO jj = 2, jpjm1
+!CDIR NOVERRCHK
+            DO ji = fs_2, fs_jpim1   ! vector opt.
                zrn2 = MAX( rn2(ji,jj,jk), rsmall )
-               zmxlm(ji,jj,jk) = MAX(  rmxl_min,  SQRT( 2._wp * en(ji,jj,jk) / zrn2 )  )
-            END DO
-            zmxld(ji,jj,mikt(ji,jj)) = zmxlm(ji,jj,mikt(ji,jj))   ! surface set to the minimum value 
+               zmxlm(ji,jj,jk) = MAX(  rmxl_min,  SQRT( 2._wp * en(ji,jj,jk) / zrn2 ) )
+            END DO
          END DO
       END DO
@@ -519,17 +543,19 @@
       !                     !* Physical limits for the mixing length
       !
-      zmxld(:,:, 1 ) = zmxlm(:,:,1)   ! surface set to the zmxlm   value
+      zmxld(:,:,1  ) = zmxlm(:,:,1)   ! surface set to the minimum value 
       zmxld(:,:,jpk) = rmxl_min       ! last level  set to the minimum value
       !
       SELECT CASE ( nn_mxl )
       !
+      ! where wmask = 0 set zmxlm == fse3w
       CASE ( 0 )           ! bounded by the distance to surface and bottom
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1   ! vector opt.
-               DO jk = mikt(ji,jj)+1, jpkm1
+         DO jk = 2, jpkm1
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zemxl = MIN( fsdepw(ji,jj,jk) - fsdepw(ji,jj,mikt(ji,jj)), zmxlm(ji,jj,jk),   &
                   &            fsdepw(ji,jj,mbkt(ji,jj)+1) - fsdepw(ji,jj,jk) )
-                  zmxlm(ji,jj,jk) = zemxl
-                  zmxld(ji,jj,jk) = zemxl
+                  ! wmask prevent zmxlm = 0 if jk = mikt(ji,jj)
+                  zmxlm(ji,jj,jk) = zemxl * wmask(ji,jj,jk) + MIN(zmxlm(ji,jj,jk),fse3w(ji,jj,jk)) * (1 - wmask(ji,jj,jk))
+                  zmxld(ji,jj,jk) = zemxl * wmask(ji,jj,jk) + MIN(zmxlm(ji,jj,jk),fse3w(ji,jj,jk)) * (1 - wmask(ji,jj,jk))
                END DO
             END DO
@@ -537,7 +563,7 @@
          !
       CASE ( 1 )           ! bounded by the vertical scale factor
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1   ! vector opt.
-               DO jk = mikt(ji,jj)+1, jpkm1
+         DO jk = 2, jpkm1
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zemxl = MIN( fse3w(ji,jj,jk), zmxlm(ji,jj,jk) )
                   zmxlm(ji,jj,jk) = zemxl
@@ -548,10 +574,14 @@
          !
       CASE ( 2 )           ! |dk[xml]| bounded by e3t :
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1   ! vector opt.
-               DO jk = mikt(ji,jj)+1, jpkm1         ! from the surface to the bottom :
+         DO jk = 2, jpkm1         ! from the surface to the bottom :
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zmxlm(ji,jj,jk) = MIN( zmxlm(ji,jj,jk-1) + fse3t(ji,jj,jk-1), zmxlm(ji,jj,jk) )
                END DO
-               DO jk = jpkm1, mikt(ji,jj)+1, -1     ! from the bottom to the surface :
+            END DO
+         END DO
+         DO jk = jpkm1, 2, -1     ! from the bottom to the surface :
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zemxl = MIN( zmxlm(ji,jj,jk+1) + fse3t(ji,jj,jk+1), zmxlm(ji,jj,jk) )
                   zmxlm(ji,jj,jk) = zemxl
@@ -562,10 +592,14 @@
          !
       CASE ( 3 )           ! lup and ldown, |dk[xml]| bounded by e3t :
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1   ! vector opt.
-               DO jk = mikt(ji,jj)+1, jpkm1         ! from the surface to the bottom : lup
+         DO jk = 2, jpkm1         ! from the surface to the bottom : lup
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zmxld(ji,jj,jk) = MIN( zmxld(ji,jj,jk-1) + fse3t(ji,jj,jk-1), zmxlm(ji,jj,jk) )
                END DO
-               DO jk = jpkm1, mikt(ji,jj)+1, -1     ! from the bottom to the surface : ldown
+            END DO
+         END DO
+         DO jk = jpkm1, 2, -1     ! from the bottom to the surface : ldown
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zmxlm(ji,jj,jk) = MIN( zmxlm(ji,jj,jk+1) + fse3t(ji,jj,jk+1), zmxlm(ji,jj,jk) )
                END DO
@@ -604,6 +638,6 @@
                zsqen = SQRT( en(ji,jj,jk) )
                zav   = rn_ediff * zmxlm(ji,jj,jk) * zsqen
-               avm  (ji,jj,jk) = MAX( zav,                  avmb(jk) ) * tmask(ji,jj,jk)
-               avt  (ji,jj,jk) = MAX( zav, avtb_2d(ji,jj) * avtb(jk) ) * tmask(ji,jj,jk)
+               avm  (ji,jj,jk) = MAX( zav,                  avmb(jk) ) * wmask(ji,jj,jk)
+               avt  (ji,jj,jk) = MAX( zav, avtb_2d(ji,jj) * avtb(jk) ) * wmask(ji,jj,jk)
                dissl(ji,jj,jk) = zsqen / zmxld(ji,jj,jk)
             END DO
@@ -612,11 +646,9 @@
       CALL lbc_lnk( avm, 'W', 1. )      ! Lateral boundary conditions (sign unchanged)
       !
-      DO jj = 2, jpjm1
-         DO ji = fs_2, fs_jpim1   ! vector opt.
-            DO jk = miku(ji,jj)+1, jpkm1            !* vertical eddy viscosity at u- and v-points
-               avmu(ji,jj,jk) = 0.5 * ( avm(ji,jj,jk) + avm(ji+1,jj  ,jk) ) * umask(ji,jj,jk)
-            END DO
-            DO jk = mikv(ji,jj)+1, jpkm1
-               avmv(ji,jj,jk) = 0.5 * ( avm(ji,jj,jk) + avm(ji  ,jj+1,jk) ) * vmask(ji,jj,jk)
+      DO jk = 2, jpkm1            !* vertical eddy viscosity at wu- and wv-points
+         DO jj = 2, jpjm1
+            DO ji = fs_2, fs_jpim1   ! vector opt.
+               avmu(ji,jj,jk) = 0.5 * ( avm(ji,jj,jk) + avm(ji+1,jj  ,jk) ) * wumask(ji,jj,jk)
+               avmv(ji,jj,jk) = 0.5 * ( avm(ji,jj,jk) + avm(ji  ,jj+1,jk) ) * wvmask(ji,jj,jk)
             END DO
          END DO
@@ -625,7 +657,7 @@
       !
       IF( nn_pdl == 1 ) THEN      !* Prandtl number case: update avt
-         DO jj = 2, jpjm1
-            DO ji = fs_2, fs_jpim1   ! vector opt.
-               DO jk = mikt(ji,jj)+1, jpkm1
+         DO jk = 2, jpkm1
+            DO jj = 2, jpjm1
+               DO ji = fs_2, fs_jpim1   ! vector opt.
                   zcoef = avm(ji,jj,jk) * 2._wp * fse3w(ji,jj,jk) * fse3w(ji,jj,jk)
                   !                                          ! shear
@@ -639,8 +671,8 @@
 !!gm and even better with the use of the "true" ri_crit=0.22222...  (this change the results!)
 !!gm              zpdlr = MAX(  0.1_wp,  ri_crit / MAX( ri_crit , zri )  )
-                  avt(ji,jj,jk)   = MAX( zpdlr * avt(ji,jj,jk), avtb_2d(ji,jj) * avtb(jk) ) * tmask(ji,jj,jk)
+                  avt(ji,jj,jk)   = MAX( zpdlr * avt(ji,jj,jk), avtb_2d(ji,jj) * avtb(jk) ) * wmask(ji,jj,jk)
 # if defined key_c1d
-                  e_pdl(ji,jj,jk) = zpdlr * tmask(ji,jj,jk)  ! c1d configuration : save masked Prandlt number
-                  e_ric(ji,jj,jk) = zri   * tmask(ji,jj,jk)  ! c1d config. : save Ri
+                  e_pdl(ji,jj,jk) = zpdlr * wmask(ji,jj,jk)  ! c1d configuration : save masked Prandlt number
+                  e_ric(ji,jj,jk) = zri   * wmask(ji,jj,jk)  ! c1d config. : save Ri
 # endif
               END DO
@@ -729,5 +761,5 @@
       IF( nn_pdl  < 0   .OR.  nn_pdl  > 1 )   CALL ctl_stop( 'bad flag: nn_pdl is  0 or 1    ' )
       IF( nn_htau < 0   .OR.  nn_htau > 1 )   CALL ctl_stop( 'bad flag: nn_htau is 0, 1 or 2 ' )
-      IF( nn_etau == 3 .AND. .NOT. lk_cpl )   CALL ctl_stop( 'nn_etau == 3 : HF taum only known in coupled mode' )
+      IF( nn_etau == 3 .AND. .NOT. ln_cpl )   CALL ctl_stop( 'nn_etau == 3 : HF taum only known in coupled mode' )
 
       IF( ln_mxl0 ) THEN
@@ -749,8 +781,8 @@
       !                               !* set vertical eddy coef. to the background value
       DO jk = 1, jpk
-         avt (:,:,jk) = avtb(jk) * tmask(:,:,jk)
-         avm (:,:,jk) = avmb(jk) * tmask(:,:,jk)
-         avmu(:,:,jk) = avmb(jk) * umask(:,:,jk)
-         avmv(:,:,jk) = avmb(jk) * vmask(:,:,jk)
+         avt (:,:,jk) = avtb(jk) * wmask (:,:,jk)
+         avm (:,:,jk) = avmb(jk) * wmask (:,:,jk)
+         avmu(:,:,jk) = avmb(jk) * wumask(:,:,jk)
+         avmv(:,:,jk) = avmb(jk) * wvmask(:,:,jk)
       END DO
       dissl(:,:,:) = 1.e-12_wp
@@ -803,4 +835,10 @@
               en (:,:,:) = rn_emin * tmask(:,:,:)
               CALL tke_avn                               ! recompute avt, avm, avmu, avmv and dissl (approximation)
+              !
+              avt_k (:,:,:) = avt (:,:,:)
+              avm_k (:,:,:) = avm (:,:,:)
+              avmu_k(:,:,:) = avmu(:,:,:)
+              avmv_k(:,:,:) = avmv(:,:,:)
+              !
               DO jit = nit000 + 1, nit000 + 10   ;   CALL zdf_tke( jit )   ;   END DO
            ENDIF
@@ -808,8 +846,8 @@
            en(:,:,:) = rn_emin * tmask(:,:,:)
            DO jk = 1, jpk                           ! set the Kz to the background value
-              avt (:,:,jk) = avtb(jk) * tmask(:,:,jk)
-              avm (:,:,jk) = avmb(jk) * tmask(:,:,jk)
-              avmu(:,:,jk) = avmb(jk) * umask(:,:,jk)
-              avmv(:,:,jk) = avmb(jk) * vmask(:,:,jk)
+              avt (:,:,jk) = avtb(jk) * wmask (:,:,jk)
+              avm (:,:,jk) = avmb(jk) * wmask (:,:,jk)
+              avmu(:,:,jk) = avmb(jk) * wumask(:,:,jk)
+              avmv(:,:,jk) = avmb(jk) * wvmask(:,:,jk)
            END DO
         ENDIF
