From 99e99c37da48f5e89420c813adb8db89f79bcefc Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 9 Oct 2020 15:51:29 +0200 Subject: [PATCH 1/7] Add bigthetao to thetao conversion --- diagonals/density.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/diagonals/density.py b/diagonals/density.py index 010a016..4510722 100644 --- a/diagonals/density.py +++ b/diagonals/density.py @@ -3,7 +3,7 @@ from numba import vectorize # import diagonals -__all__ = ['compute'] +__all__ = ['compute', 'convert'] @vectorize(['float32(float32, float32, float32)'], target='cpu') @@ -219,3 +219,29 @@ def compute(so, bigthetao, p): # b = b / ss return rho # ,a,b,r0,r + +@vectorize(['float32(float32, float32)'], target='cpu') +def convert(so, thetao): + x = np.sqrt(0.0248826675584615*so) + y = thetao*0.025e0 + enthalpy = 61.01362420681071e0 + y*(168776.46138048015e0 + + y*(-2735.2785605119625e0 + y*(2574.2164453821433e0 + + y*(-1536.6644434977543e0 + y*(545.7340497931629e0 + + (-50.91091728474331e0 - 18.30489878927802e0*y)* + y))))) + x**2*(268.5520265845071e0 + y*(-12019.028203559312e0 + + y*(3734.858026725145e0 + y*(-2046.7671145057618e0 + + y*(465.28655623826234e0 + (-0.6370820302376359e0 - + 10.650848542359153e0*y)*y)))) + + x*(937.2099110620707e0 + y*(588.1802812170108e0+ + y*(248.39476522971285e0 + (-3.871557904936333e0- + 2.6268019854268356e0*y)*y)) + + x*(-1687.914374187449e0 + x*(246.9598888781377e0 + + x*(123.59576582457964e0 - 48.5891069025409e0*x)) + + y*(936.3206544460336e0 + + y*(-942.7827304544439e0 + y*(369.4389437509002e0 + + (-33.83664947895248e0 - 9.987880382780322e0*y)*y)))))) + + # enthalpy/Specific heat + bigthetao = enthalpy/3991.86795711963 + + return bigthetao -- GitLab From 6ca6459ebfefee921bf76b6c4ab5f8d895a8feb9 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 13 Oct 2020 17:15:04 +0200 Subject: [PATCH 2/7] Adapt to dask --- diagonals/ohc.py | 4 +++- diagonals/regmean.py | 29 +++++++++++++++-------------- diagonals/regsum.py | 8 +++++--- 3 files changed, 23 insertions(+), 18 deletions(-) diff --git a/diagonals/ohc.py b/diagonals/ohc.py index e5cd6c6..ff0b04b 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -5,6 +5,8 @@ from numba import cuda import diagonals +import dask.array as da + __all__ = ['get_weights', 'compute', 'get_basin_area'] @@ -170,7 +172,7 @@ def _compute_ohc_cpu(layers, thetao, weights, area): 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 = da.sum(ohc_basin, axis=(1, 2)) ohc1d_total.append(ohc1d) return ohc, ohc1d_total diff --git a/diagonals/regmean.py b/diagonals/regmean.py index efce77b..3418922 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -1,5 +1,6 @@ import logging import numpy as np +import dask.array as da from numba import vectorize @@ -111,14 +112,14 @@ def _compute_regmean_2d_cpu(var, basins, area): regmean_total: float32 List containing regional mean for variable var. """ - times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): weights = _compute_weights_2d(mask, area) - regmean = np.empty(times) - for t in range(times): - regmean[t] = np.ma.average( - var[t, :, :], axis=(0, 1), weights=np.squeeze(weights)) + regmean = da.ma.average( + var, + axis=(1, 2), + weights=da.broadcast_to(weights, var.shape) + ) regmean_total[basin] = regmean return regmean_total @@ -142,14 +143,14 @@ def _compute_regmean_3d_cpu(var, basins, volume): regmean_total: float32 List containing the regional mean for variable var. """ - times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): weights = _compute_weights_3d(mask, volume) - regmean = np.empty(times) - for t in range(times): - regmean[t] = np.ma.average( - var[t, :, :, :], axis=(0, 1, 2), weights=np.squeeze(weights)) + regmean = da.ma.average( + var, + axis=(1, 2, 3), + weights=da.broadcast_to(weights, var.shape) + ) regmean_total[basin] = regmean return regmean_total @@ -182,10 +183,10 @@ def _compute_regmean_levels_cpu(var, basins, volume): w = _compute_weights_3d(mask, volume) for time in range(times): for lev in range(levs): - regmean[time, lev] = np.ma.average( - var[time, lev, :, :], - axis=(0, 1), - weights=np.squeeze(w)[lev, :, :] + regmean = da.ma.average( + var, + axis=(2, 3), + weights=da.broadcast_to(weights, var.shape) ) regmean_total[basin] = regmean return regmean_total diff --git a/diagonals/regsum.py b/diagonals/regsum.py index d4ea3e9..88c63cd 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -4,6 +4,8 @@ from numba import vectorize import diagonals +from dask.array import da + logger = logging.getLogger(__name__) __all__ = ['compute_regsum_2d', 'compute_regsum_3d', 'compute_regsum_levels'] @@ -83,7 +85,7 @@ def _compute_regsum_2d_cpu(var, basins, area): regsum = {} for basin, mask in basins.items(): weighted_var = _weigh_var_2d(var, mask, area) - regsum[basin] = np.sum(weighted_var, axis=(1, 2)) + regsum[basin] = da.sum(weighted_var, axis=(1, 2)) return regsum @@ -109,7 +111,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) - regsum[basin] = np.sum(weighted_var, axis=(1, 2, 3)) + regsum[basin] = da.sum(weighted_var, axis=(1, 2, 3)) return regsum @@ -134,7 +136,7 @@ def compute_regsum_levels(var, basins, volume, tmask): regsum = {} for basin, mask in basins.items(): weighted_var = _weigh_var_3d(var, mask, volume, tmask) - regsum[basin] = np.sum(weighted_var, axis=(2, 3)) + regsum[basin] = da.sum(weighted_var, axis=(2, 3)) return regsum -- GitLab From 940ed94ae9786009ddba56ae838008b586816811 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 14 Oct 2020 09:33:51 +0200 Subject: [PATCH 3/7] Adapt to dask --- diagonals/moc.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/diagonals/moc.py b/diagonals/moc.py index 211d62e..ef40514 100644 --- a/diagonals/moc.py +++ b/diagonals/moc.py @@ -5,6 +5,7 @@ from numba import vectorize from numba import guvectorize from numba import cuda from numba import float32 +import dask.array as da import diagonals @@ -69,7 +70,8 @@ def _compute_moc_cpu(vo, area): """ moc = {} for basin, mask in area.items(): - moc_basin = np.sum(_multiply_vo_basin(vo, mask), axis=3) + moc_basin = da.sum(_multiply_vo_basin(vo, mask), axis=3) + moc_basin = moc_basin.compute() _vertical_cumsum(moc_basin, moc_basin) moc[basin] = moc_basin return moc -- GitLab From 52faebf20402bc2cf67a5ad1aeef471deb7992fe Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 11 Jan 2021 09:54:41 +0100 Subject: [PATCH 4/7] Bump version --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 48cb8a4..06225e7 100644 --- a/setup.py +++ b/setup.py @@ -36,7 +36,7 @@ REQUIREMENTS = { setup(name='diagonals', - version='0.3.2', + version='0.3.3', 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 9f8853947dfef15d768469f58de8954ac77e463f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 11 Jan 2021 10:10:58 +0100 Subject: [PATCH 5/7] Fix import --- diagonals/regsum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index 88c63cd..1b63e74 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -4,7 +4,7 @@ from numba import vectorize import diagonals -from dask.array import da +import dask.array as da logger = logging.getLogger(__name__) -- GitLab From fd8896fa16f8070ab2869aa6c4d022e93a63258a Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 11 Jan 2021 10:35:09 +0100 Subject: [PATCH 6/7] Fix tests --- diagonals/regmean.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 3418922..bd79f12 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -180,7 +180,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) + weights = _compute_weights_3d(mask, volume) for time in range(times): for lev in range(levs): regmean = da.ma.average( -- GitLab From cec4357b6b9285b441784b1cda92261e19200f67 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 11 Jan 2021 11:23:38 +0100 Subject: [PATCH 7/7] Fix flake --- diagonals/density.py | 85 +++++++++++++++++++++++++------------------- diagonals/ohc.py | 2 -- diagonals/regmean.py | 8 ++--- diagonals/regsum.py | 1 - 4 files changed, 53 insertions(+), 43 deletions(-) diff --git a/diagonals/density.py b/diagonals/density.py index 4510722..44701bb 100644 --- a/diagonals/density.py +++ b/diagonals/density.py @@ -1,3 +1,5 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- import numpy as np from numba import vectorize @@ -28,24 +30,28 @@ def compute(so, bigthetao, p): """ # reduced variables - SAu = 40.*35.16504/35. + + SAu = 40. * 35.16504 / 35. CTu = 40. Zu = 1e4 deltaS = 32. - ss = np.sqrt((so+deltaS)/SAu) + ss = np.sqrt((so + deltaS) / SAu) tt = bigthetao / CTu pp = p / Zu # vertical reference profile of density + R00 = 4.6494977072e+01 R01 = -5.2099962525e+00 R02 = 2.2601900708e-01 R03 = 6.4326772569e-02 R04 = 1.5616995503e-02 R05 = -1.7243708991e-03 - r0 = (((((R05*pp + R04)*pp + R03)*pp + R02)*pp + R01)*pp + R00)*pp + r0 = (((((R05 * pp + R04) * pp + R03) * pp + R02) * pp + R01) * pp + + R00) * pp # density anomaly + R000 = 8.0189615746e+02 R100 = 8.6672408165e+02 R200 = -1.7864682637e+03 @@ -99,21 +105,24 @@ def compute(so, bigthetao, p): R103 = -1.8507636718e-02 R013 = 3.7969820455e-01 - rz3 = R013*tt + R103*ss + R003 - rz2 = (R022*tt+R112*ss+R012)*tt+(R202*ss+R102)*ss+R002 - rz1 = (((R041*tt+R131*ss+R031)*tt + - (R221*ss+R121)*ss+R021)*tt + - ((R311*ss+R211)*ss+R111)*ss+R011)*tt + \ - (((R401*ss+R301)*ss+R201)*ss+R101)*ss+R001 - rz0 = (((((R060*tt+R150*ss+R050)*tt + - (R240*ss+R140)*ss+R040)*tt + - ((R330*ss+R230)*ss+R130)*ss+R030)*tt + - (((R420*ss+R320)*ss+R220)*ss+R120)*ss+R020)*tt + - ((((R510*ss+R410)*ss+R310)*ss+R210)*ss+R110)*ss+R010)*tt + \ - (((((R600*ss+R500)*ss+R400)*ss+R300)*ss+R200)*ss+R100)*ss+R000 - r = ((rz3*pp + rz2)*pp + rz1)*pp + rz0 + rz3 = R013 * tt + R103 * ss + R003 + rz2 = (R022 * tt + R112 * ss + R012) * tt + (R202 * ss + R102) * ss \ + + R002 + rz1 = (((R041 * tt + R131 * ss + R031) * tt + (R221 * ss + R121) + * ss + R021) * tt + ((R311 * ss + R211) * ss + R111) * ss + + R011) * tt + (((R401 * ss + R301) * ss + R201) * ss + + R101) * ss + R001 + rz0 = (((((R060 * tt + R150 * ss + R050) * tt + (R240 * ss + R140) + * ss + R040) * tt + ((R330 * ss + R230) * ss + R130) * ss + + R030) * tt + (((R420 * ss + R320) * ss + R220) * ss + + R120) * ss + R020) * tt + ((((R510 * ss + R410) * ss + + R310) * ss + R210) * ss + R110) * ss + R010) * tt \ + + (((((R600 * ss + R500) * ss + R400) * ss + R300) * ss + R200) + * ss + R100) * ss + R000 + r = ((rz3 * pp + rz2) * pp + rz1) * pp + rz0 # in-situ density + rho = r + r0 # # thermal expansion a @@ -220,28 +229,32 @@ def compute(so, bigthetao, p): return rho # ,a,b,r0,r + @vectorize(['float32(float32, float32)'], target='cpu') def convert(so, thetao): - x = np.sqrt(0.0248826675584615*so) - y = thetao*0.025e0 - enthalpy = 61.01362420681071e0 + y*(168776.46138048015e0 + - y*(-2735.2785605119625e0 + y*(2574.2164453821433e0 + - y*(-1536.6644434977543e0 + y*(545.7340497931629e0 + - (-50.91091728474331e0 - 18.30489878927802e0*y)* - y))))) + x**2*(268.5520265845071e0 + y*(-12019.028203559312e0 + - y*(3734.858026725145e0 + y*(-2046.7671145057618e0 + - y*(465.28655623826234e0 + (-0.6370820302376359e0 - - 10.650848542359153e0*y)*y)))) + - x*(937.2099110620707e0 + y*(588.1802812170108e0+ - y*(248.39476522971285e0 + (-3.871557904936333e0- - 2.6268019854268356e0*y)*y)) + - x*(-1687.914374187449e0 + x*(246.9598888781377e0 + - x*(123.59576582457964e0 - 48.5891069025409e0*x)) + - y*(936.3206544460336e0 + - y*(-942.7827304544439e0 + y*(369.4389437509002e0 + - (-33.83664947895248e0 - 9.987880382780322e0*y)*y)))))) - + x = np.sqrt(0.0248826675584615 * so) + y = thetao * 0.025e0 + enthalpy = 61.01362420681071e0 + y * (168776.46138048015e0 + y + * (-2735.2785605119625e0 + y * (2574.2164453821433e0 + y + * (-1536.6644434977543e0 + y * (545.7340497931629e0 + + (-50.91091728474331e0 - 18.30489878927802e0 * y) * y))))) \ + + x ** 2 * (268.5520265845071e0 + y * (-12019.028203559312e0 + + y * (3734.858026725145e0 + y + * (-2046.7671145057618e0 + y + * (465.28655623826234e0 + (-0.6370820302376359e0 + - 10.650848542359153e0 * y) * y)))) + x + * (937.2099110620707e0 + y * (588.1802812170108e0 + + y * (248.39476522971285e0 + (-3.871557904936333e0 + - 2.6268019854268356e0 * y) * y)) + x + * (-1687.914374187449e0 + x * (246.9598888781377e0 + + x * (123.59576582457964e0 - 48.5891069025409e0 + * x)) + y * (936.3206544460336e0 + y + * (-942.7827304544439e0 + y * (369.4389437509002e0 + + (-33.83664947895248e0 - 9.987880382780322e0 * y) + * y)))))) + # enthalpy/Specific heat - bigthetao = enthalpy/3991.86795711963 + + bigthetao = enthalpy / 3991.86795711963 return bigthetao diff --git a/diagonals/ohc.py b/diagonals/ohc.py index a7e4ce3..fb83754 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -6,8 +6,6 @@ from numba import cuda import diagonals -import dask.array as da - __all__ = ['get_weights', 'compute', 'get_basin_area'] diff --git a/diagonals/regmean.py b/diagonals/regmean.py index bd79f12..0523d0b 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -116,8 +116,8 @@ def _compute_regmean_2d_cpu(var, basins, area): for basin, mask in basins.items(): weights = _compute_weights_2d(mask, area) regmean = da.ma.average( - var, - axis=(1, 2), + var, + axis=(1, 2), weights=da.broadcast_to(weights, var.shape) ) regmean_total[basin] = regmean @@ -147,8 +147,8 @@ def _compute_regmean_3d_cpu(var, basins, volume): for basin, mask in basins.items(): weights = _compute_weights_3d(mask, volume) regmean = da.ma.average( - var, - axis=(1, 2, 3), + var, + axis=(1, 2, 3), weights=da.broadcast_to(weights, var.shape) ) regmean_total[basin] = regmean diff --git a/diagonals/regsum.py b/diagonals/regsum.py index 1b63e74..d105ad3 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -1,5 +1,4 @@ import logging -import numpy as np from numba import vectorize import diagonals -- GitLab