Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-oce.xml
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-oce.xml	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-oce.xml	(revision 11145)
@@ -200,4 +200,6 @@
         <field id="zws0"                long_name="surface non-local salinity flux"          unit="psu m/s" />
         <field id="hbl"                 long_name="boundary layer depth"                     unit="m"       />
+        <field id="dh"          long_name="OBL pycnocline thickness"                                     unit="m"        />
+        <field id="hml"          long_name="OBL MLD thickness"                                     unit="m"        />
         <field id="hbli"                long_name="initial boundary layer depth"             unit="m"       />
         <field id="dstokes"             long_name="stokes drift  depth scale"                unit="m"       />
@@ -205,17 +207,39 @@
         <field id="zwstrc"              long_name="convective velocity scale"                unit="m/s"     />
         <field id="zwstrl"              long_name="langmuir velocity scale"                  unit="m/s"     />
+        <field id="zvstr"              long_name="mixed La/wind velocity"                        unit="m/s"     />
         <field id="zustar"              long_name="friction velocity"                        unit="m/s"     />
         <field id="zhbl"                long_name="boundary layer depth"                     unit="m"       />
         <field id="zhml"                long_name="mixed layer depth"                        unit="m"       />
+        <field id="zla"                long_name="Langmuir #"                        unit="#"       />
         <field id="wind_wave_abs_power" long_name="\rho |U_s| x u*^2"                        unit="mW"      />
         <field id="wind_wave_power"     long_name="U_s \dot  tau"                            unit="mW"      />
         <field id="wind_power"          long_name="\rho  u*^3"                               unit="mW"      />
 
+        <field id="ttrd_osm"          long_name="T-trend from non-local terms"         unit=" deg/s"    grid_ref="grid_T_3D"     />
+        <field id="strd_osm"          long_name="S-trend from non-local terms"         unit=" psu/s"        grid_ref="grid_T_3D"   />
+
         <!-- extra OSMOSIS diagnostics -->
+        <field id="ibld"              long_name="boundary-layer depth index"                  unit="#" />
+        <field id="imld"              long_name="mixed-layer depth index"                  unit="#" />
         <field id="zwthav"              long_name="av turb flux of T in ml"                  unit="deg m/s" />
         <field id="zt_ml"               long_name="av T in ml"                               unit="deg"     />
+        <field id="zdt_bl"               long_name="dT in ml"                               unit="deg"     />
+        <field id="zds_bl"               long_name="dS in ml"                               unit="g/kg"     />
+        <field id="zdb_bl"               long_name="db in ml"                               unit="m/s^2"     />
+        <field id="zdu_bl"               long_name="du in ml"                               unit="m/s"     />
+        <field id="zdv_bl"               long_name="dv in ml"                               unit="m/s"     />
         <field id="zwth_ent"            long_name="entrainment turb flux of T"               unit="deg m/s" />
+        <field id="zws_ent"            long_name="entrainment turb flux of S"               unit="g/kg m/s" />
+        <field id="zwb_ent"            long_name="entrainment turb flux of buoyancy"     unit="m^2/s^3" />
         <field id="zhol"                long_name="Hoenekker number"                         unit="#"       />
         <field id="zdh"                 long_name="Pycnocline  depth - grid"                 unit=" m"      />
+
+        <field id="zsc_uw_1_0"       long_name="zsc u-momentum flux on T after Stokes"                       unit="m^2/s^2" />
+        <field id="zsc_uw_1_f"       long_name="zsc u-momentum flux on T after Coriolis"                       unit="m^2/s^2" />
+        <field id="zsc_vw_1_f"       long_name="zsc v-momentum flux on T after Coriolis"                       unit="m^2/s^2" />
+        <field id="zsc_uw_2_f"       long_name="2nd zsc u-momentum flux on T after Coriolis"                       unit="m^2/s^2" />
+        <field id="zsc_vw_2_f"       long_name="2nd zsc v-momentum flux on T after Coriolis"                       unit="m^2/s^2" />
+        <field id="zuw_bse"       long_name="base u-flux T-points"                          unit="m^2/s^2" />
+        <field id="zvw_bse"       long_name="base v-flux T-points"                          unit="m^2/s^2" />
       </field_group>
 
@@ -224,4 +248,20 @@
         <field id="ghams"       long_name="non-local salinity flux"                          unit="psu m/s" />
         <field id="zdtdz_pyc"   long_name="Pycnocline temperature gradient"                  unit=" deg/m"  />
+        <field id="zdsdz_pyc"   long_name="Pycnocline salinity gradient"                  unit=" g/kg/m"  />
+        <field id="zdbdz_pyc"   long_name="Pycnocline buoyancy gradient"                  unit=" s^-2"  />
+
+        <!-- extra OSMOSIS diagnostics -->
+        <field id="ghamu_00"       long_name="non-local u-momentum flux on T before Stokes"                       unit="m^2/s^2" />
+        <field id="ghamv_00"       long_name="non-local v-momentum flux on T before Stokes"                       unit="m^2/s^2" />
+	<field id="ghamu_0"       long_name="non-local u-momentum flux on T after Stokes"                       unit="m^2/s^2" />
+        <field id="ghamu_f"       long_name="non-local u-momentum flux on T after Coriolis"                       unit="m^2/s^2" />
+        <field id="ghamv_f"       long_name="non-local v-momentum flux on T after Coriolis"                          unit="m^2/s^2" />
+        <field id="ghamu_b"       long_name="non-local u-momentum flux on T before base"                       unit="m^2/s^2" />
+        <field id="ghamv_b"       long_name="non-local v-momentum flux on T before base"                          unit="m^2/s^2" />
+        <field id="ghamu_1"       long_name="non-local u-momentum flux on T"                       unit="m^2/s^2" />
+        <field id="ghamv_1"       long_name="non-local v-momentum flux on T"                          unit="m^2/s^2" />
+        <field id="zviscos"       long_name="OSMOSIS viscosity on T-points"                          unit="m/s^2" />
+        <field id="zdudz_pyc"       long_name="pycnocline u-shear T-points"                          unit="m/s^2" />
+        <field id="zdvdz_pyc"       long_name="pycnocline v-shear T-points"                          unit="m/s^2" />
       </field_group>
 
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-waves.xml
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-waves.xml	(revision 11145)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/cfgs/SHARED/field_def_nemo-waves.xml	(revision 11145)
@@ -0,0 +1,22 @@
+<?xml version="1.0"?>
+    <!-- $id$ -->
+
+   <field_definition level="1" prec="4" operation="average" enabled=".TRUE." default_value="1.e20" detect_missing_value="true"> <!-- time step automaticaly defined -->
+
+    <!-- 
+============================================================================================================
+                                  WAVE variables
+============================================================================================================
+    -->
+
+
+        <!-- WAVE  diagnostics -->
+      <field_group id="WAVE_T" grid_ref="grid_T_2D"  >
+	<field id="hsw"          long_name="significant wave height"                                     unit="m"        />
+	<field id="wmp"          long_name="wave mean period"                                     unit="s"        />
+	<field id="hsw_NP"          long_name="significant wave height from NP"                                     unit="m"        />
+	<field id="wmp_NP"          long_name="wave mean period from NP"                                     unit="s"        />
+ 	<field id="wndm"          long_name="U_10"                                                                         unit="m/s"        />
+       </field_group>
+  
+    </field_definition>
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/C1D/step_c1d.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/C1D/step_c1d.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/C1D/step_c1d.F90	(revision 11145)
@@ -58,5 +58,8 @@
 
                              indic = 0                ! reset to no error condition
-      IF( kstp == nit000 )   CALL iom_init( "nemo")   ! iom_put initialization (must be done after nemo_init for AGRIF+XIOS+OASIS)
+                             IF( kstp == nit000 )   THEN
+                                CALL iom_init( "nemo")   ! iom_put initialization (must be done after nemo_init for AGRIF+XIOS+OASIS)
+                                CALL dia_hth_init    ! extra ML depth diagnostics, thermocline depths
+                             END IF
       IF( kstp /= nit000 )   CALL day( kstp )         ! Calendar (day was already called at nit000 in day_init)
                              CALL iom_setkt( kstp - nit000 + 1, "nemo" )   ! say to iom that we are at time step kstp
@@ -86,5 +89,5 @@
       !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
                          CALL dia_wri( kstp )       ! ocean model: outputs
-      IF( lk_diahth  )   CALL dia_hth( kstp )       ! Thermocline depth (20°C)
+      IF( ll_diahth  )   CALL dia_hth( kstp )       ! Thermocline depth (20°C)
 
 
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diahth.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diahth.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diahth.F90	(revision 11145)
@@ -5,13 +5,9 @@
    !!======================================================================
    !! History :  OPA  !  1994-09  (J.-P. Boulanger)  Original code
-   !!                 !  1996-11  (E. Guilyardi)  OPA8 
+   !!                 !  1996-11  (E. Guilyardi)  OPA8
    !!                 !  1997-08  (G. Madec)  optimization
-   !!                 !  1999-07  (E. Guilyardi)  hd28 + heat content 
+   !!                 !  1999-07  (E. Guilyardi)  hd28 + heat content
    !!   NEMO     1.0  !  2002-06  (G. Madec)  F90: Free form and module
    !!            3.2  !  2009-07  (S. Masson) hc300 bugfix + cleaning + add new diag
-   !!----------------------------------------------------------------------
-#if   defined key_diahth
-   !!----------------------------------------------------------------------
-   !!   'key_diahth' :                              thermocline depth diag.
    !!----------------------------------------------------------------------
    !!   dia_hth      : Compute varius diagnostics associated with the mixed layer
@@ -24,5 +20,4 @@
    USE lib_mpp         ! MPP library
    USE iom             ! I/O library
-   USE timing          ! preformance summary
 
    IMPLICIT NONE
@@ -30,13 +25,7 @@
 
    PUBLIC   dia_hth       ! routine called by step.F90
-   PUBLIC   dia_hth_alloc ! routine called by nemogcm.F90
-
-   LOGICAL , PUBLIC, PARAMETER ::   lk_diahth = .TRUE.    !: thermocline-20d depths flag
-   
-   ! note: following variables should move to local variables once iom_put is always used 
-   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   hth    !: depth of the max vertical temperature gradient [m]
-   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   hd20   !: depth of 20 C isotherm                         [m]
-   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   hd28   !: depth of 28 C isotherm                         [m]
-   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   htc3   !: heat content of first 300 m                    [W]
+   PUBLIC   dia_hth_init  ! routine called by nemogcm.F90
+
+   LOGICAL, PUBLIC ::   ll_diahth   !: Compute further diagnostics of ML and thermocline depth
 
    !!----------------------------------------------------------------------
@@ -47,306 +36,299 @@
 CONTAINS
 
-   FUNCTION dia_hth_alloc()
-      !!---------------------------------------------------------------------
-      INTEGER :: dia_hth_alloc
-      !!---------------------------------------------------------------------
+
+   SUBROUTINE dia_hth( kt )
+     !!---------------------------------------------------------------------
+     !!                  ***  ROUTINE dia_hth  ***
+     !!
+     !! ** Purpose : Computes
+     !!      the mixing layer depth (turbocline): avt = 5.e-4
+     !!      the depth of strongest vertical temperature gradient
+     !!      the mixed layer depth with density     criteria: rho = rho(10m or surf) + 0.03(or 0.01)
+     !!      the mixed layer depth with temperature criteria: abs( tn - tn(10m) ) = 0.2
+     !!      the top of the thermochine: tn = tn(10m) - ztem2
+     !!      the pycnocline depth with density criteria equivalent to a temperature variation
+     !!                rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC)
+     !!      the barrier layer thickness
+     !!      the maximal verical inversion of temperature and its depth max( 0, max of tn - tn(10m) )
+     !!      the depth of the 20 degree isotherm (linear interpolation)
+     !!      the depth of the 28 degree isotherm (linear interpolation)
+     !!      the heat content of first 300 m
+     !!
+     !! ** Method :
+     !!-------------------------------------------------------------------
+     INTEGER, INTENT( in ) ::   kt      ! ocean time-step index
+     !!
+     INTEGER                          ::   ji, jj, jk            ! dummy loop arguments
+     INTEGER                          ::   iid, ilevel           ! temporary integers
+     INTEGER, DIMENSION(jpi,jpj) ::   ik20, ik28  ! levels
+     REAL(wp)                         ::   zavt5 = 5.e-4_wp      ! Kz criterion for the turbocline depth
+     REAL(wp)                         ::   zrho3 = 0.03_wp       ! density     criterion for mixed layer depth
+     REAL(wp)                         ::   zrho1 = 0.01_wp       ! density     criterion for mixed layer depth
+     REAL(wp)                         ::   ztem2 = 0.2_wp        ! temperature criterion for mixed layer depth
+     REAL(wp)                         ::   zthick_0, zcoef       ! temporary scalars
+     REAL(wp)                         ::   zztmp, zzdep          ! temporary scalars inside do loop
+     REAL(wp)                         ::   zu, zv, zw, zut, zvt  ! temporary workspace
+     REAL(wp), DIMENSION(jpi,jpj) ::   zabs2      ! MLD: abs( tn - tn(10m) ) = ztem2
+     REAL(wp), DIMENSION(jpi,jpj) ::   ztm2       ! Top of thermocline: tn = tn(10m) - ztem2
+     REAL(wp), DIMENSION(jpi,jpj) ::   zrho10_3   ! MLD: rho = rho10m + zrho3
+     REAL(wp), DIMENSION(jpi,jpj) ::   zpycn      ! pycnocline: rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC)
+     REAL(wp), DIMENSION(jpi,jpj) ::   ztinv      ! max of temperature inversion
+     REAL(wp), DIMENSION(jpi,jpj) ::   zdepinv    ! depth of temperature inversion
+     REAL(wp), DIMENSION(jpi,jpj) ::   zrho0_3    ! MLD rho = rho(surf) = 0.03
+     REAL(wp), DIMENSION(jpi,jpj) ::   zrho0_1    ! MLD rho = rho(surf) = 0.01
+     REAL(wp), DIMENSION(jpi,jpj) ::   zmaxdzT    ! max of dT/dz
+     REAL(wp), DIMENSION(jpi,jpj) ::   zthick     ! vertical integration thickness
+     REAL(wp), DIMENSION(jpi,jpj) ::   zdelr      ! delta rho equivalent to deltaT = 0.2
+   ! note: following variables should move to local variables once iom_put is always used
+     REAL(wp), DIMENSION(jpi,jpj) ::   zhth    !: depth of the max vertical temperature gradient [m]
+     REAL(wp), DIMENSION(jpi,jpj) ::   zhd20   !: depth of 20 C isotherm                         [m]
+     REAL(wp), DIMENSION(jpi,jpj) ::   zhd28   !: depth of 28 C isotherm                         [m]
+     REAL(wp), DIMENSION(jpi,jpj) ::   zhtc3   !: heat content of first 300 m                    [W]
+
+     IF (iom_use("mlddzt") .OR. iom_use("mldr0_3") .OR. iom_use("mldr0_1")) THEN
+        ! ------------------------------------------------------------- !
+        ! thermocline depth: strongest vertical gradient of temperature !
+        ! turbocline depth (mixing layer depth): avt = zavt5            !
+        ! MLD: rho = rho(1) + zrho3                                     !
+        ! MLD: rho = rho(1) + zrho1                                     !
+        ! ------------------------------------------------------------- !
+        zmaxdzT(:,:) = 0._wp
+        IF( nla10 > 1 ) THEN
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 zztmp = gdepw_n(ji,jj,mbkt(ji,jj)+1)
+                 zrho0_3(ji,jj) = zztmp
+                 zrho0_1(ji,jj) = zztmp
+                 zhth(ji,jj) = zztmp
+              END DO
+           END DO
+        ELSE IF (iom_use("mlddzt")) THEN
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 zztmp = gdepw_n(ji,jj,mbkt(ji,jj)+1)
+                 zhth(ji,jj) = zztmp
+              END DO
+           END DO
+        ELSE
+           zhth(:,:) = 0._wp
+
+        ENDIF
+
+        DO jk = jpkm1, 2, -1   ! loop from bottom to 2
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 !
+                 zzdep = gdepw_n(ji,jj,jk)
+                 zztmp = ( tsn(ji,jj,jk-1,jp_tem) - tsn(ji,jj,jk,jp_tem) ) / zzdep * tmask(ji,jj,jk)   ! vertical gradient of temperature (dT/dz)
+                 zzdep = zzdep * tmask(ji,jj,1)
+
+                 IF( zztmp > zmaxdzT(ji,jj) ) THEN
+                    zmaxdzT(ji,jj) = zztmp   ;   zhth    (ji,jj) = zzdep                ! max and depth of dT/dz
+                 ENDIF
+
+                 IF( nla10 > 1 ) THEN
+                    zztmp = rhop(ji,jj,jk) - rhop(ji,jj,1)                             ! delta rho(1)
+                    IF( zztmp > zrho3 )          zrho0_3(ji,jj) = zzdep                ! > 0.03
+                    IF( zztmp > zrho1 )          zrho0_1(ji,jj) = zzdep                ! > 0.01
+                 ENDIF
+
+              END DO
+           END DO
+        END DO
+
+        IF (iom_use("mlddzt")) CALL iom_put( "mlddzt", zhth*tmask(:,:,1) )            ! depth of the thermocline
+        IF( nla10 > 1 ) THEN
+           IF (iom_use("mldr0_3")) CALL iom_put( "mldr0_3", zrho0_3*tmask(:,:,1) )   ! MLD delta rho(surf) = 0.03
+           IF (iom_use("mldr0_1")) CALL iom_put( "mldr0_1", zrho0_1*tmask(:,:,1) )   ! MLD delta rho(surf) = 0.01
+        ENDIF
+     ENDIF
+
+     IF (iom_use("mld_dt02") .OR. iom_use("topthdep") .OR. iom_use("mldr10_3") .OR.  &
+          & iom_use("pycndep") .OR. iom_use("tinv") .OR. iom_use("depti")) THEN
+        DO jj = 1, jpj
+           DO ji = 1, jpi
+              zztmp = gdepw_n(ji,jj,mbkt(ji,jj)+1)
+              zabs2   (ji,jj) = zztmp
+              ztm2    (ji,jj) = zztmp
+              zrho10_3(ji,jj) = zztmp
+              zpycn   (ji,jj) = zztmp
+           END DO
+        END DO
+        ztinv  (:,:) = 0._wp
+        zdepinv(:,:) = 0._wp
+
+        IF (iom_use("pycndep")) THEN
+           ! Preliminary computation
+           ! computation of zdelr = (dr/dT)(T,S,10m)*(-0.2 degC)
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 IF( tmask(ji,jj,nla10) == 1. ) THEN
+                    zu  =  1779.50 + 11.250 * tsn(ji,jj,nla10,jp_tem) - 3.80   * tsn(ji,jj,nla10,jp_sal)                             &
+                         &                                              - 0.0745 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_tem)   &
+                         &                                              - 0.0100 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_sal)
+                    zv  =  5891.00 + 38.000 * tsn(ji,jj,nla10,jp_tem) + 3.00   * tsn(ji,jj,nla10,jp_sal)                             &
+                         &                                              - 0.3750 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_tem)
+                    zut =    11.25 -  0.149 * tsn(ji,jj,nla10,jp_tem) - 0.01   * tsn(ji,jj,nla10,jp_sal)
+                    zvt =    38.00 -  0.750 * tsn(ji,jj,nla10,jp_tem)
+                    zw  = (zu + 0.698*zv) * (zu + 0.698*zv)
+                    zdelr(ji,jj) = ztem2 * (1000.*(zut*zv - zvt*zu)/zw)
+                 ELSE
+                    zdelr(ji,jj) = 0._wp
+                 ENDIF
+              END DO
+           END DO
+        ELSE
+           zdelr(:,:) = 0._wp
+        ENDIF
+
+        ! ------------------------------------------------------------- !
+        ! MLD: abs( tn - tn(10m) ) = ztem2                              !
+        ! Top of thermocline: tn = tn(10m) - ztem2                      !
+        ! MLD: rho = rho10m + zrho3                                     !
+        ! pycnocline: rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC)       !
+        ! temperature inversion: max( 0, max of tn - tn(10m) )          !
+        ! depth of temperature inversion                                !
+        ! ------------------------------------------------------------- !
+        DO jk = jpkm1, nlb10, -1   ! loop from bottom to nlb10
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 !
+                 zzdep = gdepw_n(ji,jj,jk) * tmask(ji,jj,1)
+                 !
+                 zztmp = tsn(ji,jj,nla10,jp_tem) - tsn(ji,jj,jk,jp_tem)  ! - delta T(10m)
+                 IF( ABS(zztmp) > ztem2 )      zabs2   (ji,jj) = zzdep   ! abs > 0.2
+                 IF(     zztmp  > ztem2 )      ztm2    (ji,jj) = zzdep   ! > 0.2
+                 zztmp = -zztmp                                          ! delta T(10m)
+                 IF( zztmp >  ztinv(ji,jj) ) THEN                        ! temperature inversion
+                    ztinv(ji,jj) = zztmp   ;   zdepinv (ji,jj) = zzdep   ! max value and depth
+                 ENDIF
+
+                 zztmp = rhop(ji,jj,jk) - rhop(ji,jj,nla10)              ! delta rho(10m)
+                 IF( zztmp > zrho3        )    zrho10_3(ji,jj) = zzdep   ! > 0.03
+                 IF( zztmp > zdelr(ji,jj) )    zpycn   (ji,jj) = zzdep   ! > equi. delta T(10m) - 0.2
+                 !
+              END DO
+           END DO
+        END DO
+
+        IF (iom_use("mld_dt02")) CALL iom_put( "mld_dt02", zabs2*tmask(:,:,1)        )   ! MLD abs(delta t) - 0.2
+        IF (iom_use("topthdep")) CALL iom_put( "topthdep", ztm2*tmask(:,:,1)         )   ! T(10) - 0.2
+        IF (iom_use("mldr10_3")) CALL iom_put( "mldr10_3", zrho10_3*tmask(:,:,1)     )   ! MLD delta rho(10m) = 0.03
+        IF (iom_use("pycndep")) CALL iom_put( "pycndep" , zpycn*tmask(:,:,1)        )   ! MLD delta rho equi. delta T(10m) = 0.2
+        IF (iom_use("tinv")) CALL iom_put( "tinv"    , ztinv*tmask(:,:,1)        )   ! max. temp. inv. (t10 ref)
+        IF (iom_use("depti")) CALL iom_put( "depti"   , zdepinv*tmask(:,:,1)      )   ! depth of max. temp. inv. (t10 ref)
+     ENDIF
+
+     IF(iom_use("20d") .OR. iom_use("28d")) THEN
+        ! ----------------------------------- !
+        ! search deepest level above 20C/28C  !
+        ! ----------------------------------- !
+        ik20(:,:) = 1
+        ik28(:,:) = 1
+        DO jk = 1, jpkm1   ! beware temperature is not always decreasing with depth => loop from top to bottom
+           DO jj = 1, jpj
+              DO ji = 1, jpi
+                 zztmp = tsn(ji,jj,jk,jp_tem)
+                 IF( zztmp >= 20. )   ik20(ji,jj) = jk
+                 IF( zztmp >= 28. )   ik28(ji,jj) = jk
+              END DO
+           END DO
+        END DO
+
+        ! --------------------------- !
+        !  Depth of 20C/28C isotherm  !
+        ! --------------------------- !
+        DO jj = 1, jpj
+           DO ji = 1, jpi
+              !
+              zzdep = gdepw_n(ji,jj,mbkt(ji,jj)+1)       ! depth of the oean bottom
+              !
+              iid = ik20(ji,jj)
+              IF( iid /= 1 ) THEN
+                 zztmp =      gdept_n(ji,jj,iid  )   &                     ! linear interpolation
+                      &  + (    gdept_n(ji,jj,iid+1) - gdept_n(ji,jj,iid)                       )   &
+                      &  * ( 20.*tmask(ji,jj,iid+1) - tsn(ji,jj,iid,jp_tem)                       )   &
+                      &  / ( tsn(ji,jj,iid+1,jp_tem) - tsn(ji,jj,iid,jp_tem) + (1.-tmask(ji,jj,1)) )
+                 zhd20(ji,jj) = MIN( zztmp , zzdep) * tmask(ji,jj,1)       ! bound by the ocean depth
+              ELSE
+                 zhd20(ji,jj) = 0._wp
+              ENDIF
+              !
+              iid = ik28(ji,jj)
+              IF( iid /= 1 ) THEN
+                 zztmp =      gdept_n(ji,jj,iid  )   &                     ! linear interpolation
+                      &  + (    gdept_n(ji,jj,iid+1) - gdept_n(ji,jj,iid)                       )   &
+                      &  * ( 28.*tmask(ji,jj,iid+1) -    tsn(ji,jj,iid,jp_tem)                       )   &
+                      &  / (  tsn(ji,jj,iid+1,jp_tem) -    tsn(ji,jj,iid,jp_tem) + (1.-tmask(ji,jj,1)) )
+                 zhd28(ji,jj) = MIN( zztmp , zzdep ) * tmask(ji,jj,1)      ! bound by the ocean depth
+              ELSE
+                 zhd28(ji,jj) = 0._wp
+              ENDIF
+
+           END DO
+        END DO
+        CALL iom_put( "20d", zhd20 )   ! depth of the 20 isotherm
+        CALL iom_put( "28d", zhd28 )   ! depth of the 28 isotherm
+     ENDIF
+
+     ! ----------------------------- !
+     !  Heat content of first 300 m  !
+     ! ----------------------------- !
+     IF (iom_use("hc300")) THEN
+        ! find ilevel with (ilevel+1) the deepest W-level above 300m (we assume we can use e3t_1d to do this search...)
+        ilevel   = 0
+        zthick_0 = 0._wp
+        DO jk = 1, jpkm1
+           zthick_0 = zthick_0 + e3t_1d(jk)
+           IF( zthick_0 < 300. )   ilevel = jk
+        END DO
+        ! surface boundary condition
+        IF( ln_linssh ) THEN   ;   zthick(:,:) = sshn(:,:)   ;   zhtc3(:,:) = tsn(:,:,1,jp_tem) * sshn(:,:) * tmask(:,:,1)
+        ELSE                   ;   zthick(:,:) = 0._wp       ;   zhtc3(:,:) = 0._wp
+        ENDIF
+        ! integration down to ilevel
+        DO jk = 1, ilevel
+           zthick(:,:) = zthick(:,:) + e3t_n(:,:,jk)
+           zhtc3  (:,:) = zhtc3  (:,:) + e3t_n(:,:,jk) * tsn(:,:,jk,jp_tem) * tmask(:,:,jk)
+        END DO
+        ! deepest layer
+        zthick(:,:) = 300. - zthick(:,:)   !   remaining thickness to reach 300m
+        DO jj = 1, jpj
+           DO ji = 1, jpi
+              zhtc3(ji,jj) = zhtc3(ji,jj) + tsn(ji,jj,ilevel+1,jp_tem)                  &
+                   &                      * MIN( e3t_n(ji,jj,ilevel+1), zthick(ji,jj) ) * tmask(ji,jj,ilevel+1)
+           END DO
+        END DO
+        ! from temperature to heat contain
+        zcoef = rau0 * rcp
+        zhtc3(:,:) = zcoef * zhtc3(:,:)
+        CALL iom_put( "hc300", zhtc3*tmask(:,:,1) )      ! first 300m heat content
+     ENDIF
+     !
+   END SUBROUTINE dia_hth
+
+
+   SUBROUTINE dia_hth_init
+      !!---------------------------------------------------------------------------
+      !!                  ***  ROUTINE dia_hth_init  ***
+      !!
+      !! ** Purpose: Initialization for ML and thermocline depths
+      !!
+      !! ** Action : If any upper ocean diagnostic required by xml file, set in dia_hth
+      !!---------------------------------------------------------------------------
+       !
+      IF(lwp) THEN
+         WRITE(numout,*)
+         WRITE(numout,*) 'dia_hth_init : heat and salt budgets diagnostics'
+         WRITE(numout,*) '~~~~~~~~~~~~ '
+      ENDIF
+      ll_diahth = iom_use("mlddzt") .OR. iom_use("mldr0_3") .OR. iom_use("mldr0_1") .OR.  &
+           & iom_use("mld_dt02") .OR. iom_use("topthdep") .OR. iom_use("mldr10_3") .OR.  &
+           & iom_use("pycndep") .OR. iom_use("tinv") .OR. iom_use("depti").OR. &
+           & iom_use("20d") .OR. iom_use("28d") .OR. iom_use("hc300")
+      IF(lwp) THEN
+         WRITE(numout,*) '   output upper ocean diagnostics (T) or not (F)       ll_diahth = ', ll_diahth
+      ENDIF
       !
-      ALLOCATE( hth(jpi,jpj), hd20(jpi,jpj), hd28(jpi,jpj), htc3(jpi,jpj), STAT=dia_hth_alloc )
-      !
-      CALL mpp_sum ( 'diahth', dia_hth_alloc )
-      IF(dia_hth_alloc /= 0)   CALL ctl_stop( 'STOP', 'dia_hth_alloc: failed to allocate arrays.' )
-      !
-   END FUNCTION dia_hth_alloc
-
-
-   SUBROUTINE dia_hth( kt )
-      !!---------------------------------------------------------------------
-      !!                  ***  ROUTINE dia_hth  ***
-      !!
-      !! ** Purpose : Computes
-      !!      the mixing layer depth (turbocline): avt = 5.e-4
-      !!      the depth of strongest vertical temperature gradient
-      !!      the mixed layer depth with density     criteria: rho = rho(10m or surf) + 0.03(or 0.01)
-      !!      the mixed layer depth with temperature criteria: abs( tn - tn(10m) ) = 0.2       
-      !!      the top of the thermochine: tn = tn(10m) - ztem2 
-      !!      the pycnocline depth with density criteria equivalent to a temperature variation 
-      !!                rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC) 
-      !!      the barrier layer thickness
-      !!      the maximal verical inversion of temperature and its depth max( 0, max of tn - tn(10m) )
-      !!      the depth of the 20 degree isotherm (linear interpolation)
-      !!      the depth of the 28 degree isotherm (linear interpolation)
-      !!      the heat content of first 300 m
-      !!
-      !! ** Method : 
-      !!-------------------------------------------------------------------
-      INTEGER, INTENT( in ) ::   kt      ! ocean time-step index
-      !!
-      INTEGER                          ::   ji, jj, jk            ! dummy loop arguments
-      INTEGER                          ::   iid, ilevel           ! temporary integers
-      INTEGER, ALLOCATABLE, SAVE, DIMENSION(:,:) ::   ik20, ik28  ! levels
-      REAL(wp)                         ::   zavt5 = 5.e-4_wp      ! Kz criterion for the turbocline depth
-      REAL(wp)                         ::   zrho3 = 0.03_wp       ! density     criterion for mixed layer depth
-      REAL(wp)                         ::   zrho1 = 0.01_wp       ! density     criterion for mixed layer depth
-      REAL(wp)                         ::   ztem2 = 0.2_wp        ! temperature criterion for mixed layer depth
-      REAL(wp)                         ::   zthick_0, zcoef       ! temporary scalars
-      REAL(wp)                         ::   zztmp, zzdep          ! temporary scalars inside do loop
-      REAL(wp)                         ::   zu, zv, zw, zut, zvt  ! temporary workspace
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zabs2      ! MLD: abs( tn - tn(10m) ) = ztem2 
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   ztm2       ! Top of thermocline: tn = tn(10m) - ztem2     
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zrho10_3   ! MLD: rho = rho10m + zrho3      
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zpycn      ! pycnocline: rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC)
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   ztinv      ! max of temperature inversion
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zdepinv    ! depth of temperature inversion
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zrho0_3    ! MLD rho = rho(surf) = 0.03
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zrho0_1    ! MLD rho = rho(surf) = 0.01
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zmaxdzT    ! max of dT/dz
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zthick     ! vertical integration thickness 
-      REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:) ::   zdelr      ! delta rho equivalent to deltaT = 0.2
-      !!----------------------------------------------------------------------
-      IF( ln_timing )   CALL timing_start('dia_hth')
-
-      IF( kt == nit000 ) THEN
-         !                                      ! allocate dia_hth array
-         IF( dia_hth_alloc() /= 0 )   CALL ctl_stop( 'STOP', 'dia_hth : unable to allocate standard arrays' )
-
-         IF(.NOT. ALLOCATED(ik20) ) THEN
-            ALLOCATE(ik20(jpi,jpj), ik28(jpi,jpj), &
-               &      zabs2(jpi,jpj),   &
-               &      ztm2(jpi,jpj),    &
-               &      zrho10_3(jpi,jpj),&
-               &      zpycn(jpi,jpj),   &
-               &      ztinv(jpi,jpj),   &
-               &      zdepinv(jpi,jpj), &
-               &      zrho0_3(jpi,jpj), &
-               &      zrho0_1(jpi,jpj), &
-               &      zmaxdzT(jpi,jpj), &
-               &      zthick(jpi,jpj),  &
-               &      zdelr(jpi,jpj), STAT=ji)
-            CALL mpp_sum('diahth', ji)
-            IF( ji /= 0 )   CALL ctl_stop( 'STOP', 'dia_hth : unable to allocate standard ocean arrays' )
-         END IF
-
-         IF(lwp) WRITE(numout,*)
-         IF(lwp) WRITE(numout,*) 'dia_hth : diagnostics of the thermocline depth'
-         IF(lwp) WRITE(numout,*) '~~~~~~~ '
-         IF(lwp) WRITE(numout,*)
-      ENDIF
-
-      ! initialization
-      ztinv  (:,:) = 0._wp  
-      zdepinv(:,:) = 0._wp  
-      zmaxdzT(:,:) = 0._wp  
-      DO jj = 1, jpj
-         DO ji = 1, jpi
-            zztmp = gdepw_n(ji,jj,mbkt(ji,jj)+1) 
-            hth     (ji,jj) = zztmp
-            zabs2   (ji,jj) = zztmp
-            ztm2    (ji,jj) = zztmp
-            zrho10_3(ji,jj) = zztmp
-            zpycn   (ji,jj) = zztmp
-        END DO
-      END DO
-      IF( nla10 > 1 ) THEN 
-         DO jj = 1, jpj
-            DO ji = 1, jpi
-               zztmp = gdepw_n(ji,jj,mbkt(ji,jj)+1) 
-               zrho0_3(ji,jj) = zztmp
-               zrho0_1(ji,jj) = zztmp
-            END DO
-         END DO
-      ENDIF
-      
-      ! Preliminary computation
-      ! computation of zdelr = (dr/dT)(T,S,10m)*(-0.2 degC)
-      DO jj = 1, jpj
-         DO ji = 1, jpi
-            IF( tmask(ji,jj,nla10) == 1. ) THEN
-               zu  =  1779.50 + 11.250 * tsn(ji,jj,nla10,jp_tem) - 3.80   * tsn(ji,jj,nla10,jp_sal)                             &
-                  &                                              - 0.0745 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_tem)   &
-                  &                                              - 0.0100 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_sal)
-               zv  =  5891.00 + 38.000 * tsn(ji,jj,nla10,jp_tem) + 3.00   * tsn(ji,jj,nla10,jp_sal)                             &
-                  &                                              - 0.3750 * tsn(ji,jj,nla10,jp_tem) * tsn(ji,jj,nla10,jp_tem)
-               zut =    11.25 -  0.149 * tsn(ji,jj,nla10,jp_tem) - 0.01   * tsn(ji,jj,nla10,jp_sal)
-               zvt =    38.00 -  0.750 * tsn(ji,jj,nla10,jp_tem)
-               zw  = (zu + 0.698*zv) * (zu + 0.698*zv)
-               zdelr(ji,jj) = ztem2 * (1000.*(zut*zv - zvt*zu)/zw)
-            ELSE
-               zdelr(ji,jj) = 0._wp
-            ENDIF
-         END DO
-      END DO
-
-      ! ------------------------------------------------------------- !
-      ! thermocline depth: strongest vertical gradient of temperature !
-      ! turbocline depth (mixing layer depth): avt = zavt5            !
-      ! MLD: rho = rho(1) + zrho3                                     !
-      ! MLD: rho = rho(1) + zrho1                                     !
-      ! ------------------------------------------------------------- !
-      DO jk = jpkm1, 2, -1   ! loop from bottom to 2
-         DO jj = 1, jpj
-            DO ji = 1, jpi
-               !
-               zzdep = gdepw_n(ji,jj,jk)
-               zztmp = ( tsn(ji,jj,jk-1,jp_tem) - tsn(ji,jj,jk,jp_tem) ) / zzdep * tmask(ji,jj,jk)   ! vertical gradient of temperature (dT/dz)
-               zzdep = zzdep * tmask(ji,jj,1)
-
-               IF( zztmp > zmaxdzT(ji,jj) ) THEN                        
-                  zmaxdzT(ji,jj) = zztmp   ;   hth    (ji,jj) = zzdep                ! max and depth of dT/dz
-               ENDIF
-               
-               IF( nla10 > 1 ) THEN 
-                  zztmp = rhop(ji,jj,jk) - rhop(ji,jj,1)                             ! delta rho(1)
-                  IF( zztmp > zrho3 )          zrho0_3(ji,jj) = zzdep                ! > 0.03
-                  IF( zztmp > zrho1 )          zrho0_1(ji,jj) = zzdep                ! > 0.01
-               ENDIF
-
-            END DO
-         END DO
-      END DO
-      
-      CALL iom_put( "mlddzt", hth )            ! depth of the thermocline
-      IF( nla10 > 1 ) THEN 
-         CALL iom_put( "mldr0_3", zrho0_3 )   ! MLD delta rho(surf) = 0.03
-         CALL iom_put( "mldr0_1", zrho0_1 )   ! MLD delta rho(surf) = 0.01
-      ENDIF
-
-      ! ------------------------------------------------------------- !
-      ! MLD: abs( tn - tn(10m) ) = ztem2                              !
-      ! Top of thermocline: tn = tn(10m) - ztem2                      !
-      ! MLD: rho = rho10m + zrho3                                     !
-      ! pycnocline: rho = rho10m + (dr/dT)(T,S,10m)*(-0.2 degC)       !
-      ! temperature inversion: max( 0, max of tn - tn(10m) )          !
-      ! depth of temperature inversion                                !
-      ! ------------------------------------------------------------- !
-      DO jk = jpkm1, nlb10, -1   ! loop from bottom to nlb10
-         DO jj = 1, jpj
-            DO ji = 1, jpi
-               !
-               zzdep = gdepw_n(ji,jj,jk) * tmask(ji,jj,1)
-               !
-               zztmp = tsn(ji,jj,nla10,jp_tem) - tsn(ji,jj,jk,jp_tem)  ! - delta T(10m)
-               IF( ABS(zztmp) > ztem2 )      zabs2   (ji,jj) = zzdep   ! abs > 0.2
-               IF(     zztmp  > ztem2 )      ztm2    (ji,jj) = zzdep   ! > 0.2
-               zztmp = -zztmp                                          ! delta T(10m)
-               IF( zztmp >  ztinv(ji,jj) ) THEN                        ! temperature inversion
-                  ztinv(ji,jj) = zztmp   ;   zdepinv (ji,jj) = zzdep   ! max value and depth
-               ENDIF
-
-               zztmp = rhop(ji,jj,jk) - rhop(ji,jj,nla10)              ! delta rho(10m)
-               IF( zztmp > zrho3        )    zrho10_3(ji,jj) = zzdep   ! > 0.03
-               IF( zztmp > zdelr(ji,jj) )    zpycn   (ji,jj) = zzdep   ! > equi. delta T(10m) - 0.2
-               !
-            END DO
-         END DO
-      END DO
-
-      CALL iom_put( "mld_dt02", zabs2        )   ! MLD abs(delta t) - 0.2
-      CALL iom_put( "topthdep", ztm2         )   ! T(10) - 0.2
-      CALL iom_put( "mldr10_3", zrho10_3     )   ! MLD delta rho(10m) = 0.03
-      CALL iom_put( "pycndep" , zpycn        )   ! MLD delta rho equi. delta T(10m) = 0.2
-      CALL iom_put( "tinv"    , ztinv        )   ! max. temp. inv. (t10 ref) 
-      CALL iom_put( "depti"   , zdepinv      )   ! depth of max. temp. inv. (t10 ref) 
-
-
-      ! ----------------------------------- !
-      ! search deepest level above 20C/28C  !
-      ! ----------------------------------- !
-      ik20(:,:) = 1
-      ik28(:,:) = 1
-      DO jk = 1, jpkm1   ! beware temperature is not always decreasing with depth => loop from top to bottom
-         DO jj = 1, jpj
-            DO ji = 1, jpi
-               zztmp = tsn(ji,jj,jk,jp_tem)
-               IF( zztmp >= 20. )   ik20(ji,jj) = jk
-               IF( zztmp >= 28. )   ik28(ji,jj) = jk
-            END DO
-         END DO
-      END DO
-
-      ! --------------------------- !
-      !  Depth of 20C/28C isotherm  !
-      ! --------------------------- !
-      DO jj = 1, jpj
-         DO ji = 1, jpi
-            !
-            zzdep = gdepw_n(ji,jj,mbkt(ji,jj)+1)       ! depth of the oean bottom
-            !
-            iid = ik20(ji,jj)
-            IF( iid /= 1 ) THEN 
-               zztmp =      gdept_n(ji,jj,iid  )   &                     ! linear interpolation
-                  &  + (    gdept_n(ji,jj,iid+1) - gdept_n(ji,jj,iid)                       )   &
-                  &  * ( 20.*tmask(ji,jj,iid+1) - tsn(ji,jj,iid,jp_tem)                       )   &
-                  &  / ( tsn(ji,jj,iid+1,jp_tem) - tsn(ji,jj,iid,jp_tem) + (1.-tmask(ji,jj,1)) )
-               hd20(ji,jj) = MIN( zztmp , zzdep) * tmask(ji,jj,1)       ! bound by the ocean depth
-            ELSE 
-               hd20(ji,jj) = 0._wp
-            ENDIF
-            !
-            iid = ik28(ji,jj)
-            IF( iid /= 1 ) THEN 
-               zztmp =      gdept_n(ji,jj,iid  )   &                     ! linear interpolation
-                  &  + (    gdept_n(ji,jj,iid+1) - gdept_n(ji,jj,iid)                       )   &
-                  &  * ( 28.*tmask(ji,jj,iid+1) -    tsn(ji,jj,iid,jp_tem)                       )   &
-                  &  / (  tsn(ji,jj,iid+1,jp_tem) -    tsn(ji,jj,iid,jp_tem) + (1.-tmask(ji,jj,1)) )
-               hd28(ji,jj) = MIN( zztmp , zzdep ) * tmask(ji,jj,1)      ! bound by the ocean depth
-            ELSE 
-               hd28(ji,jj) = 0._wp
-            ENDIF
-
-         END DO
-      END DO
-      CALL iom_put( "20d", hd20 )   ! depth of the 20 isotherm
-      CALL iom_put( "28d", hd28 )   ! depth of the 28 isotherm
-
-      ! ----------------------------- !
-      !  Heat content of first 300 m  !
-      ! ----------------------------- !
-
-      ! find ilevel with (ilevel+1) the deepest W-level above 300m (we assume we can use e3t_1d to do this search...)
-      ilevel   = 0
-      zthick_0 = 0._wp
-      DO jk = 1, jpkm1                      
-         zthick_0 = zthick_0 + e3t_1d(jk)
-         IF( zthick_0 < 300. )   ilevel = jk
-      END DO
-      ! surface boundary condition
-      IF( ln_linssh ) THEN   ;   zthick(:,:) = sshn(:,:)   ;   htc3(:,:) = tsn(:,:,1,jp_tem) * sshn(:,:) * tmask(:,:,1)  
-      ELSE                   ;   zthick(:,:) = 0._wp       ;   htc3(:,:) = 0._wp                                   
-      ENDIF
-      ! integration down to ilevel
-      DO jk = 1, ilevel
-         zthick(:,:) = zthick(:,:) + e3t_n(:,:,jk)
-         htc3  (:,:) = htc3  (:,:) + e3t_n(:,:,jk) * tsn(:,:,jk,jp_tem) * tmask(:,:,jk)
-      END DO
-      ! deepest layer
-      zthick(:,:) = 300. - zthick(:,:)   !   remaining thickness to reach 300m
-      DO jj = 1, jpj
-         DO ji = 1, jpi
-            htc3(ji,jj) = htc3(ji,jj) + tsn(ji,jj,ilevel+1,jp_tem)                  &
-               &                      * MIN( e3t_n(ji,jj,ilevel+1), zthick(ji,jj) ) * tmask(ji,jj,ilevel+1)
-         END DO
-      END DO
-      ! from temperature to heat contain
-      zcoef = rau0 * rcp
-      htc3(:,:) = zcoef * htc3(:,:)
-      CALL iom_put( "hc300", htc3 )      ! first 300m heat content
-      !
-      IF( ln_timing )   CALL timing_stop('dia_hth')
-      !
-   END SUBROUTINE dia_hth
-
-#else
-   !!----------------------------------------------------------------------
-   !!   Default option :                                       Empty module
-   !!----------------------------------------------------------------------
-   LOGICAL , PUBLIC, PARAMETER ::   lk_diahth = .FALSE.  !: thermocline-20d depths flag
-CONTAINS
-   SUBROUTINE dia_hth( kt )         ! Empty routine
-      IMPLICIT NONE
-      INTEGER, INTENT( in ) :: kt
-      WRITE(*,*) 'dia_hth: You should not have seen this print! error?', kt
-   END SUBROUTINE dia_hth
-#endif
-
-   !!======================================================================
+   END SUBROUTINE dia_hth_init
 END MODULE diahth
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diawri.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diawri.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/DIA/diawri.F90	(revision 11145)
@@ -43,4 +43,5 @@
    USE zdfdrg         ! ocean vertical physics: top/bottom friction
    USE zdfmxl         ! mixed layer
+   USE zdfosm         ! mixed layer
    !
    USE lbclnk         ! ocean lateral boundary conditions (or mpp link)
@@ -927,4 +928,22 @@
          CALL iom_rstput( 0, 0, inum, 'sdvecrtz', wsd            )    ! now StokesDrift k-velocity
       ENDIF
+
+      IF( ln_zdfosm ) THEN
+         CALL iom_rstput( 0, 0, inum, 'hbl', hbl*tmask(:,:,1)   )    ! now boundary-layer depth
+         CALL iom_rstput( 0, 0, inum, 'hml', hml*tmask(:,:,1)    )    ! now mixed-layer depth
+         CALL iom_rstput( 0, 0, inum, 'avt_k', avt_k*wmask       )    ! w-level diffusion
+         CALL iom_rstput( 0, 0, inum, 'avm_k', avm_k*wmask       )    ! now w-level viscosity
+         CALL iom_rstput( 0, 0, inum, 'ghamt', ghamt*wmask       )    ! non-local t forcing
+         CALL iom_rstput( 0, 0, inum, 'ghams', ghams*wmask       )    ! non-local s forcing
+         CALL iom_rstput( 0, 0, inum, 'ghamu', ghamu*wmask       )    ! non-local u forcing
+         CALL iom_rstput( 0, 0, inum, 'ghamv', ghamu*wmask       )    ! non-local v forcing
+      ENDIF
+     !     ! CALL histwrite( id_i, "zla", kt, zla*tmask(:,:,1)   , jpi*jpj, idex)         ! now Langmuir #
+     !     ! CALL histwrite( id_i, "zvstr", kt, zvstr*tmask(:,:,1)   , jpi*jpj, idex)     ! now mixed velocity scale
+     !     ! CALL histwrite( id_i, "zustar", kt, zustar*tmask(:,:,1)   , jpi*jpj, idex)   ! now friction velocity scale
+     !     ! CALL histwrite( id_i, "zwstrl", kt, zwstrl*tmask(:,:,1)   , jpi*jpj, idex)   ! now Langmuir velocity scale
+     !     ! CALL histwrite( id_i, "zwstrc", kt, zwstrc*tmask(:,:,1)   , jpi*jpj, idex)   ! now convective velocity scale
+     !     ! CALL histwrite( id_i, "zwb_ent", kt, zwb_ent*tmask(:,:,1)   , jpi*jpj, idex) ! now upward turb buoyancy entrainment flux
+     !     ! CALL histwrite( id_i, "zdb_bl", kt, zdb_bl*tmask(:,:,1)   , jpi*jpj, idex)   ! now db at ml base
  
 #if defined key_si3
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trd_oce.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trd_oce.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trd_oce.F90	(revision 11145)
@@ -33,5 +33,5 @@
 # endif
    !                                                  !!!* Active tracers trends indexes
-   INTEGER, PUBLIC, PARAMETER ::   jptot_tra  = 20     !: Total trend nb: change it when adding/removing one indice below
+   INTEGER, PUBLIC, PARAMETER ::   jptot_tra  = 21     !: Total trend nb: change it when adding/removing one indice below
    !                               ===============     !  
    INTEGER, PUBLIC, PARAMETER ::   jptra_xad  =  1     !: x- horizontal advection
@@ -46,4 +46,5 @@
    INTEGER, PUBLIC, PARAMETER ::   jptra_bbc  = 10     !: Bottom Boundary Condition (geoth. heating) 
    INTEGER, PUBLIC, PARAMETER ::   jptra_bbl  = 11     !: Bottom Boundary Layer (diffusive and/or advective)
+   INTEGER, PUBLIC, PARAMETER ::   jptra_osm  = 21     !: Non-local terms from OSMOSIS OBL model
    INTEGER, PUBLIC, PARAMETER ::   jptra_npc  = 12     !: non-penetrative convection treatment
    INTEGER, PUBLIC, PARAMETER ::   jptra_dmp  = 13     !: internal restoring (damping)
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trdtra.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trdtra.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/TRD/trdtra.F90	(revision 11145)
@@ -346,4 +346,6 @@
          CASE( jptra_bbl  )   ;   CALL iom_put( "ttrd_bbl"  , ptrdx )        ! bottom boundary layer
                                   CALL iom_put( "strd_bbl"  , ptrdy )
+         CASE( jptra_osm  )   ;   CALL iom_put( "ttrd_osm"  , ptrdx )        ! OSMOSIS non-local forcing
+                                  CALL iom_put( "strd_osm"  , ptrdy )
          CASE( jptra_npc  )   ;   CALL iom_put( "ttrd_npc"  , ptrdx )        ! static instability mixing
                                   CALL iom_put( "strd_npc"  , ptrdy )
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/ZDF/zdfosm.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/ZDF/zdfosm.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/ZDF/zdfosm.F90	(revision 11145)
@@ -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
@@ -51,4 +58,5 @@
    USE traqsr         ! details of solar radiation absorption
    USE zdfddm         ! double diffusion mixing (avs array)
+   USE zdfmxl         ! mixed layer depth
    USE iom            ! I/O library
    USE lib_mpp        ! MPP library
@@ -77,7 +85,7 @@
    REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   etmean   !: averaging operator for avt
    REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hbl      !: boundary layer depth
-   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hbli     !: intial boundary layer depth for stable blayer
+   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dh   ! depth of pycnocline
+   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hml  ! ML depth
    REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dstokes  !: penetration depth of the Stokes drift.
-
    !                      !!** Namelist  namzdf_osm  **
    LOGICAL  ::   ln_use_osm_la      ! Use namelist  rn_osm_la
@@ -97,5 +105,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
@@ -114,7 +123,16 @@
       !!                 ***  FUNCTION zdf_osm_alloc  ***
       !!----------------------------------------------------------------------
-     ALLOCATE( ghamu(jpi,jpj,jpk), ghamv(jpi,jpj,jpk), ghamt(jpi,jpj,jpk), ghams(jpi,jpj,jpk), &
-          &      hbl(jpi,jpj)    ,  hbli(jpi,jpj)    , dstokes(jpi, jpj) ,                     &
-          &   etmean(jpi,jpj,jpk),  STAT= zdf_osm_alloc )
+     ALLOCATE( ghamu(jpi,jpj,jpk), ghamv(jpi,jpj,jpk), ghamt(jpi,jpj,jpk),ghams(jpi,jpj,jpk), &
+          &       hbl(jpi,jpj), dh(jpi,jpj), hml(jpi,jpj), dstokes(jpi, jpj), &
+          &       etmean(jpi,jpj,jpk), STAT= zdf_osm_alloc )
+!     ALLOCATE( ghamu(jpi,jpj,jpk), ghamv(jpi,jpj,jpk), ghamt(jpi,jpj,jpk),ghams(jpi,jpj,jpk), &    ! would ths be better ?
+!          &       hbl(jpi,jpj), dh(jpi,jpj), hml(jpi,jpj), dstokes(jpi, jpj), &
+!          &       etmean(jpi,jpj,jpk), STAT= zdf_osm_alloc )
+!     IF( zdf_osm_alloc /= 0 )   CALL ctl_warn('zdf_osm_alloc: failed to allocate zdf_osm arrays')
+!     IF ( ln_osm_mle ) THEN
+!        Allocate( hmle(jpi,jpj), r1_ft(jpi,jpj), STAT= zdf_osm_alloc )
+!        IF( zdf_osm_alloc /= 0 )   CALL ctl_warn('zdf_osm_alloc: failed to allocate zdf_osm mle arrays')
+!     ENDIF
+
      IF( zdf_osm_alloc /= 0 )   CALL ctl_warn('zdf_osm_alloc: failed to allocate zdf_osm arrays')
      CALL mpp_sum ( 'zdfosm', zdf_osm_alloc )
@@ -166,5 +184,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
@@ -191,4 +210,5 @@
       REAL(wp), DIMENSION(jpi,jpj) :: zwbav     ! Buoyancy flux - bl average
       REAL(wp), DIMENSION(jpi,jpj) :: zwb_ent   ! Buoyancy entrainment flux
+
       REAL(wp), DIMENSION(jpi,jpj) :: zustke    ! Surface Stokes drift
       REAL(wp), DIMENSION(jpi,jpj) :: zla       ! Trubulent Langmuir number
@@ -196,5 +216,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
@@ -202,5 +222,4 @@
       INTEGER, DIMENSION(jpi,jpj) :: ibld ! level of boundary layer base
       INTEGER, DIMENSION(jpi,jpj) :: imld ! level of mixed-layer depth (pycnocline top)
-
       REAL(wp) :: ztgrad,zsgrad,zbgrad ! Temporary variables used to calculate pycnocline gradients
       REAL(wp) :: zugrad,zvgrad        ! temporary variables for calculating pycnocline shear
@@ -210,4 +229,6 @@
       REAL(wp), DIMENSION(jpi,jpj) :: zdh   ! pycnocline depth - grid
       REAL(wp), DIMENSION(jpi,jpj) :: zdhdt ! BL depth tendency
+      REAL(wp), DIMENSION(jpi,jpj) :: zdhdt_2                                    ! correction to dhdt due to internal structure.
+      REAL(wp), DIMENSION(jpi,jpj) :: zdtdz_ext,zdsdz_ext,zdbdz_ext              ! external temperature/salinity and buoyancy gradients
       REAL(wp), DIMENSION(jpi,jpj) :: zt_bl,zs_bl,zu_bl,zv_bl,zrh_bl  ! averages over the depth of the blayer
       REAL(wp), DIMENSION(jpi,jpj) :: zt_ml,zs_ml,zu_ml,zv_ml,zrh_ml  ! averages over the depth of the mixed layer
@@ -238,5 +259,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 +268,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 +291,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
@@ -277,4 +301,5 @@
       zdudz_pyc(:,:,:) = 0._wp ; zdvdz_pyc(:,:,:) = 0._wp
       !
+      zdtdz_ext(:,:) = 0._wp ; zdsdz_ext(:,:) = 0._wp ; zdbdz_ext(:,:) = 0._wp
       ! Flux-Gradient arrays.
       zdifml_sc(:,:)  = 0._wp ; zvisml_sc(:,:)  = 0._wp ; zdifpyc_sc(:,:) = 0._wp
@@ -287,4 +312,5 @@
       ghams(:,:,:)   = 0._wp ; ghamu(:,:,:)   = 0._wp ; ghamv(:,:,:) = 0._wp
 
+      zdhdt_2(:,:) = 0._wp
       ! hbl = MAX(hbl,epsln)
       !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
@@ -356,4 +382,5 @@
      CASE(2)
         zfac =  2.0_wp * rpi / 16.0_wp
+
         DO jj = 2, jpjm1
            DO ji = 2, jpim1
@@ -362,7 +389,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 +402,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 +410,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 +426,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,75 +439,32 @@
       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
-
+      ! Averages over well-mixed and boundary layer
+      ! Alan: do we need zb_nl?, 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)
+      CALL zdf_osm_vertical_average(imld, zt_ml, zs_ml, zu_ml, zv_ml, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml)
+! External gradient
+      CALL zdf_osm_external_gradients( zdtdz_ext, zdsdz_ext, zdbdz_ext )
+
+! Rate of change of hbl
+      CALL zdf_osm_calculate_dhdt( zdhdt, zdhdt_2 )
       ! Calculate averages over depth of boundary layer
       imld = ibld           ! use imld to hold previous blayer index
-      ibld(:,:) = 3
-
-      zhbl_t(:,:) = hbl(:,:) + (zdhdt(:,:) - wn(ji,jj,ibld(ji,jj)))* rn_rdt ! certainly need wb here, so subtract it
-      zhbl_t(:,:) = MIN(zhbl_t(:,:), ht_n(:,:))
-      zdhdt(:,:) = MIN(zdhdt(:,:), (zhbl_t(:,:) - hbl(:,:))/rn_rdt + wn(ji,jj,ibld(ji,jj))) ! adjustment to represent limiting by ocean bottom
+      ibld(:,:) = 4
+
+         DO jj = 2, jpjm1
+            DO ji = 2, jpim1
+               zhbl_t(ji,jj) = hbl(ji,jj) + (zdhdt(ji,jj) - wn(ji,jj,ibld(ji,jj)))* rn_rdt ! certainly need wn here, so subtract it
+               zhbl_t(ji,jj) = MIN(zhbl_t(ji,jj), gdepw_n(ji,jj, mbkt(ji,jj) + 1) - depth_tol)! ht_n(:,:))
+               !zdhdt(ji,jj) = MIN(zdhdt(ji,jj), (zhbl_t(ji,jj) - hbl(ji,jj))/rn_rdt + wn(ji,jj,ibld(ji,jj)))
+            END DO
+         END DO
+      ! adjustment to represent limiting by ocean bottom
 
       DO jk = 4, jpkm1
@@ -495,5 +472,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 +481,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 +525,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 +540,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,14 +561,23 @@
                 ! 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
-                zdiffut(ji,jj,ibld(ji,jj)) = zdhdt(ji,jj)* e3t_n(ji,jj,ibld(ji,jj))
+                ! Temporary fix to ensure zdiffut is +ve; won't be necessary with wn taken out
+                zdiffut(ji,jj,ibld(ji,jj)) = zdhdt(ji,jj) * e3t_n(ji,jj,ibld(ji,jj))
+                zviscos(ji,jj,ibld(ji,jj)) = zdhdt(ji,jj) * e3t_n(ji,jj,ibld(ji,jj))
                 ! could be taken out, take account of entrainment represents as a diffusivity
                 ! should remove w from here, represents entrainment
@@ -908,6 +589,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 +641,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 +698,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 +711,4 @@
           END DO   ! ji loop
        END DO      ! jj loop
-
 
        WHERE ( lconv )
@@ -1051,4 +741,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 +783,4 @@
        END DO      ! jj loop
 
-
        WHERE ( lconv )
           zsc_uw_1 = zustar**2
@@ -1134,4 +827,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 +867,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 +876,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 +893,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 +903,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
@@ -1232,4 +949,6 @@
        IF(ln_dia_osm) THEN
           IF ( iom_use("zdtdz_pyc") ) CALL iom_put( "zdtdz_pyc", wmask*zdtdz_pyc )
+          IF ( iom_use("zdsdz_pyc") ) CALL iom_put( "zdsdz_pyc", wmask*zdsdz_pyc )
+          IF ( iom_use("zdbdz_pyc") ) CALL iom_put( "zdbdz_pyc", wmask*zdbdz_pyc )
        END IF
 
@@ -1287,5 +1006,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 +1038,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 +1049,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 +1066,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,18 +1079,507 @@
          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
-         IF ( iom_use("zwth_ent") ) CALL iom_put( "zwth_ent", tmask(:,:,1)*zwth_ent )               ! ML depth internal to zdf_osm routine
-         IF ( iom_use("zt_ml") ) CALL iom_put( "zt_ml", tmask(:,:,1)*zt_ml )               ! average T in ML
+         IF ( iom_use("zwthav") ) CALL iom_put( "zwthav", tmask(:,:,1)*zwthav )         ! upward BL-avged turb temp flux
+         IF ( iom_use("zwth_ent") ) CALL iom_put( "zwth_ent", tmask(:,:,1)*zwth_ent )   ! upward turb temp entrainment flux
+         IF ( iom_use("zwb_ent") ) CALL iom_put( "zwb_ent", tmask(:,:,1)*zwb_ent )      ! upward turb buoyancy entrainment flux
+         IF ( iom_use("zws_ent") ) CALL iom_put( "zws_ent", tmask(:,:,1)*zws_ent )      ! upward turb salinity entrainment flux
+         IF ( iom_use("zt_ml") ) CALL iom_put( "zt_ml", tmask(:,:,1)*zt_ml )            ! average T in ML
       END IF
-      ! Lateral boundary conditions on p_avt  (sign unchanged)
-      CALL lbc_lnk( 'zdfosm', p_avt(:,:,:), 'W', 1. )
+
+CONTAINS
+
+
+   ! Alan: do we need zb?
+   SUBROUTINE zdf_osm_vertical_average( jnlev_av, zt, zs, zu, zv, zdt, zds, zdb, zdu, zdv )
+     !!---------------------------------------------------------------------
+     !!                   ***  ROUTINE zdf_vertical_average  ***
+     !!
+     !! ** Purpose : Determines vertical averages from surface to jnlev.
+     !!
+     !! ** Method  : Averages are calculated from the surface to jnlev.
+     !!              The external level used to calculate differences is ibld+ibld_ext
+     !!
+     !!----------------------------------------------------------------------
+
+        INTEGER, DIMENSION(jpi,jpj) :: jnlev_av  ! Number of levels to average over.
+
+        ! Alan: do we need zb?
+        REAL(wp), DIMENSION(jpi,jpj) :: zt, zs        ! Average temperature and salinity
+        REAL(wp), DIMENSION(jpi,jpj) :: zu,zv         ! Average current components
+        REAL(wp), DIMENSION(jpi,jpj) :: zdt, zds, zdb ! Difference between average and value at base of OSBL
+        REAL(wp), DIMENSION(jpi,jpj) :: zdu, zdv      ! Difference for velocity components.
+
+        INTEGER :: jk, ji, jj
+        REAL(wp) :: zthick, zthermal, zbeta
+
+
+        zt   = 0._wp
+        zs   = 0._wp
+        zu   = 0._wp
+        zv   = 0._wp
+        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)
+               ! average over depth of boundary layer
+            zthick = epsln
+            DO jk = 2, jnlev_av(ji,jj)
+               zthick = zthick + e3t_n(ji,jj,jk)
+               zt(ji,jj)   = zt(ji,jj)  + e3t_n(ji,jj,jk) * tsn(ji,jj,jk,jp_tem)
+               zs(ji,jj)   = zs(ji,jj)  + e3t_n(ji,jj,jk) * tsn(ji,jj,jk,jp_sal)
+               zu(ji,jj)   = zu(ji,jj)  + e3t_n(ji,jj,jk) &
+                     &            * ( ub(ji,jj,jk) + ub(ji - 1,jj,jk) ) &
+                     &            / MAX( 1. , umask(ji,jj,jk) + umask(ji - 1,jj,jk) )
+               zv(ji,jj)   = zv(ji,jj)  + e3t_n(ji,jj,jk) &
+                     &            * ( vb(ji,jj,jk) + vb(ji,jj - 1,jk) ) &
+                     &            / MAX( 1. , vmask(ji,jj,jk) + vmask(ji,jj - 1,jk) )
+            END DO
+            zt(ji,jj) = zt(ji,jj) / zthick
+            zs(ji,jj) = zs(ji,jj) / zthick
+            zu(ji,jj) = zu(ji,jj) / zthick
+            zv(ji,jj) = zv(ji,jj) / zthick
+            ! Alan: do we need zb?
+            zdt(ji,jj) = zt(ji,jj) - tsn(ji,jj,ibld(ji,jj)+ibld_ext,jp_tem)
+            zds(ji,jj) = zs(ji,jj) - tsn(ji,jj,ibld(ji,jj)+ibld_ext,jp_sal)
+            zdu(ji,jj) = zu(ji,jj) - ( ub(ji,jj,ibld(ji,jj)+ibld_ext) + ub(ji-1,jj,ibld(ji,jj)+ibld_ext ) ) &
+                     &    / MAX(1. , umask(ji,jj,ibld(ji,jj)+ibld_ext ) + umask(ji-1,jj,ibld(ji,jj)+ibld_ext ) )
+            zdv(ji,jj) = zv(ji,jj) - ( vb(ji,jj,ibld(ji,jj)+ibld_ext) + vb(ji,jj-1,ibld(ji,jj)+ibld_ext ) ) &
+                     &   / MAX(1. , vmask(ji,jj,ibld(ji,jj)+ibld_ext ) + vmask(ji,jj-1,ibld(ji,jj)+ibld_ext ) )
+            zdb(ji,jj) = grav * zthermal * zdt(ji,jj) - grav * zbeta * zds(ji,jj)
+         END DO
+        END DO
+   END SUBROUTINE zdf_osm_vertical_average
+
+   SUBROUTINE zdf_osm_velocity_rotation( zcos_w, zsin_w, zu, zv, zdu, zdv )
+     !!---------------------------------------------------------------------
+     !!                   ***  ROUTINE zdf_velocity_rotation  ***
+     !!
+     !! ** Purpose : Rotates frame of reference of averaged velocity components.
+     !!
+     !! ** Method  : The velocity components are rotated into frame specified by zcos_w and zsin_w
+     !!
+     !!----------------------------------------------------------------------
+
+        REAL(wp), DIMENSION(jpi,jpj) :: zcos_w, zsin_w       ! Cos and Sin of rotation angle
+        REAL(wp), DIMENSION(jpi,jpj) :: zu, zv               ! Components of current
+        REAL(wp), DIMENSION(jpi,jpj) :: zdu, zdv             ! Change in velocity components across pycnocline
+
+        INTEGER :: ji, jj
+        REAL(wp) :: ztemp
+
+        DO jj = 2, jpjm1
+           DO ji = 2, jpim1
+              ztemp = zu(ji,jj)
+              zu(ji,jj) = zu(ji,jj) * zcos_w(ji,jj) + zv(ji,jj) * zsin_w(ji,jj)
+              zv(ji,jj) = zv(ji,jj) * zcos_w(ji,jj) - ztemp * zsin_w(ji,jj)
+              ztemp = zdu(ji,jj)
+              zdu(ji,jj) = zdu(ji,jj) * zcos_w(ji,jj) + zdv(ji,jj) * zsin_w(ji,jj)
+              zdv(ji,jj) = zdv(ji,jj) * zsin_w(ji,jj) - ztemp * zsin_w(ji,jj)
+           END DO
+        END DO
+    END SUBROUTINE zdf_osm_velocity_rotation
+
+    SUBROUTINE zdf_osm_external_gradients( zdtdz, zdsdz, zdbdz )
+     !!---------------------------------------------------------------------
+     !!                   ***  ROUTINE zdf_osm_external_gradients  ***
+     !!
+     !! ** Purpose : Calculates the gradients below the OSBL
+     !!
+     !! ** Method  : Uses ibld and ibld_ext to determine levels to calculate the gradient.
+     !!
+     !!----------------------------------------------------------------------
+
+     REAL(wp), DIMENSION(jpi,jpj) :: zdtdz, zdsdz, zdbdz   ! External gradients of temperature, salinity and buoyancy.
+
+     INTEGER :: jj, ji, jkb, jkb1
+     REAL(wp) :: zthermal, zbeta
+
+
+     DO jj = 2, jpjm1
+        DO ji = 2, jpim1
+           IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) ) THEN
+              zthermal = rab_n(ji,jj,1,jp_tem) !ideally use ibld not 1??
+              zbeta    = rab_n(ji,jj,1,jp_sal)
+              jkb = ibld(ji,jj) + ibld_ext
+              jkb1 = MIN(jkb + 1, mbkt(ji,jj))
+              zdtdz(ji,jj) = - ( tsn(ji,jj,jkb1,jp_tem) - tsn(ji,jj,jkb,jp_tem ) ) &
+                   &   / e3t_n(ji,jj,ibld(ji,jj))
+              zdsdz(ji,jj) = - ( tsn(ji,jj,jkb1,jp_sal) - tsn(ji,jj,jkb,jp_sal ) ) &
+                   &   / e3t_n(ji,jj,ibld(ji,jj))
+              zdbdz(ji,jj) = grav * zthermal * zdtdz(ji,jj) - grav * zbeta * zdsdz(ji,jj)
+           END IF
+        END DO
+     END DO
+    END SUBROUTINE zdf_osm_external_gradients
+
+    SUBROUTINE zdf_osm_pycnocline_scalar_profiles( zdtdz, zdsdz, zdbdz )
+
+     REAL(wp), DIMENSION(jpi,jpj,jpk) :: zdtdz, zdsdz, zdbdz      ! gradients in the pycnocline
+
+     INTEGER :: jk, jj, ji
+     REAL(wp) :: ztgrad, zsgrad, zbgrad
+     REAL(wp) :: zgamma_b_nd, zgamma_c, znd
+     REAL(wp) :: zzeta_s=0.3
+
+     DO jj = 2, jpjm1
+        DO ji = 2, jpim1
+           IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) ) THEN
+              IF ( lconv(ji,jj) ) THEN  ! convective conditions
+                 IF ( zdb_bl(ji,jj) > 0._wp .AND. zdbdz_ext(ji,jj) > 0._wp ) THEN
+                    ztgrad = 0.5 * zdt_ml(ji,jj) / zdh(ji,jj) + zdtdz_ext(ji,jj)
+                    zsgrad = 0.5 * zds_ml(ji,jj) / zdh(ji,jj) + zdsdz_ext(ji,jj)
+                    zbgrad = 0.5 * zdb_ml(ji,jj) / zdh(ji,jj) + zdbdz_ext(ji,jj)
+                    zgamma_b_nd = zdbdz_ext(ji,jj) * zdh(ji,jj) / zdb_ml(ji,jj)
+                    zgamma_c = ( 3.14159 / 4.0 ) * ( 0.5 + zgamma_b_nd ) /&
+                         & ( 1.0 - 0.25 * SQRT( 3.14159 / 6.0 ) - 2.0 * zgamma_b_nd * zzeta_s )**2 ! check
+                    DO jk = 2, ibld(ji,jj)+ibld_ext
+                       znd = -( gdepw_n(ji,jj,jk) - zhbl(ji,jj) ) / zdh(ji,jj)
+                       IF ( znd <= zzeta_s ) THEN
+                          zdtdz(ji,jj,jk) = zdtdz_ext(ji,jj) + 0.5 * zdt_ml(ji,jj) / zdh(ji,jj) * &
+                               & EXP( -6.0 * ( znd -zzeta_s )**2 )
+                          zdsdz(ji,jj,jk) = zdsdz_ext(ji,jj) + 0.5 * zds_ml(ji,jj) / zdh(ji,jj) * &
+                               & EXP( -6.0 * ( znd -zzeta_s )**2 )
+                          zdbdz(ji,jj,jk) = zdbdz_ext(ji,jj) + 0.5 * zdb_ml(ji,jj) / zdh(ji,jj) * &
+                               & EXP( -6.0 * ( znd -zzeta_s )**2 )
+                       ELSE
+                          zdtdz(ji,jj,jk) =  ztgrad * EXP( -zgamma_c * ( znd - zzeta_s )**2 )
+                          zdsdz(ji,jj,jk) =  zsgrad * EXP( -zgamma_c * ( znd - zzeta_s )**2 )
+                          zdbdz(ji,jj,jk) =  zbgrad * EXP( -zgamma_c * ( znd - zzeta_s )**2 )
+                       ENDIF
+                    END DO
+                 ENDIF ! If condition not satisfied, no pycnocline present. Gradients have been initialised to zero, so do nothing
+              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(ji,jj,jk) =  ztgrad * EXP( -15.0 * ( znd - 0.9 )**2 )
+                             zdbdz(ji,jj,jk) =  zbgrad * EXP( -15.0 * ( znd - 0.9 )**2 )
+                             zdsdz(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(ji,jj,jk) =  ztgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
+                             zdbdz(ji,jj,jk) =  zbgrad * EXP( -1.75 * ( znd + 0.75 )**2 )
+                             zdsdz(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 and profile arrays are intialized to zero
+              ENDIF          ! IF (lconv)
+           END IF      ! IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) )
+        END DO
+     END DO
+
+    END SUBROUTINE zdf_osm_pycnocline_scalar_profiles
+
+    SUBROUTINE zdf_osm_pycnocline_shear_profiles( zdudz, zdvdz )
+      !!---------------------------------------------------------------------
+      !!                   ***  ROUTINE zdf_osm_pycnocline_shear_profiles  ***
+      !!
+      !! ** Purpose : Calculates velocity shear in the pycnocline
+      !!
+      !! ** Method  :
+      !!
+      !!----------------------------------------------------------------------
+
+      REAL(wp), DIMENSION(jpi,jpj,jpk) :: zdudz, zdvdz
+
+      INTEGER :: jk, jj, ji
+      REAL(wp) :: zugrad, zvgrad, znd
+      REAL(wp) :: zzeta_v = 0.45
       !
-   END SUBROUTINE zdf_osm
+      DO jj = 2, jpjm1
+         DO ji = 2, jpim1
+            !
+            IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) ) THEN
+               IF ( lconv (ji,jj) ) THEN
+                  ! Unstable conditions
+                  zugrad = 0.7 * zdu_ml(ji,jj) / zdh(ji,jj) + 0.3 * zustar(ji,jj)*zustar(ji,jj) / &
+                       &      ( ( ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird * zhml(ji,jj) ) * &
+                       &      MIN(zla(ji,jj)**(8.0/3.0) + epsln, 0.12 ))
+                  !Alan is this right?
+                  zvgrad = ( 0.7 * zdv_ml(ji,jj) + &
+                       &    2.0 * ff_t(ji,jj) * zustke(ji,jj) * dstokes(ji,jj) / &
+                       &          ( ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird  + epsln ) &
+                       &      )/ (zdh(ji,jj)  + epsln )
+                  DO jk = 2, ibld(ji,jj) - 1 + ibld_ext
+                     znd = -( gdepw_n(ji,jj,jk) - zhbl(ji,jj) ) / (zdh(ji,jj) + epsln ) - zzeta_v
+                     IF ( znd <= 0.0 ) THEN
+                        zdudz(ji,jj,jk) = 1.25 * zugrad * EXP( 3.0 * znd )
+                        zdvdz(ji,jj,jk) = 1.25 * zvgrad * EXP( 3.0 * znd )
+                     ELSE
+                        zdudz(ji,jj,jk) = 1.25 * zugrad * EXP( -2.0 * znd )
+                        zdvdz(ji,jj,jk) = 1.25 * zvgrad * EXP( -2.0 * znd )
+                     ENDIF
+                  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(ji,jj,jk) = zugrad * EXP( -40.0 * ( znd - 1.0 )**2 )
+                     ELSE
+                        zdudz(ji,jj,jk) = zugrad * EXP( -20.0 * ( znd - 1.0 )**2 )
+                     ENDIF
+                     zdvdz(ji,jj,jk) = zvgrad * EXP( -20.0 * ( znd - 0.85 )**2 )
+                  END DO
+               ENDIF
+               !
+            END IF      ! IF ( ibld(ji,jj) + ibld_ext < mbkt(ji,jj) )
+         END DO
+      END DO
+    END SUBROUTINE zdf_osm_pycnocline_shear_profiles
+
+    SUBROUTINE zdf_osm_calculate_dhdt( zdhdt, zdhdt_2 )
+     !!---------------------------------------------------------------------
+     !!                   ***  ROUTINE zdf_osm_calculate_dhdt  ***
+     !!
+     !! ** Purpose : Calculates the rate at which hbl changes.
+     !!
+     !! ** Method  :
+     !!
+     !!----------------------------------------------------------------------
+
+    REAL(wp), DIMENSION(jpi,jpj) :: zdhdt, zdhdt_2        ! Rate of change of hbl
+
+    INTEGER :: jj, ji
+    REAL(wp) :: zgamma_b_nd, zgamma_dh_nd, zpert
+    REAL(wp) :: zvel_max, zwb_min
+    REAL(wp) :: zwcor, zrf_conv, zrf_shear, zrf_langmuir, zr_stokes
+    REAL(wp) :: zzeta_m = 0.3
+    REAL(wp) :: zgamma_c = 2.0
+    REAL(wp) :: zdhoh = 0.1
+    REAL(wp) :: alpha_bc = 0.5
+
+    DO jj = 2, jpjm1
+       DO ji = 2, jpim1
+          IF ( lconv(ji,jj) ) THEN    ! Convective
+             ! Alan is this right?  Yes, it's a bit different from the previous relationship
+             ! zwb_ent(ji,jj) = - 2.0 * 0.2 * zwbav(ji,jj) &
+             !   &            - ( 0.15 * ( 1.0 - EXP( -1.5 * zla(ji,jj) ) ) * zustar(ji,jj)**3 + 0.03 * zwstrl(ji,jj)**3 ) / zhml(ji,jj)
+             zwcor = ABS(ff_t(ji,jj)) * zhbl(ji,jj) + epsln
+             zrf_conv = TANH( ( zwstrc(ji,jj) / zwcor )**0.69 )
+             zrf_shear = TANH( ( zustar(ji,jj) / zwcor )**0.69 )
+             zrf_langmuir = TANH( ( zwstrl(ji,jj) / zwcor )**0.69 )
+             zr_stokes = 1.0 - EXP( -25.0 * dstokes(ji,jj) / hbl(ji,jj) &
+                  &                  * ( 1.0 + 4.0 * dstokes(ji,jj) / hbl(ji,jj) ) )
+
+             zwb_ent(ji,jj) = - 2.0 * 0.2 * zrf_conv * zwbav(ji,jj) &
+                  &                  - 0.15 * zrf_shear * zustar(ji,jj)**3 /zhml(ji,jj) &
+                  &         + zr_stokes * ( 0.15 * EXP( -1.5 * zla(ji,jj) ) * zrf_shear * zustar(ji,jj)**3 &
+                  &                                         - zrf_langmuir * 0.03 * zwstrl(ji,jj)**3 ) / zhml(ji,jj)
+             !
+             zwb_min = dh(ji,jj) * zwb0(ji,jj) / zhml(ji,jj) + zwb_ent(ji,jj)
+                  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
+                  zdhdt(ji,jj) = -zwb_ent(ji,jj) / ( zvel_max + MAX(zdb_bl(ji,jj), 1.0e-15) )
+                ! added ajgn 23 July as temporay fix
+             zdhdt_2(ji,jj) = 0._wp
+
+                ! commented out ajgn 23 July as temporay fix
+!                 IF ( zdb_ml(ji,jj) > 0.0 .and. zdbdz_ext(ji,jj) > 0.0 ) THEN
+! !additional term to account for thickness of pycnocline on dhdt. Small correction, so could get rid of this if necessary.
+!                     zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
+!                     zgamma_b_nd = zdbdz_ext(ji,jj) * zhml(ji,jj) / zdb_ml(ji,jj)
+!                     zgamma_dh_nd = zdbdz_ext(ji,jj) * zdh(ji,jj) / zdb_ml(ji,jj)
+!                     zdhdt_2(ji,jj) = ( 1.0 - SQRT( 3.1415 / ( 4.0 * zgamma_c) ) * zdhoh ) * zdh(ji,jj) / zhml(ji,jj)
+!                     zdhdt_2(ji,jj) = zdhdt_2(ji,jj) * ( zwb0(ji,jj) - (1.0 + zgamma_b_nd / alpha_bc ) * zwb_min )
+!                     ! Alan no idea what this should be?
+!                     zdhdt_2(ji,jj) = alpha_bc / ( 4.0 * zgamma_c ) * zdhdt_2(ji,jj) &
+!                        &        + (alpha_bc + zgamma_dh_nd ) * ( 1.0 + SQRT( 3.1414 / ( 4.0 * zgamma_c ) ) * zdh(ji,jj) / zhbl(ji,jj) ) &
+!                        &        * (1.0 / ( 4.0 * zgamma_c * alpha_bc ) ) * zwb_min * zdh(ji,jj) / zhbl(ji,jj)
+!                     zdhdt_2(ji,jj) = zdhdt_2(ji,jj) / ( zvel_max + MAX( zdb_bl(ji,jj), 1.0e-15 ) )
+!                     IF ( zdhdt_2(ji,jj) <= 0.2 * zdhdt(ji,jj) ) THEN
+!                         zdhdt(ji,jj) = zdhdt(ji,jj) + zdhdt_2(ji,jj)
+!                 ENDIF
+            ELSE                        ! Stable
+                zdhdt(ji,jj) = ( 0.06 + 0.52 * zhol(ji,jj) / 2.0 ) * zvstr(ji,jj)**3 / hbl(ji,jj) + zwbav(ji,jj)
+                zdhdt_2(ji,jj) = 0._wp
+                IF ( zdhdt(ji,jj) < 0._wp ) THEN
+                   ! For long timsteps factor in brackets slows the rapid collapse of the OSBL
+                    zpert = 2.0 * ( 1.0 + 0.0 * 2.0 * zvstr(ji,jj) * rn_rdt / hbl(ji,jj) ) * zvstr(ji,jj)**2 / hbl(ji,jj)
+                ELSE
+                    zpert = MAX( 2.0 * ( 1.0 + 0.0 * 2.0 * zvstr(ji,jj) * rn_rdt / hbl(ji,jj) ) * zvstr(ji,jj)**2 / hbl(ji,jj), zdb_bl(ji,jj) )
+                ENDIF
+                zdhdt(ji,jj) = 2.0 * zdhdt(ji,jj) / zpert
+            ENDIF
+         END DO
+      END DO
+    END SUBROUTINE zdf_osm_calculate_dhdt
+
+    SUBROUTINE zdf_osm_timestep_hbl( zdhdt, zdhdt_2 )
+     !!---------------------------------------------------------------------
+     !!                   ***  ROUTINE zdf_osm_timestep_hbl  ***
+     !!
+     !! ** Purpose : Increments hbl.
+     !!
+     !! ** Method  : If thechange in hbl exceeds one model level the change is
+     !!              is calculated by moving down the grid, changing the buoyancy
+     !!              jump. This is to ensure that the change in hbl does not
+     !!              overshoot a stable layer.
+     !!
+     !!----------------------------------------------------------------------
+
+
+    REAL(wp), DIMENSION(jpi,jpj) :: zdhdt, zdhdt_2   ! rates of change of hbl.
+
+    INTEGER :: jk, jj, ji, jm
+    REAL(wp) :: zhbl_s, zvel_max, zdb
+    REAL(wp) :: zthermal, zbeta
+
+     DO jj = 2, jpjm1
+         DO ji = 2, jpim1
+           IF ( ibld(ji,jj) - imld(ji,jj) > 1 ) THEN
+!
+! 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( &
+                        & rn_rdt * (  -zwb_ent(ji,jj) / zdb + zdhdt_2(ji,jj) ) / FLOAT(ibld(ji,jj) - imld(ji,jj) ), &
+                        & e3w_n(ji,jj,jm) )
+                    zhbl_s = MIN(zhbl_s,  gdepw_n(ji,jj, mbkt(ji,jj) + 1) - depth_tol)
+
+                    IF ( zhbl_s >= gdepw_n(ji,jj,jm+1) ) jm = jm + 1
+                 END DO
+                 hbl(ji,jj) = zhbl_s
+                 ibld(ji,jj) = jm
+             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 * zvstr(ji,jj)**2 / zhbl_s
+
+                    ! Alan is thuis right? I have simply changed hbli to hbl
+                    zhol(ji,jj) = -zhbl_s / ( ( zvstr(ji,jj)**3 + epsln )/ zwbav(ji,jj) )
+                    zdhdt(ji,jj) = -( zwbav(ji,jj) - 0.04 / 2.0 * zwstrl(ji,jj)**3 / zhbl_s - 0.15 / 2.0 * ( 1.0 - EXP( -1.5 * zla(ji,jj) ) ) * &
+               &                  zustar(ji,jj)**3 / zhbl_s ) * ( 0.725 + 0.225 * EXP( -7.5 * zhol(ji,jj) ) )
+                    zdhdt(ji,jj) = zdhdt(ji,jj) + zwbav(ji,jj)
+                    zhbl_s = zhbl_s + MIN( zdhdt(ji,jj) / zdb * rn_rdt / FLOAT( ibld(ji,jj) - imld(ji,jj) ), e3w_n(ji,jj,jm) )
+
+                    zhbl_s = MIN(zhbl_s, gdepw_n(ji,jj, mbkt(ji,jj) + 1) - depth_tol)
+                    IF ( zhbl_s >= gdepw_n(ji,jj,jm) ) jm = jm + 1
+                 END DO
+             ENDIF   ! IF ( lconv )
+             hbl(ji,jj) = MAX(zhbl_s, gdepw_n(ji,jj,4) )
+             ibld(ji,jj) = MAX(jm, 4 )
+           ELSE
+! change zero or one model level.
+             hbl(ji,jj) = MAX(zhbl_t(ji,jj), gdepw_n(ji,jj,4) )
+           ENDIF
+           zhbl(ji,jj) = gdepw_n(ji,jj,ibld(ji,jj))
+         END DO
+      END DO
+
+    END SUBROUTINE zdf_osm_timestep_hbl
+
+    SUBROUTINE zdf_osm_pycnocline_thickness( dh, zdh )
+      !!---------------------------------------------------------------------
+      !!                   ***  ROUTINE zdf_osm_pycnocline_thickness  ***
+      !!
+      !! ** Purpose : Calculates thickness of the pycnocline
+      !!
+      !! ** Method  : The thickness is calculated from a prognostic equation
+      !!              that relaxes the pycnocine thickness to a diagnostic
+      !!              value. The time change is calculated assuming the
+      !!              thickness relaxes exponentially. This is done to deal
+      !!              with large timesteps.
+      !!
+      !!----------------------------------------------------------------------
+
+      REAL(wp), DIMENSION(jpi,jpj) :: dh, zdh     ! pycnocline thickness.
+      !
+      INTEGER :: jj, ji
+      INTEGER :: inhml
+      REAL(wp) :: zari, ztau, zddhdt
+
+
+      DO jj = 2, jpjm1
+         DO ji = 2, jpim1
+            IF ( lconv(ji,jj) ) THEN
+               IF ( zdb_bl(ji,jj) > 0._wp .and. zdbdz_ext(ji,jj) > 0._wp)THEN
+                  IF ( ( zwstrc(ji,jj) / zvstr(ji,jj) )**3 <= 0.5 ) THEN  ! near neutral stability
+                     zari = MIN( 1.5 * ( zdb_bl(ji,jj) / ( zdbdz_ext(ji,jj) * zhbl(ji,jj) ) ) / &
+                          (1.0 + zdb_bl(ji,jj)**2 / ( 4.5 * zvstr(ji,jj)**2 * zdbdz_ext(ji,jj) ) ), 0.2 )
+                  ELSE                                                     ! unstable
+                     zari = MIN( 1.5 * ( zdb_bl(ji,jj) / ( zdbdz_ext(ji,jj) * zhbl(ji,jj) ) ) / &
+                          (1.0 + zdb_bl(ji,jj)**2 / ( 4.5 * zwstrc(ji,jj)**2 * zdbdz_ext(ji,jj) ) ), 0.2 )
+                  ENDIF
+                  ztau = hbl(ji,jj) / (zvstr(ji,jj)**3 + 0.5 *zwstrc(ji,jj)**3)**pthird
+                  zddhdt = zari * hbl(ji,jj)
+               ELSE
+                  ztau = hbl(ji,jj) / ( zvstr(ji,jj)**3 + 0.5 * zwstrc(ji,jj)**3 )**pthird
+                  zddhdt = 0.2 * hbl(ji,jj)
+               ENDIF
+               dh(ji,jj) = dh(ji,jj) * EXP( -rn_rdt / ztau ) + zddhdt * ( 1.0 - EXP( -rn_rdt / ztau ) )
+               ! Alan: this hml is never defined or used
+            ELSE   ! IF (lconv)
+               ztau = hbl(ji,jj) / zvstr(ji,jj)
+               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 )
+                     zddhdt = MIN( zari, 0.2 ) * hbl(ji,jj)
+                  ELSE
+                     zddhdt = 0.2 * hbl(ji,jj)
+                  ENDIF
+               ELSE     ! IF(dhdt < 0)
+                  zddhdt = 0.2 * hbl(ji,jj)
+               ENDIF    ! IF (dhdt >= 0)
+               dh(ji,jj) = dh(ji,jj) * EXP( -rn_rdt / ztau )+ zddhdt * ( 1.0 - EXP( -rn_rdt / ztau ) )
+               IF ( zdhdt(ji,jj) < 0._wp .and. dh(ji,jj) >= hbl(ji,jj) ) dh(ji,jj) = zddhdt       ! can be a problem with dh>hbl for rapid collapse
+               ! Alan: this hml is never defined or used -- do we need it?
+            ENDIF       ! IF (lconv)
+
+            hml(ji,jj) = hbl(ji,jj) - dh(ji,jj)
+            inhml = MAX( INT( dh(ji,jj) / e3t_n(ji,jj,ibld(ji,jj)) ) , 1 )
+            imld(ji,jj) = MAX( ibld(ji,jj) - inhml, 3)
+            zhml(ji,jj) = gdepw_n(ji,jj,imld(ji,jj))
+            zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
+         END DO
+      END DO
+
+    END SUBROUTINE zdf_osm_pycnocline_thickness
+
+END SUBROUTINE zdf_osm
 
 
@@ -1378,8 +1598,9 @@
      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 &
           & ,nn_osm_wave, ln_dia_osm, rn_osm_hbl0 &
-          & ,ln_kpprimix, rn_riinfty, rn_difri, ln_convmix, rn_difconv
+          & ,ln_kpprimix, rn_riinfty, rn_difri, ln_convmix, rn_difconv, nn_osm_wave
      !!----------------------------------------------------------------------
      !
@@ -1397,9 +1618,9 @@
         WRITE(numout,*) 'zdf_osm_init : OSMOSIS Parameterisation'
         WRITE(numout,*) '~~~~~~~~~~~~'
-        WRITE(numout,*) '   Namelist namzdf_osm : set tke mixing parameters'
-        WRITE(numout,*) '     Use namelist  rn_osm_la                     ln_use_osm_la = ', ln_use_osm_la
+        WRITE(numout,*) '   Namelist namzdf_osm : set osm mixing parameters'
+        WRITE(numout,*) '     Use  rn_osm_la                                ln_use_osm_la = ', ln_use_osm_la
         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
@@ -1423,5 +1644,7 @@
      IF( zdf_osm_alloc() /= 0 )   CALL ctl_stop( 'STOP', 'zdf_osm_init : unable to allocate arrays' )
 
-     call osm_rst( nit000, 'READ' ) !* read or initialize hbl
+
+     call osm_rst( nit000, 'READ' ) !* read or initialize hbl, dh, hmle
+
 
      IF( ln_zdfddm) THEN
@@ -1536,5 +1759,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 +1774,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 +1794,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 +1802,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,12 +1822,12 @@
      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
-     hbl = MAX(hbl,epsln)
-     hbli(:,:) = hbl(:,:)
-     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 +1857,4 @@
       ENDIF
 
-      ! add non-local temperature and salinity flux
       DO jk = 1, jpkm1
          DO jj = 2, jpjm1
@@ -1648,12 +1870,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 +1944,4 @@
 
    !!======================================================================
+
 END MODULE zdfosm
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/step.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/step.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/step.F90	(revision 11145)
@@ -101,5 +101,6 @@
       IF( kstp == nit000 ) THEN                       ! initialize IOM context (must be done after nemo_init for AGRIF+XIOS+OASIS)
                              CALL iom_init(      cxios_context          )  ! for model grid (including passible AGRIF zoom)
-         IF( ln_crs      )   CALL iom_init( TRIM(cxios_context)//"_crs" )  ! for coarse grid
+                             IF( ln_crs      )   CALL iom_init( TRIM(cxios_context)//"_crs" )  ! for coarse grid
+                             CALL dia_hth_init    ! extra ML depth diagnostics, thermocline depths
       ENDIF
       IF( kstp /= nit000 )   CALL day( kstp )         ! Calendar (day was already called at nit000 in day_init)
@@ -212,9 +213,9 @@
       ! diagnostics and outputs
       !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      IF( lk_floats  )   CALL flo_stp ( kstp )        ! drifting Floats
-      IF( ln_diacfl  )   CALL dia_cfl ( kstp )        ! Courant number diagnostics
-      IF( lk_diahth  )   CALL dia_hth ( kstp )        ! Thermocline depth (20 degres isotherm depth)
-      IF( lk_diadct  )   CALL dia_dct ( kstp )        ! Transports
-                         CALL dia_ar5 ( kstp )        ! ar5 diag
+      IF( lk_floats  )   CALL flo_stp( kstp )         ! drifting Floats
+      IF( ln_diacfl  )   CALL dia_cfl( kstp )         ! Courant number diagnostics
+      IF( ll_diahth  )   CALL dia_hth( kstp )         ! Thermocline depth (20 degres isotherm depth)
+      IF( lk_diadct  )   CALL dia_dct( kstp )         ! Transports
+                         CALL dia_ar5( kstp )         ! ar5 diag
       IF( lk_diaharm )   CALL dia_harm( kstp )        ! Tidal harmonic analysis
                          CALL dia_wri ( kstp )        ! ocean model: outputs
Index: /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/stpctl.F90
===================================================================
--- /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/stpctl.F90	(revision 11144)
+++ /NEMO/branches/UKMO/NEMO_4.0_OSMOSIS/src/OCE/stpctl.F90	(revision 11145)
@@ -66,6 +66,6 @@
       INTEGER, DIMENSION(3)  ::   iu, is1, is2        ! min/max loc indices
       REAL(wp)               ::   zzz                 ! local real 
-      REAL(wp), DIMENSION(9) ::   zmax
-      LOGICAL                ::   ll_wrtstp, ll_colruns, ll_wrtruns
+      REAL(wp), DIMENSION(10) ::   zmax
+      LOGICAL                ::   ll_wrtstp, ll_colruns, ll_wrtruns, ll_isnan
       CHARACTER(len=20) :: clname
       !!----------------------------------------------------------------------
@@ -109,4 +109,6 @@
       ENDIF
       !
+      ll_isnan = ANY(ISNAN(tsn)) .OR. ANY(ISNAN(un))
+      IF (ll_isnan) nstop = nstop + 1
       !                                   !==  test of extrema  ==!
       IF( ll_wd ) THEN
@@ -124,4 +126,9 @@
          zmax(8) = MAXVAL(  ABS( wi(:,:,:) ) , mask = wmask(:,:,:) == 1._wp ) ! implicit vertical vel. max
          zmax(9) = MAXVAL(   Cu_adv(:,:,:)   , mask = tmask(:,:,:) == 1._wp ) !       cell Courant no. max
+      ENDIF
+      IF (ll_isnan) THEN
+         zmax(10) = 1._wp                                           ! stop indicator
+      ELSE
+         zmax(10) = 0._wp
       ENDIF
       !
@@ -147,5 +154,5 @@
       END IF
       !                                   !==  error handling  ==!
-      IF( ( ln_ctl .OR. lsomeoce ) .AND. (   &             ! have use mpp_max (because ln_ctl=.T.) or contains some ocean points
+      IF( ( (ln_ctl .OR. sn_cfctl%l_runstat) .OR. lsomeoce ) .AND. (   &             ! have use mpp_max (because ln_ctl=.T.) or contains some ocean points
          &  zmax(1) >   20._wp .OR.   &                    ! too large sea surface height ( > 20 m )
          &  zmax(2) >   10._wp .OR.   &                    ! too large velocity ( > 10 m/s)
@@ -153,6 +160,7 @@
          &  zmax(4) >= 100._wp .OR.   &                    ! too large sea surface salinity ( > 100 )
          &  zmax(4) <    0._wp .OR.   &                    ! too large sea surface salinity (keep this line for sea-ice)
-         &  ISNAN( zmax(1) + zmax(2) + zmax(3) ) ) ) THEN   ! NaN encounter in the tests
-         IF( lk_mpp .AND. ln_ctl ) THEN
+         &  zmax(10) >   0._wp  ) ) THEN   ! NaN encounter in the tests
+         ! &  ISNAN( zmax(1) + zmax(2) + zmax(3) ) ) ) THEN   ! NaN encounter in the tests
+         IF( lk_mpp .AND. (ln_ctl .OR. sn_cfctl%l_runstat)) THEN
             CALL mpp_maxloc( 'stpctl', ABS(sshn)        , ssmask(:,:)  , zzz, ih  )
             CALL mpp_maxloc( 'stpctl', ABS(un)          , umask (:,:,:), zzz, iu  )
