From 694f538bcfd695ad730d103d9fdfd2d526d4c76a Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 23 Jan 2020 13:05:41 +0100 Subject: [PATCH 1/3] Add psi computations --- diagonals/psi.py | 121 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 diagonals/psi.py diff --git a/diagonals/psi.py b/diagonals/psi.py new file mode 100644 index 0000000..2a4b068 --- /dev/null +++ b/diagonals/psi.py @@ -0,0 +1,121 @@ +import os + +import numpy as np + +import iris +import iris.cube +import iris.analysis + +import warnings +import datetime + +import numba +from numba import njit +from numba import vectorize +from numba import cuda +from numba import int32, int64, float32, float64 +from numba.cuda.cudadrv import driver + +import diagonals + +__all__ = ['compute'] + +def compute(basins, e2u, e1v, e3u, e3v, uo, vo): + """Function that checks device and calls computing functions. + + Checks if the computations are going performed in the CPU or the GPU: + + Parameters + ---------- + basins : list + List of masked arrays containing the mask for each basin. + e1u: float32 + Masked array containing variable e1u. + e1v: float32 + Masked array containing variable e1v. + e3u: float32 + Masked array containing variable e3u. + e3v: float32 + Masked array containing variable e3v. + uo : float32 + Masked array containing Sea Water X Velocity data. + vo : float32 + Masked array containing Sea Water Y Velocity data. + + Returns + ------- + vsftbarot: list + List of masked arrays containing the gyre strength. + """ + area_u = {} + area_v = {} + for basin in basins: + area_u[basin] = _compute_area(e2u, e3u, basins[basin]) + area_v[basin] = _compute_area(e1v, e3v, basins[basin]) + + del e2u, e1v, e3u, e3v + + vsftbarot = _compute_psi_cpu( + uo, vo, area_u, area_v, basins + ) + + return vsftbarot + +def _compute_psi_cpu(uo, vo, area_u, area_v, basins): + + psi = {} + for basin, mask in area_u.items(): + uo_lev = np.sum(uo*area_u[basin], axis=1) + vo_lev = np.sum(vo*area_v[basin], axis=1) + + del uo, vo, area_u[basin], area_v[basin] + + dpsiu = _u_cumsum(uo_lev) + dpsiv = _v_cumsum(vo_lev) + + dpsi = 0.5 * (dpsiu + dpsiv) + + ref_point = dpsi[:,-1,-1] + + psi[basin] = ( + dpsi - ref_point[:, np.newaxis, np.newaxis] + ) * basins[basin] + + return psi + +#@njit looks like it's not worth it +def _u_cumsum(uo): + dpsiu = np.zeros(uo.shape) + for j in range(1, dpsiu.shape[1]): + dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :] + return dpsiu + +#@njit looks like it's not worth it +def _v_cumsum(vo): + dpsiv = np.zeros(vo.shape) + for i in range(dpsiv.shape[2]-2,-1,-1): + dpsiv[:, : ,i] = dpsiv[:, :, i+1] - vo[: , :, i] + return dpsiv + +@vectorize(['float32(float32, float32, float32)'], target='cpu') +def _compute_area(e, e3, basin): + """Vectorized numba function executed in the CPU. + + Calculates cell area for each basin: + + Parameters + ---------- + e1: float32 + Masked array containing variable e1 or e2. + e3: float32 + Masked array containing variable e3. + basin : float32 + Masked array containing a basin mask. + + Returns + ------- + area: float32 + Masked array containing the cell area for a given basin. + """ + area = e * e3 + return area -- GitLab From 53adb891bab8589edb935aa5bb90ac3919ac5e49 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 23 Jan 2020 13:24:12 +0100 Subject: [PATCH 2/3] Fix lint --- diagonals/psi.py | 19 +++++++++++++------ diagonals/regmean.py | 2 +- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/diagonals/psi.py b/diagonals/psi.py index 2a4b068..cc4342b 100644 --- a/diagonals/psi.py +++ b/diagonals/psi.py @@ -20,6 +20,7 @@ import diagonals __all__ = ['compute'] + def compute(basins, e2u, e1v, e3u, e3v, uo, vo): """Function that checks device and calls computing functions. @@ -61,6 +62,7 @@ def compute(basins, e2u, e1v, e3u, e3v, uo, vo): return vsftbarot + def _compute_psi_cpu(uo, vo, area_u, area_v, basins): psi = {} @@ -75,7 +77,7 @@ def _compute_psi_cpu(uo, vo, area_u, area_v, basins): dpsi = 0.5 * (dpsiu + dpsiv) - ref_point = dpsi[:,-1,-1] + ref_point = dpsi[:, -1, -1] psi[basin] = ( dpsi - ref_point[:, np.newaxis, np.newaxis] @@ -83,20 +85,25 @@ def _compute_psi_cpu(uo, vo, area_u, area_v, basins): return psi -#@njit looks like it's not worth it +# @njit looks like it's not worth it + + def _u_cumsum(uo): dpsiu = np.zeros(uo.shape) for j in range(1, dpsiu.shape[1]): - dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :] + dpsiu[:, j, :] = dpsiu[:, j-1, :] - uo[:, j, :] return dpsiu -#@njit looks like it's not worth it +# @njit looks like it's not worth it + + def _v_cumsum(vo): dpsiv = np.zeros(vo.shape) - for i in range(dpsiv.shape[2]-2,-1,-1): - dpsiv[:, : ,i] = dpsiv[:, :, i+1] - vo[: , :, i] + for i in range(dpsiv.shape[2]-2, -1, -1): + dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[: ,: , i] return dpsiv + @vectorize(['float32(float32, float32, float32)'], target='cpu') def _compute_area(e, e3, basin): """Vectorized numba function executed in the CPU. diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 3bd88aa..1cdaced 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -193,7 +193,7 @@ def _compute_regmean_levels_cpu(var, basins, volume): for t in range(times): for l in range(levs): regmean[t, l] = np.ma.average(var[t, l, :, :], axis=(0, 1), - weights=np.squeeze(w)[l, :, :]) + weights=np.squeeze(w)[l, :, :]) regmean_total[basin] = regmean return regmean_total -- GitLab From 83ec0ffcf56019fd827a0097075748132e04e4fa Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 23 Jan 2020 13:30:14 +0100 Subject: [PATCH 3/3] Fix lint --- diagonals/psi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/psi.py b/diagonals/psi.py index cc4342b..d1b1f4a 100644 --- a/diagonals/psi.py +++ b/diagonals/psi.py @@ -100,7 +100,7 @@ def _u_cumsum(uo): def _v_cumsum(vo): dpsiv = np.zeros(vo.shape) for i in range(dpsiv.shape[2]-2, -1, -1): - dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[: ,: , i] + dpsiv[:, :, i] = dpsiv[:, :, i+1] - vo[:, :, i] return dpsiv -- GitLab