Index: /NEMO/branches/2020/dev_r14122_ENHANCE-14_smueller_OSMOSIS_streamlining/src/OCE/ZDF/zdfosm.F90
===================================================================
--- /NEMO/branches/2020/dev_r14122_ENHANCE-14_smueller_OSMOSIS_streamlining/src/OCE/ZDF/zdfosm.F90	(revision 14279)
+++ /NEMO/branches/2020/dev_r14122_ENHANCE-14_smueller_OSMOSIS_streamlining/src/OCE/ZDF/zdfosm.F90	(revision 14280)
@@ -39,10 +39,11 @@
    !!   'ln_zdfosm'                                             OSMOSIS scheme
    !!----------------------------------------------------------------------
-   !!   zdf_osm       : update momentum and tracer Kz from osm scheme
-   !!   zdf_osm_init  : initialization, namelist read, and parameters control
-   !!   osm_rst       : read (or initialize) and write osmosis restart fields
-   !!   tra_osm       : compute and add to the T & S trend the non-local flux
-   !!   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
+   !!   zdf_osm                  : update momentum and tracer Kz from osm scheme
+   !!   zdf_osm_vertical_average : compute vertical averages over boundary layers
+   !!   zdf_osm_init             : initialization, namelist read, and parameters control
+   !!   osm_rst                  : read (or initialize) and write osmosis restart fields
+   !!   tra_osm                  : compute and add to the T & S trend the non-local flux
+   !!   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.
@@ -271,5 +272,5 @@
       INTEGER, DIMENSION(jpi,jpj) :: ibld ! level of boundary layer base
       INTEGER, DIMENSION(jpi,jpj) :: imld ! level of mixed-layer depth (pycnocline top)
-      INTEGER, DIMENSION(jpi,jpj) :: jp_ext, jp_ext_mle ! offset for external level
+      INTEGER, DIMENSION(jpi,jpj) :: jp_ext ! offset for external level
       INTEGER, DIMENSION(jpi, jpj) :: j_ddh ! Type of shear layer
 
@@ -295,5 +296,4 @@
       REAL(wp), DIMENSION(jpi,jpj) :: zdt_bl,zds_bl,zdu_bl,zdv_bl,zdb_bl ! difference between blayer average and parameter at base of blayer
       REAL(wp), DIMENSION(jpi,jpj) :: zdt_ml,zds_ml,zdu_ml,zdv_ml,zdb_ml ! difference between mixed layer average and parameter at base of blayer
-      REAL(wp), DIMENSION(jpi,jpj) :: zdt_mle,zds_mle,zdu_mle,zdv_mle,zdb_mle ! difference between MLE layer average and parameter at base of blayer
 !      REAL(wp), DIMENSION(jpi,jpj) :: zwth_ent,zws_ent ! heat and salinity fluxes at the top of the pycnocline
       REAL(wp) :: zwth_ent,zws_ent ! heat and salinity fluxes at the top of the pycnocline
@@ -368,6 +368,4 @@
       zdt_ml(:,:)  = 0._wp ; zds_ml(:,:)  = 0._wp ; zdu_ml(:,:)   = 0._wp ; zdv_ml(:,:)  = 0._wp
       zdb_ml(:,:)  = 0._wp
-      zdt_mle(:,:)  = 0._wp ; zds_mle(:,:)  = 0._wp ; zdu_mle(:,:)   = 0._wp
-      zdv_mle(:,:)  = 0._wp ; zdb_mle(:,:)  = 0._wp
       zwth_ent = 0._wp ; zws_ent = 0._wp
       !
@@ -580,7 +578,11 @@
       ! Averages over well-mixed and boundary layer
       jp_ext(:,:) = 2
-      CALL zdf_osm_vertical_average(ibld, jp_ext, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl)
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                          &
+         &                           ibld, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl,           &
+         &                           jp_ext, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
 !      jp_ext(:,:) = ibld(:,:) - imld(:,:) + 1
-      CALL zdf_osm_vertical_average(ibld, ibld-imld+1, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml)
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                               &
+         &                           ibld, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml, ibld-imld+1,   &
+         &                           zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml )
 ! Velocity components in frame aligned with surface stress.
       CALL zdf_osm_velocity_rotation( zcos_wind, zsin_wind, zu_ml, zv_ml, zdu_ml, zdv_ml )
@@ -595,6 +597,6 @@
          IF ( hmle(ji,jj) >= gdepw(ji,jj,jk,Kmm) ) mld_prof(ji,jj) = MIN(mbkt(ji,jj), jk)
          END_3D
-         jp_ext_mle(:,:) = 2
-        CALL zdf_osm_vertical_average(mld_prof, jp_ext_mle, zt_mle, zs_mle, zb_mle, zu_mle, zv_mle, zdt_mle, zds_mle, zdb_mle, zdu_mle, zdv_mle)
+         CALL zdf_osm_vertical_average( Kbb, Kmm,                                            &
+            &                           mld_prof, zt_mle, zs_mle, zb_mle, zu_mle, zv_mle )
 
          DO_2D( 0, 0, 0, 0 )
@@ -635,7 +637,11 @@
       END_2D
 
-      CALL zdf_osm_vertical_average(ibld, jp_ext, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                          &
+         &                           ibld, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl,           &
+         &                           jp_ext, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
 !      jp_ext = ibld-imld+1
-      CALL zdf_osm_vertical_average(imld-1, ibld-imld+1, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml)
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                               &
+         &                           imld-1, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml,              &
+         &                           ibld-imld+1, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml )
 ! Rate of change of hbl
       CALL zdf_osm_calculate_dhdt( zdhdt, zddhdt )
@@ -664,5 +670,7 @@
 ! is external level in bounds?
 
-      CALL zdf_osm_vertical_average( ibld, jp_ext, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                          &
+         &                           ibld, zt_bl, zs_bl, zb_bl, zu_bl, zv_bl,           &
+         &                           jp_ext, zdt_bl, zds_bl, zdb_bl, zdu_bl, zdv_bl )
 !
 !
@@ -679,5 +687,7 @@
     ! Average over the depth of the mixed layer in the convective boundary layer
 !      jp_ext = ibld - imld +1
-      CALL zdf_osm_vertical_average( imld-1, ibld-imld+1, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml )
+      CALL zdf_osm_vertical_average( Kbb, Kmm,                                               &
+         &                           imld-1, zt_ml, zs_ml, zb_ml, zu_ml, zv_ml,              &
+         &                           ibld-imld+1, zdt_ml, zds_ml, zdb_ml, zdu_ml, zdv_ml )
     ! rotate mean currents and changes onto wind align co-ordinates
     !
@@ -1540,74 +1550,4 @@
    END SUBROUTINE zdf_osm_osbl_state
 
-
-   SUBROUTINE zdf_osm_vertical_average( jnlev_av, jp_ext, zt, zs, zb, 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.
-        INTEGER, DIMENSION(jpi,jpj) :: jp_ext
-
-        ! Alan: do we need zb?
-        REAL(wp), DIMENSION(jpi,jpj) :: zt, zs, zb        ! 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, ibld_ext
-        REAL(wp) :: zthick, zthermal, zbeta
-
-
-        IF( ln_timing ) CALL timing_start('zdf_osm_va')
-        zt   = 0._wp
-        zs   = 0._wp
-        zu   = 0._wp
-        zv   = 0._wp
-        DO_2D( 0, 0, 0, 0 )
-         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(ji,jj,jk,Kmm)
-            zt(ji,jj)   = zt(ji,jj)  + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_tem,Kmm)
-            zs(ji,jj)   = zs(ji,jj)  + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_sal,Kmm)
-            zu(ji,jj)   = zu(ji,jj)  + e3t(ji,jj,jk,Kmm) &
-                  &            * ( uu(ji,jj,jk,Kbb) + uu(ji - 1,jj,jk,Kbb) ) &
-                  &            / MAX( 1. , umask(ji,jj,jk) + umask(ji - 1,jj,jk) )
-            zv(ji,jj)   = zv(ji,jj)  + e3t(ji,jj,jk,Kmm) &
-                  &            * ( vv(ji,jj,jk,Kbb) + vv(ji,jj - 1,jk,Kbb) ) &
-                  &            / 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
-         zb(ji,jj) = grav * zthermal * zt(ji,jj) - grav * zbeta * zs(ji,jj)
-         ibld_ext = jnlev_av(ji,jj) + jp_ext(ji,jj)
-         IF ( ibld_ext < mbkt(ji,jj) ) THEN
-           zdt(ji,jj) = zt(ji,jj) - ts(ji,jj,ibld_ext,jp_tem,Kmm)
-           zds(ji,jj) = zs(ji,jj) - ts(ji,jj,ibld_ext,jp_sal,Kmm)
-           zdu(ji,jj) = zu(ji,jj) - ( uu(ji,jj,ibld_ext,Kbb) + uu(ji-1,jj,ibld_ext,Kbb ) ) &
-                  &    / MAX(1. , umask(ji,jj,ibld_ext ) + umask(ji-1,jj,ibld_ext ) )
-           zdv(ji,jj) = zv(ji,jj) - ( vv(ji,jj,ibld_ext,Kbb) + vv(ji,jj-1,ibld_ext,Kbb ) ) &
-                  &   / MAX(1. , vmask(ji,jj,ibld_ext ) + vmask(ji,jj-1,ibld_ext ) )
-           zdb(ji,jj) = grav * zthermal * zdt(ji,jj) - grav * zbeta * zds(ji,jj)
-         ELSE
-           zdt(ji,jj) = 0._wp
-           zds(ji,jj) = 0._wp
-           zdu(ji,jj) = 0._wp
-           zdv(ji,jj) = 0._wp
-           zdb(ji,jj) = 0._wp
-         ENDIF
-        END_2D
-     IF( ln_timing ) CALL timing_stop('zdf_osm_va')
-   END SUBROUTINE zdf_osm_vertical_average
 
    SUBROUTINE zdf_osm_velocity_rotation( zcos_w, zsin_w, zu, zv, zdu, zdv )
@@ -2478,4 +2418,109 @@
 END SUBROUTINE zdf_osm
 
+   SUBROUTINE zdf_osm_vertical_average( Kbb, Kmm,                           &
+      &                                 knlev, pt, ps, pb, pu, pv,          &
+      &                                 kp_ext, pdt, pds, pdb, pdu, pdv )
+      !!---------------------------------------------------------------------
+      !!                ***  ROUTINE zdf_vertical_average  ***
+      !!
+      !! ** Purpose : Determines vertical averages from surface to knlev,
+      !!              and optionally the differences between these vertical
+      !!              averages and values at an external level
+      !!
+      !! ** Method  : Averages are calculated from the surface to knlev.
+      !!              The external level used to calculate differences is
+      !!              knlev+kp_ext
+      !!----------------------------------------------------------------------
+      INTEGER,                      INTENT(in   )           ::   Kbb, Kmm   ! Ocean time-level indices
+      INTEGER,  DIMENSION(jpi,jpj), INTENT(in   )           ::   knlev      ! Number of levels to average over.
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out)           ::   pt, ps     ! Average temperature and salinity
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out)           ::   pb         ! Average buoyancy
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out)           ::   pu, pv     ! Average current components
+      INTEGER,  DIMENSION(jpi,jpj), INTENT(in   ), OPTIONAL ::   kp_ext     ! External-level offsets
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out), OPTIONAL ::   pdt, pds   ! Difference between average temperature, salinity,
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out), OPTIONAL ::   pdb        ! buoyancy,
+      REAL(wp), DIMENSION(jpi,jpj), INTENT(  out), OPTIONAL ::   pdu, pdv   ! velocity components and the OSBL
+      !
+      INTEGER                      ::   jk, jkflt, jkmax, ji, jj   ! Loop indices
+      INTEGER                      ::   ibld_ext                   ! External-layer index
+      REAL(wp), DIMENSION(jpi,jpj) ::   zthick                     ! Layer thickness
+      REAL(wp)                     ::   zthermal, zbeta            ! Thermal/haline expansion/contraction coefficients
+      !!----------------------------------------------------------------------
+      !
+      IF( ln_timing ) CALL timing_start('zdf_osm_va')
+      !
+      ! Averages over depth of boundary layer
+      pt(:,:)     = 0.0_wp
+      ps(:,:)     = 0.0_wp
+      pu(:,:)     = 0.0_wp
+      pv(:,:)     = 0.0_wp
+      zthick(:,:) = epsln
+      jkflt = jpk
+      jkmax = 0
+      DO_2D( 0, 0, 0, 0 )
+         IF ( knlev(ji,jj) < jkflt ) jkflt = knlev(ji,jj)
+         IF ( knlev(ji,jj) > jkmax ) jkmax = knlev(ji,jj)
+      END_2D
+      DO_3D( 0, 0, 0, 0, 2, jkflt )   ! Upper, flat part of layer
+         zthick(ji,jj) = zthick(ji,jj) + e3t(ji,jj,jk,Kmm)
+         pt(ji,jj)     = pt(ji,jj)     + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_tem,Kmm)
+         ps(ji,jj)     = ps(ji,jj)     + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_sal,Kmm)
+         pu(ji,jj)     = pu(ji,jj)     + e3t(ji,jj,jk,Kmm) *                                        &
+            &                               ( uu(ji,jj,jk,Kbb) + uu(ji - 1,jj,jk,Kbb) ) /           &
+            &                               MAX( 1.0_wp , umask(ji,jj,jk) + umask(ji - 1,jj,jk) )
+         pv(ji,jj)     = pv(ji,jj)     + e3t(ji,jj,jk,Kmm) *                                        &
+            &                               ( vv(ji,jj,jk,Kbb) + vv(ji,jj - 1,jk,Kbb) ) /           &
+            &                               MAX( 1.0_wp , vmask(ji,jj,jk) + vmask(ji,jj - 1,jk) )         
+      END_3D
+      DO_3D( 0, 0, 0, 0, jkflt+1, jkmax )   ! Lower, non-flat part of layer
+         IF ( knlev(ji,jj) >= jk ) THEN
+            zthick(ji,jj) = zthick(ji,jj) + e3t(ji,jj,jk,Kmm)
+            pt(ji,jj)     = pt(ji,jj)     + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_tem,Kmm)
+            ps(ji,jj)     = ps(ji,jj)     + e3t(ji,jj,jk,Kmm) * ts(ji,jj,jk,jp_sal,Kmm)
+            pu(ji,jj)     = pu(ji,jj)     + e3t(ji,jj,jk,Kmm) *                                        &
+               &                               ( uu(ji,jj,jk,Kbb) + uu(ji - 1,jj,jk,Kbb) ) /           &
+               &                               MAX( 1.0_wp , umask(ji,jj,jk) + umask(ji - 1,jj,jk) )
+            pv(ji,jj)     = pv(ji,jj)     + e3t(ji,jj,jk,Kmm) *                                        &
+               &                               ( vv(ji,jj,jk,Kbb) + vv(ji,jj - 1,jk,Kbb) ) /           &
+               &                               MAX( 1.0_wp , vmask(ji,jj,jk) + vmask(ji,jj - 1,jk) )
+         END IF
+      END_3D
+      DO_2D( 0, 0, 0, 0 )
+         pt(ji,jj) = pt(ji,jj) / zthick(ji,jj)
+         ps(ji,jj) = ps(ji,jj) / zthick(ji,jj)
+         pu(ji,jj) = pu(ji,jj) / zthick(ji,jj)
+         pv(ji,jj) = pv(ji,jj) / zthick(ji,jj)
+         zthermal  = rab_n(ji,jj,1,jp_tem)   ! ideally use ibld not 1??
+         zbeta     = rab_n(ji,jj,1,jp_sal)
+         pb(ji,jj) = grav * zthermal * pt(ji,jj) - grav * zbeta * ps(ji,jj)
+      END_2D
+      !
+      ! Differences between vertical averages and values at an external layer
+      IF ( PRESENT( kp_ext ) ) THEN
+         DO_2D( 0, 0, 0, 0 )
+            ibld_ext = knlev(ji,jj) + kp_ext(ji,jj)
+            IF ( ibld_ext < mbkt(ji,jj) ) THEN
+               pdt(ji,jj) = pt(ji,jj) - ts(ji,jj,ibld_ext,jp_tem,Kmm)
+               pds(ji,jj) = ps(ji,jj) - ts(ji,jj,ibld_ext,jp_sal,Kmm)
+               pdu(ji,jj) = pu(ji,jj) - ( uu(ji,jj,ibld_ext,Kbb) + uu(ji-1,jj,ibld_ext,Kbb ) ) /              &
+                  &                        MAX(1.0_wp , umask(ji,jj,ibld_ext ) + umask(ji-1,jj,ibld_ext ) )
+               pdv(ji,jj) = pv(ji,jj) - ( vv(ji,jj,ibld_ext,Kbb) + vv(ji,jj-1,ibld_ext,Kbb ) ) /              &
+                  &                        MAX(1.0_wp , vmask(ji,jj,ibld_ext ) + vmask(ji,jj-1,ibld_ext ) )
+               zthermal   = rab_n(ji,jj,1,jp_tem)   ! ideally use ibld not 1??
+               zbeta      = rab_n(ji,jj,1,jp_sal)
+               pdb(ji,jj) = grav * zthermal * pdt(ji,jj) - grav * zbeta * pds(ji,jj)
+            ELSE
+               pdt(ji,jj) = 0.0_wp
+               pds(ji,jj) = 0.0_wp
+               pdu(ji,jj) = 0.0_wp
+               pdv(ji,jj) = 0.0_wp
+               pdb(ji,jj) = 0.0_wp
+            ENDIF
+         END_2D
+      END IF
+      !
+      IF( ln_timing ) CALL timing_stop('zdf_osm_va')
+      !
+   END SUBROUTINE zdf_osm_vertical_average
 
    SUBROUTINE zdf_osm_init( Kmm )
