diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index bade64f6053eb78a153d5572856a19235ed0663a..73e3443bf795cc177bd167cc923caa502ff8ae02 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 867d818cde17168e2d94f575d281bbe8e97d9156..0f270a2873085411632acdadea785acc26b946ce 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 4f18cb18616b787d84dcee60a907bac6b0644eb3..7c694056a6850a011a48008b67916fa8d966e981 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 6a49975ac4ae14be0cd3446339368a7d5b33da29..6f49265d2b0ad255a4eff56026d49a25a1525d5c 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 87af4dc6489ef0e14ad68b805e4ac8d276cd4247..c55811991f230c01fc0fc05a817e0c7476140044 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 f9d4878b0d6e66af720f3394610ef0e3e68aa5ca..782a815609c4eed9fb6f0773a5ba2a7b0b1485db 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 926ca98f638446dd536214173740c0947125f26c..e8c3c95a7ebe745fea6e88b667a32218c4a2c7ad 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 ef70caa2859c50d1260fb608a137bffd1211baf6..e5cd6c63fb0951bf496a00a9c7e20a8f318cb1a3 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 1bdb91ec24e4536eaa0ed6c4e01550fdea3a6f2b..b80169b98c0db63f49adc641f1b1c8b23c086588 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 1cdacedd319dd4936fb021fd10883ca94fe117cd..ccbfe518775e9e0bb491c1dfc625296860fae3f2 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 4051c5f2425d3e50a6f597fbe3d9a4416852dd41..d4ea3e950ee26230330fff22895a0e6dde97d375 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 c0f2f730d1136f5fffac88a3556dbcbcd003a21a..b3ca5dc68ade62801ff257bf72882a910bbc6731 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 127dd7d3fd50ee6a7538d66622420234849a5b48..51833351d69ead5fbd46008c3a773e287fc688bb 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 0000000000000000000000000000000000000000..f4f8eb46cbd7a91599a23bcf96b5ea26d724f6ea
--- /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 b3259f083e7619ab1efcb8d6186413a87b139eb7..9d9f9d1916fc7d851b85f32e5130553a3bea02a9 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 55843b5ac09c6f5391d34f9d923ef0dba1d86d45..0000000000000000000000000000000000000000
--- 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 9a67a18c8e0935137f2c217e6c1ad79c8c8fb8c2..a7641dfd58804a2580ef7babd3b8132130f90779 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