Index: /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_opa_sponge.F90
===================================================================
--- /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_opa_sponge.F90	(revision 8834)
+++ /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_opa_sponge.F90	(revision 8835)
@@ -1,3 +1,3 @@
-#define SPONGE && define SPONGE_TOP
+#undef SPONGE && undef SPONGE_TOP
 
 MODULE agrif_opa_sponge
@@ -59,5 +59,4 @@
       REAL(wp) :: timecoeff
       !!---------------------------------------------
-
 #if defined SPONGE
       timecoeff = REAL(Agrif_NbStepint(),wp)/Agrif_rhot()
@@ -203,12 +202,76 @@
       INTEGER  ::   iku, ikv
       REAL(wp) :: ztsa, zabe1, zabe2, zbtr
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2) :: ztu, ztv
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2,n1:n2) ::tsbdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk) :: ztu, ztv
+      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk,n1:n2) ::tsbdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,jpk,n1:n2) ::tabres_child
+      REAL(wp), DIMENSION(k1:k2,n1:n2-1) :: tabin
+      REAL(wp) :: h_in(k1:k2)
+      REAL(wp) :: h_out(1:jpk)
+      INTEGER :: N_in, N_out
+      REAL(wp) :: h_diff, zrhoxy
       !
       IF( before ) THEN
-         tabres(i1:i2,j1:j2,k1:k2,n1:n2) = tsn(i1:i2,j1:j2,k1:k2,n1:n2)
+         DO jn = n1,n2-1
+            DO jk=k1,k2
+               DO jj=j1,j2
+                  DO ji=i1,i2
+                     tabres(ji,jj,jk,jn) = tsb(ji,jj,jk,jn) * e1e2t(ji,jj) * e3t_n(ji,jj,jk)
+                  END DO
+               END DO
+            END DO
+         END DO
+         DO jk=k1,k2
+            DO jj=j1,j2
+               DO ji=i1,i2
+                   tabres(ji,jj,jk,n2) = tmask(ji,jj,jk) * e1e2t(ji,jj) * e3t_n(ji,jj,jk) 
+               END DO
+            END DO
+         END DO
       ELSE   
          !
-         tsbdiff(:,:,:,:) = tsb(i1:i2,j1:j2,:,:) - tabres(:,:,:,:)    
+#ifdef key_vertical
+         tabres_child(:,:,:,:) = 0.
+         zrhoxy = Agrif_rhox()*Agrif_rhoy()
+         DO jj=j1,j2
+           DO ji=i1,i2
+             N_in = 0
+             DO jk=k1,k2 !k2 = jpk of parent grid
+               IF (tabres(ji,jj,jk,n2) == 0) EXIT
+               N_in = N_in + 1
+               tabin(jk,:) = tabres(ji,jj,jk,n1:n2-1)/tabres(ji,jj,jk,n2)
+               h_in(N_in) = tabres(ji,jj,jk,n2)/(e1e2t(ji,jj)*zrhoxy)
+             END DO
+             N_out = 0
+             DO jk=1,jpk ! jpk of child grid
+               IF (tmask(ji,jj,jk) == 0) EXIT ! TODO: Will not work with ISF. !This doesn't seem to work at the moment in GYRE but is OK in overflow model
+               N_out = N_out + 1
+               h_out(jk) = e3t_n(ji,jj,jk) !Child grid scale factors. Could multiply by e1e2t here instead of division above
+             ENDDO
+             IF (N_in > 0) THEN
+                h_diff = sum(h_out(1:N_out))-sum(h_in(1:N_in))
+                tabres(ji,jj,k2,:) = tabres(ji,jj,k2-1,:) !what is this line for?????
+                DO jn=1,jpts
+                   call reconstructandremap(tabin(1:N_in,jn),h_in,tabres_child(ji,jj,1:N_out,jn),h_out,N_in,N_out)
+                ENDDO
+             ENDIF
+             DO jk=1,jpk
+                tsbdiff(ji,jj,jk,n1:n2-1) = tsb(ji, jj,jk,n1:n2-1) - tabres_child(ji,jj,jk,n1:n2-1)
+             ENDDO
+           ENDDO
+         ENDDO
+#else
+         do jk=k1,k2
+            do jj=j1,j2
+               do ji=i1,i2
+                 ! This will be slow - Is there a way to speed it up and avoid divide by zero?
+                 IF (tabres(ji,jj,jk,n2) .NE. 0) THEN
+                    tsbdiff(ji,jj,jk,n1:n2-1) = tsb(ji,jj,jk,n1:n2) - tabres(ji,jj,jk,n1:n2-1)/tabres(ji,jj,jk,n2)
+                 ELSE
+                    tsbdiff(ji,jj,jk,n1:n2-1) = 0._wp
+                 ENDIF
+               enddo
+            enddo
+         enddo
+#endif
          DO jn = 1, jpts            
             DO jk = 1, jpkm1
@@ -258,25 +321,87 @@
 
 
-   SUBROUTINE interpun_sponge(tabres,i1,i2,j1,j2,k1,k2, before)
+   SUBROUTINE interpun_sponge(tabres,i1,i2,j1,j2,k1,k2,m1,m2, before)
       !!---------------------------------------------
       !!   *** ROUTINE interpun_sponge ***
       !!---------------------------------------------    
-      INTEGER, INTENT(in) :: i1,i2,j1,j2,k1,k2
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2), INTENT(inout) :: tabres
+      INTEGER, INTENT(in) :: i1,i2,j1,j2,k1,k2,m1,m2
+      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2,m1:m2), INTENT(inout) :: tabres
       LOGICAL, INTENT(in) :: before
 
-      INTEGER :: ji,jj,jk
+      INTEGER :: ji,jj,jk,N_in,N_out
 
       ! sponge parameters 
-      REAL(wp) :: ze2u, ze1v, zua, zva, zbtr
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2) :: ubdiff
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2) :: rotdiff, hdivdiff
+      REAL(wp) :: ze2u, ze1v, zua, zva, zbtr, h_diff, zrhoy
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: ubdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: rotdiff, hdivdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: tabres_child
+      REAL(wp), DIMENSION(k1:k2) :: tabin, h_in
+      REAL(wp), DIMENSION(1:jpk) :: h_out
       INTEGER :: jmax
       !!---------------------------------------------    
       !
       IF( before ) THEN
-         tabres = un(i1:i2,j1:j2,:)
+#ifdef key_vertical
+         DO jk=1,jpk
+            DO jj=j1,j2
+               DO ji=i1,i2
+                  tabres(ji,jj,jk,m1) = (e2u(ji,jj) * e3u_n(ji,jj,jk) * un(ji,jj,jk)*umask(ji,jj,jk)) 
+                  tabres(ji,jj,jk,m2) = (umask(ji,jj,jk) * e2u(ji,jj) * e3u_n(ji,jj,jk))
+               END DO
+            END DO
+         END DO
+
+#else
+         tabres(i1:i2,j1:j2,:,m1) = un(i1:i2,j1:j2,:)
+#endif
       ELSE
+#ifdef key_vertical
+         tabres_child(:,:,:) = 0._wp
+         zrhoy = Agrif_rhoy()
+         DO jj=j1,j2
+            DO ji=i1,i2
+               N_in = 0
+               DO jk=k1,k2
+                  IF (tabres(ji,jj,jk,m2) == 0) EXIT
+                  N_in = N_in + 1
+                  tabin(jk) = tabres(ji,jj,jk,m1)/tabres(ji,jj,jk,m2)
+                  h_in(N_in) = tabres(ji,jj,jk,m2)/(e2u(ji,jj)*zrhoy) 
+              ENDDO
+              !
+              IF (N_in == 0) THEN
+                 tabres_child(ji,jj,:) = 0.
+                 CYCLE
+              ENDIF
+         
+              N_out = 0
+              DO jk=1,jpk
+                 if (umask(ji,jj,jk) == 0) EXIT
+                 N_out = N_out + 1
+                 h_out(N_out) = e3u_n(ji,jj,jk)
+              ENDDO
+         
+              IF (N_out == 0) THEN
+                 tabres_child(ji,jj,:) = 0.
+                 CYCLE
+              ENDIF
+         
+              IF (N_in * N_out > 0) THEN
+                 h_diff = sum(h_out(1:N_out))-sum(h_in(1:N_in))
+                 if (h_diff < -1.e4) then
+                    print *,'CHECK YOUR BATHY ...', h_diff, sum(h_out(1:N_out)), sum(h_in(1:N_in))
+                 endif
+              ENDIF
+              call reconstructandremap(tabin(1:N_in),h_in(1:N_in),tabres_child(ji,jj,1:N_out),h_out(1:N_out),N_in,N_out)
+         
+            ENDDO
+         ENDDO
+         DO jk = 1, jpkm1
+            DO jj=j1,j2
+               ubdiff(i1:i2,jj,jk) = (ub(i1:i2,jj,jk) - tabres_child(i1:i2,jj,jk))*umask(i1:i2,jj,jk)
+            END DO
+         END DO         
+#else
          ubdiff(i1:i2,j1:j2,:) = (ub(i1:i2,j1:j2,:) - tabres(:,:,:))*umask(i1:i2,j1:j2,:)
+#endif
          !
          DO jk = 1, jpkm1                                 ! Horizontal slab
@@ -356,25 +481,80 @@
 
 
-   SUBROUTINE interpvn_sponge(tabres,i1,i2,j1,j2,k1,k2, before,nb,ndir)
+   SUBROUTINE interpvn_sponge(tabres,i1,i2,j1,j2,k1,k2,m1,m2, before,nb,ndir)
       !!---------------------------------------------
       !!   *** ROUTINE interpvn_sponge ***
       !!--------------------------------------------- 
-      INTEGER, INTENT(in) :: i1,i2,j1,j2,k1,k2
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2), INTENT(inout) :: tabres
+      INTEGER, INTENT(in) :: i1,i2,j1,j2,k1,k2,m1,m2
+      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2,m1:m2), INTENT(inout) :: tabres
       LOGICAL, INTENT(in) :: before
       INTEGER, INTENT(in) :: nb , ndir
       !
-      INTEGER  ::   ji, jj, jk
-      REAL(wp) ::   ze2u, ze1v, zua, zva, zbtr
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2) :: vbdiff
-      REAL(wp), DIMENSION(i1:i2,j1:j2,k1:k2) :: rotdiff, hdivdiff
+      INTEGER  ::   ji, jj, jk, N_in, N_out
+      REAL(wp) ::   ze2u, ze1v, zua, zva, zbtr, h_diff, zrhox
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: vbdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: rotdiff, hdivdiff
+      REAL(wp), DIMENSION(i1:i2,j1:j2,1:jpk) :: tabres_child
+      REAL(wp), DIMENSION(k1:k2) :: tabin, h_in
+      REAL(wp), DIMENSION(1:jpk) :: h_out
       INTEGER :: imax
       !!--------------------------------------------- 
-
+      
       IF( before ) THEN 
+#ifdef key_vertical
+         DO jk=1,jpk
+            DO jj=j1,j2
+               DO ji=i1,i2
+                  tabres(ji,jj,jk,m1) = (e1v(ji,jj) * e3v_n(ji,jj,jk) * vn(ji,jj,jk)*vmask(ji,jj,jk))
+                  tabres(ji,jj,jk,m2) = (vmask(ji,jj,jk) * e1v(ji,jj) * e3v_n(ji,jj,jk))
+               END DO
+            END DO
+         END DO
+#else
          tabres = vn(i1:i2,j1:j2,:)
+#endif
       ELSE
+#ifdef key_vertical
+         zrhox = Agrif_rhox()
+         tabres_child(:,:,:) = 0._wp
+         DO jj=j1,j2
+            DO ji=i1,i2
+               N_in = 0
+               DO jk=k1,k2
+                  IF (tabres(ji,jj,jk,m2) == 0) EXIT
+                  N_in = N_in + 1
+                  tabin(jk) = tabres(ji,jj,jk,m1)/tabres(ji,jj,jk,m2)
+                  h_in(N_in) = tabres(ji,jj,jk,m2)/(e1v(ji,jj)*zrhox) 
+              ENDDO
+         
+              IF (N_in == 0) THEN
+                 tabres_child(ji,jj,:) = 0.
+                 CYCLE
+              ENDIF
+         
+              N_out = 0
+              DO jk=1,jpk
+                 if (umask(ji,jj,jk) == 0) EXIT
+                 N_out = N_out + 1
+                 h_out(N_out) = e3v_n(ji,jj,jk)
+              ENDDO
+         
+              IF (N_in * N_out > 0) THEN
+                 h_diff = sum(h_out(1:N_out))-sum(h_in(1:N_in))
+                 if (h_diff < -1.e4) then
+                    print *,'CHECK YOUR BATHY ...', h_diff, sum(h_out(1:N_out)), sum(h_in(1:N_in))
+                 endif
+              ENDIF
+              call reconstructandremap(tabin(1:N_in),h_in(1:N_in),tabres_child(ji,jj,1:N_out),h_out(1:N_out),N_in,N_out)
+            ENDDO
+         ENDDO
+         DO jk = 1, jpkm1
+            DO jj=j1,j2
+               vbdiff(i1:i2,jj,jk) = (vb(i1:i2,jj,jk) - tabres_child(i1:i2,jj,jk))*vmask(i1:i2,jj,jk)
+            END DO
+         END DO         
+#else
          !
          vbdiff(i1:i2,j1:j2,:) = (vb(i1:i2,j1:j2,:) - tabres(:,:,:))*vmask(i1:i2,j1:j2,:)
+#endif
          !
          DO jk = 1, jpkm1                                 ! Horizontal slab
Index: /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_user.F90
===================================================================
--- /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_user.F90	(revision 8834)
+++ /branches/2017/dev_r8126_UKMO_AGRIF_vert_interp/NEMOGCM/NEMO/NST_SRC/agrif_user.F90	(revision 8835)
@@ -347,5 +347,5 @@
    !---------------------------------------------------------------------
    CALL agrif_declare_variable((/2,2,0,0/),(/3,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,jpts+1/),tsn_id)
-   CALL agrif_declare_variable((/2,2,0,0/),(/3,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,jpts/),tsn_sponge_id)
+   CALL agrif_declare_variable((/2,2,0,0/),(/3,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,jpts+1/),tsn_sponge_id)
 
    CALL agrif_declare_variable((/1,2,0,0/),(/2,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,2/),un_interp_id)
@@ -353,6 +353,6 @@
    CALL agrif_declare_variable((/1,2,0,0/),(/2,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,2/),un_update_id)
    CALL agrif_declare_variable((/2,1,0,0/),(/3,2,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,2/),vn_update_id)
-   CALL agrif_declare_variable((/1,2,0/),(/2,3,0/),(/'x','y','N'/),(/1,1,1/),(/nlci,nlcj,jpk/),un_sponge_id)
-   CALL agrif_declare_variable((/2,1,0/),(/3,2,0/),(/'x','y','N'/),(/1,1,1/),(/nlci,nlcj,jpk/),vn_sponge_id)
+   CALL agrif_declare_variable((/1,2,0,0/),(/2,3,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,2/),un_sponge_id)
+   CALL agrif_declare_variable((/2,1,0,0/),(/3,2,0,0/),(/'x','y','N','N'/),(/1,1,1,1/),(/nlci,nlcj,jpk,2/),vn_sponge_id)
 
    CALL agrif_declare_variable((/2,2,0/),(/3,3,0/),(/'x','y','N'/),(/1,1,1/),(/nlci,nlcj,jpk/),e3t_id)
