From 694f538bcfd695ad730d103d9fdfd2d526d4c76a Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Thu, 23 Jan 2020 13:05:41 +0100
Subject: [PATCH 1/7] Add psi computations

---
 diagonals/psi.py | 121 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 diagonals/psi.py

diff --git a/diagonals/psi.py b/diagonals/psi.py
new file mode 100644
index 0000000..2a4b068
--- /dev/null
+++ b/diagonals/psi.py
@@ -0,0 +1,121 @@
+import os
+
+import numpy as np
+
+import iris
+import iris.cube
+import iris.analysis
+
+import warnings
+import datetime
+
+import numba
+from numba import njit
+from numba import vectorize
+from numba import cuda
+from numba import int32, int64, float32, float64
+from numba.cuda.cudadrv import driver
+
+import diagonals
+
+__all__ = ['compute']
+
+def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
+    """Function that checks device and calls computing functions.
+
+    Checks if the computations are going performed in the CPU or the GPU:
+
+    Parameters
+    ----------
+    basins : list
+        List of masked arrays containing the mask for each basin.
+    e1u: float32
+        Masked array containing variable e1u.
+    e1v: float32
+        Masked array containing variable e1v.
+    e3u: float32
+        Masked array containing variable e3u.
+    e3v: float32
+        Masked array containing variable e3v.
+    uo : float32
+        Masked array containing Sea Water X Velocity data.
+    vo : float32
+        Masked array containing Sea Water Y Velocity data.
+
+    Returns
+    -------
+    vsftbarot: list
+        List of masked arrays containing the gyre strength.
+    """
+    area_u = {}
+    area_v = {}
+    for basin in basins:
+        area_u[basin] = _compute_area(e2u, e3u, basins[basin])
+        area_v[basin] = _compute_area(e1v, e3v, basins[basin])
+
+    del e2u, e1v, e3u, e3v
+
+    vsftbarot = _compute_psi_cpu(
+        uo, vo, area_u, area_v, basins
+    )
+
+    return vsftbarot
+
+def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
+
+    psi = {}
+    for basin, mask in area_u.items():
+        uo_lev = np.sum(uo*area_u[basin], axis=1)
+        vo_lev = np.sum(vo*area_v[basin], axis=1)
+
+        del uo, vo, area_u[basin], area_v[basin]
+
+        dpsiu = _u_cumsum(uo_lev)
+        dpsiv = _v_cumsum(vo_lev)
+
+        dpsi = 0.5 * (dpsiu + dpsiv)
+
+        ref_point = dpsi[:,-1,-1]
+
+        psi[basin] = (
+                       dpsi - ref_point[:, np.newaxis, np.newaxis]
+                     ) * basins[basin]
+
+    return psi
+
+#@njit looks like it's not worth it
+def _u_cumsum(uo):
+    dpsiu = np.zeros(uo.shape)
+    for j in range(1, dpsiu.shape[1]):
+       dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :]
+    return dpsiu
+
+#@njit looks like it's not worth it
+def _v_cumsum(vo):
+    dpsiv = np.zeros(vo.shape)
+    for i in range(dpsiv.shape[2]-2,-1,-1):
+        dpsiv[:, : ,i] = dpsiv[:, :, i+1] - vo[: , :, i]
+    return dpsiv
+
+@vectorize(['float32(float32, float32, float32)'], target='cpu')
+def _compute_area(e, e3, basin):
+    """Vectorized numba function executed in the CPU.
+
+    Calculates cell area for each basin:
+
+    Parameters
+    ----------
+    e1: float32
+        Masked array containing variable e1 or e2.
+    e3: float32
+        Masked array containing variable e3.
+    basin : float32
+        Masked array containing a basin mask.
+
+    Returns
+    -------
+    area: float32
+        Masked array containing the cell area for a given basin.
+    """
+    area = e * e3
+    return area
-- 
GitLab


From 53adb891bab8589edb935aa5bb90ac3919ac5e49 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Thu, 23 Jan 2020 13:24:12 +0100
Subject: [PATCH 2/7] Fix lint

---
 diagonals/psi.py     | 19 +++++++++++++------
 diagonals/regmean.py |  2 +-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/diagonals/psi.py b/diagonals/psi.py
index 2a4b068..cc4342b 100644
--- a/diagonals/psi.py
+++ b/diagonals/psi.py
@@ -20,6 +20,7 @@ import diagonals
 
 __all__ = ['compute']
 
+
 def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
     """Function that checks device and calls computing functions.
 
@@ -61,6 +62,7 @@ def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
 
     return vsftbarot
 
+
 def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
 
     psi = {}
@@ -75,7 +77,7 @@ def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
 
         dpsi = 0.5 * (dpsiu + dpsiv)
 
-        ref_point = dpsi[:,-1,-1]
+        ref_point = dpsi[:, -1, -1]
 
         psi[basin] = (
                        dpsi - ref_point[:, np.newaxis, np.newaxis]
@@ -83,20 +85,25 @@ def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
 
     return psi
 
-#@njit looks like it's not worth it
+# @njit looks like it's not worth it
+
+
 def _u_cumsum(uo):
     dpsiu = np.zeros(uo.shape)
     for j in range(1, dpsiu.shape[1]):
-       dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :]
+        dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :]
     return dpsiu
 
-#@njit looks like it's not worth it
+# @njit looks like it's not worth it
+
+
 def _v_cumsum(vo):
     dpsiv = np.zeros(vo.shape)
-    for i in range(dpsiv.shape[2]-2,-1,-1):
-        dpsiv[:, : ,i] = dpsiv[:, :, i+1] - vo[: , :, i]
+    for i in range(dpsiv.shape[2]-2, -1, -1):
+        dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[: ,: , i]
     return dpsiv
 
+
 @vectorize(['float32(float32, float32, float32)'], target='cpu')
 def _compute_area(e, e3, basin):
     """Vectorized numba function executed in the CPU.
diff --git a/diagonals/regmean.py b/diagonals/regmean.py
index 3bd88aa..1cdaced 100644
--- a/diagonals/regmean.py
+++ b/diagonals/regmean.py
@@ -193,7 +193,7 @@ def _compute_regmean_levels_cpu(var, basins, volume):
         for t in range(times):
             for l in range(levs):
                 regmean[t, l] = np.ma.average(var[t, l, :, :], axis=(0, 1),
-                                           weights=np.squeeze(w)[l, :, :])
+                                              weights=np.squeeze(w)[l, :, :])
         regmean_total[basin] = regmean
     return regmean_total
 
-- 
GitLab


From 83ec0ffcf56019fd827a0097075748132e04e4fa Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Thu, 23 Jan 2020 13:30:14 +0100
Subject: [PATCH 3/7] Fix lint

---
 diagonals/psi.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diagonals/psi.py b/diagonals/psi.py
index cc4342b..d1b1f4a 100644
--- a/diagonals/psi.py
+++ b/diagonals/psi.py
@@ -100,7 +100,7 @@ def _u_cumsum(uo):
 def _v_cumsum(vo):
     dpsiv = np.zeros(vo.shape)
     for i in range(dpsiv.shape[2]-2, -1, -1):
-        dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[: ,: , i]
+        dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[:, :, i]
     return dpsiv
 
 
-- 
GitLab


From bff9901f0ed21d0b83288c36d23eab7024825d44 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Thu, 23 Jan 2020 16:28:09 +0100
Subject: [PATCH 4/7] Remove mask from area computaton

---
 diagonals/psi.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/diagonals/psi.py b/diagonals/psi.py
index d1b1f4a..c1712e8 100644
--- a/diagonals/psi.py
+++ b/diagonals/psi.py
@@ -48,11 +48,10 @@ def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
     vsftbarot: list
         List of masked arrays containing the gyre strength.
     """
-    area_u = {}
-    area_v = {}
-    for basin in basins:
-        area_u[basin] = _compute_area(e2u, e3u, basins[basin])
-        area_v[basin] = _compute_area(e1v, e3v, basins[basin])
+
+
+    area_u = _compute_area(e2u, e3u)
+    area_v = _compute_area(e1v, e3v)
 
     del e2u, e1v, e3u, e3v
 
@@ -66,11 +65,12 @@ def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
 def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
 
     psi = {}
-    for basin, mask in area_u.items():
-        uo_lev = np.sum(uo*area_u[basin], axis=1)
-        vo_lev = np.sum(vo*area_v[basin], axis=1)
+    uo_lev = np.sum(uo*area_u, axis=1)
+    vo_lev = np.sum(vo*area_v, axis=1)
+
+    del uo, vo, area_u, area_v
 
-        del uo, vo, area_u[basin], area_v[basin]
+    for basin, mask in basins.items():
 
         dpsiu = _u_cumsum(uo_lev)
         dpsiv = _v_cumsum(vo_lev)
@@ -81,7 +81,7 @@ def _compute_psi_cpu(uo, vo, area_u, area_v, basins):
 
         psi[basin] = (
                        dpsi - ref_point[:, np.newaxis, np.newaxis]
-                     ) * basins[basin]
+                     ) * mask
 
     return psi
 
@@ -104,8 +104,8 @@ def _v_cumsum(vo):
     return dpsiv
 
 
-@vectorize(['float32(float32, float32, float32)'], target='cpu')
-def _compute_area(e, e3, basin):
+@vectorize(['float32(float32, float32)'], target='cpu')
+def _compute_area(e, e3):
     """Vectorized numba function executed in the CPU.
 
     Calculates cell area for each basin:
-- 
GitLab


From 997d2a0fcec921e4646d8cbd8b6ed46104a0d72c Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Thu, 23 Jan 2020 16:39:13 +0100
Subject: [PATCH 5/7] Fix lint issues

---
 diagonals/psi.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/diagonals/psi.py b/diagonals/psi.py
index c1712e8..1bdb91e 100644
--- a/diagonals/psi.py
+++ b/diagonals/psi.py
@@ -48,8 +48,6 @@ def compute(basins, e2u, e1v, e3u, e3v, uo, vo):
     vsftbarot: list
         List of masked arrays containing the gyre strength.
     """
-
-
     area_u = _compute_area(e2u, e3u)
     area_v = _compute_area(e1v, e3v)
 
-- 
GitLab


From a3b4ba1443092a86b2af5b02d8d587c619db7e7c Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Thu, 30 Jan 2020 13:25:00 +0100
Subject: [PATCH 6/7] Zonal mean 2D and flake8 adoption

---
 .gitlab-ci.yml                                | 12 +----
 .../examples/examples-oceanheatcontent.py     | 12 ++---
 diagonals/examples/examples-regionmean.py     | 18 +++----
 diagonals/examples/examples-regsum.py         | 19 +++----
 diagonals/examples/examples-siarea.py         |  2 +-
 diagonals/mesh_helpers/nemo.py                | 14 +++--
 diagonals/moc.py                              | 20 ++-----
 diagonals/ohc.py                              | 30 ++++-------
 diagonals/psi.py                              | 18 +------
 diagonals/regmean.py                          | 40 ++++++--------
 diagonals/regsum.py                           | 41 ++++++--------
 diagonals/siasie.py                           | 15 ------
 diagonals/zonmean.py                          | 42 +++++++++------
 setup.cfg                                     | 22 ++++++++
 setup.py                                      | 53 +++++--------------
 test/unit/test_lint.py                        | 34 ------------
 test/unit/test_ohc.py                         |  3 --
 17 files changed, 142 insertions(+), 253 deletions(-)
 create mode 100644 setup.cfg
 delete mode 100644 test/unit/test_lint.py

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index bade64f..73e3443 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -17,17 +17,7 @@ prepare:
   script:
     - conda update conda
 
-test_python2:
-  stage: test
-  script:
-    - git submodule sync --recursive
-    - git submodule update --init --recursive
-    - conda env update -f environment.yml -n diagonals2 python=2.7
-    - source activate diagonals2
-    - pip install -e .[test]
-    - python setup.py test
-
-test_python3:
+test_python:
  stage: test
  script:
    - git submodule sync --recursive
diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py
index 867d818..0f270a2 100644
--- a/diagonals/examples/examples-oceanheatcontent.py
+++ b/diagonals/examples/examples-oceanheatcontent.py
@@ -79,19 +79,19 @@ def load_thetao():
     return thetao_data
 
 
-def save_data(layers, basins, ohc_2D, ohc1D, device):
+def save_data(layers, basins, ohc_2d, ohc_1d, device):
     ohc_cube = []
-    logger.info('ohc1d length is %s', len(ohc1D))
+    logger.info('ohc1d length is %s', len(ohc_1d))
     for i, layer in enumerate(layers):
-        ohc_cube.append(iris.cube.Cube(ohc_2D[i],
+        ohc_cube.append(iris.cube.Cube(ohc_2d[i],
                                        long_name='Ocean heat content'
                                                  ' {0[0]} to {0[1]} meters'
                                        .format(layer)))
-        ohc_1D = []
+        ohc_1s = []
         for j, basin in enumerate(basins):
-            ohc_1D.append(iris.cube.Cube(ohc1D[j][:],
+            ohc_1s.append(iris.cube.Cube(ohc_1d[j][:],
                           long_name='{0}'.format(basin)))
-        iris.save(ohc_1D,
+        iris.save(ohc_1s,
                   '/esarchive/scratch/sloosvel/numba_outputs'
                   '/ohc_1D_{1}_{0[0]}_{0[1]}.nc'
                   .format(layer, device), zlib=True)
diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py
index 4f18cb1..7c69405 100644
--- a/diagonals/examples/examples-regionmean.py
+++ b/diagonals/examples/examples-regionmean.py
@@ -33,13 +33,13 @@ def main():
         areacello = mesh.get_areacello()
         tos = load_tos()
         basins = load_masks()
-        tosmean = regmean.compute_regmean_2D(tos, basins, areacello)
+        tosmean = regmean.compute_regmean_2d(tos, basins, areacello)
         volcello = mesh.get_volcello()
         thetao = load_thetao()
         thetaomean = regmean.compute_regmean_levels(thetao, basins, volcello)
-        thetaomean3D = regmean.compute_regmean_3D(thetao, basins, volcello)
+        thetaomean3d = regmean.compute_regmean_3d(thetao, basins, volcello)
         device = 'CPU'
-        save_cubes(tosmean, thetaomean, thetaomean3D, basins, device)
+        save_cubes(tosmean, thetaomean, thetaomean3d, basins, device)
     ellapsed = datetime.datetime.now() - start
     logger.info('Total ellapsed time on the %s: %s', device, ellapsed)
 
@@ -70,28 +70,28 @@ def load_thetao():
     return thetao_data
 
 
-def save_cubes(tosmean, thetaomean, thetaomean3D, basins, device):
+def save_cubes(tosmean, thetaomean, thetaomean_3d, basins, device):
     cubes_tosmean = iris.cube.CubeList()
     cubes_thetaomean = iris.cube.CubeList()
-    cubes_thetaomean3D = iris.cube.CubeList()
+    cubes_thetaomean_3d = iris.cube.CubeList()
     for basin in basins.keys():
         cube_tosmean = iris.cube.Cube(tosmean[basin])
         cube_thetaomean = iris.cube.Cube(thetaomean[basin])
-        cube_thetaomean3D = iris.cube.Cube(thetaomean3D[basin])
+        cube_thetaomean_3d = iris.cube.Cube(thetaomean_3d[basin])
 
         cube_tosmean.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
         cube_thetaomean.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
-        cube_thetaomean3D.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
+        cube_thetaomean_3d.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
 
         cubes_tosmean.append(cube_tosmean)
         cubes_thetaomean.append(cube_thetaomean)
-        cubes_thetaomean3D.append(cube_thetaomean3D)
+        cubes_thetaomean_3d.append(cube_thetaomean_3d)
 
     iris.save(cubes_tosmean.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/tosmean_{0}.nc'.format(device), zlib=True)
     iris.save(cubes_thetaomean.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/thetaomean_{0}.nc'.format(device), zlib=True)
-    iris.save(cubes_thetaomean3D.merge_cube(), '/esarchive/scratch/sloosvel/'
+    iris.save(cubes_thetaomean_3d.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/thetaomean3D_{0}.nc'.format(device), zlib=True)
 
 
diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py
index 6a49975..6f49265 100644
--- a/diagonals/examples/examples-regsum.py
+++ b/diagonals/examples/examples-regsum.py
@@ -32,15 +32,16 @@ def main():
         areacello = mesh.get_areacello()
         tos = load_tos()
         basins = load_masks()
-        tossum = regsum.compute_regsum_2D(tos, basins, areacello)
+        tossum = regsum.compute_regsum_2d(tos, basins, areacello)
         volcello = mesh.get_volcello()
         tmask = mesh.get_landsea_mask()
         thetao = load_thetao()
         thetaosum = regsum.compute_regsum_levels(thetao, basins,
                                                  volcello, tmask)
-        thetaosum3D = regsum.compute_regsum_3D(thetao, basins, volcello, tmask)
+        thetaosum_3d = regsum.compute_regsum_3d(
+            thetao, basins, volcello, tmask)
         device = 'CPU'
-        save_cubes(tossum, thetaosum, thetaosum3D, basins, device)
+        save_cubes(tossum, thetaosum, thetaosum_3d, basins, device)
     ellapsed = datetime.datetime.now() - start
     logger.info('Total ellapsed time on the %s: %s', device, ellapsed)
 
@@ -69,28 +70,28 @@ def load_thetao():
     return thetao_data
 
 
-def save_cubes(tossum, thetaosum, thetaosum3D, basins, device):
+def save_cubes(tossum, thetaosum, thetaosum3d, basins, device):
     cubes_tossum = iris.cube.CubeList()
     cubes_thetaosum = iris.cube.CubeList()
-    cubes_thetaosum3D = iris.cube.CubeList()
+    cubes_thetaosum3d = iris.cube.CubeList()
     for basin in basins.items():
         cube_tossum = iris.cube.Cube(tossum[basin])
         cube_thetaosum = iris.cube.Cube(thetaosum[basin])
-        cube_thetaosum3D = iris.cube.Cube(thetaosum3D[basin])
+        cube_thetaosum3d = iris.cube.Cube(thetaosum3d[basin])
 
         cube_tossum.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
         cube_thetaosum.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
-        cube_thetaosum3D.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
+        cube_thetaosum3d.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
 
         cubes_tossum.append(cube_tossum)
         cubes_thetaosum.append(cube_thetaosum)
-        cubes_thetaosum3D.append(cube_thetaosum)
+        cubes_thetaosum3d.append(cube_thetaosum)
 
     iris.save(cubes_tossum.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/tossum_{0}.nc'.format(device), zlib=True)
     iris.save(cubes_thetaosum.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/thetaosum_{0}.nc'.format(device), zlib=True)
-    iris.save(cubes_thetaosum3D.merge_cube(), '/esarchive/scratch/sloosvel/'
+    iris.save(cubes_thetaosum3d.merge_cube(), '/esarchive/scratch/sloosvel/'
               'numba_outputs/thetaosum3D_{0}.nc'.format(device), zlib=True)
 
 
diff --git a/diagonals/examples/examples-siarea.py b/diagonals/examples/examples-siarea.py
index 87af4dc..c558119 100644
--- a/diagonals/examples/examples-siarea.py
+++ b/diagonals/examples/examples-siarea.py
@@ -7,7 +7,7 @@ import iris.cube
 import iris.analysis
 import iris.coords
 import iris.coord_categorisation
-import iris.analysis
+from iris.experimental.equalise_cubes import equalise_attributes
 
 import diagonals
 import diagonals.siasie as siasie
diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py
index f9d4878..782a815 100644
--- a/diagonals/mesh_helpers/nemo.py
+++ b/diagonals/mesh_helpers/nemo.py
@@ -1,7 +1,7 @@
 import numpy as np
 import iris
-
-from numba import vectorize, float32
+from iris.exceptions import ConstraintMismatchError
+from numba import vectorize, float32 # noqa
 
 
 class Nemo():
@@ -29,12 +29,18 @@ class Nemo():
     def get_i_length(self, cell_point=None, dtype=np.float32):
         if cell_point is None:
             cell_point = self.default_cell_point
-        return self.get_mesh_var('e1' + cell_point.lower(), dtype)
+        try:
+            return self.get_mesh_var('e1' + cell_point.lower(), dtype)
+        except ConstraintMismatchError:
+            return self.get_mesh_var('e1' + cell_point.lower() + '_0', dtype)
 
     def get_j_length(self, cell_point=None, dtype=np.float32):
         if cell_point is None:
             cell_point = self.default_cell_point
-        return self.get_mesh_var('e2' + cell_point.lower(), dtype)
+        try:
+            return self.get_mesh_var('e2' + cell_point.lower(), dtype)
+        except ConstraintMismatchError:
+            return self.get_mesh_var('e2' + cell_point.lower() + '_0', dtype)
 
     def get_k_length(self, cell_point=None, dtype=np.float32):
         if cell_point is None:
diff --git a/diagonals/moc.py b/diagonals/moc.py
index 926ca98..e8c3c95 100644
--- a/diagonals/moc.py
+++ b/diagonals/moc.py
@@ -1,20 +1,10 @@
-import os
-
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
 
-import numba
 from numba import vectorize
 from numba import guvectorize
 from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
+from numba import float32
 
 import diagonals
 
@@ -109,16 +99,16 @@ def _compute_moc_gpu(vo, area):
 
     block = (128, 1, 1)
     grid_size = (lats // block[0]) + 1
-    grid3D = (grid_size, levels, times)
+    grid_3d = (grid_size, levels, times)
     gpu_vo = cuda.to_device(vo.astype(np.float32))
     gpu_moc = cuda.device_array((times, levels, lats), dtype=np.float32)
-    grid2D = (grid_size, times)
+    grid_2d = (grid_size, times)
 
     moc = {}
     for basin, mask in area.items():
         gpu_area = cuda.to_device(mask.astype(np.float32))
-        _horizontal_integral[grid3D, block](gpu_vo, gpu_area, gpu_moc)
-        _vertical_cumsum_gpu[grid2D, block](gpu_moc)
+        _horizontal_integral[grid_3d, block](gpu_vo, gpu_area, gpu_moc)
+        _vertical_cumsum_gpu[grid_2d, block](gpu_moc)
         moc[basin] = gpu_moc.copy_to_host()
 
     del gpu_area, gpu_moc, gpu_vo
diff --git a/diagonals/ohc.py b/diagonals/ohc.py
index ef70caa..e5cd6c6 100644
--- a/diagonals/ohc.py
+++ b/diagonals/ohc.py
@@ -1,19 +1,7 @@
-import os
-
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
-
-import numba
 from numba import vectorize
 from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
 
 import diagonals
 
@@ -179,12 +167,12 @@ def _compute_ohc_cpu(layers, thetao, weights, area):
             axis=1
         )
         ohc.append(ohc_layer)
-        ohc1D_total = []
+        ohc1d_total = []
         for i, basin in enumerate(area):
             ohc_basin = _multiply_array(ohc_layer, area[basin])
-            ohc1D = np.sum(ohc_basin, axis=(1, 2))
-            ohc1D_total.append(ohc1D)
-    return ohc, ohc1D_total
+            ohc1d = np.sum(ohc_basin, axis=(1, 2))
+            ohc1d_total.append(ohc1d)
+    return ohc, ohc1d_total
 
 
 def _compute_ohc_gpu(layers, thetao, weights, area):
@@ -236,15 +224,15 @@ def _compute_ohc_gpu(layers, thetao, weights, area):
         _compute_ohc[grid, block](gpu_thetao, weights[layer], gpu_ohc, levels)
         ohc.append(gpu_ohc.copy_to_host())  # moure al final
         _multiply_ohc_basin[grid, block](gpu_ohc, gpu_basins_area, gpu_temp)
-        ohc1D_basin = []
+        ohc1d_basin = []
         for basin in range(basins):
-            ohc_1D = np.empty(times, dtype=np.float32)
+            ohc_1d = np.empty(times, dtype=np.float32)
             for time in range(times):
-                ohc_1D[time] = _sum_red_cuda(gpu_temp[basin, time, :])
-            ohc1D_basin.append(ohc_1D)
+                ohc_1d[time] = _sum_red_cuda(gpu_temp[basin, time, :])
+            ohc1d_basin.append(ohc_1d)
     del gpu_ohc, gpu_temp, gpu_basins_area
 
-    return ohc, ohc1D_basin
+    return ohc, ohc1d_basin
 
 
 @vectorize(['float32(int32, int32, float32, float32, float32)'], target='cuda')
diff --git a/diagonals/psi.py b/diagonals/psi.py
index 1bdb91e..b80169b 100644
--- a/diagonals/psi.py
+++ b/diagonals/psi.py
@@ -1,22 +1,6 @@
-import os
-
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
-
-import numba
-from numba import njit
-from numba import vectorize
-from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
-
-import diagonals
+from numba import vectorize, float32 # noqa
 
 __all__ = ['compute']
 
diff --git a/diagonals/regmean.py b/diagonals/regmean.py
index 1cdaced..ccbfe51 100644
--- a/diagonals/regmean.py
+++ b/diagonals/regmean.py
@@ -1,27 +1,17 @@
-import os
-
+import logging
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
-
-import numba
 from numba import vectorize
-from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
 
 import diagonals
 
-__all__ = ['compute_regmean_2D', 'compute_regmean_3D',
+logger = logging.getLogger(__name__)
+
+__all__ = ['compute_regmean_2d', 'compute_regmean_3d',
            'compute_regmean_levels']
 
 
-def compute_regmean_2D(var, basins, area):
+def compute_regmean_2d(var, basins, area):
     """Function that checks device and calls computing functions.
 
     Checks if the computations are going performed in the CPU or the GPU:
@@ -43,11 +33,11 @@ def compute_regmean_2D(var, basins, area):
     if diagonals.CONFIG.use_gpu:
         logger.warning('GPU routines not implemented for regmean diagnostic.'
                        'Using CPU instead')
-    regmean = _compute_regmean_2D_cpu(var, basins, area)
+    regmean = _compute_regmean_2d_cpu(var, basins, area)
     return regmean
 
 
-def compute_regmean_3D(var, basins, volume):
+def compute_regmean_3d(var, basins, volume):
     """Function that checks device and calls computing functions.
 
     Checks if the computations are going performed in the CPU or the GPU:
@@ -69,7 +59,7 @@ def compute_regmean_3D(var, basins, volume):
     if diagonals.CONFIG.use_gpu:
         logger.warning('GPU routines not implemented for regmean diagnostic.'
                        'Using CPU instead')
-    regmean = _compute_regmean_3D_cpu(var, basins, volume)
+    regmean = _compute_regmean_3d_cpu(var, basins, volume)
     return regmean
 
 
@@ -102,7 +92,7 @@ def compute_regmean_levels(var, basins, volume):
     return regmean
 
 
-def _compute_regmean_2D_cpu(var, basins, area):
+def _compute_regmean_2d_cpu(var, basins, area):
     """Function that computes the regional mean for 2D vars in the cpu.
 
        Computes the weights for each region and performs a weighted average.
@@ -124,7 +114,7 @@ def _compute_regmean_2D_cpu(var, basins, area):
     times = var.shape[0]
     regmean_total = {}
     for basin, mask in basins.items():
-        weights = _compute_weights_2D(mask, area)
+        weights = _compute_weights_2d(mask, area)
         regmean = np.empty(times)
         for t in range(times):
             regmean[t] = np.average(var[t, :, :], axis=(0, 1),
@@ -133,7 +123,7 @@ def _compute_regmean_2D_cpu(var, basins, area):
     return regmean_total
 
 
-def _compute_regmean_3D_cpu(var, basins, volume):
+def _compute_regmean_3d_cpu(var, basins, volume):
     """Function that computes the regional mean for 3D vars in the cpu.
 
     Computes the weights for each region and performs a weighted average.
@@ -155,7 +145,7 @@ def _compute_regmean_3D_cpu(var, basins, volume):
     times = var.shape[0]
     regmean_total = {}
     for basin, mask in basins.items():
-        weights = _compute_weights_3D(mask, volume)
+        weights = _compute_weights_3d(mask, volume)
         regmean = np.empty(times)
         for t in range(times):
             regmean[t] = np.average(var[t, :, :, :], axis=(0, 1, 2),
@@ -189,7 +179,7 @@ def _compute_regmean_levels_cpu(var, basins, volume):
     regmean_total = {}
     for basin, mask in basins.items():
         regmean = np.empty((times, levs))
-        w = _compute_weights_3D(mask, volume)
+        w = _compute_weights_3d(mask, volume)
         for t in range(times):
             for l in range(levs):
                 regmean[t, l] = np.ma.average(var[t, l, :, :], axis=(0, 1),
@@ -199,7 +189,7 @@ def _compute_regmean_levels_cpu(var, basins, volume):
 
 
 @vectorize(['float32(float32, float32)'], target='cpu')
-def _compute_weights_2D(mask, area):
+def _compute_weights_2d(mask, area):
     """Function that computes the regional weights for 2D vars in the cpu.
 
     Parameters
@@ -221,7 +211,7 @@ def _compute_weights_2D(mask, area):
 
 
 @vectorize(['float32(float32, float32)'], target='cpu')
-def _compute_weights_3D(mask, volume):
+def _compute_weights_3d(mask, volume):
     """Function that computes the regional weights for 3D vars in the cpu.
 
     Parameters
diff --git a/diagonals/regsum.py b/diagonals/regsum.py
index 4051c5f..d4ea3e9 100644
--- a/diagonals/regsum.py
+++ b/diagonals/regsum.py
@@ -1,26 +1,15 @@
-import os
-
+import logging
 import numpy as np
-
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
-
-import numba
 from numba import vectorize
-from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
 
 import diagonals
 
-__all__ = ['compute_regsum_2D', 'compute_regsum_3D', 'compute_regsum_levels']
+logger = logging.getLogger(__name__)
+
+__all__ = ['compute_regsum_2d', 'compute_regsum_3d', 'compute_regsum_levels']
 
 
-def compute_regsum_2D(var, basins, area):
+def compute_regsum_2d(var, basins, area):
     """Function that checks device and calls computing functions.
 
     Checks if the computations are going performed in the CPU or the GPU:
@@ -42,11 +31,11 @@ def compute_regsum_2D(var, basins, area):
     if diagonals.CONFIG.use_gpu:
         logger.warning('GPU routines not implemented for regmean diagnostic.'
                        'Using CPU instead')
-    regsum = _compute_regsum_2D_cpu(var, basins, area)
+    regsum = _compute_regsum_2d_cpu(var, basins, area)
     return regsum
 
 
-def compute_regsum_3D(var, basins, volume, tmask):
+def compute_regsum_3d(var, basins, volume, tmask):
     """Function that checks device and calls computing functions.
 
     Checks if the computations are going performed in the CPU or the GPU:
@@ -68,11 +57,11 @@ def compute_regsum_3D(var, basins, volume, tmask):
     if diagonals.CONFIG.use_gpu:
         logger.warning('GPU routines not implemented for regmean diagnostic.'
                        'Using CPU instead')
-    regsum = _compute_regsum_3D_cpu(var, basins, volume, tmask)
+    regsum = _compute_regsum_3d_cpu(var, basins, volume, tmask)
     return regsum
 
 
-def _compute_regsum_2D_cpu(var, basins, area):
+def _compute_regsum_2d_cpu(var, basins, area):
     """Function that computes the regional sum for 2D vars in the cpu.
 
        Computes the weights for each region and performs a weighted sum.
@@ -93,12 +82,12 @@ def _compute_regsum_2D_cpu(var, basins, area):
     """
     regsum = {}
     for basin, mask in basins.items():
-        weighted_var = _weigh_var_2D(var, mask, area)
+        weighted_var = _weigh_var_2d(var, mask, area)
         regsum[basin] = np.sum(weighted_var, axis=(1, 2))
     return regsum
 
 
-def _compute_regsum_3D_cpu(var, basins, volume, tmask):
+def _compute_regsum_3d_cpu(var, basins, volume, tmask):
     """Function that computes the regional sum for 3D vars in the cpu.
 
     Computes the weights for each region and performs a weighted sum.
@@ -119,7 +108,7 @@ def _compute_regsum_3D_cpu(var, basins, volume, tmask):
     """
     regsum = {}
     for basin, mask in basins.items():
-        weighted_var = _weigh_var_3D(var, mask, volume, tmask)
+        weighted_var = _weigh_var_3d(var, mask, volume, tmask)
         regsum[basin] = np.sum(weighted_var, axis=(1, 2, 3))
     return regsum
 
@@ -144,13 +133,13 @@ def compute_regsum_levels(var, basins, volume, tmask):
     """
     regsum = {}
     for basin, mask in basins.items():
-        weighted_var = _weigh_var_3D(var, mask, volume, tmask)
+        weighted_var = _weigh_var_3d(var, mask, volume, tmask)
         regsum[basin] = np.sum(weighted_var, axis=(2, 3))
     return regsum
 
 
 @vectorize(['float32(float32, float32, float32)'], target='cpu')
-def _weigh_var_2D(var, mask, area):
+def _weigh_var_2d(var, mask, area):
     """Function that weights a 2D variable for each region.
 
     Parameters
@@ -173,7 +162,7 @@ def _weigh_var_2D(var, mask, area):
 
 @vectorize(['float32(float32, float32, float32, float32)'],
            target='cpu')
-def _weigh_var_3D(var, mask, volume, tmask):
+def _weigh_var_3d(var, mask, volume, tmask):
     """Function that weights a 3D variable for each region.
 
     Parameters
diff --git a/diagonals/siasie.py b/diagonals/siasie.py
index c0f2f73..b3ca5dc 100644
--- a/diagonals/siasie.py
+++ b/diagonals/siasie.py
@@ -1,22 +1,7 @@
-import os
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-import iris.coords
-import iris.coord_categorisation
-import iris.analysis
-from iris.experimental.equalise_cubes import equalise_attributes
-from iris.time import PartialDateTime
-
-import warnings
-
-import numba
 from numba import vectorize
 from numba import cuda
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
 
 import diagonals
 
diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 127dd7d..5183335 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -1,23 +1,13 @@
-import os
-
+import logging
 import numpy as np
 
-import iris
-import iris.cube
-import iris.analysis
-
-import warnings
-import datetime
-
 import numba
 from numba import vectorize
-from numba import cuda
-from numba import njit
-from numba import int32, int64, float32, float64
-from numba.cuda.cudadrv import driver
 
 import diagonals
 
+logger = logging.getLogger(__name__)
+
 __all__ = ['compute_zonmean', 'get_basin_area']
 
 
@@ -69,6 +59,14 @@ def _compute_zonal_mean_cpu(var, lats, area):
         List of masked arrays containing a variable's zonal mean for
         each depth level and timestep.
     """
+    if len(var.shape) == 4:
+        value = _compute_zonal_mean_cpu_3d(var, lats, area)
+    else:
+        value = _compute_zonal_mean_cpu_2d(var, lats, area)
+    return value
+
+
+def _compute_zonal_mean_cpu_3d(var, lats, area):
     times = var.shape[0]
     levs = var.shape[1]
     value = {}
@@ -77,9 +75,21 @@ def _compute_zonal_mean_cpu(var, lats, area):
         value[basin] = np.empty((times, levs, 180))
         for t in range(times):
             for lev in range(levs):
-                value[basin][t][lev][:] = _zonal_mean_cpu(var[t, lev, :, :],
-                                                          weight,
-                                                          lats)
+                value[basin][t][lev][:] = _zonal_mean_cpu(
+                    var[t, lev, :, :], weight, lats
+                )
+    return value
+
+
+def _compute_zonal_mean_cpu_2d(var, lats, area):
+    times = var.shape[0]
+    value = {}
+    for basin, mask in area.items():
+        weight = np.squeeze(area[basin])
+        value[basin] = np.empty((times, 180))
+        for t in range(times):
+            value[basin][t, :] = _zonal_mean_cpu(
+                var[t, :, :], weight, lats)
     return value
 
 
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..f4f8eb4
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,22 @@
+[aliases]
+test=pytest
+
+[tool:pytest]
+addopts =
+    --flake8
+    --doctest-modules
+    --cov=diagonals
+    --cov-report=term
+    --cov-report=xml:test/report/python3/coverage.xml
+    --cov-report=html:test/report/python3/coverage_html
+    --junit-xml=test/report/python3/report.xml
+    --html=test/report/python3/report.html
+env =
+    MPLBACKEND = Agg
+log_level = WARNING
+
+[coverage:run]
+parallel = true
+
+[pydocstyle]
+convention = numpy
diff --git a/setup.py b/setup.py
index b3259f0..9d9f9d1 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,4 @@
-import sys
-from setuptools import Command, setup, find_packages
+from setuptools import setup, find_packages
 
 
 PACKAGES = [
@@ -10,6 +9,10 @@ REQUIREMENTS = {
     # Installation script (this file) dependencies
     # Installation dependencies
     # Use with pip install . to install from source
+    "setup": [
+        'pytest-runner',
+        'setuptools_scm',
+    ],
     'install': [
         'netCDF4',
         'numba',
@@ -19,51 +22,21 @@ REQUIREMENTS = {
     'gpu': [
         'pycuda'
     ],
-    'test':[
+    'test': [
         'mock',
         'pycodestyle',
         'pytest',
         'pytest-cov',
         'pytest-html',
+        'pytest-flake8',
+        'pytest-metadata>=1.5.1',
     ]
 
 }
 
-class RunTests(Command):
-    """Class to run tests and generate reports."""
-
-    user_options = []
-
-    def initialize_options(self):
-        """Do nothing."""
-
-    def finalize_options(self):
-        """Do nothing."""
-
-    def run(self):
-        """Run tests and generate a coverage report."""
-        import pytest
-
-        version = sys.version_info[0]
-        report_dir = 'test/report/python{}'.format(version)
-        args = [
-            'test',
-            'diagonals',  # for doctests
-            '--ignore=test/report',
-            '--doctest-modules',
-            '--cov=diagonals',
-            '--cov-report=term',
-            '--cov-report=html:{}/coverage_html'.format(report_dir),
-            '--cov-report=xml:{}/coverage.xml'.format(report_dir),
-            '--junit-xml={}/report.xml'.format(report_dir),
-            '--html={}/report.html'.format(report_dir),
-        ]
-        errno = pytest.main(args)
-        sys.exit(errno)
-
 
 setup(name='diagonals',
-      version='0.1',
+      version='0.1.1',
       description='Compute diagnostics targeting the CPU or the GPU',
       url='https://earth.bsc.es/gitlab/es/diagonals',
       author='BSC-CNS Earth Sciences Department',
@@ -76,12 +49,10 @@ setup(name='diagonals',
           'Programming Language :: Python :: 3.6',
       ],
       packages=find_packages(),
-      install_requires=REQUIREMENTS['install'],
+      setup_requires=REQUIREMENTS["setup"],
+      install_requires=REQUIREMENTS["install"],
+      tests_requires=REQUIREMENTS["test"],
       extras_require={
           'gpu': REQUIREMENTS['gpu'],
-          'test': REQUIREMENTS['test'],
-      },
-      cmdclass={
-            'test': RunTests,
       },
       zip_safe=False,)
diff --git a/test/unit/test_lint.py b/test/unit/test_lint.py
deleted file mode 100644
index 55843b5..0000000
--- a/test/unit/test_lint.py
+++ /dev/null
@@ -1,34 +0,0 @@
-""" Lint tests """
-import os
-import unittest
-
-import pycodestyle  # formerly known as pep8
-
-
-class TestLint(unittest.TestCase):
-
-    def test_pep8_conformance(self):
-        """Test that we conform to PEP-8."""
-
-        check_paths = [
-            'diagonals',
-            'test',
-        ]
-        exclude_paths = [
-
-        ]
-
-        print("PEP8 check of directories: {}\n".format(', '.join(check_paths)))
-
-        # Get paths wrt package root
-        package_root = os.path.dirname(os.path.dirname(os.path.dirname(
-                                                                       __file__
-                                                                       )))
-        for paths in (check_paths, exclude_paths):
-            for i, path in enumerate(paths):
-                paths[i] = os.path.join(package_root, path)
-
-        style = pycodestyle.StyleGuide()
-        style.options.exclude.extend(exclude_paths)
-
-        self.assertEqual(style.check_files(check_paths).total_errors, 0)
diff --git a/test/unit/test_ohc.py b/test/unit/test_ohc.py
index 9a67a18..a7641df 100644
--- a/test/unit/test_ohc.py
+++ b/test/unit/test_ohc.py
@@ -1,8 +1,5 @@
-import os
 import unittest
 
-import pycodestyle  # formerly known as pep8
-
 import numpy as np
 import diagonals.ohc as ohc
 
-- 
GitLab


From 583dd7b503cfd6014b69939e30ec8845088c70f2 Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Wed, 5 Feb 2020 11:35:57 +0100
Subject: [PATCH 7/7] Bump version to 0.2

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 9d9f9d1..7efa253 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@ REQUIREMENTS = {
 
 
 setup(name='diagonals',
-      version='0.1.1',
+      version='0.2.0',
       description='Compute diagnostics targeting the CPU or the GPU',
       url='https://earth.bsc.es/gitlab/es/diagonals',
       author='BSC-CNS Earth Sciences Department',
-- 
GitLab