Index: NEMO/branches/UKMO/NEMO_4.0_momentum_trends/src/OCE/TRD/trddyn.F90
===================================================================
--- NEMO/branches/UKMO/NEMO_4.0_momentum_trends/src/OCE/TRD/trddyn.F90	(revision 11611)
+++ NEMO/branches/UKMO/NEMO_4.0_momentum_trends/src/OCE/TRD/trddyn.F90	(revision 11613)
@@ -18,5 +18,4 @@
    USE sbc_oce        ! surface boundary condition: ocean
    USE zdf_oce        ! ocean vertical physics: variables
-!!gm   USE zdfdrg         ! ocean vertical physics: bottom friction
    USE trd_oce        ! trends: ocean variables
    USE trdken         ! trends: Kinetic ENergy 
@@ -35,4 +34,11 @@
    PUBLIC trd_dyn        ! called by all dynXXX modules
 
+   INTERFACE trd_dyn
+      module procedure trd_dyn_3d, trd_dyn_2d
+   END INTERFACE
+
+   REAL(wp), ALLOCATABLE, DIMENSION(:,:,:), SAVE :: zutrd_hpg, zvtrd_hpg
+   REAL(wp), ALLOCATABLE, DIMENSION(:,:,:), SAVE :: zutrd_pvo, zvtrd_pvo
+
    !! * Substitutions
 #  include "vectopt_loop_substitute.h90"
@@ -44,7 +50,7 @@
 CONTAINS
 
-   SUBROUTINE trd_dyn( putrd, pvtrd, ktrd, kt )
+   SUBROUTINE trd_dyn_3d( putrd, pvtrd, ktrd, kt )
       !!---------------------------------------------------------------------
-      !!                  ***  ROUTINE trd_mod  ***
+      !!                  ***  ROUTINE trd_dyn_3d  ***
       !! 
       !! ** Purpose :   Dispatch momentum trend computation, e.g. 3D output, 
@@ -63,23 +69,147 @@
 !!gm NB : here a lbc_lnk should probably be added
 
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      !   3D output of momentum and/or tracers trends using IOM interface
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      IF( ln_dyn_trd )   CALL trd_dyn_iom( putrd, pvtrd, ktrd, kt )
+      SELECT CASE( ktrd )
+      CASE( jpdyn_hpg_save ) 
+         !
+         ! save 3D HPG trends to possibly have barotropic part corrected later before writing out
+         ALLOCATE( zutrd_hpg(jpi,jpj,jpk), zvtrd_hpg(jpi,jpj,jpk) )
+         zutrd_hpg(:,:,:) = putrd(:,:,:)
+         zvtrd_hpg(:,:,:) = pvtrd(:,:,:)
+
+      CASE( jpdyn_pvo_save ) 
+         !
+         ! save 3D coriolis trends to possibly have barotropic part corrected later before writing out
+         ALLOCATE( zutrd_pvo(jpi,jpj,jpk), zvtrd_pvo(jpi,jpj,jpk) )
+         zutrd_pvo(:,:,:) = putrd(:,:,:)
+         zvtrd_pvo(:,:,:) = pvtrd(:,:,:)
+
+      CASE DEFAULT
+
+         IF( ktrd == jpdyn_spg ) THEN
+            ! For explicit scheme SPG trends come here as 3D fields
+            ! Add SPG trend to 3D HPG trend and also output as 2D diagnostic in own right.
+            CALL trd_dyn_iom_2d( putrd(:,:,1), pvtrd(:,:,1), jpdyn_spg, kt ) 
+            putrd(:,:,:) = putrd(:,:,:) + zutrd_hpg(:,:,:) 
+            pvtrd(:,:,:) = pvtrd(:,:,:) + zvtrd_hpg(:,:,:) 
+            DEALLOCATE( zutrd_hpg, zvtrd_hpg )
+         ENDIF
+
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !   3D output of momentum and/or tracers trends using IOM interface
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         IF( ln_dyn_trd )   CALL trd_dyn_iom_3d( putrd, pvtrd, ktrd, kt )
+
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !  Integral Constraints Properties for momentum and/or tracers trends
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         IF( ln_glo_trd )   CALL trd_glo( putrd, pvtrd, ktrd, 'DYN', kt )
+
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !  Kinetic Energy trends
+         !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
+         IF( ln_KE_trd  )   CALL trd_ken( putrd, pvtrd, ktrd, kt )
+
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !  Vorticity trends
+         !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
+         IF( ln_vor_trd )   CALL trd_vor( putrd, pvtrd, ktrd, kt )
+
+         !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
+         !  Mixed layer trends for active tracers
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !!gm      IF( ln_dyn_mxl )   CALL trd_mxl_dyn   
+         !
+      END SELECT
+      !
+   END SUBROUTINE trd_dyn_3d
+
+
+   SUBROUTINE trd_dyn_2d( putrd, pvtrd, ktrd, kt )
+      !!---------------------------------------------------------------------
+      !!                  ***  ROUTINE trd_mod  ***
+      !! 
+      !! ** Purpose :   Dispatch momentum trend computation, e.g. 2D output, 
+      !!              integral constraints, barotropic vorticity, kinetic enrgy, 
+      !!              and/or mixed layer budget.
+      !!----------------------------------------------------------------------
+      REAL(wp), DIMENSION(:,:), INTENT(inout) ::   putrd, pvtrd   ! U and V trends 
+      INTEGER                 , INTENT(in   ) ::   ktrd           ! trend index
+      INTEGER                 , INTENT(in   ) ::   kt             ! time step
+      INTEGER                                 ::   jk
+      !!----------------------------------------------------------------------
+      !
+      putrd(:,:) = putrd(:,:) * umask(:,:,1)                       ! mask the trends
+      pvtrd(:,:) = pvtrd(:,:) * vmask(:,:,1)
+      !
+
+!!gm NB : here a lbc_lnk should probably be added
+
+      SELECT CASE(ktrd)
+
+      CASE ( jpdyn_hpg_corr )
+         !
+         ! Remove "first-guess" SPG trend from 3D HPG trend. 
+         DO jk = 1, jpkm1
+            zutrd_hpg(:,:,jk) = zutrd_hpg(:,:,jk) - putrd(:,:)
+            zvtrd_hpg(:,:,jk) = zvtrd_hpg(:,:,jk) - pvtrd(:,:)
+         ENDDO
+         CALL trd_dyn_iom_2d( putrd, pvtrd, jpdyn_hpg_corr, kt ) 
+
+      CASE( jpdyn_pvo_corr )
+         !
+         ! Remove "first-guess" barotropic coriolis trend from 3D PVO trend. 
+         DO jk = 1, jpkm1
+            zutrd_pvo(:,:,jk) = zutrd_pvo(:,:,jk) - putrd(:,:)
+            zvtrd_pvo(:,:,jk) = zvtrd_pvo(:,:,jk) - pvtrd(:,:)
+         ENDDO
+         CALL trd_dyn_iom_2d( putrd, pvtrd, jpdyn_pvo_corr, kt ) 
+
+      CASE( jpdyn_spg )
+          !
+          ! For split-explicit scheme SPG trends come here as 2D fields
+          ! Add SPG trend to 3D HPG trend and also output as 2D diagnostic in own right.
+          DO jk = 1, jpkm1
+             zutrd_hpg(:,:,jk) = zutrd_hpg(:,:,jk) + putrd(:,:)
+             zvtrd_hpg(:,:,jk) = zvtrd_hpg(:,:,jk) + pvtrd(:,:)
+          ENDDO
+          CALL trd_dyn_iom_2d( putrd(:,:), pvtrd(:,:), jpdyn_spg, kt ) 
+          CALL trd_dyn_3d( zutrd_hpg, zvtrd_hpg, jpdyn_hpg, kt )
+          DEALLOCATE( zutrd_hpg, zvtrd_hpg )
+
+      CASE( jpdyn_pvo )
+          !
+          ! Add 2D PVO trend to 3D PVO trend and also output as diagnostic in own right.
+          DO jk = 1, jpkm1
+             zutrd_pvo(:,:,jk) = zutrd_pvo(:,:,jk) + putrd(:,:)
+             zvtrd_pvo(:,:,jk) = zvtrd_pvo(:,:,jk) + pvtrd(:,:)
+          ENDDO
+          CALL trd_dyn_iom_2d( putrd, pvtrd, jpdyn_pvo, kt ) 
+          CALL trd_dyn_3d( zutrd_pvo, zvtrd_pvo, jpdyn_pvo, kt )
+          DEALLOCATE( zutrd_pvo, zvtrd_pvo )
+
+      CASE DEFAULT 
+
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         !   2D output of momentum and/or tracers trends using IOM interface
+         !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+         IF( ln_dyn_trd )   CALL trd_dyn_iom_2d( putrd, pvtrd, ktrd, kt )
          
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      !  Integral Constraints Properties for momentum and/or tracers trends
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      IF( ln_glo_trd )   CALL trd_glo( putrd, pvtrd, ktrd, 'DYN', kt )
-
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      !  Kinetic Energy trends
-      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
-      IF( ln_KE_trd  )   CALL trd_ken( putrd, pvtrd, ktrd, kt )
-
-      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-      !  Vorticity trends
-      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
-      IF( ln_vor_trd )   CALL trd_vor( putrd, pvtrd, ktrd, kt )
+      END SELECT
+
+!!$   CALLS TO THESE ROUTINES FOR 2D DIAGOSTICS NOT CODED YET
+!!$      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+!!$      !  Integral Constraints Properties for momentum and/or tracers trends
+!!$      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+!!$      IF( ln_glo_trd )   CALL trd_glo( putrd, pvtrd, ktrd, 'DYN', kt )
+!!$
+!!$      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+!!$      !  Kinetic Energy trends
+!!$      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
+!!$      IF( ln_KE_trd  )   CALL trd_ken( putrd, pvtrd, ktrd, kt )
+!!$
+!!$      !<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
+!!$      !  Vorticity trends
+!!$      !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
+!!$      IF( ln_vor_trd )   CALL trd_vor( putrd, pvtrd, ktrd, kt )
 
       !>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
@@ -88,8 +218,8 @@
 !!gm      IF( ln_dyn_mxl )   CALL trd_mxl_dyn   
       !
-   END SUBROUTINE trd_dyn
-
-
-   SUBROUTINE trd_dyn_iom( putrd, pvtrd, ktrd, kt )
+   END SUBROUTINE trd_dyn_2d
+
+
+   SUBROUTINE trd_dyn_iom_3d( putrd, pvtrd, ktrd, kt )
       !!---------------------------------------------------------------------
       !!                  ***  ROUTINE trd_dyn_iom  ***
@@ -110,6 +240,4 @@
       CASE( jpdyn_hpg )   ;   CALL iom_put( "utrd_hpg", putrd )    ! hydrostatic pressure gradient
                               CALL iom_put( "vtrd_hpg", pvtrd )
-      CASE( jpdyn_spg )   ;   CALL iom_put( "utrd_spg", putrd )    ! surface pressure gradient
-                              CALL iom_put( "vtrd_spg", pvtrd )
       CASE( jpdyn_pvo )   ;   CALL iom_put( "utrd_pvo", putrd )    ! planetary vorticity
                               CALL iom_put( "vtrd_pvo", pvtrd )
@@ -147,36 +275,45 @@
                               CALL iom_put( "vtrd_tau", z2dy )
                               DEALLOCATE( z2dx , z2dy )
-!!gm  to be changed : computation should be done in dynzdf.F90
-!!gm                + missing the top friction 
-!                              !                                    ! bottom stress tends (implicit case)
-!                              IF( ln_drgimp ) THEN
-!                                 ALLOCATE( z3dx(jpi,jpj,jpk) , z3dy(jpi,jpj,jpk) )
-!                             z3dx(:,:,:) = 0._wp   ;   z3dy(:,:,:) = 0._wp  ! after velocity known (now filed at this stage)
-!                            DO jk = 1, jpkm1
-!                                    DO jj = 2, jpjm1
-!                                       DO ji = 2, jpim1
-!                                      ikbu = mbku(ji,jj)          ! deepest ocean u- & v-levels
-!                                          ikbv = mbkv(ji,jj)
-!                                          z3dx(ji,jj,jk) = 0.5 * ( rCdU_bot(ji+1,jj) + rCdU_bot(ji,jj) ) & 
-!                                               &         * un(ji,jj,ikbu) / e3u_n(ji,jj,ikbu)
-!                                          z3dy(ji,jj,jk) = 0.5 * ( rCdU_bot(ji,jj+1) + rCdU_bot(ji,jj) ) &
-!                                               &         * vn(ji,jj,ikbv) / e3v_n(ji,jj,ikbv)
-!                                    END DO
-!                                 END DO
-!                              END DO
-!                              CALL lbc_lnk_multi( 'trddyn', z3dx, 'U', -1., z3dy, 'V', -1. )
-!                              CALL iom_put( "utrd_bfr", z3dx )
-!                              CALL iom_put( "vtrd_bfr", z3dy )
-!                                 DEALLOCATE( z3dx , z3dy )
-!                              ENDIF
-!!gm end
-      CASE( jpdyn_bfr )       ! called if ln_drgimp=F
-                              CALL iom_put( "utrd_bfr", putrd )    ! bottom friction (explicit case)
+      CASE( jpdyn_bfr )   ;   CALL iom_put( "utrd_bfr", putrd )    ! bottom friction (explicit case)
                               CALL iom_put( "vtrd_bfr", pvtrd )
+      CASE( jpdyn_bfri)   ;   CALL iom_put( "utrd_bfri", putrd )    ! bottom friction (implicit case)
+                              CALL iom_put( "vtrd_bfri", pvtrd )
       CASE( jpdyn_atf )   ;   CALL iom_put( "utrd_atf", putrd )        ! asselin filter trends 
                               CALL iom_put( "vtrd_atf", pvtrd )
       END SELECT
       !
-   END SUBROUTINE trd_dyn_iom
+   END SUBROUTINE trd_dyn_iom_3d
+
+
+   SUBROUTINE trd_dyn_iom_2d( putrd, pvtrd, ktrd, kt )
+      !!---------------------------------------------------------------------
+      !!                  ***  ROUTINE trd_dyn_iom  ***
+      !! 
+      !! ** Purpose :   output 2D trends using IOM
+      !!----------------------------------------------------------------------
+      REAL(wp), DIMENSION(:,:), INTENT(inout) ::   putrd, pvtrd   ! U and V trends
+      INTEGER                 , INTENT(in   ) ::   ktrd           ! trend index
+      INTEGER                 , INTENT(in   ) ::   kt             ! time step
+      !
+      INTEGER ::   ji, jj, jk   ! dummy loop indices
+      INTEGER ::   ikbu, ikbv   ! local integers
+      REAL(wp), ALLOCATABLE, DIMENSION(:,:)   ::   z2dx, z2dy   ! 2D workspace 
+      REAL(wp), ALLOCATABLE, DIMENSION(:,:,:) ::   z3dx, z3dy   ! 3D workspace 
+      !!----------------------------------------------------------------------
+      !
+      SELECT CASE( ktrd )
+      CASE( jpdyn_spg )      ;   CALL iom_put( "utrd_spg2d", putrd )      ! surface pressure gradient
+                                 CALL iom_put( "vtrd_spg2d", pvtrd )
+      CASE( jpdyn_pvo )      ;   CALL iom_put( "utrd_pvo2d", putrd )      ! planetary vorticity (barotropic part)
+                                 CALL iom_put( "vtrd_pvo2d", pvtrd )
+      CASE( jpdyn_hpg_corr ) ;   CALL iom_put( "utrd_hpg_corr", putrd )   ! horizontal pressure gradient correction
+                                 CALL iom_put( "vtrd_hpg_corr", pvtrd )
+      CASE( jpdyn_pvo_corr ) ;   CALL iom_put( "utrd_pvo_corr", putrd )   ! planetary vorticity correction
+                                 CALL iom_put( "vtrd_pvo_corr", pvtrd )
+      CASE( jpdyn_bfr )      ;   CALL iom_put( "utrd_bfr2d", putrd )      ! bottom friction due to barotropic currents
+                                 CALL iom_put( "vtrd_bfr2d", pvtrd )
+      END SELECT
+      !
+   END SUBROUTINE trd_dyn_iom_2d
 
    !!======================================================================
