From b41fd958953e02db11272de68cf0b6b867346413 Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Thu, 30 Jan 2020 09:38:31 +0100
Subject: [PATCH 1/4] Zonal mean working

---
 diagonals/mesh_helpers/nemo.py | 12 +++++++++---
 diagonals/zonmean.py           | 25 ++++++++++++++++++++++---
 setup.py                       |  2 +-
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py
index f9d4878..728a36d 100644
--- a/diagonals/mesh_helpers/nemo.py
+++ b/diagonals/mesh_helpers/nemo.py
@@ -1,6 +1,6 @@
 import numpy as np
 import iris
-
+from iris.exceptions import ConstraintMismatchError
 from numba import vectorize, float32
 
 
@@ -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/zonmean.py b/diagonals/zonmean.py
index 127dd7d..c46af49 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -69,6 +69,13 @@ 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 +84,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]
+    levs = var.shape[1]
+    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.py b/setup.py
index b3259f0..ce2c227 100644
--- a/setup.py
+++ b/setup.py
@@ -63,7 +63,7 @@ class RunTests(Command):
 
 
 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',
-- 
GitLab


From 180ba88ec4872cff6e7024e0bb4739aad71a334b Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Thu, 30 Jan 2020 13:07:24 +0100
Subject: [PATCH 2/4] Diagonals formatted

---
 .../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                |  2 +-
 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                          | 21 +++------
 setup.cfg                                     | 22 +++++++++
 setup.py                                      | 46 ++++---------------
 test/unit/test_lint.py                        | 34 --------------
 test/unit/test_ohc.py                         |  3 --
 16 files changed, 108 insertions(+), 235 deletions(-)
 create mode 100644 setup.cfg
 delete mode 100644 test/unit/test_lint.py

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 728a36d..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 iris.exceptions import ConstraintMismatchError
-from numba import vectorize, float32
+from numba import vectorize, float32 # noqa
 
 
 class Nemo():
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 c46af49..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']
 
 
@@ -75,6 +65,7 @@ def _compute_zonal_mean_cpu(var, lats, area):
         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]
@@ -86,12 +77,12 @@ def _compute_zonal_mean_cpu_3d(var, lats, area):
             for lev in range(levs):
                 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]
-    levs = var.shape[1]
     value = {}
     for basin, mask in area.items():
         weight = np.squeeze(area[basin])
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 ce2c227..a52fc02 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,48 +22,18 @@ 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.1',
@@ -81,7 +54,4 @@ setup(name='diagonals',
           '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 3416c3b790f187cd9768c58933f5d089e6bddba4 Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Thu, 30 Jan 2020 13:17:44 +0100
Subject: [PATCH 3/4] Remove python2 tests

---
 .gitlab-ci.yml | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

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
-- 
GitLab


From 153309076d9d6636dbbb318159b3458838ed016b Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Thu, 30 Jan 2020 13:22:14 +0100
Subject: [PATCH 4/4] Update setup.py

---
 setup.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index a52fc02..9d9f9d1 100644
--- a/setup.py
+++ b/setup.py
@@ -49,9 +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'],
       },
       zip_safe=False,)
-- 
GitLab