Index: NEMO/branches/2019/dev_r11078_OSMOSIS_IMMERSE_Nurser/src/OCE/ZDF/zdfosm.F90
===================================================================
--- NEMO/branches/2019/dev_r11078_OSMOSIS_IMMERSE_Nurser/src/OCE/ZDF/zdfosm.F90	(revision 12214)
+++ NEMO/branches/2019/dev_r11078_OSMOSIS_IMMERSE_Nurser/src/OCE/ZDF/zdfosm.F90	(revision 12216)
@@ -25,8 +25,13 @@
    !!            (12) Replace zwstrl with zvstr in calculation of eddy viscosity.
    !! 27/09/2017 (13) Calculate Stokes drift and Stokes penetration depth from wave information
-   !!            (14) Bouyancy flux due to entrainment changed to include contribution from shear turbulence (for testing commented out).
+   !!            (14) Buoyancy flux due to entrainment changed to include contribution from shear turbulence.
    !! 28/09/2017 (15) Calculation of Stokes drift moved into separate do-loops to allow for different options for the determining the Stokes drift to be added.
    !!            (16) Calculation of Stokes drift from windspeed for PM spectrum (for testing, commented out)
    !!            (17) Modification to Langmuir velocity scale to include effects due to the Stokes penetration depth (for testing, commented out)
+   !! ??/??/2018 (18) Revision to code structure, selected using key_osmldpth1. Inline code moved into subroutines. Changes to physics made,
+   !!                  (a) Pycnocline temperature and salinity profies changed for unstable layers
+   !!                  (b) The stable OSBL depth parametrization changed.
+   !! 16/05/2019 (19) Fox-Kemper parametrization of restratification through mixed layer eddies added to revised code.
+   !! 23/05/19   (20) Old code where key_osmldpth1` is *not* set removed, together with the key key_osmldpth1
    !!----------------------------------------------------------------------
 
@@ -40,4 +45,6 @@
    !!   trc_osm       : compute and add to the passive tracer trend the non-local flux (TBD)
    !!   dyn_osm       : compute and add to u & v trensd the non-local flux
+   !!
+   !! Subroutines in revised code.
    !!----------------------------------------------------------------------
    USE oce            ! ocean dynamics and active tracers
@@ -97,5 +104,6 @@
 
    !                                    !!! ** General constants  **
-   REAL(wp) ::   epsln   = 1.0e-20_wp   ! a small positive number
+   REAL(wp) ::   epsln   = 1.0e-20_wp   ! a small positive number to ensure no div by zero
+   REAL(wp) ::   depth_tol = 1.0e-6_wp  ! a small-ish positive number to give a hbl slightly shallower than gdepw
    REAL(wp) ::   pthird  = 1._wp/3._wp  ! 1/3
    REAL(wp) ::   p2third = 2._wp/3._wp  ! 2/3
@@ -166,5 +174,6 @@
       REAL(wp) ::   zbeta, zthermal                                  !
       REAL(wp) ::   zehat, zeta, zhrib, zsig, zscale, zwst, zws, zwm ! Velocity scales
-      REAL(wp) ::   zwsun, zwmun, zcons, zconm, zwcons, zwconm       !
+      REAL(wp) ::   zwsun, zwmun, zcons, zconm, zwcons, zwconm      !
+
       REAL(wp) ::   zsr, zbw, ze, zb, zd, zc, zaw, za, zb1, za1, zkw, zk0, zcomp , zrhd,zrhdr,zbvzed   ! In situ density
       INTEGER  ::   jm                          ! dummy loop indices
@@ -196,5 +205,5 @@
       REAL(wp), DIMENSION(jpi,jpj) :: zsin_wind ! Sin angle of surface stress
       REAL(wp), DIMENSION(jpi,jpj) :: zhol      ! Stability parameter for boundary layer
-      LOGICAL, DIMENSION(:,:), ALLOCATABLE :: lconv ! unstable/stable bl
+      LOGICAL, DIMENSION(jpi,jpj)  :: lconv     ! unstable/stable bl
 
       ! mixed-layer variables
@@ -238,5 +247,4 @@
       ! Temporary variables
       INTEGER :: inhml
-      INTEGER :: i_lconv_alloc
       REAL(wp) :: znd,znd_d,zznd_ml,zznd_pyc,zznd_d ! temporary non-dimensional depths used in various routines
       REAL(wp) :: ztemp, zari, zpert, zzdhdt, zdb   ! temporary variables
@@ -248,11 +256,14 @@
       REAL(wp), DIMENSION(jpi,jpj,jpk) :: zdiffut ! t-diffusivity
 
+      INTEGER :: ibld_ext=0                          ! does not have to be zero for modified scheme
+      REAL(wp) :: zwb_min, zgamma_b_nd, zgamma_b, zdhoh, ztau, zddhdt
+      REAL(wp) :: zzeta_s = 0._wp
+      REAL(wp) :: zzeta_v = 0.46
+      REAL(wp) :: zabsstke
+
       ! For debugging
       INTEGER :: ikt
       !!--------------------------------------------------------------------
       !
-      ALLOCATE( lconv(jpi,jpj),  STAT= i_lconv_alloc )
-      IF( i_lconv_alloc /= 0 )   CALL ctl_warn('zdf_osm: failed to allocate lconv')
-
       ibld(:,:)   = 0     ; imld(:,:)  = 0
       zrad0(:,:)  = 0._wp ; zradh(:,:) = 0._wp ; zradav(:,:)    = 0._wp ; zustar(:,:)    = 0._wp
@@ -268,4 +279,5 @@
       zt_bl(:,:)   = 0._wp ; zs_bl(:,:)   = 0._wp ; zu_bl(:,:)    = 0._wp ; zv_bl(:,:)   = 0._wp
       zrh_bl(:,:)  = 0._wp ; zt_ml(:,:)   = 0._wp ; zs_ml(:,:)    = 0._wp ; zu_ml(:,:)   = 0._wp
+
       zv_ml(:,:)   = 0._wp ; zrh_ml(:,:)  = 0._wp ; zdt_bl(:,:)   = 0._wp ; zds_bl(:,:)  = 0._wp
       zdu_bl(:,:)  = 0._wp ; zdv_bl(:,:)  = 0._wp ; zdrh_bl(:,:)  = 0._wp ; zdb_bl(:,:)  = 0._wp
@@ -287,4 +299,5 @@
       ghams(:,:,:)   = 0._wp ; ghamu(:,:,:)   = 0._wp ; ghamv(:,:,:) = 0._wp
 
+      zdhdt_2(:,:) = 0._wp
       ! hbl = MAX(hbl,epsln)
       !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
@@ -350,5 +363,5 @@
               ! Use wind speed wndm included in sbc_oce module
               zustke(ji,jj) = MAX ( 0.016 * wndm(ji,jj), 1.0e-8 )
-              dstokes(ji,jj) = 0.12 * wndm(ji,jj)**2 / grav
+              dstokes(ji,jj) = MAX( 0.12 * wndm(ji,jj)**2 / grav, 5.e-1)
            END DO
         END DO
@@ -362,7 +375,7 @@
               ! It could represent the effects of the spread of wave directions
               ! around the mean wind. The effect of this adjustment needs to be tested.
-              zustke(ji,jj) = MAX ( 1.0 * ( zcos_wind(ji,jj) * ut0sd(ji,jj ) + zsin_wind(ji,jj)  * vt0sd(ji,jj) ), &
-                   &                zustar(ji,jj) / ( 0.45 * 0.45 )                                                  )
-              dstokes(ji,jj) = MAX(zfac * hsw(ji,jj)*hsw(ji,jj) / ( MAX(zustke(ji,jj)*wmp(ji,jj), 1.0e-7 ) ), 5.0e-1) !rn_osm_dstokes !
+              zabsstke = SQRT(ut0sd(ji,jj)**2 + vt0sd(ji,jj)**2)
+              zustke(ji,jj) = MAX (0.8 * ( zcos_wind(ji,jj) * ut0sd(ji,jj) + zsin_wind(ji,jj)  * vt0sd(ji,jj) ), 1.0e-8)
+              dstokes(ji,jj) = MAX(zfac * hsw(ji,jj)*hsw(ji,jj) / ( MAX(zabsstke*wmp(ji,jj), 1.0e-7 ) ), 5.0e-1) !rn_osm_dstokes !
            END DO
         END DO
@@ -375,12 +388,6 @@
            ! Langmuir velocity scale (zwstrl), at T-point
            zwstrl(ji,jj) = ( zustar(ji,jj) * zustar(ji,jj) * zustke(ji,jj) )**pthird
-           ! Modify zwstrl to allow for small and large values of dstokes/hbl.
-           ! Intended as a possible test. Doesn't affect LES results for entrainment,
-           !  but hasn't been shown to be correct as dstokes/h becomes large or small.
-           zwstrl(ji,jj) = zwstrl(ji,jj) *  &
-                & (1.12 * ( 1.0 - ( 1.0 - EXP( -hbl(ji,jj) / dstokes(ji,jj) ) ) * dstokes(ji,jj) / hbl(ji,jj) ))**pthird * &
-                & ( 1.0 - EXP( -15.0 * dstokes(ji,jj) / hbl(ji,jj) ))
-           ! define La this way so effects of Stokes penetration depth on velocity scale are included
-           zla(ji,jj) = SQRT ( zustar(ji,jj) / zwstrl(ji,jj) )**3
+           zla(ji,jj) = MAX(MIN(SQRT ( zustar(ji,jj) / ( zwstrl(ji,jj) + epsln ) )**3, 4.0), 0.2)
+           IF(zla(ji,jj) > 0.45) dstokes(ji,jj) = MIN(dstokes(ji,jj), 0.5_wp*hbl(ji,jj))
            ! Velocity scale that tends to zustar for large Langmuir numbers
            zvstr(ji,jj) = ( zwstrl(ji,jj)**3  + &
@@ -389,5 +396,4 @@
            ! limit maximum value of Langmuir number as approximate treatment for shear turbulence.
            ! Note zustke and zwstrl are not amended.
-           IF ( zla(ji,jj) >= 0.45 ) zla(ji,jj) = 0.45
            !
            ! get convective velocity (zwstrc), stabilty scale (zhol) and logical conection flag lconv
@@ -406,8 +412,8 @@
      ! Mixed-layer model - calculate averages over the boundary layer, and the change in the boundary layer depth
      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-     ! BL must be always 2 levels deep.
-      hbl(:,:) = MAX(hbl(:,:), gdepw_n(:,:,3) )
-      ibld(:,:) = 3
-      DO jk = 4, jpkm1
+     ! BL must be always 4 levels deep.
+      hbl(:,:) = MAX(hbl(:,:), gdepw_n(:,:,4) )
+      ibld(:,:) = 4
+      DO jk = 5, jpkm1
          DO jj = 2, jpjm1
             DO ji = 2, jpim1
@@ -419,66 +425,10 @@
       END DO
 
-      DO jj = 2, jpjm1                                 !  Vertical slab
+      DO jj = 2, jpjm1
          DO ji = 2, jpim1
-               zthermal = rab_n(ji,jj,1,jp_tem) !ideally use ibld not 1??
-               zbeta    = rab_n(ji,jj,1,jp_sal)
-               zt   = 0._wp
-               zs   = 0._wp
-               zu   = 0._wp
-               zv   = 0._wp
-               ! average over depth of boundary layer
-               zthick=0._wp
-               DO jm = 2, ibld(ji,jj)
-                  zthick=zthick+e3t_n(ji,jj,jm)
-                  zt   = zt  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_tem)
-                  zs   = zs  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_sal)
-                  zu   = zu  + e3t_n(ji,jj,jm) &
-                     &            * ( ub(ji,jj,jm) + ub(ji - 1,jj,jm) ) &
-                     &            / MAX( 1. , umask(ji,jj,jm) + umask(ji - 1,jj,jm) )
-                  zv   = zv  + e3t_n(ji,jj,jm) &
-                     &            * ( vb(ji,jj,jm) + vb(ji,jj - 1,jm) ) &
-                     &            / MAX( 1. , vmask(ji,jj,jm) + vmask(ji,jj - 1,jm) )
-               END DO
-               zt_bl(ji,jj) = zt / zthick
-               zs_bl(ji,jj) = zs / zthick
-               zu_bl(ji,jj) = zu / zthick
-               zv_bl(ji,jj) = zv / zthick
-               zdt_bl(ji,jj) = zt_bl(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_tem)
-               zds_bl(ji,jj) = zs_bl(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_sal)
-               zdu_bl(ji,jj) = zu_bl(ji,jj) - ( ub(ji,jj,ibld(ji,jj)) + ub(ji-1,jj,ibld(ji,jj) ) ) &
-                     &    / MAX(1. , umask(ji,jj,ibld(ji,jj) ) + umask(ji-1,jj,ibld(ji,jj) ) )
-               zdv_bl(ji,jj) = zv_bl(ji,jj) - ( vb(ji,jj,ibld(ji,jj)) + vb(ji,jj-1,ibld(ji,jj) ) ) &
-                     &   / MAX(1. , vmask(ji,jj,ibld(ji,jj) ) + vmask(ji,jj-1,ibld(ji,jj) ) )
-               zdb_bl(ji,jj) = grav * zthermal * zdt_bl(ji,jj) - grav * zbeta * zds_bl(ji,jj)
-               IF ( lconv(ji,jj) ) THEN    ! Convective
-                      zwb_ent(ji,jj) = -( 2.0 * 0.2 * zwbav(ji,jj) &
-                           &            + 0.135 * zla(ji,jj) * zwstrl(ji,jj)**3/hbl(ji,jj) )
-
-                      zvel_max =  - ( 1.0 + 1.0 * ( zwstrl(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * rn_rdt / hbl(ji,jj) ) &
-                           &   * zwb_ent(ji,jj) / ( zwstrl(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird
-! Entrainment including component due to shear turbulence. Modified Langmuir component, but gives same result for La=0.3 For testing uncomment.
-!                      zwb_ent(ji,jj) = -( 2.0 * 0.2 * zwbav(ji,jj) &
-!                           &            + ( 0.15 * ( 1.0 - EXP( -0.5 * zla(ji,jj) ) ) + 0.03 / zla(ji,jj)**2 ) * zustar(ji,jj)**3/hbl(ji,jj) )
-
-!                      zvel_max = - ( 1.0 + 1.0 * ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * rn_rdt / zhbl(ji,jj) ) * zwb_ent(ji,jj) / &
-!                           &       ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird
-                      zzdhdt = - zwb_ent(ji,jj) / ( zvel_max + MAX(zdb_bl(ji,jj),0.0) )
-               ELSE                        ! Stable
-                      zzdhdt = 0.32 * ( hbli(ji,jj) / hbl(ji,jj) -1.0 ) * zwstrl(ji,jj)**3 / hbli(ji,jj) &
-                           &   + ( ( 0.32 / 3.0 ) * exp ( -2.5 * ( hbli(ji,jj) / hbl(ji,jj) - 1.0 ) ) &
-                           & - ( 0.32 / 3.0 - 0.135 * zla(ji,jj) ) * exp ( -12.5 * ( hbli(ji,jj) / hbl(ji,jj) ) ) ) &
-                           &  * zwstrl(ji,jj)**3 / hbli(ji,jj)
-                      zzdhdt = zzdhdt + zwbav(ji,jj)
-                      IF ( zzdhdt < 0._wp ) THEN
-                      ! For long timsteps factor in brackets slows the rapid collapse of the OSBL
-                         zpert   = 2.0 * ( 1.0 + 2.0 * zwstrl(ji,jj) * rn_rdt / hbl(ji,jj) ) * zwstrl(ji,jj)**2 / hbl(ji,jj)
-                      ELSE
-                         zpert   = 2.0 * ( 1.0 + 2.0 * zwstrl(ji,jj) * rn_rdt / hbl(ji,jj) ) * zwstrl(ji,jj)**2 / hbl(ji,jj) &
-                              &  + MAX( zdb_bl(ji,jj), 0.0 )
-                      ENDIF
-                      zzdhdt = 2.0 * zzdhdt / zpert
-               ENDIF
-               zdhdt(ji,jj) = zzdhdt
-           END DO
+            zhbl(ji,jj) = gdepw_n(ji,jj,ibld(ji,jj))
+            imld(ji,jj) = MAX(3,ibld(ji,jj) - MAX( INT( dh(ji,jj) / e3t_n(ji, jj, ibld(ji,jj) )) , 1 ))
+            zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
+         END DO
       END DO
 
@@ -495,5 +445,5 @@
             DO ji = 2, jpim1
                IF ( zhbl_t(ji,jj) >= gdepw_n(ji,jj,jk) ) THEN
-                  ibld(ji,jj) =  MIN(mbkt(ji,jj), jk)
+                  ibld(ji,jj) = jk
                ENDIF
             END DO
@@ -504,256 +454,40 @@
 ! Step through model levels taking account of buoyancy change to determine the effect on dhdt
 !
-      DO jj = 2, jpjm1
-         DO ji = 2, jpim1
-            IF ( ibld(ji,jj) - imld(ji,jj) > 1 ) THEN
+      CALL zdf_osm_timestep_hbl( zdhdt, zdhdt_2 )
+      ! Alan: do we need zb_ml?
+      CALL zdf_osm_vertical_average( ibld, zt_bl, zs_bl, zu_bl, zv_bl, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
 !
-! If boundary layer changes by more than one level, need to check for stable layers between initial and final depths.
 !
-               zhbl_s = hbl(ji,jj)
-               jm = imld(ji,jj)
-               zthermal = rab_n(ji,jj,1,jp_tem)
-               zbeta = rab_n(ji,jj,1,jp_sal)
-               IF ( lconv(ji,jj) ) THEN
-!unstable
-                  zvel_max =  - ( 1.0 + 1.0 * ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * rn_rdt / hbl(ji,jj) ) &
-                       &   * zwb_ent(ji,jj) / ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird
-
-                  DO jk = imld(ji,jj), ibld(ji,jj)
-                     zdb = MAX( grav * ( zthermal * ( zt_bl(ji,jj) - tsn(ji,jj,jm,jp_tem) ) &
-                          & - zbeta * ( zs_bl(ji,jj) - tsn(ji,jj,jm,jp_sal) ) ), 0.0 ) + zvel_max
-
-                     zhbl_s = zhbl_s + MIN( - zwb_ent(ji,jj) / zdb * rn_rdt / FLOAT(ibld(ji,jj)-imld(ji,jj) ), e3w_n(ji,jj,jk) )
-                     zhbl_s = MIN(zhbl_s, ht_n(ji,jj))
-
-                     IF ( zhbl_s >= gdepw_n(ji,jj,jm+1) ) jm = jm + 1
-                  END DO
-                  hbl(ji,jj) = zhbl_s
-                  ibld(ji,jj) = jm
-                  hbli(ji,jj) = hbl(ji,jj)
-               ELSE
-! stable
-                  DO jk = imld(ji,jj), ibld(ji,jj)
-                     zdb = MAX( grav * ( zthermal * ( zt_bl(ji,jj) - tsn(ji,jj,jm,jp_tem) )          &
-                          &               - zbeta * ( zs_bl(ji,jj) - tsn(ji,jj,jm,jp_sal) ) ), 0.0 ) &
-                          & + 2.0 * zwstrl(ji,jj)**2 / zhbl_s
-
-                     zhbl_s = zhbl_s +  (                                                                                &
-                          &                     0.32         *                         ( hbli(ji,jj) / zhbl_s -1.0 )     &
-                          &               * zwstrl(ji,jj)**3 / hbli(ji,jj)                                               &
-                          &               + ( ( 0.32 / 3.0 )           * EXP( -  2.5 * ( hbli(ji,jj) / zhbl_s -1.0 ) )   &
-                          &               -   ( 0.32 / 3.0  - 0.0485 ) * EXP( - 12.5 * ( hbli(ji,jj) / zhbl_s      ) ) ) &
-                          &          * zwstrl(ji,jj)**3 / hbli(ji,jj) ) / zdb * e3w_n(ji,jj,jk) / zdhdt(ji,jj)  ! ALMG to investigate whether need to include wn here
-
-                     zhbl_s = MIN(zhbl_s, ht_n(ji,jj))
-                     IF ( zhbl_s >= gdepw_n(ji,jj,jm) ) jm = jm + 1
-                  END DO
-                  hbl(ji,jj) = MAX(zhbl_s, gdepw_n(ji,jj,3) )
-                  ibld(ji,jj) = MAX(jm, 3 )
-                  IF ( hbl(ji,jj) > hbli(ji,jj) ) hbli(ji,jj) = hbl(ji,jj)
-               ENDIF   ! IF ( lconv )
-            ELSE
-! change zero or one model level.
-               hbl(ji,jj) = zhbl_t(ji,jj)
-               IF ( lconv(ji,jj) ) THEN
-                  hbli(ji,jj) = hbl(ji,jj)
-               ELSE
-                  hbl(ji,jj) = MAX(hbl(ji,jj), gdepw_n(ji,jj,3) )
-                  IF ( hbl(ji,jj) > hbli(ji,jj) ) hbli(ji,jj) = hbl(ji,jj)
-               ENDIF
-            ENDIF
-            zhbl(ji,jj) = gdepw_n(ji,jj,ibld(ji,jj))
-         END DO
-      END DO
+      CALL zdf_osm_pycnocline_thickness( dh, zdh )
       dstokes(:,:) = MIN ( dstokes(:,:), hbl(:,:)/3. )  !  Limit delta for shallow boundary layers for calculating flux-gradient terms.
-
-! Recalculate averages over boundary layer after depth updated
-     ! Consider later  combining this into the loop above and looking for columns
-     ! where the index for base of the boundary layer have changed
-      DO jj = 2, jpjm1                                 !  Vertical slab
-         DO ji = 2, jpim1
-               zthermal = rab_n(ji,jj,1,jp_tem) !ideally use ibld not 1??
-               zbeta    = rab_n(ji,jj,1,jp_sal)
-               zt   = 0._wp
-               zs   = 0._wp
-               zu   = 0._wp
-               zv   = 0._wp
-               ! average over depth of boundary layer
-               zthick=0._wp
-               DO jm = 2, ibld(ji,jj)
-                  zthick=zthick+e3t_n(ji,jj,jm)
-                  zt   = zt  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_tem)
-                  zs   = zs  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_sal)
-                  zu   = zu  + e3t_n(ji,jj,jm) &
-                     &            * ( ub(ji,jj,jm) + ub(ji - 1,jj,jm) ) &
-                     &            / MAX( 1. , umask(ji,jj,jm) + umask(ji - 1,jj,jm) )
-                  zv   = zv  + e3t_n(ji,jj,jm) &
-                     &            * ( vb(ji,jj,jm) + vb(ji,jj - 1,jm) ) &
-                     &            / MAX( 1. , vmask(ji,jj,jm) + vmask(ji,jj - 1,jm) )
-               END DO
-               zt_bl(ji,jj) = zt / zthick
-               zs_bl(ji,jj) = zs / zthick
-               zu_bl(ji,jj) = zu / zthick
-               zv_bl(ji,jj) = zv / zthick
-               zdt_bl(ji,jj) = zt_bl(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_tem)
-               zds_bl(ji,jj) = zs_bl(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_sal)
-               zdu_bl(ji,jj) = zu_bl(ji,jj) - ( ub(ji,jj,ibld(ji,jj)) + ub(ji-1,jj,ibld(ji,jj) ) ) &
-                      &   / MAX(1. , umask(ji,jj,ibld(ji,jj) ) + umask(ji-1,jj,ibld(ji,jj) ) )
-               zdv_bl(ji,jj) = zv_bl(ji,jj) - ( vb(ji,jj,ibld(ji,jj)) + vb(ji,jj-1,ibld(ji,jj) ) ) &
-                      &  / MAX(1. , vmask(ji,jj,ibld(ji,jj) ) + vmask(ji,jj-1,ibld(ji,jj) ) )
-               zdb_bl(ji,jj) = grav * zthermal * zdt_bl(ji,jj) - grav * zbeta * zds_bl(ji,jj)
-               zhbl(ji,jj) = gdepw_n(ji,jj,ibld(ji,jj))
-               IF ( lconv(ji,jj) ) THEN
-                  IF ( zdb_bl(ji,jj) > 0._wp )THEN
-                     IF ( ( zwstrc(ji,jj) / zvstr(ji,jj) )**3 <= 0.5 ) THEN  ! near neutral stability
-                           zari = 4.5 * ( zvstr(ji,jj)**2 ) &
-                             & / ( zdb_bl(ji,jj) * zhbl(ji,jj) ) + 0.01
-                     ELSE                                                     ! unstable
-                           zari = 4.5 * ( zwstrc(ji,jj)**2 ) &
-                             & / ( zdb_bl(ji,jj) * zhbl(ji,jj) ) + 0.01
-                     ENDIF
-                     IF ( zari > 0.2 ) THEN                                                ! This test checks for weakly stratified pycnocline
-                        zari = 0.2
-                        zwb_ent(ji,jj) = 0._wp
-                     ENDIF
-                     inhml = MAX( INT( zari * zhbl(ji,jj) / e3t_n(ji,jj,ibld(ji,jj)) ) , 1 )
-                     imld(ji,jj) = MAX( ibld(ji,jj) - inhml, 1)
-                     zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
-                     zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
-                  ELSE  ! IF (zdb_bl)
-                     imld(ji,jj) = ibld(ji,jj) - 1
-                     zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
-                     zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
-                  ENDIF
-               ELSE   ! IF (lconv)
-                  IF ( zdhdt(ji,jj) >= 0.0 ) THEN    ! probably shouldn't include wm here
-                  ! boundary layer deepening
-                     IF ( zdb_bl(ji,jj) > 0._wp ) THEN
-                  ! pycnocline thickness set by stratification - use same relationship as for neutral conditions.
-                        zari = MIN( 4.5 * ( zvstr(ji,jj)**2 ) &
-                          & / ( zdb_bl(ji,jj) * zhbl(ji,jj) ) + 0.01  , 0.2 )
-                        inhml = MAX( INT( zari * zhbl(ji,jj) / e3t_n(ji,jj,ibld(ji,jj)) ) , 1 )
-                        imld(ji,jj) = MAX( ibld(ji,jj) - inhml, 1)
-                        zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
-                        zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
-                     ELSE
-                        imld(ji,jj) = ibld(ji,jj) - 1
-                        zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
-                        zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
-                     ENDIF ! IF (zdb_bl > 0.0)
-                  ELSE     ! IF(dhdt >= 0)
-                  ! boundary layer collapsing.
-                     imld(ji,jj) = ibld(ji,jj)
-                     zhml(ji,jj) = zhbl(ji,jj)
-                     zdh(ji,jj) = 0._wp
-                  ENDIF    ! IF (dhdt >= 0)
-               ENDIF       ! IF (lconv)
-         END DO
-      END DO
-
-      ! Average over the depth of the mixed layer in the convective boundary layer
-      ! Also calculate entrainment fluxes for temperature and salinity
-      DO jj = 2, jpjm1                                 !  Vertical slab
-         DO ji = 2, jpim1
-            zthermal = rab_n(ji,jj,1,jp_tem) !ideally use ibld not 1??
-            zbeta    = rab_n(ji,jj,1,jp_sal)
-            IF ( lconv(ji,jj) ) THEN
-               zt   = 0._wp
-               zs   = 0._wp
-               zu   = 0._wp
-               zv   = 0._wp
-               ! average over depth of boundary layer
-               zthick=0._wp
-               DO jm = 2, imld(ji,jj)
-                  zthick=zthick+e3t_n(ji,jj,jm)
-                  zt   = zt  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_tem)
-                  zs   = zs  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_sal)
-                  zu   = zu  + e3t_n(ji,jj,jm) &
-                     &            * ( ub(ji,jj,jm) + ub(ji - 1,jj,jm) ) &
-                     &            / MAX( 1. , umask(ji,jj,jm) + umask(ji - 1,jj,jm) )
-                  zv   = zv  + e3t_n(ji,jj,jm) &
-                     &            * ( vb(ji,jj,jm) + vb(ji,jj - 1,jm) ) &
-                     &            / MAX( 1. , vmask(ji,jj,jm) + vmask(ji,jj - 1,jm) )
-               END DO
-               zt_ml(ji,jj) = zt / zthick
-               zs_ml(ji,jj) = zs / zthick
-               zu_ml(ji,jj) = zu / zthick
-               zv_ml(ji,jj) = zv / zthick
-               zdt_ml(ji,jj) = zt_ml(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_tem)
-               zds_ml(ji,jj) = zs_ml(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_sal)
-               zdu_ml(ji,jj) = zu_ml(ji,jj) - ( ub(ji,jj,ibld(ji,jj)) + ub(ji-1,jj,ibld(ji,jj) ) ) &
-                     &    / MAX(1. , umask(ji,jj,ibld(ji,jj) ) + umask(ji-1,jj,ibld(ji,jj) ) )
-               zdv_ml(ji,jj) = zv_ml(ji,jj) - ( vb(ji,jj,ibld(ji,jj)) + vb(ji,jj-1,ibld(ji,jj) ) ) &
-                     &    / MAX(1. , vmask(ji,jj,ibld(ji,jj) ) + vmask(ji,jj-1,ibld(ji,jj) ) )
-               zdb_ml(ji,jj) = grav * zthermal * zdt_ml(ji,jj) - grav * zbeta * zds_ml(ji,jj)
-            ELSE
-            ! stable, if entraining calulate average below interface layer.
-               IF ( zdhdt(ji,jj) >= 0._wp ) THEN
-                  zt   = 0._wp
-                  zs   = 0._wp
-                  zu   = 0._wp
-                  zv   = 0._wp
-                  ! average over depth of boundary layer
-                  zthick=0._wp
-                  DO jm = 2, imld(ji,jj)
-                     zthick=zthick+e3t_n(ji,jj,jm)
-                     zt   = zt  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_tem)
-                     zs   = zs  + e3t_n(ji,jj,jm) * tsn(ji,jj,jm,jp_sal)
-                     zu   = zu  + e3t_n(ji,jj,jm) &
-                        &            * ( ub(ji,jj,jm) + ub(ji - 1,jj,jm) ) &
-                        &            / MAX( 1. , umask(ji,jj,jm) + umask(ji - 1,jj,jm) )
-                     zv   = zv  + e3t_n(ji,jj,jm) &
-                        &            * ( vb(ji,jj,jm) + vb(ji,jj - 1,jm) ) &
-                        &            / MAX( 1. , vmask(ji,jj,jm) + vmask(ji,jj - 1,jm) )
-                  END DO
-                  zt_ml(ji,jj) = zt / zthick
-                  zs_ml(ji,jj) = zs / zthick
-                  zu_ml(ji,jj) = zu / zthick
-                  zv_ml(ji,jj) = zv / zthick
-                  zdt_ml(ji,jj) = zt_ml(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_tem)
-                  zds_ml(ji,jj) = zs_ml(ji,jj) - tsn(ji,jj,ibld(ji,jj),jp_sal)
-                  zdu_ml(ji,jj) = zu_ml(ji,jj) - ( ub(ji,jj,ibld(ji,jj)) + ub(ji-1,jj,ibld(ji,jj) ) ) &
-                        &    / MAX(1. , umask(ji,jj,ibld(ji,jj) ) + umask(ji-1,jj,ibld(ji,jj) ) )
-                  zdv_ml(ji,jj) = zv_ml(ji,jj) - ( vb(ji,jj,ibld(ji,jj)) + vb(ji,jj-1,ibld(ji,jj) ) ) &
-                        &    / MAX(1. , vmask(ji,jj,ibld(ji,jj) ) + vmask(ji,jj-1,ibld(ji,jj) ) )
-                  zdb_ml(ji,jj) = grav * zthermal * zdt_ml(ji,jj) - grav * zbeta * zds_ml(ji,jj)
-               ENDIF
-            ENDIF
-         END DO
-      END DO
-    !
+!
+    ! Average over the depth of the mixed layer in the convective boundary layer
+    ! Alan: do we need zb_ml?
+    CALL zdf_osm_vertical_average( imld, zt_ml, zs_ml, zu_ml, zv_ml, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml )
     ! rotate mean currents and changes onto wind align co-ordinates
     !
-
-      DO jj = 2, jpjm1
-         DO ji = 2, jpim1
-            ztemp = zu_ml(ji,jj)
-            zu_ml(ji,jj) = zu_ml(ji,jj) * zcos_wind(ji,jj) + zv_ml(ji,jj) * zsin_wind(ji,jj)
-            zv_ml(ji,jj) = zv_ml(ji,jj) * zcos_wind(ji,jj) - ztemp * zsin_wind(ji,jj)
-            ztemp = zdu_ml(ji,jj)
-            zdu_ml(ji,jj) = zdu_ml(ji,jj) * zcos_wind(ji,jj) + zdv_ml(ji,jj) * zsin_wind(ji,jj)
-            zdv_ml(ji,jj) = zdv_ml(ji,jj) * zsin_wind(ji,jj) - ztemp * zsin_wind(ji,jj)
-    !
-            ztemp = zu_bl(ji,jj)
-            zu_bl = zu_bl(ji,jj) * zcos_wind(ji,jj) + zv_bl(ji,jj) * zsin_wind(ji,jj)
-            zv_bl(ji,jj) = zv_bl(ji,jj) * zcos_wind(ji,jj) - ztemp * zsin_wind(ji,jj)
-            ztemp = zdu_bl(ji,jj)
-            zdu_bl(ji,jj) = zdu_bl(ji,jj) * zcos_wind(ji,jj) + zdv_bl(ji,jj) * zsin_wind(ji,jj)
-            zdv_bl(ji,jj) = zdv_bl(ji,jj) * zsin_wind(ji,jj) - ztemp * zsin_wind(ji,jj)
-         END DO
-      END DO
-
+    CALL zdf_osm_velocity_rotation( zcos_wind, zsin_wind, zu_ml, zv_ml, zdu_ml, zdv_ml )
+    CALL zdf_osm_velocity_rotation( zcos_wind, zsin_wind, zu_bl, zv_bl, zdu_bl, zdv_bl )
      zuw_bse = 0._wp
      zvw_bse = 0._wp
+     zwth_ent = 0._wp
+     zws_ent = 0._wp
      DO jj = 2, jpjm1
         DO ji = 2, jpim1
-
-           IF ( lconv(ji,jj) ) THEN
-              IF ( zdb_bl(ji,jj) > 0._wp ) THEN
-                 zwth_ent(ji,jj) = zwb_ent(ji,jj) * zdt_ml(ji,jj) / (zdb_ml(ji,jj) + epsln)
-                 zws_ent(ji,jj) = zwb_ent(ji,jj) * zds_ml(ji,jj) / (zdb_ml(ji,jj) + epsln)
+           IF ( ibld(ji,jj) < mbkt(ji,jj) ) THEN
+              IF ( lconv(ji,jj) ) THEN
+         	 zuw_bse(ji,jj) = -0.0075*((zvstr(ji,jj)**3+0.5*zwstrc(ji,jj)**3)**pthird*zdu_ml(ji,jj) + &
+                      &                    1.5*zustar(ji,jj)**2*(zhbl(ji,jj)-zhml(ji,jj)) )/ &
+                      &                     ( zhml(ji,jj)*MIN(zla(ji,jj)**(8./3.),1.) + epsln)
+	         zvw_bse(ji,jj) = 0.01*(-(zvstr(ji,jj)**3+0.5*zwstrc(ji,jj)**3)**pthird*zdv_ml(ji,jj)+ &
+                      &                    2.0*ff_t(ji,jj)*zustke(ji,jj)*dstokes(ji,jj)*zla(ji,jj))
+                 IF ( zdb_ml(ji,jj) > 0._wp ) THEN
+                    zwth_ent(ji,jj) = zwb_ent(ji,jj) * zdt_ml(ji,jj) / (zdb_ml(ji,jj) + epsln)
+                    zws_ent(ji,jj) = zwb_ent(ji,jj) * zds_ml(ji,jj) / (zdb_ml(ji,jj) + epsln)
+                 ENDIF
+              ELSE
+                 zwth_ent(ji,jj) = -2.0 * zwthav(ji,jj) * ( (1.0 - 0.8) - ( 1.0 - 0.8)**(3.0/2.0) )
+                 zws_ent(ji,jj) = -2.0 * zwsav(ji,jj) * ( (1.0 - 0.8 ) - ( 1.0 - 0.8 )**(3.0/2.0) )
               ENDIF
-           ELSE
-              zwth_ent(ji,jj) = -2.0 * zwthav(ji,jj) * ( (1.0 - 0.8) - ( 1.0 - 0.8)**(3.0/2.0) )
-              zws_ent(ji,jj) = -2.0 * zwsav(ji,jj) * ( (1.0 - 0.8 ) - ( 1.0 - 0.8 )**(3.0/2.0) )
            ENDIF
         END DO
@@ -764,100 +498,11 @@
       !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
 
-       DO jj = 2, jpjm1
-          DO ji = 2, jpim1
-          !
-             IF ( lconv (ji,jj) ) THEN
-             ! Unstable conditions
-                IF( zdb_bl(ji,jj) > 0._wp ) THEN
-                ! calculate pycnocline profiles, no need if zdb_bl <= 0. since profile is zero and arrays have been initialized to zero
-                   ztgrad = ( zdt_ml(ji,jj) / zdh(ji,jj) )
-                   zsgrad = ( zds_ml(ji,jj) / zdh(ji,jj) )
-                   zbgrad = ( zdb_ml(ji,jj) / zdh(ji,jj) )
-                   DO jk = 2 , ibld(ji,jj)
-                      znd = -( gdepw_n(ji,jj,jk) - zhml(ji,jj) ) / zdh(ji,jj)
-                      zdtdz_pyc(ji,jj,jk) =  ztgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                      zdbdz_pyc(ji,jj,jk) =  zbgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                      zdsdz_pyc(ji,jj,jk) =  zsgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                   END DO
-                ENDIF
-             ELSE
-             ! stable conditions
-             ! if pycnocline profile only defined when depth steady of increasing.
-                IF ( zdhdt(ji,jj) >= 0.0 ) THEN        ! Depth increasing, or steady.
-                   IF ( zdb_bl(ji,jj) > 0._wp ) THEN
-                     IF ( zhol(ji,jj) >= 0.5 ) THEN      ! Very stable - 'thick' pycnocline
-                         ztgrad = zdt_bl(ji,jj) / zhbl(ji,jj)
-                         zsgrad = zds_bl(ji,jj) / zhbl(ji,jj)
-                         zbgrad = zdb_bl(ji,jj) / zhbl(ji,jj)
-                         DO jk = 2, ibld(ji,jj)
-                            znd = gdepw_n(ji,jj,jk) / zhbl(ji,jj)
-                            zdtdz_pyc(ji,jj,jk) =  ztgrad * EXP( -15.0 * ( znd - 0.9 )**2 )
-                            zdbdz_pyc(ji,jj,jk) =  zbgrad * EXP( -15.0 * ( znd - 0.9 )**2 )
-                            zdsdz_pyc(ji,jj,jk) =  zsgrad * EXP( -15.0 * ( znd - 0.9 )**2 )
-                         END DO
-                     ELSE                                   ! Slightly stable - 'thin' pycnoline - needed when stable layer begins to form.
-                         ztgrad = zdt_bl(ji,jj) / zdh(ji,jj)
-                         zsgrad = zds_bl(ji,jj) / zdh(ji,jj)
-                         zbgrad = zdb_bl(ji,jj) / zdh(ji,jj)
-                         DO jk = 2, ibld(ji,jj)
-                            znd = -( gdepw_n(ji,jj,jk) - zhml(ji,jj) ) / zdh(ji,jj)
-                            zdtdz_pyc(ji,jj,jk) =  ztgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                            zdbdz_pyc(ji,jj,jk) =  zbgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                            zdsdz_pyc(ji,jj,jk) =  zsgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                         END DO
-                      ENDIF ! IF (zhol >=0.5)
-                   ENDIF    ! IF (zdb_bl> 0.)
-                ENDIF       ! IF (zdhdt >= 0) zdhdt < 0 not considered since pycnocline profile is zero, profile arrays are intialized to zero
-             ENDIF          ! IF (lconv)
-            !
-          END DO
-       END DO
-!
-       DO jj = 2, jpjm1
-          DO ji = 2, jpim1
-          !
-             IF ( lconv (ji,jj) ) THEN
-             ! Unstable conditions
-                 zugrad = ( zdu_ml(ji,jj) / zdh(ji,jj) ) + 0.275 * zustar(ji,jj)*zustar(ji,jj) / &
-               & (( zwstrl(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * zhml(ji,jj) ) / zla(ji,jj)**(8.0/3.0)
-                zvgrad = ( zdv_ml(ji,jj) / zdh(ji,jj) ) + 3.5 * ff_t(ji,jj) * zustke(ji,jj) / &
-              & ( zwstrl(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird
-                DO jk = 2 , ibld(ji,jj)-1
-                   znd = -( gdepw_n(ji,jj,jk) - zhml(ji,jj) ) / zdh(ji,jj)
-                   zdudz_pyc(ji,jj,jk) =  zugrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                   zdvdz_pyc(ji,jj,jk) = zvgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
-                END DO
-             ELSE
-             ! stable conditions
-                zugrad = 3.25 * zdu_bl(ji,jj) / zhbl(ji,jj)
-                zvgrad = 2.75 * zdv_bl(ji,jj) / zhbl(ji,jj)
-                DO jk = 2, ibld(ji,jj)
-                   znd = gdepw_n(ji,jj,jk) / zhbl(ji,jj)
-                   IF ( znd < 1.0 ) THEN
-                      zdudz_pyc(ji,jj,jk) = zugrad * EXP( -40.0 * ( znd - 1.0 )**2 )
-                   ELSE
-                      zdudz_pyc(ji,jj,jk) = zugrad * EXP( -20.0 * ( znd - 1.0 )**2 )
-                   ENDIF
-                   zdvdz_pyc(ji,jj,jk) = zvgrad * EXP( -20.0 * ( znd - 0.85 )**2 )
-                END DO
-             ENDIF
-            !
-          END DO
-       END DO
+      CALL zdf_osm_external_gradients( zdtdz_ext, zdsdz_ext, zdbdz_ext )
+      CALL zdf_osm_pycnocline_scalar_profiles( zdtdz_pyc, zdsdz_pyc, zdbdz_pyc )
+      CALL zdf_osm_pycnocline_shear_profiles( zdudz_pyc, zdvdz_pyc )
        !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        ! Eddy viscosity/diffusivity and non-gradient terms in the flux-gradient relationship
        !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
 
-      ! WHERE ( lconv )
-      !     zdifml_sc = zhml * ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird
-      !     zvisml_sc = zdifml_sc
-      !     zdifpyc_sc = 0.165 * ( zwstrl**3 + zwstrc**3 )**pthird * ( zhbl - zhml )
-      !     zvispyc_sc = 0.142 * ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird * ( zhbl - zhml )
-      !     zbeta_d_sc = 1.0 - (0.165 / 0.8 * ( zhbl - zhml ) / zhbl )**p2third
-      !     zbeta_v_sc = 1.0 -  2.0 * (0.142 /0.375) * (zhbl - zhml ) / zhml
-      !  ELSEWHERE
-      !     zdifml_sc = zwstrl * zhbl * EXP ( -( zhol / 0.183_wp )**2 )
-      !     zvisml_sc = zwstrl * zhbl * EXP ( -( zhol / 0.183_wp )**2 )
-      !  ENDWHERE
        DO jj = 2, jpjm1
           DO ji = 2, jpim1
@@ -868,11 +513,11 @@
                zvispyc_sc(ji,jj) = 0.142 * ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * zdh(ji,jj)
                zbeta_d_sc(ji,jj) = 1.0 - (0.165 / 0.8 * zdh(ji,jj) / zhbl(ji,jj) )**p2third
-               zbeta_v_sc(ji,jj) = 1.0 -  2.0 * (0.142 /0.375) * zdh(ji,jj) / zhml(ji,jj)
+               zbeta_v_sc(ji,jj) = 1.0 -  2.0 * (0.142 /0.375) * zdh(ji,jj) / ( zhml(ji,jj) + epsln )
              ELSE
                zdifml_sc(ji,jj) = zvstr(ji,jj) * zhbl(ji,jj) * EXP ( -( zhol(ji,jj) / 0.6_wp )**2 )
                zvisml_sc(ji,jj) = zvstr(ji,jj) * zhbl(ji,jj) * EXP ( -( zhol(ji,jj) / 0.6_wp )**2 )
-            END IF
-        END DO
-    END DO
+             END IF
+          END DO
+       END DO
 !
        DO jj = 2, jpjm1
@@ -889,11 +534,19 @@
                 ! pycnocline - if present linear profile
                 IF ( zdh(ji,jj) > 0._wp ) THEN
+                   zgamma_b = 6.0
                    DO jk = imld(ji,jj)+1 , ibld(ji,jj)
                        zznd_pyc = -( gdepw_n(ji,jj,jk) - zhml(ji,jj) ) / zdh(ji,jj)
                        !
-                       zdiffut(ji,jj,jk) = zdifpyc_sc(ji,jj) * ( 1.0 + zznd_pyc )
+                       zdiffut(ji,jj,jk) = zdifpyc_sc(ji,jj) * EXP( zgamma_b * zznd_pyc )
                        !
-                       zviscos(ji,jj,jk) = zvispyc_sc(ji,jj) * ( 1.0 + zznd_pyc )
+                       zviscos(ji,jj,jk) = zvispyc_sc(ji,jj) * EXP( zgamma_b * zznd_pyc )
                    END DO
+                   IF ( ibld_ext == 0 ) THEN
+                       zdiffut(ji,jj,ibld(ji,jj)) = 0._wp
+                       zviscos(ji,jj,ibld(ji,jj)) = 0._wp
+                   ELSE
+                       zdiffut(ji,jj,ibld(ji,jj)) = zdhdt(ji,jj) * ( hbl(ji,jj) - gdepw_n(ji, jj, ibld(ji,jj)-1) )
+                       zviscos(ji,jj,ibld(ji,jj)) = zdhdt(ji,jj) * ( hbl(ji,jj) - gdepw_n(ji, jj, ibld(ji,jj)-1) )
+                   ENDIF
                 ENDIF
                 ! Temporay fix to ensure zdiffut is +ve; won't be necessary with wn taken out
@@ -908,6 +561,14 @@
                    zviscos(ji,jj,jk) = 0.375 * zvisml_sc(ji,jj) * zznd_ml * (1.0 - zznd_ml) * ( 1.0 - zznd_ml**2 )
                 END DO
+
+                IF ( ibld_ext == 0 ) THEN
+                   zdiffut(ji,jj,ibld(ji,jj)) = 0._wp
+                   zviscos(ji,jj,ibld(ji,jj)) = 0._wp
+                ELSE
+                   zdiffut(ji,jj,ibld(ji,jj)) = MAX(zdhdt(ji,jj), 0._wp) * e3w_n(ji, jj, ibld(ji,jj))
+                   zviscos(ji,jj,ibld(ji,jj)) = MAX(zdhdt(ji,jj), 0._wp) * e3w_n(ji, jj, ibld(ji,jj))
+                ENDIF
              ENDIF   ! end if ( lconv )
-!
+             !
           END DO  ! end of ji loop
        END DO     ! end of jj loop
@@ -952,15 +613,17 @@
        END DO     ! end of jj loop
 
-
 ! Stokes term in flux-gradient relationship (note in zsc_uw_n don't use zvstr since term needs to go to zero as zwstrl goes to zero)
        WHERE ( lconv )
-          zsc_uw_1 = ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird * zustke /( 1.0 - 1.0 * 6.5 * zla**(8.0/3.0) )
-          zsc_uw_2 = ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird * zustke / ( zla**(8.0/3.0) + epsln )
-          zsc_vw_1 = ff_t * zhml * zustke**3 * zla**(8.0/3.0) / ( ( zvstr**3 + 0.5 * zwstrc**3 )**(2.0/3.0) + epsln )
+          zsc_uw_1 = ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird * zustke / MAX( ( 1.0 - 1.0 * 6.5 * zla**(8.0/3.0) ), 0.2 )
+          zsc_uw_2 = ( zwstrl**3 + 0.5 * zwstrc**3 )**pthird * zustke / MIN( zla**(8.0/3.0) + epsln, 0.12 )
+          zsc_vw_1 = ff_t * zhml * zustke**3 * MIN( zla**(8.0/3.0), 0.12 ) / ( ( zvstr**3 + 0.5 * zwstrc**3 )**(2.0/3.0) + epsln )
        ELSEWHERE
           zsc_uw_1 = zustar**2
-          zsc_vw_1 = ff_t * zhbl * zustke**3 * zla**(8.0/3.0) / (zvstr**2 + epsln)
+          zsc_vw_1 = ff_t * zhbl * zustke**3 * MIN( zla**(8.0/3.0), 0.12 ) / (zvstr**2 + epsln)
        ENDWHERE
-
+       IF(ln_dia_osm) THEN
+          IF ( iom_use("ghamu_00") ) CALL iom_put( "ghamu_00", wmask*ghamu )
+          IF ( iom_use("ghamv_00") ) CALL iom_put( "ghamv_00", wmask*ghamv )
+       END IF
        DO jj = 2, jpjm1
           DO ji = 2, jpim1
@@ -1007,5 +670,5 @@
                    zl_l = 2.0 * ( 1.0 - EXP ( - 2.0 * ( zznd_ml - zznd_ml**3 / 3.0 ) ) )                                           &
                         &     * ( 1.0 - EXP ( - 5.0 * (     1.0 - zznd_ml          ) ) ) * ( 1.0 + dstokes(ji,jj) / zhml (ji,jj) )
-                   zl_eps = zl_l + ( zl_c - zl_l ) / ( 1.0 + EXP ( 3.0 * LOG10 ( - zhol(ji,jj) ) ) ) ** (3.0/2.0)
+                   zl_eps = zl_l + ( zl_c - zl_l ) / ( 1.0 + EXP ( -3.0 * LOG10 ( - zhol(ji,jj) ) ) ) ** (3.0/2.0)
                    ! non-gradient buoyancy terms
                    ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + 0.3 * 0.5 * zsc_wth_1(ji,jj) * zl_eps * zhml(ji,jj) / ( 0.15 + zznd_ml )
@@ -1020,5 +683,4 @@
           END DO   ! ji loop
        END DO      ! jj loop
-
 
        WHERE ( lconv )
@@ -1051,4 +713,8 @@
        END DO           ! jj loop
 
+       IF(ln_dia_osm) THEN
+          IF ( iom_use("ghamu_0") ) CALL iom_put( "ghamu_0", wmask*ghamu )
+          IF ( iom_use("zsc_uw_1_0") ) CALL iom_put( "zsc_uw_1_0", tmask(:,:,1)*zsc_uw_1 )
+       END IF
 ! Transport term in flux-gradient relationship [note : includes ROI ratio (X0.3) ]
 
@@ -1089,5 +755,4 @@
        END DO      ! jj loop
 
-
        WHERE ( lconv )
           zsc_uw_1 = zustar**2
@@ -1134,4 +799,13 @@
           END DO
        END DO
+
+       IF(ln_dia_osm) THEN
+          IF ( iom_use("ghamu_f") ) CALL iom_put( "ghamu_f", wmask*ghamu )
+          IF ( iom_use("ghamv_f") ) CALL iom_put( "ghamv_f", wmask*ghamv )
+          IF ( iom_use("zsc_uw_1_f") ) CALL iom_put( "zsc_uw_1_f", tmask(:,:,1)*zsc_uw_1 )
+          IF ( iom_use("zsc_vw_1_f") ) CALL iom_put( "zsc_vw_1_f", tmask(:,:,1)*zsc_vw_1 )
+          IF ( iom_use("zsc_uw_2_f") ) CALL iom_put( "zsc_uw_2_f", tmask(:,:,1)*zsc_uw_2 )
+          IF ( iom_use("zsc_vw_2_f") ) CALL iom_put( "zsc_vw_2_f", tmask(:,:,1)*zsc_vw_2 )
+       END IF
 !
 ! Make surface forced velocity non-gradient terms go to zero at the base of the mixed layer.
@@ -1165,4 +839,8 @@
       END DO
 
+       IF(ln_dia_osm) THEN
+          IF ( iom_use("ghamu_b") ) CALL iom_put( "ghamu_b", wmask*ghamu )
+          IF ( iom_use("ghamv_b") ) CALL iom_put( "ghamv_b", wmask*ghamv )
+       END IF
       ! pynocline contributions
        ! Temporary fix to avoid instabilities when zdb_bl becomes very very small
@@ -1170,13 +848,15 @@
        DO jj = 2, jpjm1
           DO ji = 2, jpim1
-             DO jk= 2, ibld(ji,jj)
-                znd = gdepw_n(ji,jj,jk) / zhbl(ji,jj)
-                ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zdiffut(ji,jj,jk) * zdtdz_pyc(ji,jj,jk)
-                ghams(ji,jj,jk) = ghams(ji,jj,jk) + zdiffut(ji,jj,jk) * zdsdz_pyc(ji,jj,jk)
-                ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zviscos(ji,jj,jk) * zdudz_pyc(ji,jj,jk)
-                ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zsc_uw_1(ji,jj) * ( 1.0 - znd )**(7.0/4.0) * zdbdz_pyc(ji,jj,jk)
-                ghamv(ji,jj,jk) = ghamv(ji,jj,jk) + zviscos(ji,jj,jk) * zdvdz_pyc(ji,jj,jk)
-             END DO
-           END DO
+             IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) ) THEN
+                DO jk= 2, ibld(ji,jj)
+                   znd = gdepw_n(ji,jj,jk) / zhbl(ji,jj)
+                   ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zdiffut(ji,jj,jk) * zdtdz_pyc(ji,jj,jk)
+                   ghams(ji,jj,jk) = ghams(ji,jj,jk) + zdiffut(ji,jj,jk) * zdsdz_pyc(ji,jj,jk)
+                   ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zviscos(ji,jj,jk) * zdudz_pyc(ji,jj,jk)
+                   ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zsc_uw_1(ji,jj) * ( 1.0 - znd )**(7.0/4.0) * zdbdz_pyc(ji,jj,jk)
+                   ghamv(ji,jj,jk) = ghamv(ji,jj,jk) + zviscos(ji,jj,jk) * zdvdz_pyc(ji,jj,jk)
+                END DO
+             END IF
+          END DO
        END DO
 
@@ -1185,9 +865,9 @@
        DO jj=2, jpjm1
           DO ji = 2, jpim1
-             IF ( lconv(ji,jj) ) THEN
+            IF ( lconv(ji,jj) .AND. ibld(ji,jj) + ibld_ext < mbkt(ji,jj)) THEN
                DO jk = 1, imld(ji,jj) - 1
                   znd=gdepw_n(ji,jj,jk) / zhml(ji,jj)
-                  ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zwth_ent(ji,jj) * znd
-                  ghams(ji,jj,jk) = ghams(ji,jj,jk) + zws_ent(ji,jj) * znd
+                  ! ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zwth_ent(ji,jj) * znd
+                  ! ghams(ji,jj,jk) = ghams(ji,jj,jk) + zws_ent(ji,jj) * znd
                   ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zuw_bse(ji,jj) * znd
                   ghamv(ji,jj,jk) = ghamv(ji,jj,jk) + zvw_bse(ji,jj) * znd
@@ -1195,18 +875,27 @@
                DO jk = imld(ji,jj), ibld(ji,jj)
                   znd = -( gdepw_n(ji,jj,jk) - zhml(ji,jj) ) / zdh(ji,jj)
-                  ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zwth_ent(ji,jj) * ( 1.0 + znd )
-                  ghams(ji,jj,jk) = ghams(ji,jj,jk) + zws_ent(ji,jj) * ( 1.0 + znd )
+                  ! ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + zwth_ent(ji,jj) * ( 1.0 + znd )
+                  ! ghams(ji,jj,jk) = ghams(ji,jj,jk) + zws_ent(ji,jj) * ( 1.0 + znd )
                   ghamu(ji,jj,jk) = ghamu(ji,jj,jk) + zuw_bse(ji,jj) * ( 1.0 + znd )
                   ghamv(ji,jj,jk) = ghamv(ji,jj,jk) + zvw_bse(ji,jj) * ( 1.0 + znd )
                 END DO
              ENDIF
-             ghamt(ji,jj,ibld(ji,jj)) = 0._wp
-             ghams(ji,jj,ibld(ji,jj)) = 0._wp
-             ghamu(ji,jj,ibld(ji,jj)) = 0._wp
-             ghamv(ji,jj,ibld(ji,jj)) = 0._wp
+
+             ghamt(ji,jj,ibld(ji,jj)+ibld_ext) = 0._wp
+             ghams(ji,jj,ibld(ji,jj)+ibld_ext) = 0._wp
+             ghamu(ji,jj,ibld(ji,jj)+ibld_ext) = 0._wp
+             ghamv(ji,jj,ibld(ji,jj)+ibld_ext) = 0._wp
           END DO       ! ji loop
        END DO          ! jj loop
 
-
+       IF(ln_dia_osm) THEN
+          IF ( iom_use("ghamu_1") ) CALL iom_put( "ghamu_1", wmask*ghamu )
+          IF ( iom_use("ghamv_1") ) CALL iom_put( "ghamv_1", wmask*ghamv )
+          IF ( iom_use("zuw_bse") ) CALL iom_put( "zuw_bse", tmask(:,:,1)*zuw_bse )
+          IF ( iom_use("zvw_bse") ) CALL iom_put( "zvw_bse", tmask(:,:,1)*zvw_bse )
+          IF ( iom_use("zdudz_pyc") ) CALL iom_put( "zdudz_pyc", wmask*zdudz_pyc )
+          IF ( iom_use("zdvdz_pyc") ) CALL iom_put( "zdvdz_pyc", wmask*zdvdz_pyc )
+          IF ( iom_use("zviscos") ) CALL iom_put( "zviscos", wmask*zviscos )
+       END IF
        !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        ! Need to put in code for contributions that are applied explicitly to
@@ -1287,5 +976,5 @@
 
        ! Lateral boundary conditions on zvicos (sign unchanged), needed to caclulate viscosities on u and v grids
-       CALL lbc_lnk( 'zdfosm', zviscos(:,:,:), 'W', 1. )
+       !CALL lbc_lnk( zviscos(:,:,:), 'W', 1. )
 
        ! GN 25/8: need to change tmask --> wmask
@@ -1319,7 +1008,7 @@
         ! Lateral boundary conditions on final outputs for gham[uv],  on [UV]-grid  (sign unchanged)
         CALL lbc_lnk_multi( 'zdfosm', ghamt, 'W', 1. , ghams, 'W', 1.,   &
-         &                  ghamu, 'U', 1. , ghamv, 'V', 1. )
-
-       IF(ln_dia_osm) THEN
+         &                  ghamu, 'U', -1. , ghamv, 'V', -1. )
+
+      IF(ln_dia_osm) THEN
          SELECT CASE (nn_osm_wave)
          ! Stokes drift set by assumimg onstant La#=0.3(=0)  or Pierson-Moskovitz spectrum (=1).
@@ -1330,6 +1019,11 @@
          ! Stokes drift read in from sbcwave  (=2).
          CASE(2)
-            IF ( iom_use("us_x") ) CALL iom_put( "us_x", ut0sd )               ! x surface Stokes drift
-            IF ( iom_use("us_y") ) CALL iom_put( "us_y", vt0sd )               ! y surface Stokes drift
+            IF ( iom_use("us_x") ) CALL iom_put( "us_x", ut0sd*umask(:,:,1) )               ! x surface Stokes drift
+            IF ( iom_use("us_y") ) CALL iom_put( "us_y", vt0sd*vmask(:,:,1) )               ! y surface Stokes drift
+            IF ( iom_use("wmp") ) CALL iom_put( "wmp", wmp*tmask(:,:,1) )                   ! wave mean period
+            IF ( iom_use("hsw") ) CALL iom_put( "hsw", hsw*tmask(:,:,1) )                   ! significant wave height
+            IF ( iom_use("wmp_NP") ) CALL iom_put( "wmp_NP", (2.*rpi*1.026/(0.877*grav) )*wndm*tmask(:,:,1) )                  ! wave mean period from NP spectrum
+            IF ( iom_use("hsw_NP") ) CALL iom_put( "hsw_NP", (0.22/grav)*wndm**2*tmask(:,:,1) )                   ! significant wave height from NP spectrum
+            IF ( iom_use("wndm") ) CALL iom_put( "wndm", wndm*tmask(:,:,1) )                   ! U_10
             IF ( iom_use("wind_wave_abs_power") ) CALL iom_put( "wind_wave_abs_power", 1000.*rau0*tmask(:,:,1)*zustar**2* &
                  & SQRT(ut0sd**2 + vt0sd**2 ) )
@@ -1342,5 +1036,12 @@
          IF ( iom_use("zws0") ) CALL iom_put( "zws0", tmask(:,:,1)*zws0 )                ! <Sw_0>
          IF ( iom_use("hbl") ) CALL iom_put( "hbl", tmask(:,:,1)*hbl )                  ! boundary-layer depth
-         IF ( iom_use("hbli") ) CALL iom_put( "hbli", tmask(:,:,1)*hbli )               ! Initial boundary-layer depth
+         IF ( iom_use("ibld") ) CALL iom_put( "ibld", tmask(:,:,1)*ibld )               ! boundary-layer max k
+         IF ( iom_use("zdt_bl") ) CALL iom_put( "zdt_bl", tmask(:,:,1)*zdt_bl )           ! dt at ml base
+         IF ( iom_use("zds_bl") ) CALL iom_put( "zds_bl", tmask(:,:,1)*zds_bl )           ! ds at ml base
+         IF ( iom_use("zdb_bl") ) CALL iom_put( "zdb_bl", tmask(:,:,1)*zdb_bl )           ! db at ml base
+         IF ( iom_use("zdu_bl") ) CALL iom_put( "zdu_bl", tmask(:,:,1)*zdu_bl )           ! du at ml base
+         IF ( iom_use("zdv_bl") ) CALL iom_put( "zdv_bl", tmask(:,:,1)*zdv_bl )           ! dv at ml base
+         IF ( iom_use("dh") ) CALL iom_put( "dh", tmask(:,:,1)*dh )               ! Initial boundary-layer depth
+         IF ( iom_use("hml") ) CALL iom_put( "hml", tmask(:,:,1)*hml )               ! Initial boundary-layer depth
          IF ( iom_use("dstokes") ) CALL iom_put( "dstokes", tmask(:,:,1)*dstokes )      ! Stokes drift penetration depth
          IF ( iom_use("zustke") ) CALL iom_put( "zustke", tmask(:,:,1)*zustke )            ! Stokes drift magnitude at T-points
@@ -1348,9 +1049,12 @@
          IF ( iom_use("zwstrl") ) CALL iom_put( "zwstrl", tmask(:,:,1)*zwstrl )         ! Langmuir velocity scale
          IF ( iom_use("zustar") ) CALL iom_put( "zustar", tmask(:,:,1)*zustar )         ! friction velocity scale
+         IF ( iom_use("zvstr") ) CALL iom_put( "zvstr", tmask(:,:,1)*zvstr )         ! mixed velocity scale
+         IF ( iom_use("zla") ) CALL iom_put( "zla", tmask(:,:,1)*zla )         ! langmuir #
          IF ( iom_use("wind_power") ) CALL iom_put( "wind_power", 1000.*rau0*tmask(:,:,1)*zustar**3 ) ! BL depth internal to zdf_osm routine
          IF ( iom_use("wind_wave_power") ) CALL iom_put( "wind_wave_power", 1000.*rau0*tmask(:,:,1)*zustar**2*zustke )
          IF ( iom_use("zhbl") ) CALL iom_put( "zhbl", tmask(:,:,1)*zhbl )               ! BL depth internal to zdf_osm routine
          IF ( iom_use("zhml") ) CALL iom_put( "zhml", tmask(:,:,1)*zhml )               ! ML depth internal to zdf_osm routine
-         IF ( iom_use("zdh") ) CALL iom_put( "zdh", tmask(:,:,1)*zdh )               ! ML depth internal to zdf_osm routine
+         IF ( iom_use("imld") ) CALL iom_put( "imld", tmask(:,:,1)*imld )               ! index for ML depth internal to zdf_osm routine
+         IF ( iom_use("zdh") ) CALL iom_put( "zdh", tmask(:,:,1)*zdh )                  ! pyc thicknessh internal to zdf_osm routine
          IF ( iom_use("zhol") ) CALL iom_put( "zhol", tmask(:,:,1)*zhol )               ! ML depth internal to zdf_osm routine
          IF ( iom_use("zwthav") ) CALL iom_put( "zwthav", tmask(:,:,1)*zwthav )               ! ML depth internal to zdf_osm routine
@@ -1378,4 +1082,5 @@
      INTEGER  ::   ios            ! local integer
      INTEGER  ::   ji, jj, jk     ! dummy loop indices
+     REAL z1_t2
      !!
      NAMELIST/namzdf_osm/ ln_use_osm_la, rn_osm_la, rn_osm_dstokes, nn_ave &
@@ -1401,5 +1106,5 @@
         WRITE(numout,*) '     Turbulent Langmuir number                     rn_osm_la   = ', rn_osm_la
         WRITE(numout,*) '     Initial hbl for 1D runs                       rn_osm_hbl0   = ', rn_osm_hbl0
-        WRITE(numout,*) '     Depth scale of Stokes drift                rn_osm_dstokes = ', rn_osm_dstokes
+        WRITE(numout,*) '     Depth scale of Stokes drift                   rn_osm_dstokes = ', rn_osm_dstokes
         WRITE(numout,*) '     horizontal average flag                       nn_ave      = ', nn_ave
         WRITE(numout,*) '     Stokes drift                                  nn_osm_wave = ', nn_osm_wave
@@ -1536,5 +1241,5 @@
      REAL(wp) ::   zN2_c           ! local scalar
      REAL(wp) ::   rho_c = 0.01_wp    !: density criterion for mixed layer depth
-     INTEGER, DIMENSION(:,:), ALLOCATABLE :: imld_rst ! level of mixed-layer depth (pycnocline top)
+     INTEGER, DIMENSION(jpi,jpj) :: imld_rst ! level of mixed-layer depth (pycnocline top)
      !!----------------------------------------------------------------------
      !
@@ -1551,12 +1256,13 @@
            WRITE(numout,*) ' ===>>>> :  wn not in restart file, set to zero initially'
         END IF
+
         id1 = iom_varid( numror, 'hbl'   , ldstop = .FALSE. )
-        id2 = iom_varid( numror, 'hbli'   , ldstop = .FALSE. )
+        id2 = iom_varid( numror, 'dh'   , ldstop = .FALSE. )
         IF( id1 > 0 .AND. id2 > 0) THEN                       ! 'hbl' exists; read and return
            CALL iom_get( numror, jpdom_autoglo, 'hbl' , hbl , ldxios = lrxios )
-           CALL iom_get( numror, jpdom_autoglo, 'hbli', hbli, ldxios = lrxios  )
-           WRITE(numout,*) ' ===>>>> :  hbl & hbli read from restart file'
+           CALL iom_get( numror, jpdom_autoglo, 'dh', dh, ldxios = lrxios  )
+           WRITE(numout,*) ' ===>>>> :  hbl & dh read from restart file'
            RETURN
-        ELSE                      ! 'hbl' & 'hbli' not in restart file, recalculate
+        ELSE                      ! 'hbl' & 'dh' not in restart file, recalculate
            WRITE(numout,*) ' ===>>>> : previous run without osmosis scheme, hbl computed from stratification'
         END IF
@@ -1570,5 +1276,5 @@
          CALL iom_rstput( kt, nitrst, numrow, 'wn'     , wn  , ldxios = lwxios )
          CALL iom_rstput( kt, nitrst, numrow, 'hbl'    , hbl , ldxios = lwxios )
-         CALL iom_rstput( kt, nitrst, numrow, 'hbli'   , hbli, ldxios = lwxios )
+         CALL iom_rstput( kt, nitrst, numrow, 'dh'     , dh, ldxios = lwxios )
         RETURN
      END IF
@@ -1578,5 +1284,4 @@
      !!-----------------------------------------------------------------------------
      IF( lwp ) WRITE(numout,*) ' ===>>>> : calculating hbl computed from stratification'
-     ALLOCATE( imld_rst(jpi,jpj) )
      ! w-level of the mixing and mixed layers
      CALL eos_rab( tsn, rab_n )
@@ -1599,6 +1304,7 @@
      DO jj = 1, jpj
         DO ji = 1, jpi
-           iiki = imld_rst(ji,jj)
-           hbl (ji,jj) = gdepw_n(ji,jj,iiki  ) * ssmask(ji,jj)    ! Turbocline depth
+           iiki = MAX(4,imld_rst(ji,jj))
+           hbl (ji,jj) = gdepw_n(ji,jj,iiki  )    ! Turbocline depth
+           dh (ji,jj) = e3t_n(ji,jj,iiki-1  )     ! Turbocline depth
         END DO
      END DO
@@ -1607,4 +1313,6 @@
      DEALLOCATE( imld_rst )
      WRITE(numout,*) ' ===>>>> : hbl computed from stratification'
+     wn(:,:,:) = 0._wp
+     WRITE(numout,*) ' ===>>>> :  wn not in restart file, set to zero initially'
    END SUBROUTINE osm_rst
 
@@ -1634,5 +1342,4 @@
       ENDIF
 
-      ! add non-local temperature and salinity flux
       DO jk = 1, jpkm1
          DO jj = 2, jpjm1
@@ -1648,12 +1355,11 @@
       END DO
 
-
-      ! save the non-local tracer flux trends for diagnostic
+      ! save the non-local tracer flux trends for diagnostics
       IF( l_trdtra )   THEN
          ztrdt(:,:,:) = tsa(:,:,:,jp_tem) - ztrdt(:,:,:)
          ztrds(:,:,:) = tsa(:,:,:,jp_sal) - ztrds(:,:,:)
-!!bug gm jpttdzdf ==> jpttosm
-         CALL trd_tra( kt, 'TRA', jp_tem, jptra_zdf, ztrdt )
-         CALL trd_tra( kt, 'TRA', jp_sal, jptra_zdf, ztrds )
+
+         CALL trd_tra( kt, 'TRA', jp_tem, jptra_osm, ztrdt )
+         CALL trd_tra( kt, 'TRA', jp_sal, jptra_osm, ztrds )
          DEALLOCATE( ztrdt )      ;     DEALLOCATE( ztrds )
       ENDIF
@@ -1723,3 +1429,4 @@
 
    !!======================================================================
+
 END MODULE zdfosm
