Index: codes/icosagcm/trunk/src/diagnostics/diagflux.F90
===================================================================
--- codes/icosagcm/trunk/src/diagnostics/diagflux.F90	(revision 599)
+++ codes/icosagcm/trunk/src/diagnostics/diagflux.F90	(revision 604)
@@ -10,6 +10,6 @@
        f_massfluxt(:), f_qfluxt(:), & ! time-integrated mass flux and tracer flux
        f_qfluxt_lon(:), f_qfluxt_lat(:), & ! scalar flux reconstructed at cell centers
-       f_epot(:), f_ekin(:), f_enthalpy(:), & ! time-averaged potential E, kinetic E and enthalpy
-       f_epotfluxt(:), f_ekinfluxt(:), f_enthalpyfluxt(:) ! time averaged 'fluxes' of epot, ekin and enthalpy
+       f_ulont(:), f_thetat(:), f_epot(:), f_ekin(:), f_enthalpy(:), & ! time-averaged potential E, kinetic E and enthalpy
+       f_ulonfluxt(:), f_thetafluxt(:), f_epotfluxt(:), f_ekinfluxt(:), f_enthalpyfluxt(:) ! time averaged 'fluxes' of epot, ekin and enthalpy
   LOGICAL :: diagflux_on
   !$OMP THREADPRIVATE(diagflux_on)
@@ -26,4 +26,6 @@
     ll = MERGE(llm,1,diagflux_on)
     CALL allocate_field(f_masst,         field_t,type_real,ll,       name="masst")
+    CALL allocate_field(f_ulont,         field_t,type_real,ll,       name="ulont")
+    CALL allocate_field(f_thetat,        field_t,type_real,ll,       name="thetat")
     CALL allocate_field(f_epot,          field_t,type_real,ll,       name="epot")
     CALL allocate_field(f_ekin,          field_t,type_real,ll,       name="ekin")
@@ -31,4 +33,6 @@
     CALL allocate_field(f_qmasst,        field_t,type_real,ll,nqtot, name="qmasst")
     CALL allocate_field(f_massfluxt,     field_u,type_real,ll,       name="massfluxt")
+    CALL allocate_field(f_ulonfluxt,     field_u,type_real,ll,       name="ulonfluxt")
+    CALL allocate_field(f_thetafluxt,    field_u,type_real,ll,       name="thetafluxt")
     CALL allocate_field(f_epotfluxt,     field_u,type_real,ll,       name="epotfluxt")
     CALL allocate_field(f_ekinfluxt,     field_u,type_real,ll,       name="ekinfluxt")
@@ -50,4 +54,6 @@
        CALL swap_dimensions(ind)
        ZERO2(f_masst)
+       ZERO2(f_ulont)
+       ZERO2(f_thetat)
        ZERO2(f_epot)
        ZERO2(f_ekin)
@@ -55,4 +61,6 @@
        ZERO3(f_qmasst)
        ZERO2(f_massfluxt)
+       ZERO2(f_ulonfluxt)
+       ZERO2(f_thetafluxt)
        ZERO2(f_epotfluxt)
        ZERO2(f_ekinfluxt)
@@ -110,11 +118,11 @@
 !------------------------------------ Compute energy fluxes ---------------------------------------
 
-  SUBROUTINE diagflux_energy(frac, f_phis,f_rhodz,f_theta_rhodz,f_u, f_geopot,f_theta, f_hfluxt)
+  SUBROUTINE diagflux_energy(frac, f_phis,f_rhodz,f_theta_rhodz,f_u, f_geopot,f_theta,f_pk, f_hfluxt)
     REAL(rstd), INTENT(IN) :: frac
-    TYPE(t_field),POINTER :: f_phis(:),f_rhodz(:),f_theta_rhodz(:),f_u(:), f_geopot(:), f_theta(:), f_hfluxt(:)
+    TYPE(t_field),POINTER :: f_phis(:),f_rhodz(:),f_theta_rhodz(:),f_u(:), f_geopot(:), f_theta(:), f_pk(:), f_hfluxt(:)
     REAL(rstd), POINTER :: phis(:), rhodz(:,:), theta_rhodz(:,:,:), u(:,:), &
-         geopot(:,:), pk(:,:,:), hfluxt(:,:), &
-         epot(:,:), ekin(:,:), enthalpy(:,:), &
-         epotflux(:,:), ekinflux(:,:), enthalpyflux(:,:)
+         geopot(:,:), theta(:,:,:), pk(:,:), hfluxt(:,:), &
+         ulont(:,:), thetat(:,:), epot(:,:), ekin(:,:), enthalpy(:,:), &
+         thetaflux(:,:), ulonflux(:,:), epotflux(:,:), ekinflux(:,:), enthalpyflux(:,:)
     INTEGER :: ind
     DO ind=1,ndomain
@@ -128,31 +136,38 @@
        u = f_u(ind)
        geopot = f_geopot(ind)
-       pk = f_theta(ind) ! buffer
+       theta = f_theta(ind) ! buffer
+       pk = f_pk(ind) ! buffer
+       ulont = f_ulont(ind)
+       thetat = f_thetat(ind)
        epot = f_epot(ind)
        ekin = f_ekin(ind)
        enthalpy = f_enthalpy(ind)
+       ulonflux = f_ulonfluxt(ind)
+       thetaflux = f_thetafluxt(ind)
        epotflux = f_epotfluxt(ind)
        ekinflux = f_ekinfluxt(ind)
        enthalpyflux = f_enthalpyfluxt(ind)
-       CALL compute_diagflux_energy(frac,hfluxt, phis,rhodz,theta_rhodz,u, geopot,pk, epot,ekin,enthalpy, epotflux, ekinflux, enthalpyflux)
+       CALL compute_diagflux_energy(frac,hfluxt, phis,rhodz,theta_rhodz,u, geopot,theta,pk, &
+            ulont, thetat, epot, ekin, enthalpy, &
+            ulonflux, thetaflux, epotflux, ekinflux, enthalpyflux)
     END DO
   END SUBROUTINE diagflux_energy
 
-  SUBROUTINE compute_diagflux_energy(frac, massflux, phis,rhodz,theta_rhodz,u, geopot,pk, epot,ekin,enthalpy, epot_flux, ekin_flux, enthalpy_flux)
+  SUBROUTINE compute_diagflux_energy(frac, massflux, phis,rhodz,theta_rhodz,ue, geopot,theta,pk, &
+       ulon, thetat, epot, ekin, enthalpy, &
+       ulon_flux, thetat_flux, epot_flux, ekin_flux, enthalpy_flux)
     USE disvert_mod, ONLY : ptop
     REAL(rstd), INTENT(IN) :: frac
-    REAL(rstd), INTENT(IN) :: massflux(3*iim*jjm,llm), u(3*iim*jjm,llm),&
+    REAL(rstd), INTENT(IN) :: massflux(3*iim*jjm,llm), ue(3*iim*jjm,llm),&
                               phis(iim*jjm), rhodz(iim*jjm,llm), theta_rhodz(iim*jjm,llm,nqtot)
-    REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm,llm+1), pk(iim*jjm,llm) ! pk = buffer
-    REAL(rstd), INTENT(INOUT), DIMENSION(iim*jjm, llm)   ::  epot, ekin, enthalpy
-    REAL(rstd), INTENT(INOUT), DIMENSION(3*iim*jjm, llm) ::  epot_flux, ekin_flux, enthalpy_flux    
-    REAL(rstd) :: energy, p_ik, theta_ik, temp_ik, gv, Rd 
+    REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm,llm+1), theta(iim*jjm,llm), pk(iim*jjm,llm) ! theta,pk = buffers
+    REAL(rstd), INTENT(INOUT), DIMENSION(iim*jjm, llm)   ::  ulon, thetat, epot, ekin, enthalpy
+    REAL(rstd), INTENT(INOUT), DIMENSION(3*iim*jjm, llm) ::  ulon_flux, thetat_flux, epot_flux, ekin_flux, enthalpy_flux    
+    REAL(rstd) :: energy, p_ik, theta_ik, temp_ik, gv, Rd, cx,cy,cz, ux,uy,uz, ue_le,ulon_i
     INTEGER :: ij, l, ij_omp_begin_ext, ij_omp_end_ext
     Rd = kappa*cpp
     ! even if loops are of the _ext variant, we still need halo exchanges before reconstructing fluxes at cell centers
     ! => loop over interior region
-    CALL distrib_level(ij_end-ij_begin+1,ij_omp_begin_ext,ij_omp_end_ext)
-    ij_omp_begin_ext = ij_omp_begin_ext+ij_begin-1
-    ij_omp_end_ext = ij_omp_end_ext+ij_begin-1
+    CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)
 #include "../kernels/energy_fluxes.k90"
   END SUBROUTINE compute_diagflux_energy
