From 9afc4f0775ce8bab756bfd2c434218f5dc40fdff Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Fri, 21 Jun 2019 15:16:44 +0200
Subject: [PATCH 01/19] Add sivol functions from DIAGONALS

---
 diagonals/siasie.py | 181 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 176 insertions(+), 5 deletions(-)

diff --git a/diagonals/siasie.py b/diagonals/siasie.py
index d9b5bf6..95f3832 100644
--- a/diagonals/siasie.py
+++ b/diagonals/siasie.py
@@ -23,7 +23,7 @@ import diagonals
 __all__ = ['compute']
 
 
-def compute(gphit, area, sic_slices, basins):
+def compute(gphit, area, sic_slices, basins, sit):
     """Function that checks device and calls computing functions.
 
     Checks if the computations are going performed in the CPU or the GPU:
@@ -58,17 +58,27 @@ def compute(gphit, area, sic_slices, basins):
     """
     if diagonals.CONFIG.use_gpu:
         extn, exts, arean, areas = _compute_sic_gpu(
-            gphit, area, sic_slices, basins
+            gphit, area, sic_slices, basins, sit
         )
+        if sit:
+            extn, exts, arean, areas, voln, vols = _compute_sic_gpu(
+                gphit, area, sic_slices, basins, sit
+            )
+            return extn, exts, arean, areas, voln, vols
     else:
         extn, exts, arean, areas = _compute_sic_cpu(
-            gphit, area, sic_slices, basins
+            gphit, area, sic_slices, basins, sit
         )
+        if sit:
+            extn, exts, arean, areas, voln, vols = _compute_sic_cpu(
+                gphit, area, sic_slices, basins, sit
+            )
+            return extn, exts, arean, areas, voln, vols
 
     return extn, exts, arean, areas
 
 
-def _compute_sic_cpu(gphit, area, sic_slices, basins):
+def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick):
     """Function that computes sea ice are and extension on the CPU.
 
     Loops over time and basins and reduces over lat and lon:
@@ -104,6 +114,11 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins):
     arean = np.empty((len(basins), len(sic_slices)))
     areas = np.empty((len(basins), len(sic_slices)))
 
+    if sithick:
+        sit = sithick.data.astype(np.float32)
+        voln = np.empty((len(basins), len(sic_slices)))
+        vols = np.empty((len(basins), len(sic_slices)))
+
     for time, sic_data in enumerate(sic_slices):
         sic = sic_data.data
         b = 0
@@ -120,14 +135,24 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins):
             temp = _areas_cpu(gphit, area, tmask, sic)
             areas[b][time] = np.sum(temp, axis=(1, 2))
 
+            if sithick:
+                temp = _voln_cpu(gphit, area, tmask, sic, sit)
+                voln[b][time] = np.sum(temp, axis=(1,2))
+
+                temp = _vols_cpu(gphit, area, tmask, sic, sit)
+                vols[b][time] = np.sum(temp, axis=(1,2))
+
             b += 1
 
     del gphit, area, sic_slices, temp
 
+    if sithick:
+        return extn, exts, arean, areas, voln, vols
+
     return extn, exts, arean, areas
 
 
-def _compute_sic_gpu(gphit, area, sic_slices, basins):
+def _compute_sic_gpu(gphit, area, sic_slices, basins, sithick):
     """Function that computes sea ice are and extension on the GPU.
 
     Loops over time and basins and reduces over lat and lon:
@@ -172,6 +197,12 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins):
     grid = (grid_size, area.shape[1])
     del gphit, area
 
+    if sithick:
+        sit = sithick.data.astype(np.float32)
+        gpu_sit = cuda.to_device(sit.astype(np.float32))
+        voln = np.empty((len(basins), len(sic_slices)))
+        vols = np.empty((len(basins), len(sic_slices)))
+
     for time, sic_data in enumerate(sic_slices):
         extn[time] = {}
         exts[time] = {}
@@ -199,8 +230,20 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins):
                                     gpu_sic, gpu_temp)
             areas[b][time] = _sum_red_cuda(gpu_temp)
 
+            if sithick:
+                _voln_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask,
+                                       gpu_sic, gpu_sit, gpu_temp)
+                voln[b][time] = _sum_red_cuda(gpu_temp)
+
+                _vols_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask,
+                                       gpu_sic, gpu_sit, gpu_temp)
+                vols[b][time] = _sum_red_cuda(gpu_temp)
+
             b += 1
 
+    if sithick:
+        return extn, exts, arean, areas, voln, vols
+
     del gpu_gphit, gpu_area, gpu_tmask, gpu_sic, gpu_temp
     return extn, exts, arean, areas
 
@@ -450,6 +493,134 @@ def _exts_gpu(gphit, area, tmask, sic, temp):
         temp[x] = area[0, j, i] * tmask[0, j, i]
 
 
+@vectorize(['float32(float32, float32, float32, float32, float32)'],
+           target='cpu')
+def _voln_cpu(gphit, area, tmask, sic, sit):
+    """Vectorized numba function executed on the cpu that computes the
+       percentage of sea ice volume in the northern hemisphere for a given
+       timestep and basin:
+
+    Parameters
+    ----------
+    gphit : float32
+        Masked array containing variable gphit.
+    area : float32
+        Masked array containing cell area.
+    sic: float32
+        Masked array containing variable sic for a given timestep.
+    tmask: float32
+        Masked array for a given basin
+    sit: float32
+        Masked array containing variable sit for a given timestep.
+    Returns
+    -------
+    temp : float32
+        Masked array to be summed over lat and lon to obtain sea ice volume
+    """
+    if gphit > 0:
+        temp = area * tmask * sic * sit / 100
+    return temp
+
+
+@cuda.jit()
+def _voln_gpu(gphit, area, tmask, sic, sit, temp):
+    """Numba kernel executed on the gpu that computes the percentage of sea ice
+    volume in the northern hemisphere for a given timestep and basin:
+
+     Parameters
+     ----------
+     gphit : float32
+         Masked array containing variable gphit.
+     area : float32
+        Masked array containing cell area.
+     sic: float32
+         Masked array containing variable sic for a given timestep.
+     tmask: float32
+         Masked array for a given basin
+     sit: float32
+        Masked array containing variable sit for a given timestep.
+
+     Returns
+     -------
+     temp : float32
+         Masked array to be summed over lat and lon to obtain the sea ice area
+         value. Needs to be a one dimensional array in order to be the input
+         for the reduction kernel.
+     """
+    i, j = cuda.grid(2)
+    x = i + j * sic.shape[1]
+    if(i >= sic.shape[1]):
+        return
+    temp[x] = 0.0
+    if(gphit[0, j, i] > 0):
+        temp[x] = (area[0, j, i] * tmask[0, j, i] * sic[j, i] * sit[j, i]
+                   / 100)
+
+
+@vectorize(['float32(float32, float32, float32, float32, float32)'],
+           target='cpu')
+def _vols_cpu(gphit, area, tmask, sic, sit):
+    """Vectorized numba function executed on the cpu that computes the
+       percentage of sea ice volume in the southern hemisphere for a given
+       timestep and basin:
+
+    Parameters
+    ----------
+    gphit : float32
+        Masked array containing variable gphit.
+    area : float32
+        Masked array containing cell area.
+    sic: float32
+        Masked array containing variable sic for a given timestep.
+    tmask: float32
+        Masked array for a given basin
+    sit: float32
+        Masked array containing variable sit for a given timestep.
+    Returns
+    -------
+    temp : float32
+        Masked array to be summed over lat and lon to obtain sea ice volume
+    """
+    if gphit < 0:
+        temp = area * tmask * sic * sit / 100
+    return temp
+
+
+@cuda.jit()
+def _vols_gpu(gphit, area, tmask, sic, sit, temp):
+    """Numba kernel executed on the gpu that computes the percentage of sea ice
+    volume in the southern hemisphere for a given timestep and basin:
+
+     Parameters
+     ----------
+     gphit : float32
+         Masked array containing variable gphit.
+     area : float32
+        Masked array containing cell area.
+     sic: float32
+         Masked array containing variable sic for a given timestep.
+     tmask: float32
+         Masked array for a given basin
+     sit: float32
+        Masked array containing variable sit for a given timestep.
+
+     Returns
+     -------
+     temp : float32
+         Masked array to be summed over lat and lon to obtain the sea ice area
+         value. Needs to be a one dimensional array in order to be the input
+         for the reduction kernel.
+     """
+    i, j = cuda.grid(2)
+    x = i + j * sic.shape[1]
+    if(i >= sic.shape[1]):
+        return
+    temp[x] = 0.0
+    if(gphit[0, j, i] < 0):
+        temp[x] = (area[0, j, i] * tmask[0, j, i] * sic[j, i] * sit[j, i]
+                   / 100)
+
+
 @vectorize(['float64(float32, float32)', 'float64(float64, float64)'],
            target='cpu')
 def _sum_red(x, y):
-- 
GitLab


From ef8ba9b28cff552fd097ef88595fe8207add8d09 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 25 Jun 2019 17:06:03 +0200
Subject: [PATCH 02/19] Add first version

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

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
new file mode 100644
index 0000000..3030d92
--- /dev/null
+++ b/diagonals/zonmean.py
@@ -0,0 +1,137 @@
+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 njit
+from numba import int32, int64, float32, float64
+from numba.cuda.cudadrv import driver
+
+import diagonals
+
+__all__ = ['compute_zonmean', 'get_basin_area']
+
+
+def compute_zonmean(var, lats, area):
+    """Function that checks device and calls computing functions.
+
+    Checks if the computations are going performed in the CPU or the GPU:
+
+    Parameters
+    ----------
+    var : float32
+        Masked array containing variable data.
+    basins : float32
+        List containing basin names and masks.
+    area : float32
+        Masked array containing cell area.
+
+    Returns
+    -------
+    regmean: float32
+        List containing regional mean for variable var
+    """
+    if diagonals.CONFIG.use_gpu:
+        logger.warning('GPU routines not implemented for regmean diagnostic.'
+                       'Using CPU instead')
+    else:
+        zonmean = _compute_zonmean_cpu(var, lats, area)
+    return zonmean
+
+
+def _compute_zonal_mean_cpu(var, lats, area):
+    times = var.shape[0]
+    levs = var.shape[1]
+    value = {}
+    for basin, mask in weight.items():
+        for time in times:
+            for lev in levs:
+                value[basin]= _zonal_mean_cpu(var[t, l, :, :],
+                                              area[basin],
+                                              lats)
+    return value
+
+
+def _compute_zonal_mean_gpu(var, lats, area):
+    times = var.shape[0]
+    levs = var.shape[1]
+    lats = var.shape[2]
+    lons = var.shape[3]
+    value = {}
+
+    block = (128, 1, 1)
+    grid_size = (lons // block[0]) + 1
+    grid = (grid_size, lats)
+    gpu_var = cuda.to_device(var.astype(np.float32))
+    total = cuda.device_array(180, dtype=np.float32)
+    weights = cuda.device_array(180, dtype=np.float32)
+    for basin, mask in area.items():
+       gpu_area = cuda.to_device(area[basin].astype(np.float32))
+       gpu_value = _zonal_mean_gpu(gpu_var, gpu_area, lats)
+       value[basin] = gpu_value.copy_to_host()
+    del gpu_area, gpu_value, gpu_var
+
+    return value
+
+@numba.njit()
+def _zonal_mean_cpu(variable, weight, latitude):
+    total = np.zeros(180, np.float64)
+    weights = np.zeros(180, np.float64)
+    for i in range(variable.shape[0]):
+        for j in range(variable.shape[1]):
+            if weight[i, j] == 0:
+                continue
+            bin_value = int(round(latitude[i, j]) + 90)
+            weights[bin_value] += weight[i, j]
+            total[bin_value] += variable[i, j] * weight[i, j]
+    return total / weights
+
+
+@cuda.jit()
+def _zonal_mean_gpu(variable, weight, latitude):
+    total[i, j] = 0.0
+    weights[i, j] = 0.0
+    if(i >= total.shape[0]):
+        return
+    if weight[i, j] == 0:
+        continue
+    bin_value = int(round(latitude[i, j]) + 90)
+    weights[bin_value] += weight[i, j]
+    total[bin_value] += variable[i, j] * weight[i, j]
+    return total / weights
+
+
+def get_basin_area(areacello, basins):
+    basin_areas = {}
+    for basin in basins:
+        basin_areas[basin] = _compute_basin_area(areacello, basins[basin])
+    return basin_areas
+
+
+@vectorize(['float32(float32, float32)'], target='cpu')
+def _compute_basin_area(areacello, basin):
+    """Vectorized numba function executed on the cpu that computes the area
+       for each basin:
+
+    Parameters
+    ----------
+    areacello : float32
+        Masked array containing areacello.
+    basin: float32
+        Masked array containing the mask for a given basin.
+
+    Returns
+    -------
+    areacello*basin : float32
+        Masked array containing the area for a given basin.
+    """
+    return areacello*basin
\ No newline at end of file
-- 
GitLab


From 6d957aabd206f774b356d4db0900f8646e04ae0e Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 25 Jun 2019 17:14:02 +0200
Subject: [PATCH 03/19] Add first version

---
 diagonals/examples/examples-zonalmean.py | 59 ++++++++++++++++++++++++
 1 file changed, 59 insertions(+)
 create mode 100644 diagonals/examples/examples-zonalmean.py

diff --git a/diagonals/examples/examples-zonalmean.py b/diagonals/examples/examples-zonalmean.py
new file mode 100644
index 0000000..c7d9fed
--- /dev/null
+++ b/diagonals/examples/examples-zonalmean.py
@@ -0,0 +1,59 @@
+import logging
+import datetime
+import numpy as np
+
+import iris
+import iris.cube
+import iris.analysis
+import iris.coords
+import iris.coord_categorisation
+import iris.analysis
+
+import diagonals
+import diagonals.zonmean as zonmean
+from diagonals.mesh_helpers.nemo import Nemo
+
+MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc"
+REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc"
+THETAO_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium"\
+              "/EC-Earth3-LR/historical/r1i1p1f1/Omon/thetao/gn/v20180711/"\
+              "thetao_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_"\
+              "196101-196112.nc"
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+    start = datetime.datetime.now()
+    logger.info('Starting at %s', start)
+    mesh = Nemo(MESH_FILE, REGIONS_FILE)
+    with diagonals.CONFIG.context(use_gpu=False):
+        lats = mesh.get_mesh_var('latitude', THETAO_FILE)
+        areacello = mesh.get_areacello()
+        var = load_thetao()
+        basins = load_masks()
+        area_basin = zonmean.get_basin_area(areacello)
+        zonmean = zonmean.compute_zonmean(var, lats, area)
+
+
+def load_thetao():
+    thetao = iris.load_cube(THETAO_FILE, 'sea_water_potential_temperature')
+    thetao_data = thetao.data.astype(np.float32)
+    return thetao_data
+
+
+def load_masks():
+    basins = {}
+    cubes = iris.load(REGIONS_FILE)
+    for cube in cubes:
+        name = cube.name()
+        if name in ('nav_lat', 'nav_lon', 'Caspian_Sea'):
+            continue
+        cube = cube.extract(iris.Constraint(z=1))
+        basins[name] = cube.data.astype(np.float32)
+    return basins
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
-- 
GitLab


From 988278d58d060d277da9bedaed4dcad759ca4763 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Fri, 28 Jun 2019 15:18:44 +0200
Subject: [PATCH 04/19] First version of zonal mean, pending to test

---
 diagonals/examples/examples-zonalmean.py |  2 +-
 diagonals/zonmean.py                     | 24 ++++++++++++------------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/diagonals/examples/examples-zonalmean.py b/diagonals/examples/examples-zonalmean.py
index c7d9fed..9a90df0 100644
--- a/diagonals/examples/examples-zonalmean.py
+++ b/diagonals/examples/examples-zonalmean.py
@@ -56,4 +56,4 @@ def load_masks():
 
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()
diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 3030d92..bba8d3d 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -55,9 +55,9 @@ def _compute_zonal_mean_cpu(var, lats, area):
     for basin, mask in weight.items():
         for time in times:
             for lev in levs:
-                value[basin]= _zonal_mean_cpu(var[t, l, :, :],
-                                              area[basin],
-                                              lats)
+                value[basin] = _zonal_mean_cpu(var[t, l, :, :],
+                                               area[basin],
+                                               lats)
     return value
 
 
@@ -75,13 +75,14 @@ def _compute_zonal_mean_gpu(var, lats, area):
     total = cuda.device_array(180, dtype=np.float32)
     weights = cuda.device_array(180, dtype=np.float32)
     for basin, mask in area.items():
-       gpu_area = cuda.to_device(area[basin].astype(np.float32))
-       gpu_value = _zonal_mean_gpu(gpu_var, gpu_area, lats)
-       value[basin] = gpu_value.copy_to_host()
+        gpu_area = cuda.to_device(area[basin].astype(np.float32))
+        gpu_value = _zonal_mean_gpu(gpu_var, gpu_area, lats)
+        value[basin] = gpu_value.copy_to_host()
     del gpu_area, gpu_value, gpu_var
 
     return value
 
+
 @numba.njit()
 def _zonal_mean_cpu(variable, weight, latitude):
     total = np.zeros(180, np.float64)
@@ -102,11 +103,10 @@ def _zonal_mean_gpu(variable, weight, latitude):
     weights[i, j] = 0.0
     if(i >= total.shape[0]):
         return
-    if weight[i, j] == 0:
-        continue
-    bin_value = int(round(latitude[i, j]) + 90)
-    weights[bin_value] += weight[i, j]
-    total[bin_value] += variable[i, j] * weight[i, j]
+    if(weight[i, j] != 0.0):
+        bin_value = int(round(latitude[i, j]) + 90)
+        weights[bin_value] += weight[i, j]
+        total[bin_value] += variable[i, j] * weight[i, j]
     return total / weights
 
 
@@ -134,4 +134,4 @@ def _compute_basin_area(areacello, basin):
     areacello*basin : float32
         Masked array containing the area for a given basin.
     """
-    return areacello*basin
\ No newline at end of file
+    return areacello*basin
-- 
GitLab


From 988144620774b5526aba4336c6b0b2df4fc7a992 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Fri, 28 Jun 2019 16:32:26 +0200
Subject: [PATCH 05/19] Fix volume bugs

---
 diagonals/siasie.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/diagonals/siasie.py b/diagonals/siasie.py
index 95f3832..36635eb 100644
--- a/diagonals/siasie.py
+++ b/diagonals/siasie.py
@@ -57,23 +57,25 @@ def compute(gphit, area, sic_slices, basins, sit):
         Device where the computations are being performed (CPU or GPU)
     """
     if diagonals.CONFIG.use_gpu:
-        extn, exts, arean, areas = _compute_sic_gpu(
-            gphit, area, sic_slices, basins, sit
-        )
         if sit:
             extn, exts, arean, areas, voln, vols = _compute_sic_gpu(
                 gphit, area, sic_slices, basins, sit
             )
             return extn, exts, arean, areas, voln, vols
+        else:
+            extn, exts, arean, areas = _compute_sic_gpu(
+                gphit, area, sic_slices, basins, sit
+            )
     else:
-        extn, exts, arean, areas = _compute_sic_cpu(
-            gphit, area, sic_slices, basins, sit
-        )
         if sit:
             extn, exts, arean, areas, voln, vols = _compute_sic_cpu(
                 gphit, area, sic_slices, basins, sit
             )
             return extn, exts, arean, areas, voln, vols
+        else:
+            extn, exts, arean, areas = _compute_sic_cpu(
+                gphit, area, sic_slices, basins, sit
+            )
 
     return extn, exts, arean, areas
 
@@ -115,7 +117,6 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick):
     areas = np.empty((len(basins), len(sic_slices)))
 
     if sithick:
-        sit = sithick.data.astype(np.float32)
         voln = np.empty((len(basins), len(sic_slices)))
         vols = np.empty((len(basins), len(sic_slices)))
 
@@ -136,6 +137,7 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick):
             areas[b][time] = np.sum(temp, axis=(1, 2))
 
             if sithick:
+                sit = sithick[time].data
                 temp = _voln_cpu(gphit, area, tmask, sic, sit)
                 voln[b][time] = np.sum(temp, axis=(1,2))
 
@@ -198,8 +200,6 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins, sithick):
     del gphit, area
 
     if sithick:
-        sit = sithick.data.astype(np.float32)
-        gpu_sit = cuda.to_device(sit.astype(np.float32))
         voln = np.empty((len(basins), len(sic_slices)))
         vols = np.empty((len(basins), len(sic_slices)))
 
@@ -231,6 +231,8 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins, sithick):
             areas[b][time] = _sum_red_cuda(gpu_temp)
 
             if sithick:
+                sit = sithick[time].data.astype(np.float32)
+                gpu_sit = cuda.to_device(sit.astype(np.float32))
                 _voln_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask,
                                        gpu_sic, gpu_sit, gpu_temp)
                 voln[b][time] = _sum_red_cuda(gpu_temp)
-- 
GitLab


From 0407def824a2c8b414b99b6ab639d20bfd08fcf5 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Fri, 28 Jun 2019 16:37:18 +0200
Subject: [PATCH 06/19] Fix style issues

---
 diagonals/siasie.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/diagonals/siasie.py b/diagonals/siasie.py
index 36635eb..c0f2f73 100644
--- a/diagonals/siasie.py
+++ b/diagonals/siasie.py
@@ -139,10 +139,10 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick):
             if sithick:
                 sit = sithick[time].data
                 temp = _voln_cpu(gphit, area, tmask, sic, sit)
-                voln[b][time] = np.sum(temp, axis=(1,2))
+                voln[b][time] = np.sum(temp, axis=(1, 2))
 
                 temp = _vols_cpu(gphit, area, tmask, sic, sit)
-                vols[b][time] = np.sum(temp, axis=(1,2))
+                vols[b][time] = np.sum(temp, axis=(1, 2))
 
             b += 1
 
-- 
GitLab


From acc5e95e5cb9714ec023a8e31f2836acca06c678 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 12:34:45 +0200
Subject: [PATCH 07/19] Fix bugs in cpu routine

---
 diagonals/zonmean.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index bba8d3d..5094423 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -44,7 +44,7 @@ def compute_zonmean(var, lats, area):
         logger.warning('GPU routines not implemented for regmean diagnostic.'
                        'Using CPU instead')
     else:
-        zonmean = _compute_zonmean_cpu(var, lats, area)
+        zonmean = _compute_zonal_mean_cpu(var, lats, area)
     return zonmean
 
 
@@ -52,11 +52,13 @@ def _compute_zonal_mean_cpu(var, lats, area):
     times = var.shape[0]
     levs = var.shape[1]
     value = {}
-    for basin, mask in weight.items():
-        for time in times:
-            for lev in levs:
-                value[basin] = _zonal_mean_cpu(var[t, l, :, :],
-                                               area[basin],
+    for basin, mask in area.items():
+        weight = np.squeeze(area[basin])
+        value[basin] = np.empty((times, levs, 180))
+        for time in range(times):
+            for lev in range(levs):
+                value[basin][time][lev][:] = _zonal_mean_cpu(var[time, lev, :, :],
+                                               weight,
                                                lats)
     return value
 
@@ -85,8 +87,8 @@ def _compute_zonal_mean_gpu(var, lats, area):
 
 @numba.njit()
 def _zonal_mean_cpu(variable, weight, latitude):
-    total = np.zeros(180, np.float64)
-    weights = np.zeros(180, np.float64)
+    total = np.zeros(180, np.float32)
+    weights = np.zeros(180, np.float32)
     for i in range(variable.shape[0]):
         for j in range(variable.shape[1]):
             if weight[i, j] == 0:
-- 
GitLab


From 07348b042918c187dfd2229ab9538be804f06ac3 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 13:19:39 +0200
Subject: [PATCH 08/19] Add save function

---
 diagonals/examples/examples-zonalmean.py | 31 ++++++++++++++++++------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/diagonals/examples/examples-zonalmean.py b/diagonals/examples/examples-zonalmean.py
index 9a90df0..b3fcd28 100644
--- a/diagonals/examples/examples-zonalmean.py
+++ b/diagonals/examples/examples-zonalmean.py
@@ -15,10 +15,10 @@ from diagonals.mesh_helpers.nemo import Nemo
 
 MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc"
 REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc"
-THETAO_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium"\
-              "/EC-Earth3-LR/historical/r1i1p1f1/Omon/thetao/gn/v20180711/"\
-              "thetao_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_"\
-              "196101-196112.nc"
+THETAO_FILE = "/esarchive/exp/ecearth/a1tr/cmorfiles/CMIP/EC-Earth-Consortium"\
+              "/EC-Earth3/historical/r24i1p1f1/Omon/thetao/gn/v20190312/"\
+              "thetao_Omon_EC-Earth3_historical_r24i1p1f1_gn_"\
+              "185001-185012.nc"
 
 logger = logging.getLogger(__name__)
 
@@ -29,12 +29,19 @@ def main():
     logger.info('Starting at %s', start)
     mesh = Nemo(MESH_FILE, REGIONS_FILE)
     with diagonals.CONFIG.context(use_gpu=False):
-        lats = mesh.get_mesh_var('latitude', THETAO_FILE)
+        lats = mesh.get_mesh_var('nav_lat', dtype=np.float32)
         areacello = mesh.get_areacello()
         var = load_thetao()
         basins = load_masks()
-        area_basin = zonmean.get_basin_area(areacello)
-        zonmean = zonmean.compute_zonmean(var, lats, area)
+        area_basin = zonmean.get_basin_area(areacello, basins)
+        zonalmean = zonmean.compute_zonmean(var, lats, area_basin)
+        if diagonals.CONFIG.use_gpu:
+            device = 'GPU'
+        else:
+            device = 'CPU'
+        save_data(zonalmean, basins, device)
+    ellapsed = datetime.datetime.now() - start
+    logger.info('Total ellapsed time on the %s: %s', device, ellapsed)
 
 
 def load_thetao():
@@ -55,5 +62,15 @@ def load_masks():
     return basins
 
 
+def save_data(zonalmean, basins, device):
+    cubes_zonmean = iris.cube.CubeList()
+    for basin in basins.keys():
+        cube_zonmean = iris.cube.Cube(zonalmean[basin])
+        cube_zonmean.add_aux_coord(iris.coords.AuxCoord(basin, 'region'))
+        cubes_zonmean.append(cube_zonmean)
+    iris.save(cubes_zonmean.merge_cube(), '/esarchive/scratch/sloosvel/'
+              'numba_outputs/zonmean_{0}.nc'.format(device), zlib=True)
+
+
 if __name__ == '__main__':
     main()
-- 
GitLab


From 81c05bac1b7daf370ef223be6359771a3f59fe45 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 13:24:28 +0200
Subject: [PATCH 09/19] Fix style issues

---
 diagonals/siasie.py  |  4 ++--
 diagonals/zonmean.py | 11 +++++------
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/diagonals/siasie.py b/diagonals/siasie.py
index 95f3832..b4280d3 100644
--- a/diagonals/siasie.py
+++ b/diagonals/siasie.py
@@ -137,10 +137,10 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick):
 
             if sithick:
                 temp = _voln_cpu(gphit, area, tmask, sic, sit)
-                voln[b][time] = np.sum(temp, axis=(1,2))
+                voln[b][time] = np.sum(temp, axis=(1, 2))
 
                 temp = _vols_cpu(gphit, area, tmask, sic, sit)
-                vols[b][time] = np.sum(temp, axis=(1,2))
+                vols[b][time] = np.sum(temp, axis=(1, 2))
 
             b += 1
 
diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 5094423..082b3c4 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -41,8 +41,7 @@ def compute_zonmean(var, lats, area):
         List containing regional mean for variable var
     """
     if diagonals.CONFIG.use_gpu:
-        logger.warning('GPU routines not implemented for regmean diagnostic.'
-                       'Using CPU instead')
+        zonmean = _compute_zonal_mean_gpu(var, lats, area)
     else:
         zonmean = _compute_zonal_mean_cpu(var, lats, area)
     return zonmean
@@ -55,11 +54,11 @@ def _compute_zonal_mean_cpu(var, lats, area):
     for basin, mask in area.items():
         weight = np.squeeze(area[basin])
         value[basin] = np.empty((times, levs, 180))
-        for time in range(times):
+        for t in range(times):
             for lev in range(levs):
-                value[basin][time][lev][:] = _zonal_mean_cpu(var[time, lev, :, :],
-                                               weight,
-                                               lats)
+                value[basin][t][lev][:] = _zonal_mean_cpu(var[t, lev, :, :],
+                                                          weight,
+                                                          lats)
     return value
 
 
-- 
GitLab


From 6c7c63eb5c4d6e2131cf3dd0b3d1fc7e8d29d8ed Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 16:17:18 +0200
Subject: [PATCH 10/19] Test gpu routine

---
 diagonals/zonmean.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 082b3c4..74f23fd 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -62,7 +62,7 @@ def _compute_zonal_mean_cpu(var, lats, area):
     return value
 
 
-def _compute_zonal_mean_gpu(var, lats, area):
+def _compute_zonal_mean_gpu(var, latitudes, area):
     times = var.shape[0]
     levs = var.shape[1]
     lats = var.shape[2]
@@ -73,12 +73,17 @@ def _compute_zonal_mean_gpu(var, lats, area):
     grid_size = (lons // block[0]) + 1
     grid = (grid_size, lats)
     gpu_var = cuda.to_device(var.astype(np.float32))
-    total = cuda.device_array(180, dtype=np.float32)
-    weights = cuda.device_array(180, dtype=np.float32)
+    gpu_lats = cuda.to_device(latitudes.astype(np.float32))
+    gpu_total = cuda.device_array(180, dtype=np.float32)
+    gpu_weights = cuda.device_array(180, dtype=np.float32)
     for basin, mask in area.items():
-        gpu_area = cuda.to_device(area[basin].astype(np.float32))
-        gpu_value = _zonal_mean_gpu(gpu_var, gpu_area, lats)
-        value[basin] = gpu_value.copy_to_host()
+        value[basin] = np.empty((times, levs, 180))
+        gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
+        for t in range(times):
+            for lev in range(levels):
+                gpu_value = _zonal_mean_gpu(gpu_var[t, lev, :, :], gpu_area,
+                                            gpu_lats, gpu_total, gpu_weights)
+                value[basin][t][lev][:] = gpu_value.copy_to_host()
     del gpu_area, gpu_value, gpu_var
 
     return value
@@ -99,7 +104,7 @@ def _zonal_mean_cpu(variable, weight, latitude):
 
 
 @cuda.jit()
-def _zonal_mean_gpu(variable, weight, latitude):
+def _zonal_mean_gpu(variable, weight, latitude, total, weights):
     total[i, j] = 0.0
     weights[i, j] = 0.0
     if(i >= total.shape[0]):
@@ -108,7 +113,7 @@ def _zonal_mean_gpu(variable, weight, latitude):
         bin_value = int(round(latitude[i, j]) + 90)
         weights[bin_value] += weight[i, j]
         total[bin_value] += variable[i, j] * weight[i, j]
-    return total / weights
+    return total[i, j] / weights[i, j]
 
 
 def get_basin_area(areacello, basins):
-- 
GitLab


From 4b8257bc8eead98fc6ec111df1a747fa2900e304 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 16:19:37 +0200
Subject: [PATCH 11/19] Fix wrong var name

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

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 74f23fd..05758bd 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -80,7 +80,7 @@ def _compute_zonal_mean_gpu(var, latitudes, area):
         value[basin] = np.empty((times, levs, 180))
         gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
         for t in range(times):
-            for lev in range(levels):
+            for lev in range(levs):
                 gpu_value = _zonal_mean_gpu(gpu_var[t, lev, :, :], gpu_area,
                                             gpu_lats, gpu_total, gpu_weights)
                 value[basin][t][lev][:] = gpu_value.copy_to_host()
-- 
GitLab


From adaa8caaca9b60e66d394e58acfd36b8af82f840 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 16:22:41 +0200
Subject: [PATCH 12/19] Add missing indices

---
 diagonals/zonmean.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 05758bd..51c7aa6 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -105,6 +105,7 @@ def _zonal_mean_cpu(variable, weight, latitude):
 
 @cuda.jit()
 def _zonal_mean_gpu(variable, weight, latitude, total, weights):
+    i, j = cuda.grid(2)
     total[i, j] = 0.0
     weights[i, j] = 0.0
     if(i >= total.shape[0]):
-- 
GitLab


From db1b28785a37b8f22dc26bade29bb94bc61424c4 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 16:26:44 +0200
Subject: [PATCH 13/19] Delete wrong indexed array

---
 diagonals/zonmean.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 51c7aa6..a59bfd9 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -106,15 +106,13 @@ def _zonal_mean_cpu(variable, weight, latitude):
 @cuda.jit()
 def _zonal_mean_gpu(variable, weight, latitude, total, weights):
     i, j = cuda.grid(2)
-    total[i, j] = 0.0
-    weights[i, j] = 0.0
     if(i >= total.shape[0]):
         return
     if(weight[i, j] != 0.0):
         bin_value = int(round(latitude[i, j]) + 90)
         weights[bin_value] += weight[i, j]
         total[bin_value] += variable[i, j] * weight[i, j]
-    return total[i, j] / weights[i, j]
+    return total / weights
 
 
 def get_basin_area(areacello, basins):
-- 
GitLab


From 32fc9d6eff9c48f57f77137e37af4753627352b4 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Mon, 1 Jul 2019 16:49:30 +0200
Subject: [PATCH 14/19] segmentation fault

---
 diagonals/examples/examples-zonalmean.py | 2 +-
 diagonals/zonmean.py                     | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/diagonals/examples/examples-zonalmean.py b/diagonals/examples/examples-zonalmean.py
index b3fcd28..5b5b3a7 100644
--- a/diagonals/examples/examples-zonalmean.py
+++ b/diagonals/examples/examples-zonalmean.py
@@ -28,7 +28,7 @@ def main():
     start = datetime.datetime.now()
     logger.info('Starting at %s', start)
     mesh = Nemo(MESH_FILE, REGIONS_FILE)
-    with diagonals.CONFIG.context(use_gpu=False):
+    with diagonals.CONFIG.context(use_gpu=True):
         lats = mesh.get_mesh_var('nav_lat', dtype=np.float32)
         areacello = mesh.get_areacello()
         var = load_thetao()
diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index a59bfd9..9ac5296 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -81,9 +81,9 @@ def _compute_zonal_mean_gpu(var, latitudes, area):
         gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
         for t in range(times):
             for lev in range(levs):
-                gpu_value = _zonal_mean_gpu(gpu_var[t, lev, :, :], gpu_area,
-                                            gpu_lats, gpu_total, gpu_weights)
-                value[basin][t][lev][:] = gpu_value.copy_to_host()
+                 _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area, gpu_lats, gpu_total, gpu_weights)
+                 value[basin][t][lev][:] = (gpu_total.copy_to_host() /
+                                            gpu_weights.copy_to_host())
     del gpu_area, gpu_value, gpu_var
 
     return value
@@ -109,10 +109,10 @@ def _zonal_mean_gpu(variable, weight, latitude, total, weights):
     if(i >= total.shape[0]):
         return
     if(weight[i, j] != 0.0):
+      
         bin_value = int(round(latitude[i, j]) + 90)
         weights[bin_value] += weight[i, j]
         total[bin_value] += variable[i, j] * weight[i, j]
-    return total / weights
 
 
 def get_basin_area(areacello, basins):
-- 
GitLab


From 6cabfa8ba915d150b71331098fb8de61838ef8ad Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 2 Jul 2019 13:16:42 +0200
Subject: [PATCH 15/19] Add warning due to numba bug

---
 diagonals/zonmean.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 9ac5296..9bc5397 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -41,9 +41,12 @@ def compute_zonmean(var, lats, area):
         List containing regional mean for variable var
     """
     if diagonals.CONFIG.use_gpu:
-        zonmean = _compute_zonal_mean_gpu(var, lats, area)
-    else:
-        zonmean = _compute_zonal_mean_cpu(var, lats, area)
+            if diagonals.CONFIG.use_gpu:
+        logger.warning('GPU routines not implemented yet'
+                       'for zonmean diagnostic due to a compiler bug.'
+                       'Using CPU instead until bug has been fixed')
+
+    zonmean = _compute_zonal_mean_cpu(var, lats, area)
     return zonmean
 
 
@@ -81,7 +84,8 @@ def _compute_zonal_mean_gpu(var, latitudes, area):
         gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
         for t in range(times):
             for lev in range(levs):
-                 _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area, gpu_lats, gpu_total, gpu_weights)
+                 _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area,
+                                              gpu_lats, gpu_total, gpu_weights)
                  value[basin][t][lev][:] = (gpu_total.copy_to_host() /
                                             gpu_weights.copy_to_host())
     del gpu_area, gpu_value, gpu_var
@@ -109,7 +113,7 @@ def _zonal_mean_gpu(variable, weight, latitude, total, weights):
     if(i >= total.shape[0]):
         return
     if(weight[i, j] != 0.0):
-      
+
         bin_value = int(round(latitude[i, j]) + 90)
         weights[bin_value] += weight[i, j]
         total[bin_value] += variable[i, j] * weight[i, j]
-- 
GitLab


From 748d9a48757a6810c74a6566f7c442d867c3ae83 Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 2 Jul 2019 13:21:27 +0200
Subject: [PATCH 16/19] Fix style issues

---
 diagonals/zonmean.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 9bc5397..ccffc24 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -41,7 +41,6 @@ def compute_zonmean(var, lats, area):
         List containing regional mean for variable var
     """
     if diagonals.CONFIG.use_gpu:
-            if diagonals.CONFIG.use_gpu:
         logger.warning('GPU routines not implemented yet'
                        'for zonmean diagnostic due to a compiler bug.'
                        'Using CPU instead until bug has been fixed')
@@ -84,10 +83,10 @@ def _compute_zonal_mean_gpu(var, latitudes, area):
         gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
         for t in range(times):
             for lev in range(levs):
-                 _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area,
-                                              gpu_lats, gpu_total, gpu_weights)
-                 value[basin][t][lev][:] = (gpu_total.copy_to_host() /
-                                            gpu_weights.copy_to_host())
+                _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area,
+                                             gpu_lats, gpu_total, gpu_weights)
+                value[basin][t][lev][:] = (gpu_total.copy_to_host() /
+                                           gpu_weights.copy_to_host())
     del gpu_area, gpu_value, gpu_var
 
     return value
-- 
GitLab


From 0fd60fbabd60799bc12989c6976021e6002c28db Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 2 Jul 2019 16:37:05 +0200
Subject: [PATCH 17/19] Add documentation and comment GPU functions for now

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

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
new file mode 100644
index 0000000..33747e9
--- /dev/null
+++ b/diagonals/zonmean.py
@@ -0,0 +1,182 @@
+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 njit
+from numba import int32, int64, float32, float64
+from numba.cuda.cudadrv import driver
+
+import diagonals
+
+__all__ = ['compute_zonmean', 'get_basin_area']
+
+
+def compute_zonmean(var, lats, area):
+    """Function that checks device and calls computing functions.
+
+    Checks if the computations are going performed in the CPU or the GPU:
+
+    Parameters
+    ----------
+    var : float32
+        Masked array containing variable data.
+    lats : float32
+        Masked array containing latitude points.
+    area : float32
+        Masked array containing cell area.
+
+    Returns
+    -------
+    zonmean: float32
+        List containing zonal mean for variable var
+    """
+    if diagonals.CONFIG.use_gpu:
+        logger.warning('GPU routines not implemented yet'
+                       'for zonmean diagnostic due to a compiler bug.'
+                       'Using CPU instead until bug gets fixed')
+
+    zonmean = _compute_zonal_mean_cpu(var, lats, area)
+    return zonmean
+
+
+def _compute_zonal_mean_cpu(var, lats, area):
+    """Function that calls computing functions in the CPU.
+
+    Loops over time-steps and depth levels to compute the zonal mean:
+
+    Parameters
+    ----------
+    var : float32
+        Masked array containing variable to average.
+    lats: float32
+        Masked array containing latitude points.
+    area: float32
+        List of masked arrays containing the area for each basin.
+
+    Returns
+    -------
+    value: float32
+        List of masked arrays containing a variable's zonal mean for
+        each depth level and timestep.
+    """
+    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, levs, 180))
+        for t in range(times):
+            for lev in range(levs):
+                value[basin][t][lev][:] = _zonal_mean_cpu(var[t, lev, :, :],
+                                                          weight,
+                                                          lats)
+    return value
+
+
+#def _compute_zonal_mean_gpu(var, latitudes, area):
+#    times = var.shape[0]
+#    levs = var.shape[1]
+#    lats = var.shape[2]
+#    lons = var.shape[3]
+#    value = {}
+
+#    block = (128, 1, 1)
+#    grid_size = (lons // block[0]) + 1
+#    grid = (grid_size, lats)
+#    gpu_var = cuda.to_device(var.astype(np.float32))
+#    gpu_lats = cuda.to_device(latitudes.astype(np.float32))
+#    gpu_total = cuda.device_array(180, dtype=np.float32)
+#    gpu_weights = cuda.device_array(180, dtype=np.float32)
+#    for basin, mask in area.items():
+#        value[basin] = np.empty((times, levs, 180))
+#        gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32))
+#        for t in range(times):
+#            for lev in range(levs):
+#                _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area,
+#                                             gpu_lats, gpu_total, gpu_weights)
+#                value[basin][t][lev][:] = (gpu_total.copy_to_host() /
+#                                           gpu_weights.copy_to_host())
+#    del gpu_area, gpu_value, gpu_var
+
+#    return value
+
+
+@numba.njit()
+def _zonal_mean_cpu(variable, weight, latitude):
+    """Compiled function in executed on the CPU.
+
+    Computes the zonal mean.
+
+    Parameters
+    ----------
+    var : float32
+        Masked array containing variable data.
+    latitude : float32
+        Masked array containing latitude points.
+    weight : float32
+        Masked array containing cell weights.
+
+    Returns
+    -------
+    zonmean: float32
+        List containing zonal mean for variable var
+    """
+    total = np.zeros(180, np.float32)
+    weights = np.zeros(180, np.float32)
+    for i in range(variable.shape[0]):
+        for j in range(variable.shape[1]):
+            if weight[i, j] == 0:
+                continue
+            bin_value = int(round(latitude[i, j]) + 90)
+            weights[bin_value] += weight[i, j]
+            total[bin_value] += variable[i, j] * weight[i, j]
+    return total / weights
+
+
+#@cuda.jit()
+#def _zonal_mean_gpu(variable, weight, latitude, total, weights):
+#    i, j = cuda.grid(2)
+#    if(i >= total.shape[0]):
+#        return
+#    if(weight[i, j] != 0.0):
+
+#        bin_value = int(round(latitude[i, j]) + 90)
+#        weights[bin_value] += weight[i, j]
+#        total[bin_value] += variable[i, j] * weight[i, j]
+
+
+def get_basin_area(areacello, basins):
+    basin_areas = {}
+    for basin in basins:
+        basin_areas[basin] = _compute_basin_area(areacello, basins[basin])
+    return basin_areas
+
+
+@vectorize(['float32(float32, float32)'], target='cpu')
+def _compute_basin_area(areacello, basin):
+    """Vectorized numba function executed on the cpu that computes the area
+       for each basin:
+
+    Parameters
+    ----------
+    areacello : float32
+        Masked array containing areacello.
+    basin: float32
+        Masked array containing the mask for a given basin.
+
+    Returns
+    -------
+    areacello*basin : float32
+        Masked array containing the area for a given basin.
+    """
+    return areacello*basin
-- 
GitLab


From 33bd2235d00b1267f372bcb77ce91714cefb3c5c Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Tue, 2 Jul 2019 16:40:04 +0200
Subject: [PATCH 18/19] Solve merge conflict

---
 diagonals/zonmean.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index d774257..33747e9 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -143,7 +143,6 @@ def _zonal_mean_cpu(variable, weight, latitude):
     return total / weights
 
 
-<<<<<<< HEAD
 #@cuda.jit()
 #def _zonal_mean_gpu(variable, weight, latitude, total, weights):
 #    i, j = cuda.grid(2)
@@ -154,18 +153,6 @@ def _zonal_mean_cpu(variable, weight, latitude):
 #        bin_value = int(round(latitude[i, j]) + 90)
 #        weights[bin_value] += weight[i, j]
 #        total[bin_value] += variable[i, j] * weight[i, j]
-=======
-@cuda.jit()
-def _zonal_mean_gpu(variable, weight, latitude, total, weights):
-    i, j = cuda.grid(2)
-    if(i >= total.shape[0]):
-        return
-    if(weight[i, j] != 0.0):
-
-        bin_value = int(round(latitude[i, j]) + 90)
-        weights[bin_value] += weight[i, j]
-        total[bin_value] += variable[i, j] * weight[i, j]
->>>>>>> 466e36e43edad7f4bc2998ef1c0ee5d0d46a532d
 
 
 def get_basin_area(areacello, basins):
-- 
GitLab


From e9194d681c7a9f6a03ce500f30b265757898721f Mon Sep 17 00:00:00 2001
From: sloosvel <saskia.loosveldt@bsc.es>
Date: Wed, 24 Jul 2019 10:51:56 +0200
Subject: [PATCH 19/19] Fix lint issues

---
 diagonals/zonmean.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py
index 33747e9..127dd7d 100644
--- a/diagonals/zonmean.py
+++ b/diagonals/zonmean.py
@@ -83,7 +83,7 @@ def _compute_zonal_mean_cpu(var, lats, area):
     return value
 
 
-#def _compute_zonal_mean_gpu(var, latitudes, area):
+# def _compute_zonal_mean_gpu(var, latitudes, area):
 #    times = var.shape[0]
 #    levs = var.shape[1]
 #    lats = var.shape[2]
@@ -143,8 +143,8 @@ def _zonal_mean_cpu(variable, weight, latitude):
     return total / weights
 
 
-#@cuda.jit()
-#def _zonal_mean_gpu(variable, weight, latitude, total, weights):
+# @cuda.jit()
+# def _zonal_mean_gpu(variable, weight, latitude, total, weights):
 #    i, j = cuda.grid(2)
 #    if(i >= total.shape[0]):
 #        return
-- 
GitLab