Skip to content
Snippets Groups Projects
zdfosm.F90 212 KiB
Newer Older
Guillaume Samson's avatar
Guillaume Samson committed
MODULE zdfosm
   !!======================================================================
   !!                       ***  MODULE  zdfosm  ***
   !! Ocean physics:  vertical mixing coefficient compute from the OSMOSIS
   !!                 turbulent closure parameterization
   !!=====================================================================
   !!  History : NEMO 4.0  ! A. Grant, G. Nurser
   !! 15/03/2017  Changed calculation of pycnocline thickness in unstable conditions and stable conditions AG
   !! 15/03/2017  Calculation of pycnocline gradients for stable conditions changed. Pycnocline gradients now depend on stability of the OSBL. A.G
   !! 06/06/2017  (1) Checks on sign of buoyancy jump in calculation of  OSBL depth.  A.G.
   !!             (2) Removed variable zbrad0, zbradh and zbradav since they are not used.
   !!             (3) Approximate treatment for shear turbulence.
   !!                        Minimum values for zustar and zustke.
   !!                        Add velocity scale, zvstr, that tends to zustar for large Langmuir numbers.
   !!                        Limit maximum value for Langmuir number.
   !!                        Use zvstr in definition of stability parameter zhol.
   !!             (4) Modified parametrization of entrainment flux, changing original coefficient 0.0485 for Langmuir contribution to 0.135 * zla
   !!             (5) For stable boundary layer add factor that depends on length of timestep to 'slow' collapse and growth. Make sure buoyancy jump not negative.
   !!             (6) For unstable conditions when growth is over multiple levels, limit change to maximum of one level per cycle through loop.
   !!             (7) Change lower limits for loops that calculate OSBL averages from 1 to 2. Large gradients between levels 1 and 2 can cause problems.
   !!             (8) Change upper limits from ibld-1 to ibld.
   !!             (9) Calculation of pycnocline thickness in unstable conditions. Check added to ensure that buoyancy jump is positive before calculating Ri.
   !!            (10) Thickness of interface layer at base of the stable OSBL set by Richardson number. Gives continuity in transition from unstable OSBL.
   !!            (11) Checks that buoyancy jump is poitive when calculating pycnocline profiles.
   !!            (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) 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 fortran macro. Inline code moved into subroutines. Changes to physics made,
Guillaume Samson's avatar
Guillaume Samson committed
   !!                  (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) Remove old code excluded by fortran macro along with the fortran macro that is not needed
Guillaume Samson's avatar
Guillaume Samson committed
36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922
   !!             4.2  !  2021-05  (S. Mueller)  Efficiency improvements, source-code clarity enhancements, and adaptation to tiling
   !!----------------------------------------------------------------------

   !!----------------------------------------------------------------------
   !!   'ln_zdfosm'                                          OSMOSIS scheme
   !!----------------------------------------------------------------------
   !!   zdf_osm        : update momentum and tracer Kz from osm scheme
   !!      zdf_osm_vertical_average             : compute vertical averages over boundary layers
   !!      zdf_osm_velocity_rotation            : rotate velocity components
   !!         zdf_osm_velocity_rotation_2d      :    rotation of 2d fields
   !!         zdf_osm_velocity_rotation_3d      :    rotation of 3d fields
   !!      zdf_osm_osbl_state                   : determine the state of the OSBL
   !!      zdf_osm_external_gradients           : calculate gradients below the OSBL
   !!      zdf_osm_calculate_dhdt               : calculate rate of change of hbl
   !!      zdf_osm_timestep_hbl                 : hbl timestep
   !!      zdf_osm_pycnocline_thickness         : calculate thickness of pycnocline
   !!      zdf_osm_diffusivity_viscosity        : compute eddy diffusivity and viscosity profiles
   !!      zdf_osm_fgr_terms                    : compute flux-gradient relationship terms
   !!         zdf_osm_pycnocline_buoyancy_profiles : calculate pycnocline buoyancy profiles
   !!      zdf_osm_zmld_horizontal_gradients    : calculate horizontal buoyancy gradients for use with Fox-Kemper parametrization
   !!      zdf_osm_osbl_state_fk                : determine state of OSBL and MLE layers
   !!      zdf_osm_mle_parameters               : timestep MLE depth and calculate MLE fluxes
   !!   zdf_osm_init   : initialization, namelist read, and parameters control
   !!      zdf_osm_alloc                        : memory allocation
   !!   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_iomput : iom_put wrapper that accepts arrays without halo
   !!      zdf_osm_iomput_2d                    : iom_put wrapper for 2D fields
   !!      zdf_osm_iomput_3d                    : iom_put wrapper for 3D fields
   !!----------------------------------------------------------------------
   USE oce                       ! Ocean dynamics and active tracers
   !                             ! Uses ww from previous time step (which is now wb) to calculate hbl
   USE dom_oce                   ! Ocean space and time domain
   USE zdf_oce                   ! Ocean vertical physics
   USE sbc_oce                   ! Surface boundary condition: ocean
   USE sbcwave                   ! Surface wave parameters
   USE phycst                    ! Physical constants
   USE eosbn2                    ! Equation of state
   USE traqsr                    ! Details of solar radiation absorption
   USE zdfdrg, ONLY : rCdU_bot   ! Bottom friction velocity
   USE zdfddm                    ! Double diffusion mixing (avs array)
   USE iom                       ! I/O library
   USE lib_mpp                   ! MPP library
   USE trd_oce                   ! Ocean trends definition
   USE trdtra                    ! Tracers trends
   USE in_out_manager            ! I/O manager
   USE lbclnk                    ! Ocean lateral boundary conditions (or mpp link)
   USE prtctl                    ! Print control
   USE lib_fortran               ! Fortran utilities (allows no signed zero when 'key_nosignedzero' defined)

   IMPLICIT NONE
   PRIVATE

   ! Public subroutines
   PUBLIC zdf_osm        ! Routine called by step.F90
   PUBLIC zdf_osm_init   ! Routine called by nemogcm.F90
   PUBLIC osm_rst        ! Routine called by step.F90
   PUBLIC tra_osm        ! Routine called by step.F90
   PUBLIC trc_osm        ! Routine called by trcstp.F90
   PUBLIC dyn_osm        ! Routine called by step.F90

   ! Public variables
   LOGICAL,  PUBLIC                                      ::   ln_osm_mle   !: Flag to activate the Mixed Layer Eddy (MLE)
   !                                                                       !     parameterisation, needed by tra_mle_init in
   !                                                                       !     tramle.F90
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   ghamu        !: Non-local u-momentum flux
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   ghamv        !: Non-local v-momentum flux
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   ghamt        !: Non-local temperature flux (gamma/<ws>o)
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   ghams        !: Non-local salinity flux (gamma/<ws>o)
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hbl          !: Boundary layer depth
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hml          !: ML depth
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   hmle         !: Depth of layer affexted by mixed layer eddies in Fox-Kemper parametrization
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dbdx_mle     !: Zonal buoyancy gradient in ML
   REAL(wp), PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dbdy_mle     !: Meridional buoyancy gradient in ML
   INTEGER,  PUBLIC, ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   mld_prof     !: Level of base of MLE layer

   INTERFACE zdf_osm_velocity_rotation
      !!---------------------------------------------------------------------
      !!              ***  INTERFACE zdf_velocity_rotation  ***
      !!---------------------------------------------------------------------
      MODULE PROCEDURE zdf_osm_velocity_rotation_2d
      MODULE PROCEDURE zdf_osm_velocity_rotation_3d
   END INTERFACE
   !
   INTERFACE zdf_osm_iomput
      !!---------------------------------------------------------------------
      !!                 ***  INTERFACE zdf_osm_iomput  ***
      !!---------------------------------------------------------------------
      MODULE PROCEDURE zdf_osm_iomput_2d
      MODULE PROCEDURE zdf_osm_iomput_3d
   END INTERFACE

   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   etmean      ! Averaging operator for avt
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dh          ! Depth of pycnocline
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   r1_ft       ! Inverse of the modified Coriolis parameter at t-pts
   ! Layer indices
   INTEGER,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   nbld        ! Level of boundary layer base
   INTEGER,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   nmld        ! Level of mixed-layer depth (pycnocline top)
   ! Layer type
   INTEGER,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   n_ddh       ! Type of shear layer
   !                                                              !    n_ddh=0: active shear layer
   !                                                              !    n_ddh=1: shear layer not active
   !                                                              !    n_ddh=2: shear production low
   ! Layer flags
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_conv      ! Unstable/stable bl
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_shear     ! Shear layers
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_coup      ! Coupling to bottom
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_pyc       ! OSBL pycnocline present
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_flux      ! Surface flux extends below OSBL into MLE layer
   LOGICAL,  ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   l_mle       ! MLE layer increases in hickness.
   ! Scales
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swth0       ! Surface heat flux (Kinematic)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   sws0        ! Surface freshwater flux
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swb0        ! Surface buoyancy flux
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   suw0        ! Surface u-momentum flux
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   sustar      ! Friction velocity
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   scos_wind   ! Cos angle of surface stress
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   ssin_wind   ! Sin angle of surface stress
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swthav      ! Heat flux - bl average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swsav       ! Freshwater flux - bl average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swbav       ! Buoyancy flux - bl average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   sustke      ! Surface Stokes drift
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   dstokes     ! Penetration depth of the Stokes drift
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swstrl      ! Langmuir velocity scale
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   swstrc      ! Convective velocity scale
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   sla         ! Trubulent Langmuir number
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   svstr       ! Velocity scale that tends to sustar for large Langmuir number
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   shol        ! Stability parameter for boundary layer
   ! Layer averages: BL
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_t_bl     ! Temperature average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_s_bl     ! Salinity average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_u_bl     ! Velocity average (u)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_v_bl     ! Velocity average (v)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_b_bl     ! Buoyancy average
   ! Difference between layer average and parameter at the base of the layer: BL
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_dt_bl    ! Temperature difference
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_ds_bl    ! Salinity difference
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_du_bl    ! Velocity difference (u)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_dv_bl    ! Velocity difference (v)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_db_bl    ! Buoyancy difference
   ! Layer averages: ML
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_t_ml     ! Temperature average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_s_ml     ! Salinity average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_u_ml     ! Velocity average (u)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_v_ml     ! Velocity average (v)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_b_ml     ! Buoyancy average
   ! Difference between layer average and parameter at the base of the layer: ML
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_dt_ml    ! Temperature difference
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_ds_ml    ! Salinity difference
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_du_ml    ! Velocity difference (u)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_dv_ml    ! Velocity difference (v)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_db_ml    ! Buoyancy difference
   ! Layer averages: MLE
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_t_mle    ! Temperature average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_s_mle    ! Salinity average
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_u_mle    ! Velocity average (u)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_v_mle    ! Velocity average (v)
   REAL(wp), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   av_b_mle    ! Buoyancy average
   ! Diagnostic output
   REAL(WP), ALLOCATABLE, SAVE, DIMENSION(:,:)   ::   osmdia2d    ! Auxiliary array for diagnostic output
   REAL(WP), ALLOCATABLE, SAVE, DIMENSION(:,:,:) ::   osmdia3d    ! Auxiliary array for diagnostic output
   LOGICAL  ::   ln_dia_pyc_scl = .FALSE.                         ! Output of pycnocline scalar-gradient profiles
   LOGICAL  ::   ln_dia_pyc_shr = .FALSE.                         ! Output of pycnocline velocity-shear  profiles

   !                                               !!* namelist namzdf_osm *
   LOGICAL  ::   ln_use_osm_la                      ! Use namelist rn_osm_la
   REAL(wp) ::   rn_osm_la                          ! Turbulent Langmuir number
   REAL(wp) ::   rn_osm_dstokes                     ! Depth scale of Stokes drift
   REAL(wp) ::   rn_zdfosm_adjust_sd   = 1.0_wp     ! Factor to reduce Stokes drift by
   REAL(wp) ::   rn_osm_hblfrac        = 0.1_wp     ! For nn_osm_wave = 3/4 specify fraction in top of hbl
   LOGICAL  ::   ln_zdfosm_ice_shelter              ! Flag to activate ice sheltering
   REAL(wp) ::   rn_osm_hbl0           = 10.0_wp    ! Initial value of hbl for 1D runs
   INTEGER  ::   nn_ave                             ! = 0/1 flag for horizontal average on avt
   INTEGER  ::   nn_osm_wave = 0                    ! = 0/1/2 flag for getting stokes drift from La# / PM wind-waves/Inputs into
   !                                                !    sbcwave
   INTEGER  ::   nn_osm_SD_reduce                   ! = 0/1/2 flag for getting effective stokes drift from surface value
   LOGICAL  ::   ln_dia_osm                         ! Use namelist  rn_osm_la
   LOGICAL  ::   ln_kpprimix           = .TRUE.     ! Shear instability mixing
   REAL(wp) ::   rn_riinfty            = 0.7_wp     ! Local Richardson Number limit for shear instability
   REAL(wp) ::   rn_difri              = 0.005_wp   ! Maximum shear mixing at Rig = 0    (m2/s)
   LOGICAL  ::   ln_convmix            = .TRUE.     ! Convective instability mixing
   REAL(wp) ::   rn_difconv            = 1.0_wp     ! Diffusivity when unstable below BL  (m2/s)
   ! OSMOSIS mixed layer eddy parametrization constants
   INTEGER  ::   nn_osm_mle                         ! = 0/1 flag for horizontal average on avt
   REAL(wp) ::   rn_osm_mle_ce                      ! MLE coefficient
   !    Parameters used in nn_osm_mle = 0 case
   REAL(wp) ::   rn_osm_mle_lf                      ! Typical scale of mixed layer front
   REAL(wp) ::   rn_osm_mle_time                    ! Time scale for mixing momentum across the mixed layer
   !    Parameters used in nn_osm_mle = 1 case
   REAL(wp) ::   rn_osm_mle_lat                     ! Reference latitude for a 5 km scale of ML front
   LOGICAL  ::   ln_osm_hmle_limit                  ! If true arbitrarily restrict hmle to rn_osm_hmle_limit*zmld
   REAL(wp) ::   rn_osm_hmle_limit                  ! If ln_osm_hmle_limit true arbitrarily restrict hmle to rn_osm_hmle_limit*zmld
   REAL(wp) ::   rn_osm_mle_rho_c                   ! Density criterion for definition of MLD used by FK
   REAL(wp) ::   rb_c                               ! ML buoyancy criteria = g rho_c /rho0 where rho_c is defined in zdfmld
   REAL(wp) ::   rc_f                               ! MLE coefficient (= rn_ce / (5 km * fo) ) in nn_osm_mle=1 case
   REAL(wp) ::   rn_osm_mle_thresh                  ! Threshold buoyancy for deepening of MLE layer below OSBL base
   REAL(wp) ::   rn_osm_bl_thresh                   ! Threshold buoyancy for deepening of OSBL base
   REAL(wp) ::   rn_osm_mle_tau                     ! Adjustment timescale for MLE

   ! General constants
   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.0_wp/3.0_wp   ! 1/3
   REAL(wp) ::   p2third   = 2.0_wp/3.0_wp   ! 2/3

   !! * Substitutions
#  include "do_loop_substitute.h90"
#  include "domzgr_substitute.h90"
   !!----------------------------------------------------------------------
   !! NEMO/OCE 4.0 , NEMO Consortium (2018)
   !! $Id: zdfosm.F90 14921 2021-05-28 12:19:26Z smueller $
   !! Software governed by the CeCILL license (see ./LICENSE)
   !!----------------------------------------------------------------------
CONTAINS

   INTEGER FUNCTION zdf_osm_alloc()
      !!----------------------------------------------------------------------
      !!                 ***  FUNCTION zdf_osm_alloc  ***
      !!----------------------------------------------------------------------
      INTEGER ::   ierr
      !!----------------------------------------------------------------------
      !
      zdf_osm_alloc = 0
      !
      ALLOCATE( ghamu(jpi,jpj,jpk), ghamv(jpi,jpj,jpk), ghamt(jpi,jpj,jpk), ghams(jpi,jpj,jpk), hbl(jpi,jpj), hml(jpi,jpj),   &
         &      hmle(jpi,jpj),      dbdx_mle(jpi,jpj),  dbdy_mle(jpi,jpj),  mld_prof(jpi,jpj),  STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( etmean(A2D(nn_hls-1),jpk), dh(jpi,jpj), r1_ft(A2D(nn_hls-1)), STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( nbld(jpi,jpj), nmld(A2D(nn_hls-1)), STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( n_ddh(A2D(nn_hls-1)), STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( l_conv(A2D(nn_hls-1)), l_shear(A2D(nn_hls-1)), l_coup(A2D(nn_hls-1)), l_pyc(A2D(nn_hls-1)),   &
         &      l_flux(A2D(nn_hls-1)), l_mle(A2D(nn_hls-1)),   STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( swth0(A2D(nn_hls-1)),  sws0(A2D(nn_hls-1)),      swb0(A2D(nn_hls-1)),      suw0(A2D(nn_hls-1)),      &
         &      sustar(A2D(nn_hls-1)), scos_wind(A2D(nn_hls-1)), ssin_wind(A2D(nn_hls-1)), swthav(A2D(nn_hls-1)),    &
         &      swsav(A2D(nn_hls-1)),  swbav(A2D(nn_hls-1)),     sustke(A2D(nn_hls-1)),    dstokes(A2D(nn_hls-1)),   &
         &      swstrl(A2D(nn_hls-1)), swstrc(A2D(nn_hls-1)),    sla(A2D(nn_hls-1)),       svstr(A2D(nn_hls-1)),     &
         &      shol(A2D(nn_hls-1)),   STAT=ierr )
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( av_t_bl(jpi,jpj), av_s_bl(jpi,jpj), av_u_bl(jpi,jpj), av_v_bl(jpi,jpj),   &
         &      av_b_bl(jpi,jpj), STAT=ierr)
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( av_dt_bl(jpi,jpj), av_ds_bl(jpi,jpj), av_du_bl(jpi,jpj), av_dv_bl(jpi,jpj),   &
         &      av_db_bl(jpi,jpj), STAT=ierr)
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( av_t_ml(jpi,jpj), av_s_ml(jpi,jpj), av_u_ml(jpi,jpj), av_v_ml(jpi,jpj),   &
         &      av_b_ml(jpi,jpj), STAT=ierr)
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( av_dt_ml(jpi,jpj), av_ds_ml(jpi,jpj), av_du_ml(jpi,jpj), av_dv_ml(jpi,jpj),   &
         &      av_db_ml(jpi,jpj), STAT=ierr)
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      ALLOCATE( av_t_mle(jpi,jpj), av_s_mle(jpi,jpj), av_u_mle(jpi,jpj), av_v_mle(jpi,jpj),   &
         &      av_b_mle(jpi,jpj), STAT=ierr)
      zdf_osm_alloc = zdf_osm_alloc + ierr
      !
      IF ( ln_dia_osm ) THEN
         ALLOCATE( osmdia2d(jpi,jpj), osmdia3d(jpi,jpj,jpk), STAT=ierr )
         zdf_osm_alloc = zdf_osm_alloc + ierr
      END IF
      !
      CALL mpp_sum ( 'zdfosm', zdf_osm_alloc )
      IF( zdf_osm_alloc /= 0 ) CALL ctl_warn( 'zdf_osm_alloc: failed to allocate zdf_osm arrays' )
      !
   END FUNCTION zdf_osm_alloc

   SUBROUTINE zdf_osm( kt, Kbb, Kmm, Krhs, p_avm,   &
      &                p_avt )
      !!----------------------------------------------------------------------
      !!                   ***  ROUTINE zdf_osm  ***
      !!
      !! ** Purpose :   Compute the vertical eddy viscosity and diffusivity
      !!      coefficients and non local mixing using the OSMOSIS scheme
      !!
      !! ** Method :   The boundary layer depth hosm is diagnosed at tracer points
      !!      from profiles of buoyancy, and shear, and the surface forcing.
      !!      Above hbl (sigma=-z/hbl <1) the mixing coefficients are computed from
      !!
      !!                      Kx =  hosm  Wx(sigma) G(sigma)
      !!
      !!             and the non local term ghamt = Cs / Ws(sigma) / hosm
      !!      Below hosm  the coefficients are the sum of mixing due to internal waves
      !!      shear instability and double diffusion.
      !!
      !!      -1- Compute the now interior vertical mixing coefficients at all depths.
      !!      -2- Diagnose the boundary layer depth.
      !!      -3- Compute the now boundary layer vertical mixing coefficients.
      !!      -4- Compute the now vertical eddy vicosity and diffusivity.
      !!      -5- Smoothing
      !!
      !!        N.B. The computation is done from jk=2 to jpkm1
      !!             Surface value of avt are set once a time to zero
      !!             in routine zdf_osm_init.
      !!
      !! ** Action  :   update the non-local terms ghamts
      !!                update avt (before vertical eddy coef.)
      !!
      !! References : Large W.G., Mc Williams J.C. and Doney S.C.
      !!         Reviews of Geophysics, 32, 4, November 1994
      !!         Comments in the code refer to this paper, particularly
      !!         the equation number. (LMD94, here after)
      !!----------------------------------------------------------------------
      INTEGER                   , INTENT(in   ) ::  kt               ! Ocean time step
      INTEGER                   , INTENT(in   ) ::  Kbb, Kmm, Krhs   ! Ocean time level indices
      REAL(wp), DIMENSION(:,:,:), INTENT(inout) ::  p_avm, p_avt     ! Momentum and tracer Kz (w-points)
      !!
      INTEGER ::   ji, jj, jk, jl, jm, jkflt   ! Dummy loop indices
      !!
      REAL(wp) ::   zthermal, zbeta
      REAL(wp) ::   zesh2, zri, zfri   ! Interior Richardson mixing
      !! Scales
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zrad0       ! Surface solar temperature flux (deg m/s)
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zradh       ! Radiative flux at bl base (Buoyancy units)
      REAL(wp)                           ::   zradav      ! Radiative flux, bl average (Buoyancy Units)
      REAL(wp)                           ::   zvw0        ! Surface v-momentum flux
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zwb0tot     ! Total surface buoyancy flux including insolation
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zwb_ent     ! Buoyancy entrainment flux
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zwb_min
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zwb_fk_b    ! MLE buoyancy flux averaged over OSBL
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zwb_fk      ! Max MLE buoyancy flux
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdiff_mle   ! Extra MLE vertical diff
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zvel_mle    ! Velocity scale for dhdt with stable ML and FK
      !! Mixed-layer variables
      INTEGER,  DIMENSION(A2D(nn_hls-1)) ::   jk_nlev  ! Number of levels
      INTEGER,  DIMENSION(A2D(nn_hls-1)) ::   jk_ext   ! Offset for external level
      !!
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zhbl   ! BL depth - grid
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zhml   ! ML depth - grid
      !!
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zhmle   ! MLE depth - grid
      REAL(wp), DIMENSION(A2D(nn_hls))   ::   zmld    ! ML depth on grid
      !!
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdh                          ! Pycnocline depth - grid
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdhdt                        ! BL depth tendency
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdtdz_bl_ext, zdsdz_bl_ext   ! External temperature/salinity gradients
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdbdz_bl_ext                 ! External buoyancy gradients
      REAL(wp), DIMENSION(A2D(nn_hls))   ::   zdtdx, zdtdy, zdsdx, zdsdy   ! Horizontal gradients for Fox-Kemper parametrization
      !!
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zdbds_mle   ! Magnitude of horizontal buoyancy gradient
      !! Flux-gradient relationship variables
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zshear   ! Shear production
      !!
      REAL(wp), DIMENSION(A2D(nn_hls-1)) ::   zhbl_t   ! Holds boundary layer depth updated by full timestep
      !! For calculating Ri#-dependent mixing
      REAL(wp), DIMENSION(A2D(nn_hls)) ::   z2du     ! u-shear^2
      REAL(wp), DIMENSION(A2D(nn_hls)) ::   z2dv     ! v-shear^2
      REAL(wp)                         ::   zrimix   ! Spatial form of ri#-induced diffusion
      !! Temporary variables
      REAL(wp)                                 ::   znd              ! Temporary non-dimensional depth
      REAL(wp)                                 ::   zz0, zz1, zfac
      REAL(wp)                                 ::   zus_x, zus_y     ! Temporary Stokes drift
      REAL(wp), DIMENSION(A2D(nn_hls-1),jpk)   ::   zviscos          ! Viscosity
      REAL(wp), DIMENSION(A2D(nn_hls-1),jpk)   ::   zdiffut          ! t-diffusivity
      REAL(wp)                                 ::   zabsstke
      REAL(wp)                                 ::   zsqrtpi, z_two_thirds, zthickness
      REAL(wp)                                 ::   z2k_times_thickness, zsqrt_depth, zexp_depth, zf, zexperfc
      !! For debugging
      REAL(wp), PARAMETER ::   pp_large = -1e10_wp
      !!----------------------------------------------------------------------
      !
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         nmld(ji,jj)   = 0
         sustke(ji,jj) = pp_large
         l_pyc(ji,jj)  = .FALSE.
         l_flux(ji,jj) = .FALSE.
         l_mle(ji,jj)  = .FALSE.
      END_2D
      ! Mixed layer
      ! No initialization of zhbl or zhml (or zdh?)
      zhbl(:,:) = pp_large
      zhml(:,:) = pp_large
      zdh(:,:)  = pp_large
      !
      IF ( ln_osm_mle ) THEN   ! Only initialise arrays if needed
         zdtdx(:,:)  = pp_large ; zdtdy(:,:)    = pp_large ; zdsdx(:,:)     = pp_large
         zdsdy(:,:)  = pp_large
         zwb_fk(:,:) = pp_large ; zvel_mle(:,:) = pp_large
         zhmle(:,:)  = pp_large ; zmld(:,:)     = pp_large
         DO_2D_OVR( nn_hls, nn_hls, nn_hls, nn_hls )
            dbdx_mle(ji,jj) = pp_large
            dbdy_mle(ji,jj) = pp_large
         END_2D
      ENDIF
      zhbl_t(:,:)   = pp_large
      !
      zdiffut(:,:,:) = 0.0_wp
      zviscos(:,:,:) = 0.0_wp
      !
      DO_3D_OVR( nn_hls, nn_hls, nn_hls, nn_hls, 1, jpk )
         ghamt(ji,jj,jk) = pp_large
         ghams(ji,jj,jk) = pp_large
         ghamu(ji,jj,jk) = pp_large
         ghamv(ji,jj,jk) = pp_large
      END_3D
      DO_3D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1, 1, jpk )
         ghamt(ji,jj,jk) = 0.0_wp
         ghams(ji,jj,jk) = 0.0_wp
         ghamu(ji,jj,jk) = 0.0_wp
         ghamv(ji,jj,jk) = 0.0_wp
      END_3D
      !
      zdiff_mle(:,:) = 0.0_wp
      !
      ! Ensure only positive hbl values are accessed when using extended halo
      ! (nn_hls==2)
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         hbl(ji,jj) = MAX( hbl(ji,jj), epsln )
      END_2D
      !
      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
      ! Calculate boundary layer scales
      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
      !
      ! Turbulent surface fluxes and fluxes averaged over depth of the OSBL
      zz0 =           rn_abs   ! Assume two-band radiation model for depth of OSBL - surface equi-partition in 2-bands
      zz1 =  1.0_wp - rn_abs
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         zrad0(ji,jj)  = qsr(ji,jj) * r1_rho0_rcp   ! Surface downward irradiance (so always +ve)
         zradh(ji,jj)  = zrad0(ji,jj) *                                &   ! Downwards irradiance at base of boundary layer
            &            ( zz0 * EXP( -1.0_wp * hbl(ji,jj) / rn_si0 ) + zz1 * EXP( -1.0_wp * hbl(ji,jj) / rn_si1 ) )
         zradav        = zrad0(ji,jj) *                                              &            ! Downwards irradiance averaged
            &            ( zz0 * ( 1.0_wp - EXP( -hbl(ji,jj)/rn_si0 ) ) * rn_si0 +   &            !    over depth of the OSBL
            &              zz1 * ( 1.0_wp - EXP( -hbl(ji,jj)/rn_si1 ) ) * rn_si1 ) / hbl(ji,jj)
         swth0(ji,jj)  = - qns(ji,jj) * r1_rho0_rcp * tmask(ji,jj,1)   ! Upwards surface Temperature flux for non-local term
         swthav(ji,jj) = 0.5_wp * swth0(ji,jj) - ( 0.5_wp * ( zrad0(ji,jj) + zradh(ji,jj) ) -   &   ! Turbulent heat flux averaged
            &                                                 zradav )                              !    over depth of OSBL
      END_2D
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         sws0(ji,jj)    = -1.0_wp * ( ( emp(ji,jj) - rnf(ji,jj) ) * ts(ji,jj,1,jp_sal,Kmm) +   &   ! Upwards surface salinity flux
            &                         sfx(ji,jj) ) * r1_rho0 * tmask(ji,jj,1)                      !    for non-local term
         zthermal       = rab_n(ji,jj,1,jp_tem)
         zbeta          = rab_n(ji,jj,1,jp_sal)
         swb0(ji,jj)    = grav * zthermal * swth0(ji,jj) - grav * zbeta * sws0(ji,jj)   ! Non radiative upwards surface buoyancy flux
         zwb0tot(ji,jj) = swb0(ji,jj) - grav * zthermal * ( zrad0(ji,jj) - zradh(ji,jj) )   ! Total upwards surface buoyancy flux
         swsav(ji,jj)   = 0.5_wp * sws0(ji,jj)                              ! Turbulent salinity flux averaged over depth of the OBSL
         swbav(ji,jj)   = grav  * zthermal * swthav(ji,jj) -            &   ! Turbulent buoyancy flux averaged over the depth of the
            &             grav  * zbeta * swsav(ji,jj)                      ! OBSBL
      END_2D
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         suw0(ji,jj)    = -0.5_wp * (utau(ji-1,jj) + utau(ji,jj)) * r1_rho0 * tmask(ji,jj,1)   ! Surface upward velocity fluxes
         zvw0           = -0.5_wp * (vtau(ji,jj-1) + vtau(ji,jj)) * r1_rho0 * tmask(ji,jj,1)
         sustar(ji,jj)  = MAX( SQRT( SQRT( suw0(ji,jj) * suw0(ji,jj) + zvw0 * zvw0 ) ),   &   ! Friction velocity (sustar), at
            &                  1e-8_wp )                                                      !    T-point : LMD94 eq. 2
         scos_wind(ji,jj) = -1.0_wp * suw0(ji,jj) / ( sustar(ji,jj) * sustar(ji,jj) )
         ssin_wind(ji,jj) = -1.0_wp * zvw0        / ( sustar(ji,jj) * sustar(ji,jj) )
      END_2D
      ! Calculate Stokes drift in direction of wind (sustke) and Stokes penetration depth (dstokes)
      SELECT CASE (nn_osm_wave)
         ! Assume constant La#=0.3
      CASE(0)
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            zus_x = scos_wind(ji,jj) * sustar(ji,jj) / 0.3_wp**2
            zus_y = ssin_wind(ji,jj) * sustar(ji,jj) / 0.3_wp**2
            ! Linearly
            sustke(ji,jj)  = MAX( SQRT( zus_x * zus_x + zus_y * zus_y ), 1e-8_wp )
            dstokes(ji,jj) = rn_osm_dstokes
         END_2D
         ! Assume Pierson-Moskovitz wind-wave spectrum
      CASE(1)
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            ! Use wind speed wndm included in sbc_oce module
            sustke(ji,jj)  = MAX ( 0.016_wp * wndm(ji,jj), 1e-8_wp )
            dstokes(ji,jj) = MAX ( 0.12_wp * wndm(ji,jj)**2 / grav, 5e-1_wp )
         END_2D
         ! Use ECMWF wave fields as output from SBCWAVE
      CASE(2)
         zfac =  2.0_wp * rpi / 16.0_wp
         !
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            IF ( hsw(ji,jj) > 1e-4_wp ) THEN
               ! Use  wave fields
               zabsstke       = SQRT( ut0sd(ji,jj)**2 + vt0sd(ji,jj)**2 )
               sustke(ji,jj)  = MAX( ( scos_wind(ji,jj) * ut0sd(ji,jj) + ssin_wind(ji,jj)  * vt0sd(ji,jj) ), 1e-8_wp )
               dstokes(ji,jj) = MAX( zfac * hsw(ji,jj) * hsw(ji,jj) / ( MAX( zabsstke * wmp(ji,jj), 1e-7 ) ), 5e-1_wp )
            ELSE
               ! Assume masking issue (e.g. ice in ECMWF reanalysis but not in model run)
               ! .. so default to Pierson-Moskowitz
               sustke(ji,jj)  = MAX( 0.016_wp * wndm(ji,jj), 1e-8_wp )
               dstokes(ji,jj) = MAX( 0.12_wp * wndm(ji,jj)**2 / grav, 5e-1_wp )
            END IF
         END_2D
      END SELECT
      !
      IF (ln_zdfosm_ice_shelter) THEN
         ! Reduce both Stokes drift and its depth scale by ocean fraction to represent sheltering by ice
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            sustke(ji,jj)  = sustke(ji,jj)  * ( 1.0_wp - fr_i(ji,jj) )
            dstokes(ji,jj) = dstokes(ji,jj) * ( 1.0_wp - fr_i(ji,jj) )
         END_2D
      END IF
      !
      SELECT CASE (nn_osm_SD_reduce)
         ! Reduce surface Stokes drift by a constant factor or following Breivik (2016) + van Roekel (2012) or Grant (2020).
      CASE(0)
         ! The Langmur number from the ECMWF model (or from PM) appears to give La<0.3 for wind-driven seas.
         ! The coefficient rn_zdfosm_adjust_sd = 0.8 gives La=0.3 in this situation.
         ! It could represent the effects of the spread of wave directions around the mean wind. The effect of this adjustment needs to be tested.
         IF(nn_osm_wave > 0) THEN
            DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
               sustke(ji,jj) = rn_zdfosm_adjust_sd * sustke(ji,jj)
            END_2D
         END IF
      CASE(1)
         ! Van Roekel (2012): consider average SD over top 10% of boundary layer
         ! Assumes approximate depth profile of SD from Breivik (2016)
         zsqrtpi = SQRT(rpi)
         z_two_thirds = 2.0_wp / 3.0_wp
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            zthickness = rn_osm_hblfrac*hbl(ji,jj)
            z2k_times_thickness =  zthickness * 2.0_wp / MAX( ABS( 5.97_wp * dstokes(ji,jj) ), 1e-7_wp )
            zsqrt_depth = SQRT( z2k_times_thickness )
            zexp_depth  = EXP( -1.0_wp * z2k_times_thickness )
            sustke(ji,jj) = sustke(ji,jj) * ( 1.0_wp - zexp_depth -   &
               &                              z_two_thirds * ( zsqrtpi * zsqrt_depth * z2k_times_thickness * ERFC(zsqrt_depth) +   &
               &                                               1.0_wp - ( 1.0_wp + z2k_times_thickness ) * zexp_depth ) ) /        &
               &            z2k_times_thickness
         END_2D
      CASE(2)
         ! Grant (2020): Match to exponential with same SD and d/dz(Sd) at depth 10% of boundary layer
         ! Assumes approximate depth profile of SD from Breivik (2016)
         zsqrtpi = SQRT(rpi)
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            zthickness = rn_osm_hblfrac*hbl(ji,jj)
            z2k_times_thickness =  zthickness * 2.0_wp / MAX( ABS( 5.97_wp * dstokes(ji,jj) ), 1e-7_wp )
            IF( z2k_times_thickness < 50.0_wp ) THEN
               zsqrt_depth = SQRT( z2k_times_thickness )
               zexperfc    = zsqrtpi * zsqrt_depth * ERFC(zsqrt_depth) * EXP( z2k_times_thickness )
            ELSE
               ! Asymptotic expansion of sqrt(pi)*zsqrt_depth*EXP(z2k_times_thickness)*ERFC(zsqrt_depth) for large
               !    z2k_times_thickness
               ! See Abramowitz and Stegun, Eq. 7.1.23
               ! zexperfc = 1._wp - (1/2)/(z2k_times_thickness) + (3/4)/(z2k_times_thickness**2) - (15/8)/(z2k_times_thickness**3)
               zexperfc = ( ( -1.875_wp / z2k_times_thickness + 0.75_wp ) / z2k_times_thickness - 0.5_wp ) /   &
                  &       z2k_times_thickness + 1.0_wp
            END IF
            zf = z2k_times_thickness * ( 1.0_wp / zexperfc - 1.0_wp )
            dstokes(ji,jj) = 5.97_wp * zf * dstokes(ji,jj)
            sustke(ji,jj)  = sustke(ji,jj) * EXP( z2k_times_thickness * ( 1.0_wp / ( 2.0_wp * zf ) - 1.0_wp ) ) *   &
               &             ( 1.0_wp - zexperfc )
         END_2D
      END SELECT
      !
      ! Langmuir velocity scale (swstrl), La # (sla)
      ! Mixed scale (svstr), convective velocity scale (swstrc)
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         ! Langmuir velocity scale (swstrl), at T-point
         swstrl(ji,jj) = ( sustar(ji,jj) * sustar(ji,jj) * sustke(ji,jj) )**pthird
         sla(ji,jj)    = MAX( MIN( SQRT( sustar(ji,jj) / ( swstrl(ji,jj) + epsln ) )**3, 4.0_wp ), 0.2_wp )
         IF ( sla(ji,jj) > 0.45_wp ) dstokes(ji,jj) = MIN( dstokes(ji,jj), 0.5_wp * hbl(ji,jj) )
         ! Velocity scale that tends to sustar for large Langmuir numbers
         svstr(ji,jj)  = ( swstrl(ji,jj)**3 + ( 1.0_wp - EXP( -0.5_wp * sla(ji,jj)**2 ) ) * sustar(ji,jj) * sustar(ji,jj) *   &
            &                                 sustar(ji,jj) )**pthird
         !
         ! Limit maximum value of Langmuir number as approximate treatment for shear turbulence
         ! Note sustke and swstrl are not amended
         !
         ! Get convective velocity (swstrc), stabilty scale (shol) and logical conection flag l_conv
         IF ( swbav(ji,jj) > 0.0_wp ) THEN
            swstrc(ji,jj) = ( 2.0_wp * swbav(ji,jj) * 0.9_wp * hbl(ji,jj) )**pthird
            shol(ji,jj)   = -0.9_wp * hbl(ji,jj) * 2.0_wp * swbav(ji,jj) / ( svstr(ji,jj)**3 + epsln )
         ELSE
            swstrc(ji,jj) = 0.0_wp
            shol(ji,jj)   = -1.0_wp * hbl(ji,jj) * 2.0_wp * swbav(ji,jj) / ( svstr(ji,jj)**3  + epsln )
         ENDIF
      END_2D
      !
      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
      ! Mixed-layer model - calculate averages over the boundary layer, and the change in the boundary layer depth
      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
      ! BL must be always 4 levels deep.
      ! For calculation of lateral buoyancy gradients for FK in
      ! zdf_osm_zmld_horizontal_gradients need halo values for nbld
      !
      ! agn 23/6/20: not clear all this is needed, as hbl checked after it is re-calculated anyway
      ! ##########################################################################
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         hbl(ji,jj) = MAX(hbl(ji,jj), gdepw(ji,jj,4,Kmm) )
      END_2D
      DO_2D_OVR( nn_hls, nn_hls, nn_hls, nn_hls )
         nbld(ji,jj) = 4
      END_2D
      DO_3D_OVR( nn_hls, nn_hls, nn_hls, nn_hls, 5, jpkm1 )
         IF ( MAX( hbl(ji,jj), gdepw(ji,jj,4,Kmm) ) >= gdepw(ji,jj,jk,Kmm) ) THEN
            nbld(ji,jj) = MIN(mbkt(ji,jj)-2, jk)
         ENDIF
      END_3D
      ! ##########################################################################
      !
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         zhbl(ji,jj) = gdepw(ji,jj,nbld(ji,jj),Kmm)
         nmld(ji,jj) = MAX( 3, nbld(ji,jj) - MAX( INT( dh(ji,jj) / e3t(ji,jj,nbld(ji,jj)-1,Kmm) ), 1 ) )
         zhml(ji,jj) = gdepw(ji,jj,nmld(ji,jj),Kmm)
         zdh(ji,jj) = zhbl(ji,jj) - zhml(ji,jj)
      END_2D
      !
      ! Averages over well-mixed and boundary layer, note BL averages use jk_ext=2 everywhere
      jk_nlev(:,:) = nbld(A2D(nn_hls-1))
      jk_ext(:,:) = 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_bl,  av_s_bl,    &
         &                           av_b_bl,  av_u_bl,  av_v_bl,  jk_ext,   av_dt_bl,   &
         &                           av_ds_bl, av_db_bl, av_du_bl, av_dv_bl )
      jk_nlev(:,:) = nmld(A2D(nn_hls-1)) - 1
      jk_ext(:,:) = nbld(A2D(nn_hls-1)) - nmld(A2D(nn_hls-1)) + jk_ext(:,:) + 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_ml,  av_s_ml,    &
         &                           av_b_ml,  av_u_ml,  av_v_ml,  jk_ext,   av_dt_ml,   &
         &                           av_ds_ml, av_db_ml, av_du_ml, av_dv_ml )
      ! Velocity components in frame aligned with surface stress
      CALL zdf_osm_velocity_rotation( av_u_ml,  av_v_ml  )
      CALL zdf_osm_velocity_rotation( av_du_ml, av_dv_ml )
      CALL zdf_osm_velocity_rotation( av_u_bl,  av_v_bl  )
      CALL zdf_osm_velocity_rotation( av_du_bl, av_dv_bl )
      !
      ! Determine the state of the OSBL, stable/unstable, shear/no shear
      CALL zdf_osm_osbl_state( Kmm, zwb_ent, zwb_min, zshear, zhbl,     &
         &                     zhml, zdh )
      !
      IF ( ln_osm_mle ) THEN
         ! Fox-Kemper Scheme
         DO_2D_OVR( nn_hls, nn_hls, nn_hls, nn_hls )
            mld_prof(ji,jj) = 4
         END_2D
         DO_3D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1, 5, jpkm1 )
            IF ( hmle(ji,jj) >= gdepw(ji,jj,jk,Kmm) ) mld_prof(ji,jj) = MIN( mbkt(ji,jj), jk)
         END_3D
         jk_nlev(:,:) = mld_prof(A2D(nn_hls-1))
         CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_mle, av_s_mle,   &
            &                           av_b_mle, av_u_mle, av_v_mle )
         !
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            zhmle(ji,jj) = gdepw(ji,jj,mld_prof(ji,jj),Kmm)
         END_2D
         !
         ! Calculate fairly-well-mixed depth zmld & its index mld_prof + lateral zmld-averaged gradients
         CALL zdf_osm_zmld_horizontal_gradients( Kmm, zmld, zdtdx, zdtdy, zdsdx,   &
            &                                    zdsdy, zdbds_mle )
         ! Calculate max vertical FK flux zwb_fk & set logical descriptors
         CALL zdf_osm_osbl_state_fk( Kmm, zwb_fk, zhbl, zhmle, zwb_ent,   &
            &                        zdbds_mle )
         ! Recalculate hmle, zmle, zvel_mle, zdiff_mle & redefine mld_proc to be index for new hmle
         CALL zdf_osm_mle_parameters( Kmm, zmld, zhmle, zvel_mle, zdiff_mle,   &
            &                         zdbds_mle, zhbl, zwb0tot )
      ELSE    ! ln_osm_mle
         ! FK not selected, Boundary Layer only.
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            l_pyc(ji,jj)  = .TRUE.
            l_flux(ji,jj) = .FALSE.
            l_mle(ji,jj)  = .FALSE.
            IF ( l_conv(ji,jj) .AND. av_db_bl(ji,jj) < rn_osm_bl_thresh ) l_pyc(ji,jj) = .FALSE.
         END_2D
      ENDIF   ! ln_osm_mle
      !
      !! External gradient below BL needed both with and w/o FK
      jk_ext(:,:) = nbld(A2D(nn_hls-1)) + 1
      CALL zdf_osm_external_gradients( Kmm, jk_ext, zdtdz_bl_ext, zdsdz_bl_ext, zdbdz_bl_ext )   ! ag 19/03
      !
      ! Test if pycnocline well resolved
      !      DO_2D( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )                                         Removed with ag 19/03 changes. A change in eddy diffusivity/viscosity
      !         IF (l_conv(ji,jj) ) THEN                                  should account for this.
      !            ztmp = 0.2 * zhbl(ji,jj) / e3w(ji,jj,nbld(ji,jj),Kmm)
      !            IF ( ztmp > 6 ) THEN
      !               ! pycnocline well resolved
      !               jk_ext(ji,jj) = 1
      !            ELSE
      !               ! pycnocline poorly resolved
      !               jk_ext(ji,jj) = 0
      !            ENDIF
      !         ELSE
      !            ! Stable conditions
      !            jk_ext(ji,jj) = 0
      !         ENDIF
      !      END_2D
      !
      ! Recalculate bl averages using jk_ext & ml averages .... note no rotation of u & v here..
      jk_nlev(:,:) = nbld(A2D(nn_hls-1))
      jk_ext(:,:) = 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_bl,  av_s_bl,    &
         &                           av_b_bl,  av_u_bl,  av_v_bl,  jk_ext,   av_dt_bl,   &
         &                           av_ds_bl, av_db_bl, av_du_bl, av_dv_bl )
      jk_nlev(:,:) = nmld(A2D(nn_hls-1)) - 1
      jk_ext(:,:) = nbld(A2D(nn_hls-1)) - nmld(A2D(nn_hls-1)) + jk_ext(:,:) + 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_ml,  av_s_ml,    &
         &                           av_b_ml,  av_u_ml,  av_v_ml,  jk_ext,   av_dt_ml,   &
         &                           av_ds_ml, av_db_ml, av_du_ml, av_dv_ml )   ! ag 19/03
      !
      ! Rate of change of hbl
      CALL zdf_osm_calculate_dhdt( zdhdt, zhbl, zdh, zwb_ent, zwb_min,   &
         &                         zdbdz_bl_ext, zwb_fk_b, zwb_fk, zvel_mle )
      ! Test if surface boundary layer coupled to bottom
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         l_coup(ji,jj) = .FALSE.   ! ag 19/03
         zhbl_t(ji,jj) = hbl(ji,jj) + ( zdhdt(ji,jj) - ww(ji,jj,nbld(ji,jj)) ) * rn_Dt   ! Certainly need ww here, so subtract it
         ! Adjustment to represent limiting by ocean bottom
         IF ( mbkt(ji,jj) > 2 ) THEN   ! To ensure mbkt(ji,jj) - 2 > 0 so no incorrect array access
            IF ( zhbl_t(ji,jj) > gdepw(ji, jj,mbkt(ji,jj)-2,Kmm) ) THEN
               zhbl_t(ji,jj) = MIN( zhbl_t(ji,jj), gdepw(ji,jj,mbkt(ji,jj)-2,Kmm) )   ! ht(:,:))
               l_pyc(ji,jj)  = .FALSE.
               l_coup(ji,jj) = .TRUE.   ! ag 19/03
            END IF
         END IF
      END_2D
      !
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         nmld(ji,jj) = nbld(ji,jj)           ! use nmld to hold previous blayer index
         nbld(ji,jj) = 4
      END_2D
      !
      DO_3D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1, 4, jpkm1 )
         IF ( zhbl_t(ji,jj) >= gdepw(ji,jj,jk,Kmm) ) THEN
            nbld(ji,jj) = jk
         END IF
      END_3D
      !
      !
      ! Step through model levels taking account of buoyancy change to determine the effect on dhdt
      !
      CALL zdf_osm_timestep_hbl( Kmm, zdhdt, zhbl, zhbl_t, zwb_ent,   &
         &                       zwb_fk_b )
      ! Is external level in bounds?
      !
      ! Recalculate BL averages and differences using new BL depth
      jk_nlev(:,:) = nbld(A2D(nn_hls-1))
      jk_ext(:,:) = 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_bl,  av_s_bl,    &
         &                           av_b_bl,  av_u_bl,  av_v_bl,  jk_ext,   av_dt_bl,   &
         &                           av_ds_bl, av_db_bl, av_du_bl, av_dv_bl )
      !
      CALL zdf_osm_pycnocline_thickness( Kmm, zdh, zhml, zdhdt, zhbl,   &
         &                               zwb_ent, zdbdz_bl_ext, zwb_fk_b )
      !
      ! Reset l_pyc before calculating terms in the flux-gradient relationship
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
         IF ( av_db_bl(ji,jj) < rn_osm_bl_thresh .OR. nbld(ji,jj) >= mbkt(ji,jj) - 2 .OR.   &
            & nbld(ji,jj) - nmld(ji,jj) == 1   .OR. zdhdt(ji,jj) < 0.0_wp ) THEN   ! ag 19/03
            l_pyc(ji,jj) = .FALSE.   ! ag 19/03
            IF ( nbld(ji,jj) >= mbkt(ji,jj) -2 ) THEN
               nmld(ji,jj) = nbld(ji,jj) - 1                                               ! ag 19/03
               zdh(ji,jj)  = gdepw(ji,jj,nbld(ji,jj),Kmm) - gdepw(ji,jj,nmld(ji,jj),Kmm)   ! ag 19/03
               zhml(ji,jj) = gdepw(ji,jj,nmld(ji,jj),Kmm)                                  ! ag 19/03
               dh(ji,jj)   = zdh(ji,jj)                                                    ! ag 19/03  
               hml(ji,jj)  = hbl(ji,jj) - dh(ji,jj)                                        ! ag 19/03
            ENDIF
         ENDIF                                              ! ag 19/03
      END_2D
      !
      DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )               ! Limit delta for shallow boundary layers for calculating
         dstokes(ji,jj) = MIN ( dstokes(ji,jj), hbl(ji,jj) / 3.0_wp )   !    flux-gradient terms
      END_2D
      !                                                       
      !
      ! Average over the depth of the mixed layer in the convective boundary layer
      !      jk_ext = nbld - nmld + 1
      ! Recalculate ML averages and differences using new ML depth
      jk_nlev(:,:) = nmld(A2D(nn_hls-1)) - 1
      jk_ext(:,:) = nbld(A2D(nn_hls-1)) - nmld(A2D(nn_hls-1)) + jk_ext(:,:) + 1   ! ag 19/03
      CALL zdf_osm_vertical_average( Kbb,      Kmm,      jk_nlev,  av_t_ml,  av_s_ml,    &
         &                           av_b_ml,  av_u_ml,  av_v_ml,  jk_ext,   av_dt_ml,   &
         &                           av_ds_ml, av_db_ml, av_du_ml, av_dv_ml )
      !
      jk_ext(:,:) = nbld(A2D(nn_hls-1)) + 1
      CALL zdf_osm_external_gradients( Kmm, jk_ext, zdtdz_bl_ext, zdsdz_bl_ext, zdbdz_bl_ext )
      ! Rotate mean currents and changes onto wind aligned co-ordinates
      CALL zdf_osm_velocity_rotation( av_u_ml,  av_v_ml  )
      CALL zdf_osm_velocity_rotation( av_du_ml, av_dv_ml )
      CALL zdf_osm_velocity_rotation( av_u_bl,  av_v_bl  )
      CALL zdf_osm_velocity_rotation( av_du_bl, av_dv_bl )
      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
      ! Eddy viscosity/diffusivity and non-gradient terms in the flux-gradient relationship
      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
      CALL zdf_osm_diffusivity_viscosity( Kbb, Kmm, zdiffut, zviscos, zhbl,    &
         &                                zhml, zdh, zdhdt, zshear, zwb_ent,   &
         &                                zwb_min )
      !
      ! Calculate non-gradient components of the flux-gradient relationships
      ! --------------------------------------------------------------------
      jk_ext(:,:) = 1   ! ag 19/03
      CALL zdf_osm_fgr_terms( Kmm, jk_ext, zhbl, zhml, zdh,                              &
         &                    zdhdt, zshear, zdtdz_bl_ext, zdsdz_bl_ext, zdbdz_bl_ext,   &
         &                    zdiffut, zviscos )
      !
      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
      ! Need to put in code for contributions that are applied explicitly to
      ! the prognostic variables
      !  1. Entrainment flux
      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
      !
      ! Rotate non-gradient velocity terms back to model reference frame
      jk_nlev(:,:) = nbld(A2D(nn_hls-1))
      CALL zdf_osm_velocity_rotation( ghamu, ghamv, .FALSE.,  2, jk_nlev )
      !
      ! KPP-style Ri# mixing
      IF ( ln_kpprimix ) THEN
         jkflt = jpk
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            IF ( nbld(ji,jj) < jkflt ) jkflt = nbld(ji,jj)
         END_2D
         DO jk = jkflt+1, jpkm1
            ! Shear production at uw- and vw-points (energy conserving form)
            DO_2D( nn_hls, nn_hls-1, nn_hls, nn_hls-1 )
               z2du(ji,jj) = 0.5_wp * ( uu(ji,jj,jk-1,Kmm) - uu(ji,jj,jk,Kmm) ) * ( uu(ji,jj,jk-1,Kbb) - uu(ji,jj,jk,Kbb) ) *   &
                  &          wumask(ji,jj,jk) / ( e3uw(ji,jj,jk,Kmm) * e3uw(ji,jj,jk,Kbb) )
               z2dv(ji,jj) = 0.5_wp * ( vv(ji,jj,jk-1,Kmm) - vv(ji,jj,jk,Kmm) ) * ( vv(ji,jj,jk-1,Kbb) - vv(ji,jj,jk,Kbb) ) *   &
                  &          wvmask(ji,jj,jk) / ( e3vw(ji,jj,jk,Kmm) * e3vw(ji,jj,jk,Kbb) )
            END_2D
            DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
               IF ( jk > nbld(ji,jj) ) THEN
                  ! Shear prod. at w-point weightened by mask
                  zesh2 = ( z2du(ji-1,jj) + z2du(ji,jj) ) / MAX( 1.0_wp , umask(ji-1,jj,jk) + umask(ji,jj,jk) ) +   &
                     &    ( z2dv(ji,jj-1) + z2dv(ji,jj) ) / MAX( 1.0_wp , vmask(ji,jj-1,jk) + vmask(ji,jj,jk) )
                  ! Local Richardson number
                  zri     = MAX( rn2b(ji,jj,jk), 0.0_wp ) / MAX( zesh2, epsln )
                  zfri    = MIN( zri / rn_riinfty, 1.0_wp )
                  zfri    = ( 1.0_wp - zfri * zfri )
                  zrimix  =  zfri * zfri  * zfri * wmask(ji, jj, jk)
                  zdiffut(ji,jj,jk) = MAX( zdiffut(ji,jj,jk), zrimix*rn_difri )
                  zviscos(ji,jj,jk) = MAX( zviscos(ji,jj,jk), zrimix*rn_difri )
               END IF
            END_2D
         END DO
      END IF   ! ln_kpprimix = .true.
      !
      ! KPP-style set diffusivity large if unstable below BL
      IF ( ln_convmix) THEN
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            DO jk = nbld(ji,jj) + 1, jpkm1
               IF ( MIN( rn2(ji,jj,jk), rn2b(ji,jj,jk) ) <= -1e-12_wp ) zdiffut(ji,jj,jk) = MAX( rn_difconv, zdiffut(ji,jj,jk) )
            END DO
         END_2D
      END IF   ! ln_convmix = .true.
      !
      IF ( ln_osm_mle ) THEN   ! Set up diffusivity and non-gradient mixing
         DO_2D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1 )
            IF ( l_flux(ji,jj) ) THEN   ! MLE mixing extends below boundary layer
               ! Calculate MLE flux contribution from surface fluxes
               DO jk = 1, nbld(ji,jj)
                  znd = gdepw(ji,jj,jk,Kmm) / MAX( zhbl(ji,jj), epsln )
                  ghamt(ji,jj,jk) = ghamt(ji,jj,jk) - ( swth0(ji,jj) - zrad0(ji,jj) + zradh(ji,jj) ) * ( 1.0_wp - znd )
                  ghams(ji,jj,jk) = ghams(ji,jj,jk) - sws0(ji,jj) * ( 1.0_wp - znd )
               END DO
               DO jk = 1, mld_prof(ji,jj)
                  znd = gdepw(ji,jj,jk,Kmm) / MAX( zhmle(ji,jj), epsln )
                  ghamt(ji,jj,jk) = ghamt(ji,jj,jk) + ( swth0(ji,jj) - zrad0(ji,jj) + zradh(ji,jj) ) * ( 1.0_wp - znd )
                  ghams(ji,jj,jk) = ghams(ji,jj,jk) + sws0(ji,jj) * ( 1.0_wp -znd )
               END DO
               ! Viscosity for MLEs
               DO jk = 1, mld_prof(ji,jj)
                  znd = -1.0_wp * gdepw(ji,jj,jk,Kmm) / MAX( zhmle(ji,jj), epsln )
                  zdiffut(ji,jj,jk) = zdiffut(ji,jj,jk) + zdiff_mle(ji,jj) * ( 1.0_wp - ( 2.0_wp * znd + 1.0_wp )**2 ) *   &
                     &                                    ( 1.0_wp + 5.0_wp / 21.0_wp * ( 2.0_wp * znd + 1.0_wp )**2 )
               END DO
            ELSE   ! Surface transports limited to OSBL
               ! Viscosity for MLEs
               DO jk = 1, mld_prof(ji,jj)
                  znd = -1.0_wp * gdepw(ji,jj,jk,Kmm) / MAX( zhmle(ji,jj), epsln )
                  zdiffut(ji,jj,jk) = zdiffut(ji,jj,jk) + zdiff_mle(ji,jj) * ( 1.0_wp - ( 2.0_wp * znd + 1.0_wp )**2 ) *   &
                     &                                    ( 1.0_wp + 5.0_wp / 21.0_wp * ( 2.0_wp * znd + 1.0_wp )**2 )
               END DO
            END IF
         END_2D
      ENDIF
      !
      ! Lateral boundary conditions on zvicos (sign unchanged), needed to caclulate viscosities on u and v grids
      ! CALL lbc_lnk( 'zdfosm', zviscos(:,:,:), 'W', 1.0_wp )
      ! GN 25/8: need to change tmask --> wmask
      DO_3D_OVR( nn_hls-1, nn_hls-1, nn_hls-1, nn_hls-1, 2, jpkm1 )
         p_avt(ji,jj,jk) = MAX( zdiffut(ji,jj,jk), avtb(jk) ) * tmask(ji,jj,jk)
         p_avm(ji,jj,jk) = MAX( zviscos(ji,jj,jk), avmb(jk) ) * tmask(ji,jj,jk)
      END_3D
      !
      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)
         CASE(0:1)
            CALL zdf_osm_iomput( "us_x", tmask(A2D(0),1) * sustke(A2D(0)) * scos_wind(A2D(0)) )   ! x surface Stokes drift
            CALL zdf_osm_iomput( "us_y", tmask(A2D(0),1) * sustke(A2D(0)) * ssin_wind(A2D(0)) )   ! y surface Stokes drift
Guillaume Samson's avatar
Guillaume Samson committed
            CALL zdf_osm_iomput( "wind_wave_abs_power", 1000.0_wp * rho0 * tmask(A2D(0),1) * sustar(A2D(0))**2 * sustke(A2D(0)) )
            ! Stokes drift read in from sbcwave  (=2).
         CASE(2:3)
            CALL zdf_osm_iomput( "us_x",   ut0sd(A2D(0)) * umask(A2D(0),1) )                         ! x surface Stokes drift
            CALL zdf_osm_iomput( "us_y",   vt0sd(A2D(0)) * vmask(A2D(0),1) )                         ! y surface Stokes drift
            CALL zdf_osm_iomput( "wmp",    wmp(A2D(0)) * tmask(A2D(0),1) )                           ! Wave mean period
            CALL zdf_osm_iomput( "hsw",    hsw(A2D(0)) * tmask(A2D(0),1) )                           ! Significant wave height
            CALL zdf_osm_iomput( "wmp_NP", ( 2.0_wp * rpi * 1.026_wp / ( 0.877_wp * grav ) ) *   &   ! Wave mean period from NP
               &                           wndm(A2D(0)) * tmask(A2D(0),1) )                          !    spectrum
            CALL zdf_osm_iomput( "hsw_NP", ( 0.22_wp / grav ) * wndm(A2D(0))**2 * tmask(A2D(0),1) )  ! Significant wave height from
            !                                                                                        !    NP spectrum
            CALL zdf_osm_iomput( "wndm",   wndm(A2D(0)) * tmask(A2D(0),1) )                          ! U_10
            CALL zdf_osm_iomput( "wind_wave_abs_power", 1000.0_wp * rho0 * tmask(A2D(0),1) * sustar(A2D(0))**2 *   &
               &                                        SQRT( ut0sd(A2D(0))**2 + vt0sd(A2D(0))**2 ) )
         END SELECT
         CALL zdf_osm_iomput( "zwth0",           tmask(A2D(0),1) * swth0(A2D(0))     )      ! <Tw_0>
         CALL zdf_osm_iomput( "zws0",            tmask(A2D(0),1) * sws0(A2D(0))      )      ! <Sw_0>
         CALL zdf_osm_iomput( "zwb0",            tmask(A2D(0),1) * swb0(A2D(0))      )      ! <bw_0>
         CALL zdf_osm_iomput( "zwbav",           tmask(A2D(0),1) * swbav(A2D(0))     )      ! Upward BL-avged turb buoyancy flux
Guillaume Samson's avatar
Guillaume Samson committed
         CALL zdf_osm_iomput( "ibld",            tmask(A2D(0),1) * nbld(A2D(0))      )      ! Boundary-layer max k
         CALL zdf_osm_iomput( "zdt_bl",          tmask(A2D(0),1) * av_dt_bl(A2D(0))  )      ! dt at ml base
         CALL zdf_osm_iomput( "zds_bl",          tmask(A2D(0),1) * av_ds_bl(A2D(0))  )      ! ds at ml base
         CALL zdf_osm_iomput( "zdb_bl",          tmask(A2D(0),1) * av_db_bl(A2D(0))  )      ! db at ml base
         CALL zdf_osm_iomput( "zdu_bl",          tmask(A2D(0),1) * av_du_bl(A2D(0))  )      ! du at ml base
         CALL zdf_osm_iomput( "zdv_bl",          tmask(A2D(0),1) * av_dv_bl(A2D(0))  )      ! dv at ml base
         CALL zdf_osm_iomput( "dh",              tmask(A2D(0),1) * dh(A2D(0))        )      ! Initial boundary-layer depth
         CALL zdf_osm_iomput( "hml",             tmask(A2D(0),1) * hml(A2D(0))       )      ! Initial boundary-layer depth
         CALL zdf_osm_iomput( "zdt_ml",          tmask(A2D(0),1) * av_dt_ml(A2D(0))  )      ! dt at ml base
         CALL zdf_osm_iomput( "zds_ml",          tmask(A2D(0),1) * av_ds_ml(A2D(0))  )      ! ds at ml base
         CALL zdf_osm_iomput( "zdb_ml",          tmask(A2D(0),1) * av_db_ml(A2D(0))  )      ! db at ml base
         CALL zdf_osm_iomput( "dstokes",         tmask(A2D(0),1) * dstokes(A2D(0))   )      ! Stokes drift penetration depth
         CALL zdf_osm_iomput( "zustke",          tmask(A2D(0),1) * sustke(A2D(0))    )      ! Stokes drift magnitude at T-points
         CALL zdf_osm_iomput( "zwstrc",          tmask(A2D(0),1) * swstrc(A2D(0))    )      ! Convective velocity scale
         CALL zdf_osm_iomput( "zwstrl",          tmask(A2D(0),1) * swstrl(A2D(0))    )      ! Langmuir velocity scale
         CALL zdf_osm_iomput( "zustar",          tmask(A2D(0),1) * sustar(A2D(0))    )      ! Friction velocity scale
         CALL zdf_osm_iomput( "zvstr",           tmask(A2D(0),1) * svstr(A2D(0))     )      ! Mixed velocity scale
         CALL zdf_osm_iomput( "zla",             tmask(A2D(0),1) * sla(A2D(0))       )      ! Langmuir #
         CALL zdf_osm_iomput( "wind_power",      1000.0_wp * rho0 * tmask(A2D(0),1) *   &   ! BL depth internal to zdf_osm routine
            &                                    sustar(A2D(0))**3 )
         CALL zdf_osm_iomput( "wind_wave_power", 1000.0_wp * rho0 * tmask(A2D(0),1) *   &
            &                                    sustar(A2D(0))**2 * sustke(A2D(0))  )
         CALL zdf_osm_iomput( "zhbl",            tmask(A2D(0),1) * zhbl(A2D(0))      )      ! BL depth internal to zdf_osm routine
         CALL zdf_osm_iomput( "zhml",            tmask(A2D(0),1) * zhml(A2D(0))      )      ! ML depth internal to zdf_osm routine
         CALL zdf_osm_iomput( "imld",            tmask(A2D(0),1) * nmld(A2D(0))      )      ! Index for ML depth internal to zdf_osm
         !                                                                                  !    routine
         CALL zdf_osm_iomput( "jp_ext",          tmask(A2D(0),1) * jk_ext(A2D(0))    )      ! =1 if pycnocline resolved internal to
         !                                                                                  !    zdf_osm routine
         CALL zdf_osm_iomput( "j_ddh",           tmask(A2D(0),1) * n_ddh(A2D(0))     )      ! Index forpyc thicknessh internal to
         !                                                                                  !    zdf_osm routine
         CALL zdf_osm_iomput( "zshear",          tmask(A2D(0),1) * zshear(A2D(0))    )      ! Shear production of TKE internal to
         !                                                                                  !    zdf_osm routine
         CALL zdf_osm_iomput( "zdh",             tmask(A2D(0),1) * zdh(A2D(0))       )      ! Pyc thicknessh internal to zdf_osm
         !                                                                                  !    routine
         CALL zdf_osm_iomput( "zhol",            tmask(A2D(0),1) * shol(A2D(0))      )      ! ML depth internal to zdf_osm routine
         CALL zdf_osm_iomput( "zwb_ent",         tmask(A2D(0),1) * zwb_ent(A2D(0))   )      ! Upward turb buoyancy entrainment flux
         CALL zdf_osm_iomput( "zt_ml",           tmask(A2D(0),1) * av_t_ml(A2D(0))   )      ! Average T in ML
         CALL zdf_osm_iomput( "zmld",            tmask(A2D(0),1) * zmld(A2D(0))      )      ! FK target layer depth
         CALL zdf_osm_iomput( "zwb_fk",          tmask(A2D(0),1) * zwb_fk(A2D(0))    )      ! FK b flux
         CALL zdf_osm_iomput( "zwb_fk_b",        tmask(A2D(0),1) * zwb_fk_b(A2D(0))  )      ! FK b flux averaged over ML
         CALL zdf_osm_iomput( "mld_prof",        tmask(A2D(0),1) * mld_prof(A2D(0))  )      ! FK layer max k
         CALL zdf_osm_iomput( "zdtdx",           umask(A2D(0),1) * zdtdx(A2D(0))     )      ! FK dtdx at u-pt
         CALL zdf_osm_iomput( "zdtdy",           vmask(A2D(0),1) * zdtdy(A2D(0))     )      ! FK dtdy at v-pt
         CALL zdf_osm_iomput( "zdsdx",           umask(A2D(0),1) * zdsdx(A2D(0))     )      ! FK dtdx at u-pt
         CALL zdf_osm_iomput( "zdsdy",           vmask(A2D(0),1) * zdsdy(A2D(0))     )      ! FK dsdy at v-pt
         CALL zdf_osm_iomput( "dbdx_mle",        umask(A2D(0),1) * dbdx_mle(A2D(0))  )      ! FK dbdx at u-pt
         CALL zdf_osm_iomput( "dbdy_mle",        vmask(A2D(0),1) * dbdy_mle(A2D(0))  )      ! FK dbdy at v-pt
         CALL zdf_osm_iomput( "zdiff_mle",       tmask(A2D(0),1) * zdiff_mle(A2D(0)) )      ! FK diff in MLE at t-pt
         CALL zdf_osm_iomput( "zvel_mle",        tmask(A2D(0),1) * zvel_mle(A2D(0)) )       ! FK velocity in MLE at t-pt
Guillaume Samson's avatar
Guillaume Samson committed
      END IF
      !
      ! Lateral boundary conditions on ghamu and ghamv, currently on W-grid (sign unchanged), needed to caclulate gham[uv] on u and
      !    v grids
      IF ( .NOT. l_istiled .OR. ntile == nijtile ) THEN   ! Finalise ghamu, ghamv, hbl, and hmle only after full domain has been
         !                                                !    processed
         IF ( nn_hls == 1 ) CALL lbc_lnk( 'zdfosm', ghamu, 'W', 1.0_wp,   &
            &                                       ghamv, 'W', 1.0_wp )
         DO jk = 2, jpkm1