From c9ce8f51f6f6b5b606809f388afb1ca1f282ac0c Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Mon, 6 May 2019 17:51:29 +0200 Subject: [PATCH 01/54] ohc refactored --- .gitignore | 4 +- diagonals/__init__.py | 70 +++++ .../{diagnostics => examples}/__init__.py | 0 .../examples/examples-oceanheatcontent.py | 256 +++--------------- .../oceanheatcontent.py => ohc.py} | 72 +++-- diagonals/{diagnostics => }/regmean.py | 0 diagonals/{diagnostics => }/regsum.py | 0 diagonals/{diagnostics => }/siarea.py | 0 8 files changed, 135 insertions(+), 267 deletions(-) rename diagonals/{diagnostics => examples}/__init__.py (100%) rename diagonals/{diagnostics/oceanheatcontent.py => ohc.py} (86%) rename diagonals/{diagnostics => }/regmean.py (100%) rename diagonals/{diagnostics => }/regsum.py (100%) rename diagonals/{diagnostics => }/siarea.py (100%) diff --git a/.gitignore b/.gitignore index 2e06392..b319c4e 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,6 @@ test/report/* htmlcov .pytest_cache prof -/.vscode +.vscode .eggs - +*.pyc diff --git a/diagonals/__init__.py b/diagonals/__init__.py index e69de29..3c732db 100644 --- a/diagonals/__init__.py +++ b/diagonals/__init__.py @@ -0,0 +1,70 @@ +import threading +import contextlib +import logging +import six + +import numba.cuda + +logger = logging.getLogger(__name__) + +class Config(threading.local): + """Run-time configuration controller.""" + + def __init__(self, use_gpu=None): + """ + A container for run-time options controls. + To adjust the values simply update the relevant attribute from + within your code. For example:: + diagonals.CONFIG.use_gpu = False + If Iris code is executed with multiple threads, note the values of + these options are thread-specific. + + """ + if use_gpu is None: + use_gpu = numba.cuda.is_available() + self.__dict__['use_gpu'] = use_gpu + self.use_gpu = use_gpu + + def __repr__(self): + return 'Config(use_gpu={0})'.format(self.use_gpu) + + def __setattr__(self, name, value): + if name not in self.__dict__: + msg = "'Config' object has no attribute {!r}".format(name) + raise AttributeError(msg) + if name == 'use_gpu': + if value and not numba.cuda.is_available(): + logger.warning( + 'CUDA is not available. Usage of GPU is disabled' + ) + self.__dict__[name] = False + return + self.__dict__[name] = value + + @contextlib.contextmanager + def context(self, **kwargs): + """ + Return a context manager which allows temporary modification of + the option values for the active thread. + On entry to the `with` statement, all keyword arguments are + applied to the Config object. On exit from the `with` + statement, the previous state is restored. + For example:: + with diagonals.CONFIG.context(cell_datetime_objects=False): + # ... code that expects numbers and not datetimes + """ + # Save the current context + current_state = self.__dict__.copy() + # Update the state + for name, value in six.iteritems(kwargs): + setattr(self, name, value) + try: + yield + finally: + # Return the state + self.__dict__.clear() + self.__dict__.update(current_state) + + +#: Object containing all the Iris run-time options. +CONFIG = Config() diff --git a/diagonals/diagnostics/__init__.py b/diagonals/examples/__init__.py similarity index 100% rename from diagonals/diagnostics/__init__.py rename to diagonals/examples/__init__.py diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py index f0f2deb..e5507b7 100644 --- a/diagonals/examples/examples-oceanheatcontent.py +++ b/diagonals/examples/examples-oceanheatcontent.py @@ -1,17 +1,13 @@ -import os +import logging +import datetime import numpy as np + import iris import iris.cube import iris.analysis -import warnings -import numba -import datetime -from numba import jit -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 + +import diagonals +import diagonals.ohc as ohc MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" DATA_FILE = "/esarchive/exp/ecearth/t05l/cmorfiles/CMIP/EC-Earth-Consortium" \ @@ -19,61 +15,38 @@ DATA_FILE = "/esarchive/exp/ecearth/t05l/cmorfiles/CMIP/EC-Earth-Consortium" \ "/thetao_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_199001-199001.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" +logger = logging.getLogger(__name__) + def main(): + logging.basicConfig(level=logging.INFO) start = datetime.datetime.now() + logger.info('Starting at {}'.format(start)) one_layer = ((0, 300),) # multiple_layers = ((0, 200), (200, 700), (700, 2000), # (0, 100), (100, 200),(200, 300), (300, 400), # (400, 500), (500, 1000),) - - mask, e3t, depth = load_mesh() - weights = check_device_and_compute_weights(one_layer, mask, e3t, depth) - del mask, depth, e3t - basin_name, basins = load_basins() - e1t, e2t = load_areas() - area = {} - for basin in basins: - area[basin] = compute_area_basin(e1t, e2t, basins[basin]) - del e1t, e2t, basins - thetao = load_thetao() - thetao += 273.15 - ohc, ohc_1D = compute_OHC(one_layer, weights, thetao, area) - del weights, thetao, area - if numba.cuda.is_available() is False: - device = 'CPU' - else: - device = 'GPU' - save_data(one_layer, basin_name, ohc, ohc_1D, device) - seconds_total = datetime.datetime.now() - start - print('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) - - -def check_device_and_compute_weights(layers, mask, e3t, depth): - check_gpu = numba.cuda.is_available() - if check_gpu is False: - device = 'CPU' - weights = compute_weights_cpu(layers, mask, e3t, depth) - else: - use_gpu = True - if use_gpu is True: + with diagonals.CONFIG.context(use_gpu=True): + mask, e3t, depth = load_mesh() + weights = ohc.get_weights(one_layer, mask, e3t, depth) + del mask, depth, e3t + basin_name, basins = load_basins() + e1t, e2t = load_areas() + area = {} + for basin in basins: + area[basin] = ohc.get_basin_area(e1t, e2t, basins[basin]) + del e1t, e2t, basins + thetao = load_thetao() + thetao += 273.15 + ohc_2D, ohc_1D = ohc.compute(one_layer, weights, thetao, area) + del weights, thetao, area + if diagonals.CONFIG.use_gpu: device = 'GPU' - weights = compute_weights_gpu(layers, mask, e3t, depth) else: device = 'CPU' - weights = compute_weights_cpu(layers, mask, e3t, depth) - print('Running on {}'.format(device)) - return weights - - -def compute_OHC(layers, weights, thetao, area): - check_gpu = numba.cuda.is_available() - if check_gpu is False: - ohc, ohc_1D = compute_ohc_cpu(layers, thetao, weights, area) - else: - ohc, ohc_1D = compute_ohc_gpu(layers, thetao, weights, area) - return ohc, ohc_1D - + save_data(one_layer, basin_name, ohc_2D, ohc_1D, device) + seconds_total = datetime.datetime.now() - start + logger.info('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) def load_mesh(): mask = iris.load_cube(MESH_FILE, 'tmask').data.astype(np.float32) @@ -81,27 +54,6 @@ def load_mesh(): depth = iris.load_cube(MESH_FILE, 'gdepw_0').data.astype(np.float32) return mask, e3t, depth - -def compute_weights_cpu(layers, mask, e3t, depth): - weights = [] - for min_depth, max_depth in layers: - weights.append(calculate_weight_numba(min_depth, max_depth, e3t, depth, - mask)) - return weights - - -def compute_weights_gpu(layers, mask, e3t, depth): - gpu_mask = cuda.to_device(mask.data.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.data.astype(np.float32)) - gpu_depth = cuda.to_device(depth.data.astype(np.float32)) - weights = [] - for min_depth, max_depth in layers: - weights.append(calculate_weight_numba_cuda(min_depth, max_depth, - gpu_e3t, gpu_depth, gpu_mask)) - del gpu_depth, gpu_mask, gpu_e3t - return weights - - def load_basins(): regions = iris.load(REGIONS_FILE) basins = {} @@ -127,61 +79,11 @@ def load_thetao(): thetao_data = thetao.data.astype(np.float32) return thetao_data - -def compute_ohc_cpu(layers, thetao, weights, area): - ohc = [] - for layer in range(len(layers)): - ohc_layer = sum_red.reduce(multiply_thetao_weight(thetao, - weights[layer]), axis=1) - ohc.append(ohc_layer) - ohc1D_total = [] - for basin in area: - ohc_basin = multiply_array(ohc_layer, area[basin]) - temp = sum_red.reduce(ohc_basin, axis=1) - ohc1D = sum_red.reduce(temp, axis=1) - ohc1D_total.append(ohc1D) - return ohc, ohc1D_total - - -def compute_ohc_gpu(layers, thetao, weights, area): - levels = thetao.shape[1] - times = thetao.shape[0] - lats = thetao.shape[2] - lons = thetao.shape[3] - basins = len(area) - - area_basin = np.empty((basins, lats, lons)) - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - ohc = [] - gpu_ohc = cuda.device_array((times, lats, lons), dtype=np.float32) - gpu_temp = cuda.device_array((basins, times, lats*lons), dtype=np.float32) - for i, basin in enumerate(area): - area_basin[i, :, :] = area[basin] - gpu_basins_area = cuda.to_device(area_basin.astype(np.float32)) - gpu_thetao = cuda.to_device(thetao.astype(np.float32)) - - for layer in range(len(layers)): - 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 = [] - for basin in range(basins): - 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) - del gpu_ohc, gpu_temp, gpu_basins_area - - return ohc, ohc1D_basin - - -def save_data(layers, basins, ohc, ohc1D, device): +def save_data(layers, basins, ohc_2D, ohc1D, device): ohc_cube = [] - print(len(ohc1D)) + logger.info('ohc1d length is %s', len(ohc1D)) for i, layer in enumerate(layers): - ohc_cube.append(iris.cube.Cube(ohc[i], + ohc_cube.append(iris.cube.Cube(ohc_2D[i], long_name='Ocean heat content' ' {0[0]} to {0[1]} meters' .format(layer))) @@ -197,101 +99,5 @@ def save_data(layers, basins, ohc, ohc1D, device): '/esarchive/scratch/sloosvel/numba_outputs/ohc_{0}.nc' .format(device), zlib=True) - -@vectorize(['float32(int32, int32, float32, float32, float32)'], target='cuda') -def calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask): - """ - Calculates the weight for each cell - """ - if not mask: - return 0 - top = depth - bottom = top + e3t - if bottom < min_depth or top > max_depth: - return 0 - else: - if top < min_depth: - top = min_depth - if bottom > max_depth: - bottom = max_depth - - return (bottom - top) * 1020 * 4000 - - -@cuda.jit() -def compute_ohc(thetao, weight, ohc, levels): - i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x - j = cuda.blockIdx.y - t = cuda.blockIdx.z - # i, j, t = grid(3) - ohc[t, j, i] = 0.0 - temp = 0.0 - if(i >= thetao.shape[3]): - return - for lev in range(levels): - thetao_value = thetao[t, lev, j, i] - temp += thetao_value * weight[0, lev, j, i] - ohc[t, j, i] = temp - - -@cuda.jit() -def multiply_ohc_basin(ohc, area, temp): - i, j, t = cuda.grid(3) - basins = area.shape[0] - x = i+j*ohc.shape[2] - if(i >= ohc.shape[2]): - return - for basin in range(basins): - temp[basin, t, x] = 0.0 - temp[basin, t, x] = ohc[t, j, i] * area[basin, j, i] - - -@vectorize(['float32(int32, int32, float32, float32, float32)'], target='cpu') -def calculate_weight_numba(min_depth, max_depth, e3t, depth, mask): - """ - Calculates the weight for each cell - """ - if not mask: - return 0 - top = depth - bottom = top + e3t - if bottom < min_depth or top > max_depth: - return 0 - else: - if top < min_depth: - top = min_depth - if bottom > max_depth: - bottom = max_depth - - return (bottom - top) * 1020 * 4000 - - -@vectorize(['float32(float32, float32)'], target='cpu') -def multiply_thetao_weight(thetao, weight): - return thetao*weight - - -@vectorize(['float64(float32, float32)', 'float64(float64, float64)'], - target='cpu') -def sum_red(x, y): - return x+y - - -@vectorize(['float32(float32, float32)', 'float64(float64, float64)'], - target='cpu') -def multiply_array(a, b): - return a*b - - -@vectorize(['float32(float32, float32, float32)'], target='cpu') -def compute_area_basin(e1t, e2t, basin): - return e1t*e2t*basin - - -@cuda.reduce -def sum_red_cuda(x, y): - return x+y - - if __name__ == '__main__': main() diff --git a/diagonals/diagnostics/oceanheatcontent.py b/diagonals/ohc.py similarity index 86% rename from diagonals/diagnostics/oceanheatcontent.py rename to diagonals/ohc.py index c8ff761..60d289c 100644 --- a/diagonals/diagnostics/oceanheatcontent.py +++ b/diagonals/ohc.py @@ -7,7 +7,6 @@ import iris.cube import iris.analysis import warnings - import datetime import numba @@ -18,8 +17,11 @@ from numba import cuda from numba import int32, int64, float32, float64 from numba.cuda.cudadrv import driver +import diagonals + +__all__ = ['get_weights', 'compute', 'get_basin_area'] -def check_device_and_compute_weights(layers, mask, e3t, depth): +def get_weights(layers, mask, e3t, depth): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -40,23 +42,14 @@ def check_device_and_compute_weights(layers, mask, e3t, depth): weights: float32 List of masked arrays containing the weights for each layer. """ - check_gpu = numba.cuda.is_available() - if check_gpu is False: - device = 'CPU' - weights = compute_weights_cpu(layers, mask, e3t, depth) + if diagonals.CONFIG.use_gpu: + weights = _compute_weights_gpu(layers, mask, e3t, depth) else: - use_gpu = True - if use_gpu is True: - device = 'GPU' - weights = compute_weights_gpu(layers, mask, e3t, depth) - else: - device = 'CPU' - weights = compute_weights_cpu(layers, mask, e3t, depth) - print('Running on {}'.format(device)) + weights = _compute_weights_cpu(layers, mask, e3t, depth) return weights -def compute_OHC(layers, weights, thetao, area): +def compute(layers, weights, thetao, area): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -81,15 +74,14 @@ def compute_OHC(layers, weights, thetao, area): List of masked arrays containing the 1D ocean heat content for each layer, basin and timestep. """ - check_gpu = numba.cuda.is_available() - if check_gpu is False: - ohc, ohc_1D = compute_ohc_cpu(layers, thetao, weights, area) + if diagonals.CONFIG.use_gpu: + ohc, ohc_1D = _compute_ohc_gpu(layers, thetao, weights, area) else: - ohc, ohc_1D = compute_ohc_gpu(layers, thetao, weights, area) + ohc, ohc_1D = _compute_ohc_cpu(layers, thetao, weights, area) return ohc, ohc_1D -def compute_weights_cpu(layers, mask, e3t, depth): +def _compute_weights_cpu(layers, mask, e3t, depth): """Function that calls computing functions in the CPU. Computes weights for each layer and appends them to a list: @@ -112,12 +104,12 @@ def compute_weights_cpu(layers, mask, e3t, depth): """ weights = [] for min_depth, max_depth in layers: - weights.append(calculate_weight_numba(min_depth, max_depth, e3t, depth, + weights.append(_calculate_weight_numba(min_depth, max_depth, e3t, depth, mask)) return weights -def compute_weights_gpu(layers, mask, e3t, depth): +def _compute_weights_gpu(layers, mask, e3t, depth): """Function that calls computing functions in the GPU. Computes weights for each layer and appends them to a list: @@ -143,13 +135,13 @@ def compute_weights_gpu(layers, mask, e3t, depth): gpu_depth = cuda.to_device(depth.data.astype(np.float32)) weights = [] for min_depth, max_depth in layers: - weights.append(calculate_weight_numba_cuda(min_depth, max_depth, + weights.append(_calculate_weight_numba_cuda(min_depth, max_depth, gpu_e3t, gpu_depth, gpu_mask)) del gpu_depth, gpu_mask, gpu_e3t return weights -def compute_ohc_cpu(layers, thetao, weights, area): +def _compute_ohc_cpu(layers, thetao, weights, area): """Function that calls computing functions in the CPU. Loops over layers to compute the global ocean heat content in 2D. This @@ -178,18 +170,18 @@ def compute_ohc_cpu(layers, thetao, weights, area): """ ohc = [] for layer in range(len(layers)): - ohc_layer = sum_red.reduce(multiply_array(thetao, + ohc_layer = _sum_red.reduce(_multiply_array(thetao, weights[layer]), axis=1) ohc.append(ohc_layer) ohc1D_total = [] for i, basin in enumerate(area): - ohc_basin = multiply_array(ohc_layer, area[basin]) + ohc_basin = _multiply_array(ohc_layer, area[basin]) ohc1D = np.sum(ohc_basin, axis=(1, 2)) ohc1D_total.append(ohc1D) return ohc, ohc1D_total -def compute_ohc_gpu(layers, thetao, weights, area): +def _compute_ohc_gpu(layers, thetao, weights, area): """Function that calls computing functions in the GPU. Loops over layers to compute the global ocean heat content in 2D. This @@ -235,14 +227,14 @@ def compute_ohc_gpu(layers, thetao, weights, area): gpu_thetao = cuda.to_device(thetao.astype(np.float32)) for layer in range(len(layers)): - compute_ohc[grid, block](gpu_thetao, weights[layer], gpu_ohc, levels) + _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) + _multiply_ohc_basin[grid, block](gpu_ohc, gpu_basins_area, gpu_temp) ohc1D_basin = [] for basin in range(basins): ohc_1D = np.empty(times, dtype=np.float32) for time in range(times): - ohc_1D[time] = sum_red_cuda(gpu_temp[basin, time, :]) + ohc_1D[time] = _sum_red_cuda(gpu_temp[basin, time, :]) ohc1D_basin.append(ohc_1D) del gpu_ohc, gpu_temp, gpu_basins_area @@ -250,7 +242,7 @@ def compute_ohc_gpu(layers, thetao, weights, area): @vectorize(['float32(int32, int32, float32, float32, float32)'], target='cuda') -def calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask): +def _calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask): """Vectorized numba function executed on the gpu. Calculates the weight for each cell: @@ -289,7 +281,7 @@ def calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask): @cuda.jit() -def compute_ohc(thetao, weight, ohc, levels): +def _compute_ohc(thetao, weight, ohc, levels): """Numba kernel executed on the gpu that computes the global ocean heat content in 2D for a given layer and timestep: @@ -321,7 +313,7 @@ def compute_ohc(thetao, weight, ohc, levels): @cuda.jit() -def multiply_ohc_basin(ohc, area, temp): +def _multiply_ohc_basin(ohc, area, temp): """Numba kernel executed on the gpu that weights the global ocean heat content over the area for each basin for and timestep: @@ -343,8 +335,8 @@ def multiply_ohc_basin(ohc, area, temp): """ i, j, t = cuda.grid(3) basins = area.shape[0] - x = i+j*ohc.shape[2] - if(i >= ohc.shape[2]): + x = i + j * ohc.shape[2] + if i >= ohc.shape[2]: return for basin in range(basins): temp[basin, t, x] = 0.0 @@ -352,7 +344,7 @@ def multiply_ohc_basin(ohc, area, temp): @vectorize(['float32(int32, int32, float32, float32, float32)'], target='cpu') -def calculate_weight_numba(min_depth, max_depth, e3t, depth, mask): +def _calculate_weight_numba(min_depth, max_depth, e3t, depth, mask): """Vectorized numba function executed on the cpu. Calculates the weight for each cell: @@ -392,7 +384,7 @@ def calculate_weight_numba(min_depth, max_depth, e3t, depth, mask): @vectorize(['float64(float32, float32)', 'float64(float64, float64)'], target='cpu') -def sum_red(x, y): +def _sum_red(x, y): """Vectorized numba function executed on the cpu that performs a sum reduction: @@ -413,7 +405,7 @@ def sum_red(x, y): @vectorize(['float32(float32, float32)', 'float64(float64, float64)'], target='cpu') -def multiply_array(a, b): +def _multiply_array(a, b): """Vectorized numba function executed on the cpu that performs an element-wise multiplication of two arrays: @@ -434,7 +426,7 @@ def multiply_array(a, b): @vectorize(['float32(float32, float32, float32)'], target='cpu') -def compute_area_basin(e1t, e2t, basin): +def get_basin_area(e1t, e2t, basin): """Vectorized numba function executed on the cpu that computes the area for each basin: @@ -456,7 +448,7 @@ def compute_area_basin(e1t, e2t, basin): @cuda.reduce -def sum_red_cuda(x, y): +def _sum_red_cuda(x, y): """Numba kernel executed on the gpu that performs a sum reduction: Parameters diff --git a/diagonals/diagnostics/regmean.py b/diagonals/regmean.py similarity index 100% rename from diagonals/diagnostics/regmean.py rename to diagonals/regmean.py diff --git a/diagonals/diagnostics/regsum.py b/diagonals/regsum.py similarity index 100% rename from diagonals/diagnostics/regsum.py rename to diagonals/regsum.py diff --git a/diagonals/diagnostics/siarea.py b/diagonals/siarea.py similarity index 100% rename from diagonals/diagnostics/siarea.py rename to diagonals/siarea.py -- GitLab From a7c26f8b7ccba7b3ea481fed733ecd185a06afd5 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Tue, 7 May 2019 11:02:16 +0200 Subject: [PATCH 02/54] Added mesh manager --- .gitignore | 1 - .../examples/examples-oceanheatcontent.py | 24 ++--- diagonals/mesh_helpers/__init__.py | 0 diagonals/mesh_helpers/nemo.py | 88 +++++++++++++++++++ diagonals/ohc.py | 21 +++-- 5 files changed, 115 insertions(+), 19 deletions(-) create mode 100644 diagonals/mesh_helpers/__init__.py create mode 100644 diagonals/mesh_helpers/nemo.py diff --git a/.gitignore b/.gitignore index b319c4e..f705c88 100644 --- a/.gitignore +++ b/.gitignore @@ -19,4 +19,3 @@ htmlcov prof .vscode .eggs -*.pyc diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py index e5507b7..ffbe739 100644 --- a/diagonals/examples/examples-oceanheatcontent.py +++ b/diagonals/examples/examples-oceanheatcontent.py @@ -8,6 +8,7 @@ import iris.analysis import diagonals import diagonals.ohc as ohc +from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" DATA_FILE = "/esarchive/exp/ecearth/t05l/cmorfiles/CMIP/EC-Earth-Consortium" \ @@ -21,21 +22,25 @@ logger = logging.getLogger(__name__) def main(): logging.basicConfig(level=logging.INFO) start = datetime.datetime.now() - logger.info('Starting at {}'.format(start)) + logger.info('Starting at %s', start) one_layer = ((0, 300),) + mesh = Nemo(MESH_FILE, REGIONS_FILE) # multiple_layers = ((0, 200), (200, 700), (700, 2000), # (0, 100), (100, 200),(200, 300), (300, 400), # (400, 500), (500, 1000),) with diagonals.CONFIG.context(use_gpu=True): - mask, e3t, depth = load_mesh() + + e3t = mesh.get_k_length() + mask = mesh.get_landsea_mask() + depth = mesh.get_depth(cell_point='W') weights = ohc.get_weights(one_layer, mask, e3t, depth) del mask, depth, e3t + basin_name, basins = load_basins() - e1t, e2t = load_areas() - area = {} - for basin in basins: - area[basin] = ohc.get_basin_area(e1t, e2t, basins[basin]) - del e1t, e2t, basins + areacello = mesh.get_areacello() + area = ohc.get_basin_area(areacello, basins) + del areacello, basins + thetao = load_thetao() thetao += 273.15 ohc_2D, ohc_1D = ohc.compute(one_layer, weights, thetao, area) @@ -45,14 +50,13 @@ def main(): else: device = 'CPU' save_data(one_layer, basin_name, ohc_2D, ohc_1D, device) + seconds_total = datetime.datetime.now() - start logger.info('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) def load_mesh(): - mask = iris.load_cube(MESH_FILE, 'tmask').data.astype(np.float32) - e3t = iris.load_cube(MESH_FILE, 'e3t_0').data.astype(np.float32) depth = iris.load_cube(MESH_FILE, 'gdepw_0').data.astype(np.float32) - return mask, e3t, depth + return depth def load_basins(): regions = iris.load(REGIONS_FILE) diff --git a/diagonals/mesh_helpers/__init__.py b/diagonals/mesh_helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py new file mode 100644 index 0000000..6328e15 --- /dev/null +++ b/diagonals/mesh_helpers/nemo.py @@ -0,0 +1,88 @@ +import numpy as np +import iris + +from numba import vectorize + +class Nemo(): + + def __init__(self, mesh_file, regions_file, default_cell_point='T'): + self.mesh_file = mesh_file + self.regions_file = regions_file + self.default_cell_point = default_cell_point + + def get_areacello(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + e1 = self.get_i_length(cell_point, dtype) + e2 = self.get_j_length(cell_point, dtype) + return _get_area(e1, e2) + + def get_volcello(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + e1 = self.get_i_length(cell_point, dtype) + e2 = self.get_j_length(cell_point, dtype) + e3 = self.get_j_length(cell_point, dtype) + return _get_volume(e1, e2, e3) + + def get_i_length(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return iris.load_cube(self.mesh_file, 'e1' + cell_point.lower()).data.astype(dtype) + + def get_j_length(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return iris.load_cube(self.mesh_file, 'e2' + cell_point.lower()).data.astype(dtype) + + def get_k_length(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return iris.load_cube(self.mesh_file, 'e3' + cell_point.lower() + '_0').data.astype(dtype) + + def get_depth(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return iris.load_cube(self.mesh_file, 'gdep' + cell_point.lower() + '_0').data.astype(dtype) + + def get_landsea_mask(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return iris.load_cube(self.mesh_file, cell_point.lower() + 'mask').data.astype(dtype) + + +@vectorize(['float32(float32, float32)'], target='cpu') +def _get_area(e1, e2): + """Vectorized numba function executed on the cpu that computes the area + for each basin: + + Parameters + ---------- + e1t : float32 + Masked array containing variable e1t. + e2t : float32 + Masked array containing variable e2t. + Returns + ------- + e1t*e2t : float32 + Masked array containing the area for the whole grid. + """ + return e1*e2 + +@vectorize(['float32(float32, float32, float32)'], target='cpu') +def _get_volume(e1, e2, e3): + """Vectorized numba function executed on the cpu that computes the area + for each basin: + + Parameters + ---------- + e1t : float32 + Masked array containing variable e1t. + e2t : float32 + Masked array containing variable e2t. + Returns + ------- + e1t*e2t : float32 + Masked array containing the area for the whole grid. + """ + return e1 * e2 * e3 \ No newline at end of file diff --git a/diagonals/ohc.py b/diagonals/ohc.py index 60d289c..75dd092 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -425,26 +425,31 @@ def _multiply_array(a, b): return a*b -@vectorize(['float32(float32, float32, float32)'], target='cpu') -def get_basin_area(e1t, e2t, basin): +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 ---------- - e1t : float32 - Masked array containing variable e1t. - e2t : float32 - Masked array containing variable e2t. + areacello : float32 + Masked array containing areacello. basin: float32 Masked array containing the mask for a given basin. Returns ------- - e1t*e2t*basin : float32 + areacello*basin : float32 Masked array containing the area for a given basin. """ - return e1t*e2t*basin + return areacello*basin @cuda.reduce -- GitLab From 64f7f291f15b14ec1c32bb58693bcb17eede0eeb Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Tue, 7 May 2019 11:25:48 +0200 Subject: [PATCH 03/54] FIxed tests --- diagonals/__init__.py | 1 + diagonals/examples/examples-areamoc.py | 25 +++++++++++-------- diagonals/examples/examples-moc.py | 2 -- .../examples/examples-oceanheatcontent.py | 10 +++++--- diagonals/mesh_helpers/nemo.py | 17 ++++++++----- diagonals/ohc.py | 12 ++++++--- 6 files changed, 40 insertions(+), 27 deletions(-) diff --git a/diagonals/__init__.py b/diagonals/__init__.py index 3c732db..15bd90f 100644 --- a/diagonals/__init__.py +++ b/diagonals/__init__.py @@ -7,6 +7,7 @@ import numba.cuda logger = logging.getLogger(__name__) + class Config(threading.local): """Run-time configuration controller.""" diff --git a/diagonals/examples/examples-areamoc.py b/diagonals/examples/examples-areamoc.py index 3497831..ff54767 100644 --- a/diagonals/examples/examples-areamoc.py +++ b/diagonals/examples/examples-areamoc.py @@ -5,7 +5,7 @@ import iris import iris.cube import iris.analysis import iris.util -import iris.coords +from iris.coords import DimCoord import warnings import numba import datetime @@ -21,6 +21,7 @@ VSFTMYZ_FILE = '/esarchive/exp/ecearth/a16l/cmorfiles/CMIP'\ '/vsftmyz/gn/v20180711/vsftmyz_Omon_EC-Earth3-LR_historical_'\ 'r1i1p1f1_gn_196101-196112.nc' + def main(): one_layer = ((0, 300),) min_lat = 40 @@ -34,7 +35,6 @@ def main(): basins = vsftmyz.shape[3] - def load_vsftmyz(min_lat, max_lat, min_depth, max_depth): long_name = 'Ocean meridional overturning volume streamfunction' vsftmyz = iris.load_cube(VSFTMYZ_FILE, long_name) @@ -47,15 +47,18 @@ def load_vsftmyz(min_lat, max_lat, min_depth, max_depth): def add_coord_names(cube): dims = len(cube.shape) - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), - var_name='basin'), - dims - 1) - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), - var_name='latitude'), - dims - 2) - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 3]), - var_name='longitude'), - dims - 3) + cube.add_dim_coord( + DimCoord(np.arange(cube.shape[dims - 1]), var_name='basin'), + dims - 1 + ) + cube.add_dim_coord( + DimCoord(np.arange(cube.shape[dims - 2]), var_name='latitude'), + dims - 2 + ) + cube.add_dim_coord( + DimCoord(np.arange(cube.shape[dims - 3]), var_name='longitude'), + dims - 3 + ) return cube diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index 6b35ed0..cc875a5 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -135,7 +135,5 @@ def horizontal_integral(vo, area, moc): moc[t, lev, j] = temp - - if __name__ == '__main__': main() diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py index ffbe739..0d18414 100644 --- a/diagonals/examples/examples-oceanheatcontent.py +++ b/diagonals/examples/examples-oceanheatcontent.py @@ -29,7 +29,6 @@ def main(): # (0, 100), (100, 200),(200, 300), (300, 400), # (400, 500), (500, 1000),) with diagonals.CONFIG.context(use_gpu=True): - e3t = mesh.get_k_length() mask = mesh.get_landsea_mask() depth = mesh.get_depth(cell_point='W') @@ -50,14 +49,15 @@ def main(): else: device = 'CPU' save_data(one_layer, basin_name, ohc_2D, ohc_1D, device) - - seconds_total = datetime.datetime.now() - start - logger.info('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) + def load_mesh(): depth = iris.load_cube(MESH_FILE, 'gdepw_0').data.astype(np.float32) return depth + def load_basins(): regions = iris.load(REGIONS_FILE) basins = {} @@ -83,6 +83,7 @@ def load_thetao(): thetao_data = thetao.data.astype(np.float32) return thetao_data + def save_data(layers, basins, ohc_2D, ohc1D, device): ohc_cube = [] logger.info('ohc1d length is %s', len(ohc1D)) @@ -103,5 +104,6 @@ def save_data(layers, basins, ohc_2D, ohc1D, device): '/esarchive/scratch/sloosvel/numba_outputs/ohc_{0}.nc' .format(device), zlib=True) + if __name__ == '__main__': main() diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py index 6328e15..cd2359c 100644 --- a/diagonals/mesh_helpers/nemo.py +++ b/diagonals/mesh_helpers/nemo.py @@ -3,6 +3,7 @@ import iris from numba import vectorize + class Nemo(): def __init__(self, mesh_file, regions_file, default_cell_point='T'): @@ -28,27 +29,30 @@ 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 iris.load_cube(self.mesh_file, 'e1' + cell_point.lower()).data.astype(dtype) + return self.get_mesh_var('e1' + cell_point.lower(), dtype) def get_j_length(self, cell_point=None, dtype=np.float32): if cell_point is None: cell_point = self.default_cell_point - return iris.load_cube(self.mesh_file, 'e2' + cell_point.lower()).data.astype(dtype) + return self.get_mesh_var('e2' + cell_point.lower(), dtype) def get_k_length(self, cell_point=None, dtype=np.float32): if cell_point is None: cell_point = self.default_cell_point - return iris.load_cube(self.mesh_file, 'e3' + cell_point.lower() + '_0').data.astype(dtype) + return self.get_mesh_var('e3' + cell_point.lower() + '_0', dtype) def get_depth(self, cell_point=None, dtype=np.float32): if cell_point is None: cell_point = self.default_cell_point - return iris.load_cube(self.mesh_file, 'gdep' + cell_point.lower() + '_0').data.astype(dtype) + return self.get_mesh_var('gdep' + cell_point.lower() + '_0', dtype) def get_landsea_mask(self, cell_point=None, dtype=np.float32): if cell_point is None: cell_point = self.default_cell_point - return iris.load_cube(self.mesh_file, cell_point.lower() + 'mask').data.astype(dtype) + return self.get_mesh_var(cell_point.lower() + 'mask', dtype) + + def get_mesh_var(self, var, dtype): + return iris.load_cube(self.mesh_file, var).data.astype(dtype) @vectorize(['float32(float32, float32)'], target='cpu') @@ -69,6 +73,7 @@ def _get_area(e1, e2): """ return e1*e2 + @vectorize(['float32(float32, float32, float32)'], target='cpu') def _get_volume(e1, e2, e3): """Vectorized numba function executed on the cpu that computes the area @@ -85,4 +90,4 @@ def _get_volume(e1, e2, e3): e1t*e2t : float32 Masked array containing the area for the whole grid. """ - return e1 * e2 * e3 \ No newline at end of file + return e1 * e2 * e3 diff --git a/diagonals/ohc.py b/diagonals/ohc.py index 75dd092..59c0ca8 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -21,6 +21,7 @@ import diagonals __all__ = ['get_weights', 'compute', 'get_basin_area'] + def get_weights(layers, mask, e3t, depth): """Function that checks device and calls computing functions. @@ -104,8 +105,9 @@ def _compute_weights_cpu(layers, mask, e3t, depth): """ weights = [] for min_depth, max_depth in layers: - weights.append(_calculate_weight_numba(min_depth, max_depth, e3t, depth, - mask)) + weights.append( + _calculate_weight_numba(min_depth, max_depth, e3t, depth, mask) + ) return weights @@ -170,8 +172,10 @@ def _compute_ohc_cpu(layers, thetao, weights, area): """ ohc = [] for layer in range(len(layers)): - ohc_layer = _sum_red.reduce(_multiply_array(thetao, - weights[layer]), axis=1) + ohc_layer = _sum_red.reduce( + _multiply_array(thetao, weights[layer]), + axis=1 + ) ohc.append(ohc_layer) ohc1D_total = [] for i, basin in enumerate(area): -- GitLab From 3ade130e4f9eee12c81b5549bb693b4180220184 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Thu, 9 May 2019 14:36:44 +0200 Subject: [PATCH 04/54] Added test for ohc weights computation --- __init__.py | 0 diagonals/ohc.py | 10 +++++++--- setup.py | 3 +++ test/unit/test_ohc.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 3 deletions(-) delete mode 100644 __init__.py create mode 100644 test/unit/test_ohc.py diff --git a/__init__.py b/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/diagonals/ohc.py b/diagonals/ohc.py index 59c0ca8..5d7460a 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -22,7 +22,7 @@ import diagonals __all__ = ['get_weights', 'compute', 'get_basin_area'] -def get_weights(layers, mask, e3t, depth): +def get_weights(layers, mask, cell_height, cell_top_depth): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -44,9 +44,13 @@ def get_weights(layers, mask, e3t, depth): List of masked arrays containing the weights for each layer. """ if diagonals.CONFIG.use_gpu: - weights = _compute_weights_gpu(layers, mask, e3t, depth) + weights = _compute_weights_gpu( + layers, mask, cell_height, cell_top_depth + ) else: - weights = _compute_weights_cpu(layers, mask, e3t, depth) + weights = _compute_weights_cpu( + layers, mask, cell_height, cell_top_depth + ) return weights diff --git a/setup.py b/setup.py index 6a478da..758e1af 100644 --- a/setup.py +++ b/setup.py @@ -81,4 +81,7 @@ setup(name='diagonals', 'gpu': REQUIREMENTS['gpu'], 'test': REQUIREMENTS['test'], }, + cmdclass={ + 'test': RunTests, + }, zip_safe=False,) diff --git a/test/unit/test_ohc.py b/test/unit/test_ohc.py new file mode 100644 index 0000000..9a67a18 --- /dev/null +++ b/test/unit/test_ohc.py @@ -0,0 +1,43 @@ +import os +import unittest + +import pycodestyle # formerly known as pep8 + +import numpy as np +import diagonals.ohc as ohc + + +class TestOhc(unittest.TestCase): + + def test_weights(self): + weights = ohc.get_weights( + layers=((0, 2),), + mask=np.ones((2, 2), dtype=np.float32), + cell_height=np.array(((1., 2.), (4., 3.)), dtype=np.float32), + cell_top_depth=np.array(((0., 0.), (1., 2.)), dtype=np.float32) + ) + expected_result = 1020 * 4000 * np.array(((1., 2.), (1., 0.))) + self.assertTrue(np.all(weights == expected_result)) + + def test_weights_all_land(self): + weights = ohc.get_weights( + layers=((0, 2),), + mask=np.zeros((2, 2), dtype=np.float32), + cell_height=np.array(((1., 2.), (4., 3.)), dtype=np.float32), + cell_top_depth=np.array(((0., 0.), (1., 2.)), dtype=np.float32) + ) + self.assertTrue(np.all(weights == np.zeros((2, 2), dtype=np.float32))) + + def test_weights_mul_layers(self): + weights = ohc.get_weights( + layers=((0, 2), (2, 4)), + mask=np.ones((2, 2), dtype=np.float32), + cell_height=np.array(((1., 2.), (4., 3.)), dtype=np.float32), + cell_top_depth=np.array(((0., 0.), (1., 2.)), dtype=np.float32) + ) + print(weights) + expected_result = 1020 * 4000 * np.array(( + ((1., 2.), (1., 0.)), + ((0., 0.), (2., 2.)) + )) + self.assertTrue(np.all(weights == expected_result)) -- GitLab From d01b35013da6f0499d931aab7a3fa058fa7fc571 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Thu, 9 May 2019 14:43:34 +0200 Subject: [PATCH 05/54] Remove unneeded files from repository --- .vscode/settings.json | 3 --- build/lib/diagonals/__init__.py | 0 diagonals/__pycache__/__init__.cpython-37.pyc | Bin 141 -> 0 bytes test/__pycache__/__init__.cpython-37.pyc | Bin 134 -> 0 bytes test/unit/__pycache__/__init__.cpython-37.pyc | Bin 139 -> 0 bytes test/unit/__pycache__/test_lint.cpython-37.pyc | Bin 1009 -> 0 bytes 6 files changed, 3 deletions(-) delete mode 100644 .vscode/settings.json delete mode 100644 build/lib/diagonals/__init__.py delete mode 100644 diagonals/__pycache__/__init__.cpython-37.pyc delete mode 100644 test/__pycache__/__init__.cpython-37.pyc delete mode 100644 test/unit/__pycache__/__init__.cpython-37.pyc delete mode 100644 test/unit/__pycache__/test_lint.cpython-37.pyc diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index 5ea03ca..0000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python.pythonPath": "/home/Earth/sloosvel/.conda/envs/diagonals/bin/python" -} \ No newline at end of file diff --git a/build/lib/diagonals/__init__.py b/build/lib/diagonals/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/diagonals/__pycache__/__init__.cpython-37.pyc b/diagonals/__pycache__/__init__.cpython-37.pyc deleted file mode 100644 index 3ac23b1d693a09d0e9c06847643641ed5ee32a9b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 141 zcmZ?b<>g`k0uzVI7!ds!M8E(ekl_Ht#VkM~g&~+hlhJP_LlH*Y7a+j^ diff --git a/test/__pycache__/__init__.cpython-37.pyc b/test/__pycache__/__init__.cpython-37.pyc deleted file mode 100644 index f9dc2c9fee10bb904ca5e86f7f30655c7c164fd3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 134 zcmZ?b<>g`k0uzVI7=9rA7{q}AMj*ohh>KZ(L<&PNgC?WjN`@kkFoO7{ub+{ho2u`c zSX7dsU!0SlUtE@&qo0zQn4X`Pm{Y7@l3HA%A0MBYmst`YuUAlci^C>2KczG$)edA( IF%UBV03F#K@c;k- diff --git a/test/unit/__pycache__/__init__.cpython-37.pyc b/test/unit/__pycache__/__init__.cpython-37.pyc deleted file mode 100644 index c5fca930321cdf105aa08f8c5572894b24be475a..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 139 zcmZ?b<>g`k0uzVI7=9rA7{q}AMj*ohh>KZ(L<&PNgC?WjN`@kkFoO7HqMwnUo2u`c zSX7dsU!0SlUtE@&qo0zQn4X`Pm{Y7@l3HA%Uz(R$q8}fh38doV^$IF)aoFVMr#f^byU<95?rdbv$^H`>rW zj3fz>RI-95v|d@l>NQA$eIf(NZ;0fI4fu7CG$cDB>y4{+{5CIL?38sj{s}1LAVYgA z44fTMy8=cLKqhn48+nbH*xO zHMYt5@jG(<<{fd-bYVs&VHFNqgT-mPVl&89k!hz2r%RZ9a8EX7j1t*|PQt4->?C&m z)Wx4voawTs%`kR4KHNKey0z&SWS*YtGA*pp?!_`UDs$T8%D#v{ zPg+-OBQ{9N!TNpQ7#ZNmHxV;+o)2_hdZqztl)AnT;c=M`m5+qz<%JT$FQ{@nR3>%G zFODuVEg|J)p_0}SF8kwLDj({R%XMjes4lW%3_DGA;Z!O8Vrs22Ztr}Y7Rh2=V-0MR zw$m;xgfd1OyB>LNRnfb>g2mXR^9N1L!}ci6&eBsQjMff1!8@MjW&wKhjk>C@#5{!P z*Yj>)4^?+BHLl;ah1T|?D!PA(>mqO6ahbc?69sUxIlA-@a9@nnXiNOzJuNeZqi-LB zA)L`B<9sDprY*Kam+2~t=n|WMG913m{{kG~dyD6}bkM1BMK?Y%cv^_86B~FJ$ L5Kbc+u_*WrQB4*A -- GitLab From b13567ca1e94cd6abff9ae2f3f602649229c32a5 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 11:48:08 +0200 Subject: [PATCH 06/54] Correct descriptions --- diagonals/ohc.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/diagonals/ohc.py b/diagonals/ohc.py index 5d7460a..ef70caa 100644 --- a/diagonals/ohc.py +++ b/diagonals/ohc.py @@ -10,9 +10,7 @@ import warnings import datetime import numba -from numba import jit 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 @@ -33,9 +31,9 @@ def get_weights(layers, mask, cell_height, cell_top_depth): List containing the minimum and maximum depth values for the layer. mask : float32 Masked array containing the Global basin mask. - e3t: float32 + cell_height: float32 Masked array containing variable e3t. - depth: float32 + cell_top_depth: float32 Masked array containing variable gwdep. Returns -- GitLab From 688c053ca0d444588f374894a0427ba08488421e Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 12:00:42 +0200 Subject: [PATCH 07/54] Modify regmean2D for the user interface --- diagonals/regmean.py | 181 ++++++++++++++++++++++--------------------- 1 file changed, 91 insertions(+), 90 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index af7d667..8597c6f 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -1,20 +1,26 @@ import os + import numpy as np + import iris import iris.cube import iris.analysis + import warnings -import numba import datetime -from numba import jit, njit + +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 +import diagonals + +__all__ = ['compute_regmean_2D', ] + -def check_device_and_compute_regmean_2D(var, basins, e1t, e2t): +def compute_regmean_2D(var, basins, e1, e2): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -25,33 +31,28 @@ def check_device_and_compute_regmean_2D(var, basins, e1t, e2t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. Returns ------- regmean: float32 List containing regional mean for variable var """ - check_gpu = numba.cuda.is_available() - if check_gpu is False: - device = 'CPU' - regmean = compute_regmean_2D_cpu(var, basins, e1t, e2t) + if diagonals.CONFIG.use_gpu: + regmean = _compute_regmean_2D_gpu( + var, basins, e1, e2 + ) else: - use_gpu = True - if use_gpu is True: - device = 'GPU' - regmean = compute_regmean_2D_gpu(var, basins, e1t, e2t) - else: - device = 'CPU' - regmean = compute_regmean_2D_cpu(var, basins, e1t, e2t) - print('Running on {}'.format(device)) + regmean = _compute_regmean_2D_cpu( + var, basins, e1, e2 + ) return regmean -def compute_regmean_2D_cpu(var, basins, e1t, e2t): +def _compute_regmean_2D_cpu(var, basins, e1, e2): """Function that computes the regional mean for 2D vars in the cpu. Computes the weights for each region and performs a weighted average. @@ -62,10 +63,10 @@ def compute_regmean_2D_cpu(var, basins, e1t, e2t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. Returns ------- @@ -75,7 +76,7 @@ def compute_regmean_2D_cpu(var, basins, e1t, e2t): times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): - weights = compute_weights_2D(mask, e1t, e2t) + weights = compute_weights_2D(mask, e1, e2) regmean = np.empty(times) for t in range(times): regmean[t] = np.average(var[t, :, :], axis=(0, 1), @@ -84,6 +85,63 @@ def compute_regmean_2D_cpu(var, basins, e1t, e2t): return regmean_total +def _compute_regmean_2D_gpu(var, basins, e1, e2): + """Function that computes the regional mean for 2D vars in the gpu. + + Computes the weights for each region and performs a weighted average. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. + + Returns + ------- + regmean_total: float32 + List containing regional mean for variable var. + """ + times = var.shape[0] + lats = var.shape[1] + lons = var.shape[2] + mask = len(basins) + + masks = np.empty((mask, lats, lons)) + for i, basin in enumerate(basins): + masks[i, :, :] = basins[basin] + return masks + + block = (128, 1, 1) + grid_size = (lons // block[0]) + 1 + grid = (grid_size, lats, times) + + gpu_var = cuda.to_device(var.astype(np.float32)) + gpu_e1 = cuda.to_device(e1.astype(np.float32)) + gpu_e2 = cuda.to_device(e2.astype(np.float32)) + gpu_masks = cuda.to_device(masks) + gpu_weighted_var = cuda.device_array((mask, times, lats*lons), + dtype=np.float32) + gpu_vector_weight = cuda.device_array((mask, lats*lons), dtype=np.float32) + + _weigh_2D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_masks, + gpu_weighted_var, gpu_vector_weight) + + regmean = {} + for m in range(mask): + reduced_weight = _sum_red_cuda(gpu_vector_weight[m, :]) + regmean[m] = {} + for t in range(times): + reduced_var = _sum_red_cuda(gpu_weighted_var[m, t, :]) + regmean[m][t] = reduced_var / reduced_weight + + return regmean + + def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -157,63 +215,6 @@ def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): return regmean_total -def compute_regmean_2D_gpu(var, basins, e1t, e2t): - """Function that computes the regional mean for 2D vars in the gpu. - - Computes the weights for each region and performs a weighted average. - - Parameters - ---------- - var : float32 - Masked array containing variable data. - basins : float32 - List containing basin names and masks. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - - Returns - ------- - regmean_total: float32 - List containing regional mean for variable var. - """ - times = var.shape[0] - lats = var.shape[1] - lons = var.shape[2] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - return masks - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, lats*lons), dtype=np.float32) - - variable_weight_cuda[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_masks, - gpu_weighted_var, gpu_vector_weight) - - regmean = {} - for m in range(mask): - reduced_weight = sum_red_cuda(gpu_vector_weight[m, :]) - regmean[m] = {} - for t in range(times): - reduced_var = sum_red_cuda(gpu_weighted_var[m, t, :]) - regmean[m][t] = reduced_var / reduced_weight - - return regmean - - def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): """Function that computes the regional mean for 3D vars in the gpu. @@ -393,16 +394,16 @@ def compute_weights_3D(mask, e1t, e2t, e3t): @cuda.jit -def variable_weight_cuda(var, e1t, e2t, masks, weighted_var, vector_weight): +def _weigh_2D_var_cuda(var, e1, e2, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 2D variable for each region Parameters ---------- - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -425,7 +426,7 @@ def variable_weight_cuda(var, e1t, e2t, masks, weighted_var, vector_weight): for bas in range(basins): vector_weight[bas, x] = 0.0 weighted_var[bas, t, x] = 0.0 - vector_weight[bas, x] = e1t[j, i]*e2t[j, i]*masks[bas, j, i] + vector_weight[bas, x] = e1[j, i]*e2[j, i]*masks[bas, j, i] weighted_var[bas, t, x] = var[t, j, i] * vector_weight[bas, x] @@ -516,7 +517,7 @@ def variable_weight_cuda_level_3D(var, e1t, e2t, e3t, masks, weighted_var, @cuda.reduce -def sum_red_cuda(x, y): +def _sum_red_cuda(x, y): """Numba kernel executed on the gpu that performs a sum reduction: Parameters -- GitLab From a8c69595bce35031dcd4d4d7d08367d884bca96e Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 12:16:43 +0200 Subject: [PATCH 08/54] Correct user interface for regmean 3D --- diagonals/regmean.py | 203 ++++++++++++++++++++++++++++--------------- 1 file changed, 133 insertions(+), 70 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 8597c6f..6c41695 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -17,7 +17,7 @@ from numba.cuda.cudadrv import driver import diagonals -__all__ = ['compute_regmean_2D', ] +__all__ = ['compute_regmean_2D', 'compute_regmean_3D'] def compute_regmean_2D(var, basins, e1, e2): @@ -52,6 +52,70 @@ def compute_regmean_2D(var, basins, e1, e2): return regmean +def compute_regmean_2D(var, basins, e1, e2): + """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. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. + + Returns + ------- + regmean: float32 + List containing regional mean for variable var + """ + if diagonals.CONFIG.use_gpu: + regmean = _compute_regmean_2D_gpu( + var, basins, e1, e2 + ) + else: + regmean = _compute_regmean_2D_cpu( + var, basins, e1, e2 + ) + return regmean + + +def compute_regmean_3D(var, basins, e1, e2, e3): + """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. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. + + Returns + ------- + regmean: float32 + List containing regional mean for variable var + """ + if diagonals.CONFIG.use_gpu: + regmean = _compute_regmean_3D_gpu( + var, basins, e1, e2, e3 + ) + else: + regmean = _compute_regmean_3D_cpu( + var, basins, e1, e2, e3 + ) + return regmean + + def _compute_regmean_2D_cpu(var, basins, e1, e2): """Function that computes the regional mean for 2D vars in the cpu. @@ -142,45 +206,7 @@ def _compute_regmean_2D_gpu(var, basins, e1, e2): return regmean -def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): - """Function that computes the regional mean at every depth level - for 3D vars in the cpu. - - Computes the weights for each region and performs a weighted average. - - Parameters - ---------- - var : float32 - Masked array containing variable data. - basins : float32 - List containing basin names and masks. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - e3t: float32 - Masked array containing variable e3t. - - Returns - ------- - regmean_total: float32 - List containing regional mean at every depth level for variable var. - """ - times = var.shape[0] - levs = var.shape[1] - regmean_total = {} - for basin, mask in basins.items(): - regmean = np.empty((times, levs)) - w = compute_weights_3D(mask, e1t, e2t, e3t) - for t in range(times): - for l in range(levs): - regmean[t, l] = np.average(var[t, l, :, :], axis=(0, 1), - weights=np.squeeze(w)[l, :, :]) - regmean_total[basin] = regmean - return regmean_total - - -def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): +def _compute_regmean_3D_cpu(var, basins, e1, e2, e3): """Function that computes the regional mean for 3D vars in the cpu. Computes the weights for each region and performs a weighted average. @@ -191,11 +217,11 @@ def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 + e1 : float32 Masked array containing variable e1t. - e2t: float32 + e2: float32 Masked array containing variable e2t. - e3t: float32 + e3: float32 Masked array containing variable e3t. Returns @@ -206,7 +232,7 @@ def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): - weights = compute_weights_3D(mask, e1t, e2t, e3t) + weights = _compute_weights_3D(mask, e1, e2, e3) regmean = np.empty(times) for t in range(times): regmean[t] = np.average(var[t, :, :, :], axis=(0, 1, 2), @@ -215,7 +241,7 @@ def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): return regmean_total -def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): +def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): """Function that computes the regional mean for 3D vars in the gpu. Computes the weights for each region and performs a weighted average. @@ -226,11 +252,11 @@ def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 + e1 : float32 Masked array containing variable e1t. - e2t: float32 + e2: float32 Masked array containing variable e2t. - e3t: float32 + e3: float32 Masked array containing variable e3t. Returns @@ -253,9 +279,9 @@ def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): grid = (grid_size, lats, times) gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.astype(np.float32)) + gpu_e1 = cuda.to_device(e1.astype(np.float32)) + gpu_e2 = cuda.to_device(e2.astype(np.float32)) + gpu_e3 = cuda.to_device(e3.astype(np.float32)) gpu_masks = cuda.to_device(masks) gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), @@ -263,7 +289,7 @@ def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): gpu_vector_weight = cuda.device_array((mask, levs*lats*lons), dtype=np.float32) - variable_weight_cuda_3D[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_e3t, + _weight_3D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, gpu_masks, gpu_weighted_var, gpu_vector_weight) @@ -278,6 +304,44 @@ def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): return regmean +def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): + """Function that computes the regional mean at every depth level + for 3D vars in the cpu. + + Computes the weights for each region and performs a weighted average. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + e1t : float32 + Masked array containing variable e1t. + e2t: float32 + Masked array containing variable e2t. + e3t: float32 + Masked array containing variable e3t. + + Returns + ------- + regmean_total: float32 + List containing regional mean at every depth level for variable var. + """ + times = var.shape[0] + levs = var.shape[1] + regmean_total = {} + for basin, mask in basins.items(): + regmean = np.empty((times, levs)) + w = compute_weights_3D(mask, e1t, e2t, e3t) + for t in range(times): + for l in range(levs): + regmean[t, l] = np.average(var[t, l, :, :], axis=(0, 1), + weights=np.squeeze(w)[l, :, :]) + regmean_total[basin] = regmean + return regmean_total + + def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -370,26 +434,26 @@ def compute_weights_2D(mask, e1t, e2t): @vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def compute_weights_3D(mask, e1t, e2t, e3t): +def _compute_weights_3D(mask, e1, e2, e3): """Function that computes the regional weights for 3D vars in the cpu. Parameters ---------- mask : float32 Mask array containing a region mask. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - e3t: float32 - Masked array containing variable e2t + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. + e3 : float32 + Masked array containing variable e3t Returns ------- weights: float32 Masked array containing the weights for a given region. """ - weights = mask*e1t*e2t*e3t + weights = mask*e1*e2*e3 return weights @@ -431,19 +495,18 @@ def _weigh_2D_var_cuda(var, e1, e2, masks, weighted_var, vector_weight): @cuda.jit -def variable_weight_cuda_3D(var, e1t, e2t, e3t, masks, weighted_var, - vector_weight): +def _weigh_3D_var_cuda(var, e1, e2, e3, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 3D variable for each region. Parameters ---------- - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - e3t: float32 - Masked array containing variable e3t. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. + e3: float32 + Masked array containing variable e3. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -467,8 +530,8 @@ def variable_weight_cuda_3D(var, e1t, e2t, e3t, masks, weighted_var, for lev in range(levels): x = i+var.shape[3]*(j+var.shape[2]*lev) vector_weight[bas, x] = 0.0 - vector_weight[bas, x] = (e1t[j, i] * e2t[j, i] * - e3t[lev, j, i] * masks[bas, j, i]) + vector_weight[bas, x] = (e1[j, i] * e2[j, i] * + e3[lev, j, i] * masks[bas, j, i]) weighted_var[bas, t, x] = 0.0 weighted_var[bas, t, x] = var[t, lev, j, i] * vector_weight[bas, x] -- GitLab From 94f9a12981ac879c40461d1e63433def98edef7d Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 12:32:41 +0200 Subject: [PATCH 09/54] Add user interface for regmean levels --- diagonals/regmean.py | 80 ++++++++++++++++++++++---------------------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 6c41695..b069529 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -52,7 +52,7 @@ def compute_regmean_2D(var, basins, e1, e2): return regmean -def compute_regmean_2D(var, basins, e1, e2): +def compute_regmean_3D(var, basins, e1, e2, e3): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -74,17 +74,17 @@ def compute_regmean_2D(var, basins, e1, e2): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - regmean = _compute_regmean_2D_gpu( - var, basins, e1, e2 + regmean = _compute_regmean_3D_gpu( + var, basins, e1, e2, e3 ) else: - regmean = _compute_regmean_2D_cpu( - var, basins, e1, e2 + regmean = _compute_regmean_3D_cpu( + var, basins, e1, e2, e3 ) return regmean -def compute_regmean_3D(var, basins, e1, e2, e3): +def compute_regmean_levels(var, basins, e1, e2, e3): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -103,14 +103,14 @@ def compute_regmean_3D(var, basins, e1, e2, e3): Returns ------- regmean: float32 - List containing regional mean for variable var + List containing regional mean for variable var at each level. """ if diagonals.CONFIG.use_gpu: - regmean = _compute_regmean_3D_gpu( + regmean = _compute_regmean_levels_gpu( var, basins, e1, e2, e3 ) else: - regmean = _compute_regmean_3D_cpu( + regmean = _compute_regmean_levels_cpu( var, basins, e1, e2, e3 ) return regmean @@ -289,22 +289,22 @@ def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): gpu_vector_weight = cuda.device_array((mask, levs*lats*lons), dtype=np.float32) - _weight_3D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, + _weigh_3D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, gpu_masks, gpu_weighted_var, gpu_vector_weight) regmean = {} for m in range(mask): - reduced_weight = sum_red_cuda(gpu_vector_weight[m, :]) + reduced_weight = _sum_red_cuda(gpu_vector_weight[m, :]) regmean[m] = {} for t in range(times): - reduced_var = sum_red_cuda(gpu_weighted_var[m, t, :]) + reduced_var = _sum_red_cuda(gpu_weighted_var[m, t, :]) regmean[m][t] = reduced_var / reduced_weight return regmean -def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): +def _compute_regmean_levels_cpu(var, basins, e1, e2, e3): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -316,11 +316,11 @@ def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 + e1 : float32 Masked array containing variable e1t. - e2t: float32 + e2: float32 Masked array containing variable e2t. - e3t: float32 + e3: float32 Masked array containing variable e3t. Returns @@ -333,7 +333,7 @@ def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): regmean_total = {} for basin, mask in basins.items(): regmean = np.empty((times, levs)) - w = compute_weights_3D(mask, e1t, e2t, e3t) + w = _compute_weights_3D(mask, e1, e2, e3) for t in range(times): for l in range(levs): regmean[t, l] = np.average(var[t, l, :, :], axis=(0, 1), @@ -342,7 +342,7 @@ def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): return regmean_total -def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): +def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -354,12 +354,12 @@ def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - e3t: float32 - Masked array containing variable e3t. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. + e3: float32 + Masked array containing variable e3. Returns ------- @@ -381,9 +381,9 @@ def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): grid = (grid_size, lats, times) gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.astype(np.float32)) + gpu_e1 = cuda.to_device(e1.astype(np.float32)) + gpu_e2 = cuda.to_device(e2.astype(np.float32)) + gpu_e3 = cuda.to_device(e3.astype(np.float32)) gpu_masks = cuda.to_device(masks) gpu_weighted_var = cuda.device_array((mask, times, levs, lats*lons), @@ -391,10 +391,10 @@ def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): gpu_vector_weight = cuda.device_array((mask, levs, lats*lons), dtype=np.float32) - variable_weight_cuda_level_3D[grid, block](gpu_var, gpu_weights, + _weigh_var_levels_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, gpu_weighted_var, gpu_vector_weight) - del gpu_var, gpu_weights + del gpu_var, gpu_e1, gpu_e2, gpu_e3 regmean = {} for m in range(mask): @@ -537,19 +537,19 @@ def _weigh_3D_var_cuda(var, e1, e2, e3, masks, weighted_var, vector_weight): @cuda.jit -def variable_weight_cuda_level_3D(var, e1t, e2t, e3t, masks, weighted_var, - vector_weight): +def _weigh_var_levels_cuda(var, e1, e2, e3, masks, weighted_var, + vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 3D variable for each region. Parameters ---------- - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - e3t: float32 - Masked array containing variable e3t. + e1 : float32 + Masked array containing variable e1. + e2: float32 + Masked array containing variable e2. + e3: float32 + Masked array containing variable e3. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -571,10 +571,10 @@ def variable_weight_cuda_level_3D(var, e1t, e2t, e3t, masks, weighted_var, if(i >= var.shape[3]): return for bas in range(basins): - for lev in levels: + for lev in range(levels): vector_weight[bas, lev, x] = 0.0 - vector_weight[bas, lev, x] = (e1t[j, i] * e2t[j, i] * - e3t[lev, j, i] * masks[bas, j, i]) + vector_weight[bas, lev, x] = (e1[j, i] * e2[j, i] * + e3[lev, j, i] * masks[bas, j, i]) weighted_var[bas, t, lev, x] = (var[t, lev, j, i] * vector_weight[bas, lev, x]) -- GitLab From 3341c65028d65156fc7eb38d96c7330f1882572f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 12:36:36 +0200 Subject: [PATCH 10/54] Fix call to function --- diagonals/regmean.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index b069529..433b0e6 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -402,8 +402,8 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): for t in range(times): regmean[m][t] = {} for l in range(levs): - reduced_weight = sum_red_cuda(gpu_vector_weight[mask, l, :]) - reduced_weighted_var = sum_red_cuda(gpu_weighted_var[mask, t, + reduced_weight = _sum_red_cuda(gpu_vector_weight[mask, l, :]) + reduced_weighted_var = _sum_red_cuda(gpu_weighted_var[mask, t, l, :]) regmean[m][t][l] = reduced_weighted_var / reduced_weight del gpu_weighted_var, gpu_vector_weight -- GitLab From fcdb11b2e08ce0414f60410662149403e6e97765 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 12:37:40 +0200 Subject: [PATCH 11/54] Add compute levels to all --- diagonals/regmean.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 433b0e6..f2b663e 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -17,7 +17,8 @@ from numba.cuda.cudadrv import driver import diagonals -__all__ = ['compute_regmean_2D', 'compute_regmean_3D'] +__all__ = ['compute_regmean_2D', 'compute_regmean_3D', + 'compute_regmean_levels'] def compute_regmean_2D(var, basins, e1, e2): -- GitLab From 95a1f00162b875e6677f34bc765d33e6f25d11b4 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 13:24:21 +0200 Subject: [PATCH 12/54] Add user interface for regsum --- diagonals/regsum.py | 73 ++++++++++++++++++++++++++++++++++++--------- 1 file changed, 59 insertions(+), 14 deletions(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index 39a72f1..d6558bd 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -1,32 +1,77 @@ import os + import numpy as np + import iris import iris.cube import iris.analysis + import warnings -import numba import datetime -from numba import jit, njit + +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 +import diagonals + +__all__ = ['compute_regsum_2D', 'compute_regsum_3D'] + +def compute_regsum_2D(var, basins, e1, e2): + """Function that computes the regional sum for 2D vars in the cpu. + + Computes the weights for each region and performs a weighted sum. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. -def compute_regsum_2D_cpu(var, basins, e1t, e2t): + Returns + ------- + regmean_total: float32 + List containing regional mean for variable var. + """ regsum = {} for basin, mask in basins.items(): - weighted_var = compute_weighted_var_2D(var, mask, e1t, e2t) + weighted_var = _weigh_var_2D(var, mask, e1, e2) regsum[basin] = np.sum(weighted_var, axis=(1, 2)) return regsum -def compute_regsum_3D_cpu(var, basins, e1t, e2t, e3t, save3d): +def compute_regsum_3D_cpu(var, basins, e1, e2, e3, save3d): + """Function that computes the regional sum for 3D vars in the cpu. + + Computes the weights for each region and performs a weighted sum. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. + + Returns + ------- + regmean_total: float32 + List containing regional mean for variable var. + """ regsum = {} regsum_3D = {} for basin, mask in basins.items(): - weighted_var = compute_weighted_var_3D(var, mask, e1t, e2t, e3t) + weighted_var = _weigh_var_3D(var, mask, e1, e2, e3) regsum[basin] = np.sum(weighted_var, axis=(2, 3)) if save3d: regsum_3D[basin] = np.sum(regsum[basin], axis=(1,)) @@ -35,14 +80,14 @@ def compute_regsum_3D_cpu(var, basins, e1t, e2t, e3t, save3d): return regsum, regsum_3D -@vectorize(['float32(float32, float32, float32, float32, float32)'], - target='cpu') -def compute_weighted_var_3D(var, mask, e1t, e2t, e3t): - weighted_var = var*mask*e1t*e2t*e3t +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') +def _weigh_var_2D(var, mask, e1, e2): + weighted_var = var*mask*e1*e2 return weighted_var -@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def compute_weighted_var_2D(var, mask, e1t, e2t): - weighted_var = var*mask*e1t*e2t +@vectorize(['float32(float32, float32, float32, float32, float32)'], + target='cpu') +def _weigh_var_3D(var, mask, e1, e2, e3): + weighted_var = var*mask*e1*e2*e3 return weighted_var -- GitLab From 561e242f535737f2cdf90bfc1330ab802dfe4d01 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 15:29:31 +0200 Subject: [PATCH 13/54] Add user interface siasie --- diagonals/{siarea.py => siasie.py} | 98 ++++++++++++++---------------- 1 file changed, 47 insertions(+), 51 deletions(-) rename diagonals/{siarea.py => siasie.py} (85%) diff --git a/diagonals/siarea.py b/diagonals/siasie.py similarity index 85% rename from diagonals/siarea.py rename to diagonals/siasie.py index b7dae33..3de4512 100644 --- a/diagonals/siarea.py +++ b/diagonals/siasie.py @@ -16,15 +16,16 @@ from iris.time import PartialDateTime import warnings import numba -from numba import jit 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 +import diagonals -def compute_siarea(cube_gphit, cube_e1t, cube_e2t, cube_sic, masks): +__all__ = ['compute_siasie', 'compute', 'load_data'] + +def compute_siasie(cube_gphit, cube_e1t, cube_e2t, cube_sic, masks): """Function that calls loading and computing functions for the siarea diagnostic. @@ -64,17 +65,17 @@ def compute_siarea(cube_gphit, cube_e1t, cube_e2t, cube_sic, masks): cube_e1t, cube_e2t, cube_sic) - extn, exts, arean, areas, device = check_device_and_compute(gphit, - e1t, - e2t, - sic_slices, - masks) + extn, exts, arean, areas, device = compute(gphit, + e1t, + e2t, + sic_slices, + masks) seconds_total = datetime.datetime.now() - start print('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) return extn, exts, arean, areas -def check_device_and_compute(gphit, e1t, e2t, sic_slices, basins): +def compute(gphit, e1t, e2t, sic_slices, basins): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -109,26 +110,21 @@ def check_device_and_compute(gphit, e1t, e2t, sic_slices, basins): device: string Device where the computations are being performed (CPU or GPU) """ - check_gpu = numba.cuda.is_available() - if check_gpu is False: - device = 'CPU' - extn, exts, arean, areas = compute_sic_cpu(gphit, e1t, e2t, sic_slices, - basins) + if diagonals.CONFIG.use_gpu: + device = 'GPU' + extn, exts, arean, areas = _compute_sic_gpu( + gphit, e1t, e2t, sic_slices, basins + ) else: - use_gpu = True - if use_gpu is True: - device = 'GPU' - extn, exts, arean, areas = compute_sic_gpu(gphit, e1t, e2t, - sic_slices, basins) - else: - device = 'CPU' - extn, exts, arean, areas = compute_sic_cpu(gphit, e1t, e2t, - sic_slices, basins) + device = 'CPU' + extn, exts, arean, areas = _compute_sic_cpu( + gphit, e1t, e2t, sic_slices, basins + ) return extn, exts, arean, areas, device -def compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): +def _compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): """Function that computes sea ice are and extension on the CPU. Loops over time and basins and reduces over lat and lon: @@ -174,16 +170,16 @@ def compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): areas[time] = {} sic = sic_data.data for basin, tmask in basins.items(): - temp = extn_numba(gphit, e1t, e2t, tmask, sic) + temp = _extn_cpu(gphit, e1t, e2t, tmask, sic) extn[time][basin] = np.sum(temp, axis=(1, 2)) - temp = exts_numba(gphit, e1t, e2t, tmask, sic) + temp = _exts_cpu(gphit, e1t, e2t, tmask, sic) exts[time][basin] = np.sum(temp, axis=(1, 2)) - temp = arean_numba(gphit, e1t, e2t, tmask, sic) + temp = _arean_cpu(gphit, e1t, e2t, tmask, sic) arean[time][basin] = np.sum(temp, axis=(1, 2)) - temp = areas_numba(gphit, e1t, e2t, tmask, sic) + temp = _areas_cpu(gphit, e1t, e2t, tmask, sic) areas[time][basin] = np.sum(temp, axis=(1, 2)) del gphit, e1t, e2t, sic_slices, temp @@ -191,7 +187,7 @@ def compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): return extn, exts, arean, areas -def compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): +def _compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): """Function that computes sea ice are and extension on the GPU. Loops over time and basins and reduces over lat and lon: @@ -249,21 +245,21 @@ def compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): for basin, tmask in basins.items(): gpu_tmask = cuda.to_device(tmask.astype(np.float32)) - extn_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - extn[time][basin] = sum_red_cuda(gpu_temp) + _extn_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, + gpu_tmask, gpu_sic, gpu_temp) + extn[time][basin] = _sum_red_cuda(gpu_temp) - exts_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - exts[time][basin] = sum_red_cuda(gpu_temp) + _exts_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, + gpu_tmask, gpu_sic, gpu_temp) + exts[time][basin] = _sum_red_cuda(gpu_temp) - arean_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - arean[time][basin] = sum_red_cuda(gpu_temp) + _arean_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, + gpu_tmask, gpu_sic, gpu_temp) + arean[time][basin] = _sum_red_cuda(gpu_temp) - areas_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - areas[time][basin] = sum_red_cuda(gpu_temp) + _areas_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, + gpu_tmask, gpu_sic, gpu_temp) + areas[time][basin] = _sum_red_cuda(gpu_temp) del gphit, e1t, e2t, sic_slices del gpu_gphit, gpu_e1t, gpu_e2t, gpu_tmask, gpu_sic, gpu_temp @@ -315,7 +311,7 @@ def load_data(cube_gphit, cube_e1t, cube_e2t, cube_sic): @vectorize(['float32(float32, float32, float32, float32, float32)'], target='cpu') -def arean_numba(gphit, e1t, e2t, tmask, sic): +def _arean_cpu(gphit, e1t, e2t, tmask, sic): """Vectorized numba function executed on the cpu that computes the percentage of sea ice in the northern hemisphere for a given timestep and basin: @@ -344,7 +340,7 @@ def arean_numba(gphit, e1t, e2t, tmask, sic): @cuda.jit() -def arean_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): +def _arean_gpu(gphit, e1t, e2t, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice in the northern hemisphere for a given timestep and basin: @@ -380,7 +376,7 @@ def arean_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): @vectorize(['float32(float32, float32, float32, float32, float32)'], target='cpu') -def extn_numba(gphit, e1t, e2t, tmask, sic): +def _extn_cpu(gphit, e1t, e2t, tmask, sic): """Vectorized numba function executed on the cpu that computes the sea ice extent per unit of area in the northern hemisphere for a given timestep and basin: @@ -410,7 +406,7 @@ def extn_numba(gphit, e1t, e2t, tmask, sic): @cuda.jit() -def extn_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): +def _extn_gpu(gphit, e1t, e2t, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice extent value per unit of area in the northern hemisphere for a given timestep and basin: @@ -446,7 +442,7 @@ def extn_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): @vectorize(['float32(float32, float32, float32, float32, float32)'], target='cpu') -def areas_numba(gphit, e1t, e2t, tmask, sic): +def _areas_cpu(gphit, e1t, e2t, tmask, sic): """Vectorized numba function executed on the cpu that computes the percentage of sea ice in the southern hemisphere for a given timestep and basin: @@ -475,7 +471,7 @@ def areas_numba(gphit, e1t, e2t, tmask, sic): @cuda.jit() -def areas_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): +def _areas_gpu(gphit, e1t, e2t, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice in the southern hemisphere for a given timestep and basin: @@ -511,7 +507,7 @@ def areas_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): @vectorize(['float32(float32, float32, float32, float32, float32)'], target='cpu') -def exts_numba(gphit, e1t, e2t, tmask, sic): +def _exts_cpu(gphit, e1t, e2t, tmask, sic): """Vectorized numba function executed on the cpu that computes the sea ice extent per unit of area in the southern hemisphere for a given timestep and basin: @@ -541,7 +537,7 @@ def exts_numba(gphit, e1t, e2t, tmask, sic): @cuda.jit() -def exts_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): +def _exts_gpu(gphit, e1t, e2t, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice extent value per unit of area in the southern hemisphere for a given timestep and basin: @@ -577,7 +573,7 @@ def exts_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): @vectorize(['float64(float32, float32)', 'float64(float64, float64)'], target='cpu') -def sum_red(x, y): +def _sum_red(x, y): """Vectorized numba function executed on the cpu that performs a sum reduction: @@ -597,7 +593,7 @@ def sum_red(x, y): @cuda.reduce -def sum_red_cuda(x, y): +def _sum_red_cuda(x, y): """Numba kernel executed on the gpu that performs a sum reduction: Parameters -- GitLab From deccf5ee5870fa469a3e9160a89fa5d803f0a880 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 16:11:23 +0200 Subject: [PATCH 14/54] Remove unused function --- diagonals/siasie.py | 51 +-------------------------------------------- 1 file changed, 1 insertion(+), 50 deletions(-) diff --git a/diagonals/siasie.py b/diagonals/siasie.py index 3de4512..41b7b31 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -23,56 +23,7 @@ from numba.cuda.cudadrv import driver import diagonals -__all__ = ['compute_siasie', 'compute', 'load_data'] - -def compute_siasie(cube_gphit, cube_e1t, cube_e2t, cube_sic, masks): - """Function that calls loading and computing functions for the siarea - diagnostic. - - Loads the data from Iris cubes and computes the sea ice area and extent for - the northern and southern hemisphere: - - Parameters - ---------- - cube_gphit : Cube - Iris cube containing variable gphit. - cube_e1t : Cube - Iris cube containing variable e1t. - cube_e2t: Cube - Iris cube containing variable e2t. - cube_sic: Cube - Iris cube containing variable sic (sea_ice_area_fraction). - masks: float32 - List containing basin names and masks. - - Returns - ------- - extn: float64 - List containing sea ice extent in the northern hemisphere at each - timestep and basin. - exts: float64 - List containing sea ice extent in the southern hemisphere at each - timestep and basin. - arean: float64 - List containing sea ice area in the northern hemisphere at each - timestep and basin. - areas: float64 - List containing sea ice area in the southern hemisphere at each - timestep and basin. - """ - start = datetime.datetime.now() - gphit, e1t, e2t, sic_slices = load_data(cube_gphit, - cube_e1t, - cube_e2t, - cube_sic) - extn, exts, arean, areas, device = compute(gphit, - e1t, - e2t, - sic_slices, - masks) - seconds_total = datetime.datetime.now() - start - print('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) - return extn, exts, arean, areas +__all__ = ['compute', 'load_data'] def compute(gphit, e1t, e2t, sic_slices, basins): -- GitLab From 0f4ede3fb08e1cb0e92f7d4c441c4f3dab5ca22b Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 17:15:52 +0200 Subject: [PATCH 15/54] Working moc example --- diagonals/examples/examples-moc.py | 100 ++++++++++++++++++----------- 1 file changed, 63 insertions(+), 37 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index cc875a5..00573a0 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -16,11 +16,11 @@ from numba import cuda from numba import int32, int64, float32, float64 from numba.cuda.cudadrv import driver -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" -VO_FILE = '/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium' \ - '/EC-Earth3-LR/historical/r1i1p1f1/Omon/vo/gn/v20180711' \ - '/vo_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.nc' +MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O25L75.nc" +REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O25L75.nc" +VO_FILE = "/esarchive/exp/ecearth/a0qp/cmorfiles/CMIP/EC-Earth-Consortium/" \ + "EC-Earth3-HR/historical/r1i1p1f1/Omon/vo/gn/v20190411/" \ + "vo_Omon_EC-Earth3-HR_historical_r1i1p1f1_gn_201301-201312.nc" RES = '/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium' \ '/EC-Earth3-LR/historical/r1i1p1f1/Omon/vsftmyz/gn/v20180711' \ '/vsftmyz_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.nc' @@ -34,41 +34,47 @@ def main(): area[basin] = compute_area_basin(e1v, e3v, basins[basin]) del e3v, e1v vo = load_vo() - start = datetime.datetime.now() - vsftmyz = compute_vsftmyz_cpu(vo, area) - print('np {0}'.format(datetime.datetime.now()-start)) - cube = iris.cube.Cube(vsftmyz) - iris.save(cube, '/esarchive/scratch/sloosvel/numba_outputs' - '/vsftmyz_global_numpy.nc', zlib=True) + for i in range(1): + start = datetime.datetime.now() + vsftmyz = compute_vsftmyz_cpu(vo, area) + print('np {0}'.format(datetime.datetime.now()-start)) + start = datetime.datetime.now() + moc = compute_vsftmyz_gpu(vo, area) + print('gp {0}'.format(datetime.datetime.now()-start)) + save_data(moc) + save_data_cpu(vsftmyz) def compute_vsftmyz_cpu(vo, area): - vsftmyx = [] + vsftmyz = [] for basin in area: vsftmyz_basin = np.sum(multiply_vo_basin(vo, area[basin]), axis=3) compute_moc(vsftmyz_basin, vsftmyz_basin) - vsftmyx.append(vsftmyz_basin) - return vsftmyz_basin + vsftmyz.append(vsftmyz_basin) + return vsftmyz def compute_vsftmyz_gpu(vo, area): - levels = vo.shape[1] times = vo.shape[0] + levels = vo.shape[1] lats = vo.shape[2] lons = vo.shape[3] block = (128, 1, 1) grid_size = (lats // block[0]) + 1 - grid = (grid_size, levels, times) + grid3D = (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) + + vsftmyz = [] for basin in area: gpu_area = cuda.to_device(area[basin].astype(np.float32)) - gpu_moc = cuda.device_array((times, levels, lats), dtype=np.float32) - horizontal_integral[grid, block](gpu_vo, gpu_area, gpu_moc) - compute_moc_gpu(gpu_moc, gpu_moc) - vsftmyz_basin = gpu_moc.copy_to_host() + horizontal_integral[grid3D, block](gpu_vo, gpu_area, gpu_moc) + compute_moc_gpu[grid2D, block](gpu_moc) + vsftmyz.append(gpu_moc.copy_to_host()) - return vsftmyz_basin + return vsftmyz def load_vo(): @@ -79,23 +85,44 @@ def load_vo(): def load_mesh(): e3v = iris.load_cube(MESH_FILE, 'e3v_0').data.astype(np.float32) - e3v = np.squeeze(e3v) e1v = iris.load_cube(MESH_FILE, 'e1v').data.astype(np.float32) - e1v = np.squeeze(e1v) return e3v, e1v def load_masks(): basins = {} - cubes = iris.load_cube(MESH_FILE, 'vmask') + global_mask = iris.load_cube(MESH_FILE, 'vmask') basin_name = [] - basin_name.append('a') - cubes.data[..., 0] = 0.0 - cubes.data[..., -1] = 0.0 - basins['a'] = cubes.data.astype(np.float32) + basin_name.append('Global_Ocean') + global_mask.data[..., 0] = 0.0 + global_mask.data[..., -1] = 0.0 + basins['Global_Ocean'] = global_mask.data.astype(np.float32) + cubes = iris.load(REGIONS_FILE) + for cube in cubes: + name = cube.name() + if name in ('nav_lat', 'nav_lon', 'Caspian_Sea', 'Global_Ocean'): + continue + cube = cube.extract(iris.Constraint(z=1)) + basins[name] = cube.data.astype(np.float32) return basin_name, basins +def save_data(vsftmyz): + moc_cube = [] + for basin in range(len(vsftmyz)): + moc_cube.append(iris.cube.Cube(vsftmyz[basin])) + iris.save(moc_cube, '/esarchive/scratch/sloosvel/numba_outputs/moc.nc', + zlib=True) + + +def save_data_cpu(vsftmyz): + moc_cube = [] + for basin in range(len(vsftmyz)): + moc_cube.append(iris.cube.Cube(vsftmyz[basin])) + iris.save(moc_cube, '/esarchive/scratch/sloosvel/numba_outputs/vsftmyz.nc', + zlib=True) + + @vectorize(['float32(float32, float32, float32)'], target='cpu') def compute_area_basin(e1v, e3v, basin): return - e1v * e3v * basin / 1e6 @@ -113,19 +140,18 @@ def compute_moc(moc, out): moc[:, lev, :] = moc[:, lev, :] + moc[:, lev+1, :] -@guvectorize([(float32[:, :, :], float32[:, :, :])], '(t, l, j)->(t, l, j)', - target='cuda') -def compute_moc_gpu(moc, out): - for lev in reversed(range(moc.shape[1]-1)): - moc[:, lev, :] = moc[:, lev, :] + moc[:, lev+1, :] +@cuda.jit() +def compute_moc_gpu(moc): + j, t = cuda.grid(2) + if(j >= moc.shape[2]): + return + for lev in range(moc.shape[1]-2, -1, -1): + moc[t, lev, j] = moc[t, lev, j] + moc[t, lev+1, j] @cuda.jit() def horizontal_integral(vo, area, moc): - j = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x - lev = cuda.blockIdx.y - t = cuda.blockIdx.z - # i, j, t = grid(3) + j, lev, t = cuda.grid(3) moc[t, lev, j] = 0.0 temp = 0.0 if(j >= vo.shape[2]): -- GitLab From 3e10513a05848b346916a41e7b0f41e6beadc195 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 17:16:02 +0200 Subject: [PATCH 16/54] Add user interface for moc --- diagonals/moc.py | 108 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 diagonals/moc.py diff --git a/diagonals/moc.py b/diagonals/moc.py new file mode 100644 index 0000000..8a607c4 --- /dev/null +++ b/diagonals/moc.py @@ -0,0 +1,108 @@ +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 + +import diagonals + +__all__ = ['compute'] + + +def compute(basins, vo): + area = {} + for basin in basins: + area[basin] = _compute_area(e1v, e3v, basins[basin]) + + if diagonals.CONFIG.use_gpu: + moc = _compute_moc_gpu( + vo, area + ) + else: + moc = _compute_moc_cpu( + vo, area + ) + return moc + + +def _compute_moc_cpu(vo, area): + moc = [] + for basin in area: + moc_basin = np.sum(_multiply_vo_basin(vo, area[basin]), axis=3) + _vertical_cumsum(moc_basin, moc_basin) + moc.append(moc_basin) + return moc + + +def _compute_moc_gpu(vo, area): + times = vo.shape[0] + levels = vo.shape[1] + lats = vo.shape[2] + + block = (128, 1, 1) + grid_size = (lats // block[0]) + 1 + grid3D = (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) + + moc = [] + for basin in area: + gpu_area = cuda.to_device(area[basin].astype(np.float32)) + _horizontal_integral[grid3D, block](gpu_vo, gpu_area, gpu_moc) + _vertical_cumsum_gpu[grid2D, block](gpu_moc) + moc.append(gpu_moc.copy_to_host()) + + del gpu_area, gpu_moc, gpu_vo + + return moc + + +@vectorize(['float32(float32, float32, float32)'], target='cpu') +def _compute_area(e1v, e3v, basin): + return - e1v * e3v * basin / 1e6 + + +@vectorize(['float32(float32, float32)'], target='cpu') +def _multiply_vo_basin(vo, basin): + return vo*basin + + +@guvectorize([(float32[:, :, :], float32[:, :, :])], '(t, l, j)->(t, l, j)', + target='cpu') +def _vertical_cumsum(moc, out): + for lev in reversed(range(moc.shape[1]-1)): + moc[:, lev, :] = moc[:, lev, :] + moc[:, lev+1, :] + + +@cuda.jit() +def _vertical_cumsum_gpu(moc): + j, t = cuda.grid(2) + if(j >= moc.shape[2]): + return + for lev in range(moc.shape[1]-2, -1, -1): + moc[t, lev, j] = moc[t, lev, j] + moc[t, lev+1, j] + + +@cuda.jit() +def horizontal_integral(vo, area, moc): + j, lev, t = cuda.grid(3) + moc[t, lev, j] = 0.0 + temp = 0.0 + if(j >= vo.shape[2]): + return + for i in range(vo.shape[3]): + temp += vo[t, lev, j, i] * area[0, lev, j, i] + moc[t, lev, j] = temp \ No newline at end of file -- GitLab From 043520474690b639b4709e181935f6cfde40bb93 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 17:31:38 +0200 Subject: [PATCH 17/54] Add missing variable in function --- diagonals/regmean.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index f2b663e..7db1f99 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -393,8 +393,8 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): dtype=np.float32) _weigh_var_levels_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, - gpu_weighted_var, - gpu_vector_weight) + gpu_masks, gpu_weighted_var, + gpu_vector_weight) del gpu_var, gpu_e1, gpu_e2, gpu_e3 regmean = {} -- GitLab From 9d911aff0d328169cc438bc037cdf1509245abac Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 10 May 2019 18:25:19 +0200 Subject: [PATCH 18/54] Add documentation --- diagonals/moc.py | 144 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 141 insertions(+), 3 deletions(-) diff --git a/diagonals/moc.py b/diagonals/moc.py index 8a607c4..28a884c 100644 --- a/diagonals/moc.py +++ b/diagonals/moc.py @@ -21,11 +21,33 @@ import diagonals __all__ = ['compute'] -def compute(basins, vo): +def compute(basins, e1v, e3v, 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. + e1v: float32 + Masked array containing variable e1v + e3v: float32 + Masked array containing variable e3v + vo : float32 + Masked array containing Sea Water Y Velocity data. + + Returns + ------- + moc: list + List of masked arrays containing the moc index for each basin. + """ area = {} for basin in basins: area[basin] = _compute_area(e1v, e3v, basins[basin]) + del e1v, e3v + if diagonals.CONFIG.use_gpu: moc = _compute_moc_gpu( vo, area @@ -38,6 +60,23 @@ def compute(basins, vo): def _compute_moc_cpu(vo, area): + """Function that calls computing functions in the CPU. + + Computes moc index for each basin and appends them to a list: + + Parameters + ---------- + area: list + List containing the areas for every basin. + + vo : float32 + Masked array containing Sea Water Y Velocity data. + + Returns + ------- + moc: list + List of masked arrays containing the moc index for each basin. + """ moc = [] for basin in area: moc_basin = np.sum(_multiply_vo_basin(vo, area[basin]), axis=3) @@ -47,6 +86,23 @@ def _compute_moc_cpu(vo, area): def _compute_moc_gpu(vo, area): + """Function that calls computing functions in the GPU. + + Computes moc index for each basin and appends them to a list: + + Parameters + ---------- + area: list + List containing the areas for every basin. + + vo : float32 + Masked array containing Sea Water Y Velocity data. + + Returns + ------- + moc: list + List of masked arrays containing the moc index for each basin. + """ times = vo.shape[0] levels = vo.shape[1] lats = vo.shape[2] @@ -72,23 +128,88 @@ def _compute_moc_gpu(vo, area): @vectorize(['float32(float32, float32, float32)'], target='cpu') def _compute_area(e1v, e3v, basin): - return - e1v * e3v * basin / 1e6 + """Vectorized numba function executed in the CPU. + + Calculates cell area for each basin: + + Parameters + ---------- + e1v: float32 + Masked array containing variable e1v. + e3v: float32 + Masked array containing variable e3v. + basin : float32 + Masked array containing a basin mask. + + Returns + ------- + area: float32 + Masked array containing the cell area for a given basin. + """ + area = - e1v * e3v * basin / 1e6 + return area @vectorize(['float32(float32, float32)'], target='cpu') def _multiply_vo_basin(vo, basin): + """Vectorized numba function executed in the CPU. + + Weights vo with a given basin area: + + Parameters + ---------- + e1v: float32 + Masked array containing variable e1v. + e3v: float32 + Masked array containing variable e3v. + basin : float32 + Masked array containing a basin mask. + + Returns + ------- + area: float32 + Masked array containing the cell area for a given basin. + """ return vo*basin @guvectorize([(float32[:, :, :], float32[:, :, :])], '(t, l, j)->(t, l, j)', target='cpu') def _vertical_cumsum(moc, out): + """Numba gu-function executed in the CPU. + + Performs vertical cummulative sum in order to compute the moc index: + + Parameters + ---------- + moc: float32 + Masked array containing horizontal integration of vo. + + Returns + ------- + moc: float32 + Masked array containing the moc index for a given basin. + """ for lev in reversed(range(moc.shape[1]-1)): moc[:, lev, :] = moc[:, lev, :] + moc[:, lev+1, :] @cuda.jit() def _vertical_cumsum_gpu(moc): + """Numba kernel executed in the GPU. + + Performs vertical cummulative sum in order to compute the moc index: + + Parameters + ---------- + moc: float32 + Masked array containing horizontal integration of vo. + + Returns + ------- + moc: float32 + Masked array containing the moc indes for a given basin. + """ j, t = cuda.grid(2) if(j >= moc.shape[2]): return @@ -97,7 +218,24 @@ def _vertical_cumsum_gpu(moc): @cuda.jit() -def horizontal_integral(vo, area, moc): +def _horizontal_integral(vo, area, moc): + """Numba kernel executed in the GPU. + + Integrates along the longitudes in order to compute the moc index: + + Parameters + ---------- + vo : float32 + Masked array containing Sea Water Y Velocity data. + + area: float32 + List containing the areas for every basin. + + Returns + ------- + moc: float32 + Masked array containing vo horizontally integrated for a given basin. + """ j, lev, t = cuda.grid(3) moc[t, lev, j] = 0.0 temp = 0.0 -- GitLab From b4d1efac66312af1f65cae5afda3d357c464918c Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 11:59:19 +0200 Subject: [PATCH 19/54] Update documentation --- diagonals/mesh_helpers/nemo.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py index cd2359c..ce338b5 100644 --- a/diagonals/mesh_helpers/nemo.py +++ b/diagonals/mesh_helpers/nemo.py @@ -62,13 +62,13 @@ def _get_area(e1, e2): Parameters ---------- - e1t : float32 - Masked array containing variable e1t. - e2t : float32 - Masked array containing variable e2t. + e1 : float32 + Masked array containing variable e1. + e2 : float32 + Masked array containing variable e2. Returns ------- - e1t*e2t : float32 + e1*e2 : float32 Masked array containing the area for the whole grid. """ return e1*e2 @@ -87,7 +87,7 @@ def _get_volume(e1, e2, e3): Masked array containing variable e2t. Returns ------- - e1t*e2t : float32 - Masked array containing the area for the whole grid. + e1*e2*e3 : float32 + Masked array containing the volume for the whole grid. """ return e1 * e2 * e3 -- GitLab From d53a7a6b24537471fe8e7fbfb1f4704791236fc5 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 12:01:11 +0200 Subject: [PATCH 20/54] Adapt to user interface --- diagonals/regmean.py | 187 +++++++++++++++++-------------------------- 1 file changed, 72 insertions(+), 115 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 7db1f99..1451d21 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -21,7 +21,7 @@ __all__ = ['compute_regmean_2D', 'compute_regmean_3D', 'compute_regmean_levels'] -def compute_regmean_2D(var, basins, e1, e2): +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: @@ -32,10 +32,8 @@ def compute_regmean_2D(var, basins, e1, e2): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. + area : float32 + Masked array containing cell area. Returns ------- @@ -44,16 +42,16 @@ def compute_regmean_2D(var, basins, e1, e2): """ if diagonals.CONFIG.use_gpu: regmean = _compute_regmean_2D_gpu( - var, basins, e1, e2 + var, basins, area ) else: regmean = _compute_regmean_2D_cpu( - var, basins, e1, e2 + var, basins, area ) return regmean -def compute_regmean_3D(var, basins, e1, e2, e3): +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: @@ -64,10 +62,8 @@ def compute_regmean_3D(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -76,16 +72,16 @@ def compute_regmean_3D(var, basins, e1, e2, e3): """ if diagonals.CONFIG.use_gpu: regmean = _compute_regmean_3D_gpu( - var, basins, e1, e2, e3 + var, basins, volume ) else: regmean = _compute_regmean_3D_cpu( - var, basins, e1, e2, e3 + var, basins, volume ) return regmean -def compute_regmean_levels(var, basins, e1, e2, e3): +def compute_regmean_levels(var, basins, volume): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -96,10 +92,8 @@ def compute_regmean_levels(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -108,16 +102,16 @@ def compute_regmean_levels(var, basins, e1, e2, e3): """ if diagonals.CONFIG.use_gpu: regmean = _compute_regmean_levels_gpu( - var, basins, e1, e2, e3 + var, basins, volume ) else: regmean = _compute_regmean_levels_cpu( - var, basins, e1, e2, e3 + var, basins, volume ) return regmean -def _compute_regmean_2D_cpu(var, basins, e1, e2): +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. @@ -128,10 +122,8 @@ def _compute_regmean_2D_cpu(var, basins, e1, e2): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2 : float32 - Masked array containing variable e2. + area : float32 + Masked array containing cell area. Returns ------- @@ -141,7 +133,7 @@ def _compute_regmean_2D_cpu(var, basins, e1, e2): times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): - weights = compute_weights_2D(mask, e1, e2) + weights = _compute_weights_2D(mask, area) regmean = np.empty(times) for t in range(times): regmean[t] = np.average(var[t, :, :], axis=(0, 1), @@ -150,7 +142,7 @@ def _compute_regmean_2D_cpu(var, basins, e1, e2): return regmean_total -def _compute_regmean_2D_gpu(var, basins, e1, e2): +def _compute_regmean_2D_gpu(var, basins, area): """Function that computes the regional mean for 2D vars in the gpu. Computes the weights for each region and performs a weighted average. @@ -161,10 +153,8 @@ def _compute_regmean_2D_gpu(var, basins, e1, e2): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2 : float32 - Masked array containing variable e2. + area : float32 + Masked array containing cell area. Returns ------- @@ -186,15 +176,14 @@ def _compute_regmean_2D_gpu(var, basins, e1, e2): grid = (grid_size, lats, times) gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1 = cuda.to_device(e1.astype(np.float32)) - gpu_e2 = cuda.to_device(e2.astype(np.float32)) + gpu_area = cuda.to_device(area.astype(np.float32)) gpu_masks = cuda.to_device(masks) gpu_weighted_var = cuda.device_array((mask, times, lats*lons), dtype=np.float32) gpu_vector_weight = cuda.device_array((mask, lats*lons), dtype=np.float32) - _weigh_2D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_masks, - gpu_weighted_var, gpu_vector_weight) + _weigh_2D_var_cuda[grid, block](gpu_var, gpu_area, gpu_masks, + gpu_weighted_var, gpu_vector_weight) regmean = {} for m in range(mask): @@ -207,7 +196,7 @@ def _compute_regmean_2D_gpu(var, basins, e1, e2): return regmean -def _compute_regmean_3D_cpu(var, basins, e1, e2, e3): +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. @@ -218,12 +207,8 @@ def _compute_regmean_3D_cpu(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1t. - e2: float32 - Masked array containing variable e2t. - e3: float32 - Masked array containing variable e3t. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -233,7 +218,7 @@ def _compute_regmean_3D_cpu(var, basins, e1, e2, e3): times = var.shape[0] regmean_total = {} for basin, mask in basins.items(): - weights = _compute_weights_3D(mask, e1, e2, e3) + 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), @@ -242,7 +227,7 @@ def _compute_regmean_3D_cpu(var, basins, e1, e2, e3): return regmean_total -def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): +def _compute_regmean_3D_gpu(var, basins, volume): """Function that computes the regional mean for 3D vars in the gpu. Computes the weights for each region and performs a weighted average. @@ -253,12 +238,8 @@ def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1t. - e2: float32 - Masked array containing variable e2t. - e3: float32 - Masked array containing variable e3t. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -280,9 +261,7 @@ def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): grid = (grid_size, lats, times) gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1 = cuda.to_device(e1.astype(np.float32)) - gpu_e2 = cuda.to_device(e2.astype(np.float32)) - gpu_e3 = cuda.to_device(e3.astype(np.float32)) + gpu_volume = cuda.to_device(volume.astype(np.float32)) gpu_masks = cuda.to_device(masks) gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), @@ -290,9 +269,9 @@ def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): gpu_vector_weight = cuda.device_array((mask, levs*lats*lons), dtype=np.float32) - _weigh_3D_var_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, - gpu_masks, gpu_weighted_var, - gpu_vector_weight) + _weigh_3D_var_cuda[grid, block](gpu_var, gpu_volume, gpu_masks, + gpu_weighted_var, gpu_vector_weight) + del gpu_var, gpu_volume, gpu_masks regmean = {} for m in range(mask): @@ -305,7 +284,7 @@ def _compute_regmean_3D_gpu(var, basins, e1, e2, e3): return regmean -def _compute_regmean_levels_cpu(var, basins, e1, e2, e3): +def _compute_regmean_levels_cpu(var, basins, volume): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -317,12 +296,8 @@ def _compute_regmean_levels_cpu(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1t. - e2: float32 - Masked array containing variable e2t. - e3: float32 - Masked array containing variable e3t. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -334,7 +309,7 @@ def _compute_regmean_levels_cpu(var, basins, e1, e2, e3): regmean_total = {} for basin, mask in basins.items(): regmean = np.empty((times, levs)) - w = _compute_weights_3D(mask, e1, e2, e3) + w = _compute_weights_3D(mask, volume) for t in range(times): for l in range(levs): regmean[t, l] = np.average(var[t, l, :, :], axis=(0, 1), @@ -343,7 +318,7 @@ def _compute_regmean_levels_cpu(var, basins, e1, e2, e3): return regmean_total -def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): +def _compute_regmean_levels_gpu(var, basins, volume): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -355,12 +330,8 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. - e3: float32 - Masked array containing variable e3. + volume : float32 + Masked array containing cell volume. Returns ------- @@ -382,9 +353,7 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): grid = (grid_size, lats, times) gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1 = cuda.to_device(e1.astype(np.float32)) - gpu_e2 = cuda.to_device(e2.astype(np.float32)) - gpu_e3 = cuda.to_device(e3.astype(np.float32)) + gpu_volume = cuda.to_device(volume.astype(np.float32)) gpu_masks = cuda.to_device(masks) gpu_weighted_var = cuda.device_array((mask, times, levs, lats*lons), @@ -392,10 +361,10 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): gpu_vector_weight = cuda.device_array((mask, levs, lats*lons), dtype=np.float32) - _weigh_var_levels_cuda[grid, block](gpu_var, gpu_e1, gpu_e2, gpu_e3, + _weigh_var_levels_cuda[grid, block](gpu_var, gpu_volume, gpu_masks, gpu_weighted_var, gpu_vector_weight) - del gpu_var, gpu_e1, gpu_e2, gpu_e3 + del gpu_var, gpu_volume, gpu_masks regmean = {} for m in range(mask): @@ -413,62 +382,56 @@ def _compute_regmean_levels_gpu(var, basins, e1, e2, e3): @vectorize(['float32(float32, float32, float32)'], target='cpu') -def compute_weights_2D(mask, e1t, e2t): +def _compute_weights_2D(mask, area): """Function that computes the regional weights for 2D vars in the cpu. Parameters ---------- mask : float32 Mask array containing a region mask. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. Returns ------- weights: float32 Masked array containing the weights for a given region. """ - weights = mask*e1t*e2t + weights = mask*area return weights @vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _compute_weights_3D(mask, e1, e2, e3): +def _compute_weights_3D(mask, volume): """Function that computes the regional weights for 3D vars in the cpu. Parameters ---------- mask : float32 Mask array containing a region mask. - e1 : float32 - Masked array containing variable e1. - e2 : float32 - Masked array containing variable e2. - e3 : float32 - Masked array containing variable e3t + volume: float32 + Masked array containing cell volume. Returns ------- weights: float32 Masked array containing the weights for a given region. """ - weights = mask*e1*e2*e3 + weights = mask*volume return weights @cuda.jit -def _weigh_2D_var_cuda(var, e1, e2, masks, weighted_var, vector_weight): +def _weigh_2D_var_cuda(var, area, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 2D variable for each region Parameters ---------- - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. + var: float32 + Masked array containing variable to average. + area: float32 + Masked array containing cell area. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -491,23 +454,21 @@ def _weigh_2D_var_cuda(var, e1, e2, masks, weighted_var, vector_weight): for bas in range(basins): vector_weight[bas, x] = 0.0 weighted_var[bas, t, x] = 0.0 - vector_weight[bas, x] = e1[j, i]*e2[j, i]*masks[bas, j, i] + vector_weight[bas, x] = area[j, i]*masks[bas, j, i] weighted_var[bas, t, x] = var[t, j, i] * vector_weight[bas, x] @cuda.jit -def _weigh_3D_var_cuda(var, e1, e2, e3, masks, weighted_var, vector_weight): +def _weigh_3D_var_cuda(var, volume, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 3D variable for each region. Parameters ---------- - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. - e3: float32 - Masked array containing variable e3. + var: float32 + Masked array containing variable to average. + volume: float32 + Masked array containing cell volume. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -531,26 +492,23 @@ def _weigh_3D_var_cuda(var, e1, e2, e3, masks, weighted_var, vector_weight): for lev in range(levels): x = i+var.shape[3]*(j+var.shape[2]*lev) vector_weight[bas, x] = 0.0 - vector_weight[bas, x] = (e1[j, i] * e2[j, i] * - e3[lev, j, i] * masks[bas, j, i]) + vector_weight[bas, x] = volume[lev, j, i] * masks[bas, j, i] weighted_var[bas, t, x] = 0.0 weighted_var[bas, t, x] = var[t, lev, j, i] * vector_weight[bas, x] @cuda.jit -def _weigh_var_levels_cuda(var, e1, e2, e3, masks, weighted_var, +def _weigh_var_levels_cuda(var, volume, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and weights a 3D variable for each region. Parameters ---------- - e1 : float32 - Masked array containing variable e1. - e2: float32 - Masked array containing variable e2. - e3: float32 - Masked array containing variable e3. + var: float32 + Masked array containing variable to average. + volume: float32 + Masked array containing cell volume. masks : float32 Masked array containing region masks. weighted_var: float32 @@ -574,8 +532,7 @@ def _weigh_var_levels_cuda(var, e1, e2, e3, masks, weighted_var, for bas in range(basins): for lev in range(levels): vector_weight[bas, lev, x] = 0.0 - vector_weight[bas, lev, x] = (e1[j, i] * e2[j, i] * - e3[lev, j, i] * masks[bas, j, i]) + vector_weight[bas, lev, x] = volume[lev, j, i]* masks[bas, j, i] weighted_var[bas, t, lev, x] = (var[t, lev, j, i] * vector_weight[bas, lev, x]) -- GitLab From 15b25050711bf468dcf23f81b9b6038ce7c3c498 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 12:20:56 +0200 Subject: [PATCH 21/54] Adapt to user interface --- diagonals/examples/examples-regionmean.py | 318 ++-------------------- 1 file changed, 29 insertions(+), 289 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 86896c7..43e05f5 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -1,17 +1,14 @@ -import os +import logging +import datetime import numpy as np + import iris import iris.cube import iris.analysis -import warnings -import numba -import datetime -from numba import jit, njit -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 + +import diagonals +import diagonals.regmean as regmean +from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" TOS_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ @@ -23,46 +20,32 @@ THETAO_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ "195001-195012.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" +logger = logging.getLogger(__name__) def main(): - e1t, e2t, e3t_0 = load_mesh() - basins = load_masks() - tos = load_tos() - for i in range(1): - start_gpu = datetime.datetime.now() - regmean_tos_gpu = compute_regmean_2D_gpu(tos, basins, e1t, e2t) - end_gpu = datetime.datetime.now() - start_gpu - start_numpy = datetime.datetime.now() - regmean_tos_numpy = compute_regmean_2D_numpy(tos, basins, e1t, e2t) - end_numpy = datetime.datetime.now() - start_numpy - print('Computing time 2D var GPU: {0} s'.format(end_gpu)) - print('Computing time 2D var numpy: {0} s'.format(end_numpy)) + 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=True): + areacello = mesh.get_areacello() + tos = load_tos() + basin_name, basins = load_masks() + 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) + if diagonals.CONFIG.use_gpu: + device = 'GPU' + else: + device = 'CPU' + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) - del tos - thetao = load_thetao() - for i in range(1): - start_gpu_3D = datetime.datetime.now() - regmean_thetao_gpu = compute_regmean_gpu_3D(thetao, basins, e1t, e2t, - e3t_0) - regmean_thetao_lev_gpu = compute_regmean_lev_gpu_3D(thetao, basins, - e1t, e2t, e3t_0) - end_gpu_3D = datetime.datetime.now() - start_gpu_3D - start_numpy_3D = datetime.datetime.now() - regmean_3D = compute_regmean_3D_numpy(thetao, basins, e1t, e2t, e3t_0) - reg_lev = compute_regmean_3D_levels_numpy(thetao, e1t, e2t, e3t_0) - end_numpy_3D = datetime.datetime.now() - start_numpy_3D - print('Computing time 3D var GPU: {0} s'.format(end_gpu_3D)) - print('Computing time 3D var numpy: {0} s'.format(end_numpy_3D)) -def load_mesh(): - e1t = iris.load_cube(MESH_FILE, 'e1t').data.astype(np.float32) - e1t = np.squeeze(e1t) - e2t = iris.load_cube(MESH_FILE, 'e2t').data.astype(np.float32) - e2t = np.squeeze(e2t) - e3t = iris.load_cube(MESH_FILE, 'e3t_0').data.astype(np.float32) - e3t = np.squeeze(e3t) - return e1t, e2t, e3t def load_masks(): @@ -85,252 +68,9 @@ def load_tos(): def load_thetao(): thetao = iris.load_cube(THETAO_FILE, 'sea_water_potential_temperature') - thetao_data = np.ma.filled(thetao.data, 0.0).astype(np.float32) + thetao_data = thetao.data.astype(np.float32) return thetao_data -def compute_regmean_2D_numpy(var, basins, e1t, e2t): - times = var.shape[0] - regmean = {} - for basin, mask in basins.items(): - weights = compute_weights_2D(mask, e1t, e2t) - regmean[basin] = {} - for t in range(times): - regmean[basin][t] = np.average(var[t, :, :], axis=(0, 1), - weights=np.squeeze(weights)) - return regmean - - -def compute_regmean_3D_levels(var, basins, e1t, e2t, e3t): - times = var.shape[0] - levs = var.shape[1] - regmean_total = {} - for basin, mask in basins.items(): - regmean = np.empty((times, levs)) - w = compute_weights_3D(mask, e1t, e2t, e3t) - for t in range(times): - for l in range(levs): - regmean[t, l] = np.average(var[t, l, :, :], axis=(0, 1), - weights=np.squeeze(w)[l, :, :]) - regmean_total[basin] = regmean - return regmean_total - - -def compute_regmean_3D_cpu(var, basins, e1t, e2t, e3t): - times = var.shape[0] - regmean_total = {} - for basin, mask in basins.items(): - weights = compute_weights_3D(mask, e1t, e2t, e3t) - regmean = np.empty(times) - for t in range(times): - regmean[t] = np.average(var[t, :, :, :], axis=(0, 1, 2), - weights=np.squeeze(weights)) - regmean_total[basin] = regmean - return regmean_total - - -def compute_regmean_2D_gpu(var, basins, e1t, e2t): - times = var.shape[0] - lats = var.shape[1] - lons = var.shape[2] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, lats*lons), dtype=np.float32) - - variable_weight_cuda[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_masks, - gpu_weighted_var, gpu_vector_weight) - - regmean = {} - for m in range(mask): - reduced_weight = sum_red_cuda(gpu_vector_weight[m, :]) - regmean[m] = {} - for t in range(times): - reduced_var = sum_red_cuda(gpu_weighted_var[m, t, :]) - regmean[m][t] = reduced_var / reduced_weight - - return regmean - - -def compute_regmean_gpu_3D(var, basins, e1t, e2t, e3t): - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - - gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, levs*lats*lons), - dtype=np.float32) - - variable_weight_cuda_3D[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_e3t, - gpu_masks, gpu_weighted_var, - gpu_vector_weight) - - regmean = {} - for m in range(mask): - reduced_weight = sum_red_cuda(gpu_vector_weight[m, :]) - regmean[m] = {} - for t in range(times): - reduced_var = sum_red_cuda(gpu_weighted_var[m, t, :]) - regmean[m][t] = reduced_var / reduced_weight - - return regmean - - -def compute_regmean_lev_gpu_3D(var, basins, e1t, e2t, e3t): - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - - gpu_weighted_var = cuda.device_array((mask, times, levs, lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, levs, lats*lons), - dtype=np.float32) - - variable_weight_cuda_level_3D[grid, block](gpu_var, gpu_weights, - gpu_weighted_var, - gpu_vector_weight) - del gpu_var, gpu_weights - - regmean = {} - for m in range(mask): - regmean[m] = {} - for t in range(times): - regmean[m][t] = {} - for l in range(levs): - reduced_weight = sum_red_cuda(gpu_vector_weight[mask, l, :]) - reduced_weighted_var = sum_red_cuda(gpu_weighted_var[mask, t, - l, :]) - regmean[m][t][l] = reduced_weighted_var / reduced_weight - del gpu_weighted_var, gpu_vector_weight - - return regmean - - -def basin_weights_2D(basins, e1t, e2t): - weights = {} - for basin, mask in basins.items(): - weights[basin] = compute_weights_2D(mask, e1t, e2t) - return weights - - -@vectorize(['float32(float32, float32, float32)'], target='cpu') -def compute_weights_2D(mask, e1t, e2t): - weights = mask*e1t*e2t - return weights - - -@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def compute_weights_3D(mask, e1t, e2t, e3t): - weights = mask*e1t*e2t*e3t - return weights - - -@cuda.jit -def variable_weight_cuda(var, e1t, e2t, masks, weighted_var, vector_weight): - i, j, t = cuda.grid(3) - basins = masks.shape[0] - x = i+j*var.shape[2] - if(i >= var.shape[2]): - return - for bas in range(basins): - vector_weight[bas, x] = 0.0 - weighted_var[bas, t, x] = 0.0 - vector_weight[bas, x] = e1t[j, i]*e2t[j, i]*masks[bas, j, i] - weighted_var[bas, t, x] = var[t, j, i] * vector_weight[bas, x] - - -@cuda.jit -def variable_weight_cuda_3D(var, e1t, e2t, e3t, masks, weighted_var, - vector_weight): - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - if(i >= var.shape[3]): - return - for bas in range(basins): - for lev in range(levels): - x = i+var.shape[3]*(j+var.shape[2]*lev) - vector_weight[bas, x] = 0.0 - vector_weight[bas, x] = (e1t[j, i] * e2t[j, i] * - e3t[lev, j, i] * masks[bas, j, i]) - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = var[t, lev, j, i] * vector_weight[bas, x] - - -@cuda.jit -def variable_weight_cuda_level_3D(var, e1t, e2t, e3t, masks, weighted_var, - vector_weight): - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - x = i+j*var.shape[3] - if(i >= var.shape[3]): - return - for bas in range(basins): - for lev in levels: - vector_weight[bas, lev, x] = 0.0 - vector_weight[bas, lev, x] = (e1t[j, i] * e2t[j, i] * - e3t[lev, j, i] * masks[bas, j, i]) - weighted_var[bas, t, lev, x] = (var[t, lev, j, i] * - vector_weight[bas, lev, x]) - - -@vectorize(['float32(float64, float64)'], target='cuda') -def regmean_cuda(weighted_var, weights): - regmean = weighted_var/weights - return regmean - - -@cuda.reduce -def sum_red_cuda(x, y): - return x+y - - if __name__ == '__main__': main() -- GitLab From 815eb9a46db0fd39688eacee29e545436b8dba0f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 12:21:39 +0200 Subject: [PATCH 22/54] Adapt to user interface --- diagonals/regsum.py | 68 +++++++++++------- diagonals/siasie.py | 172 +++++++++++++++++++------------------------- 2 files changed, 117 insertions(+), 123 deletions(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index d6558bd..e080c7f 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -17,9 +17,9 @@ from numba.cuda.cudadrv import driver import diagonals -__all__ = ['compute_regsum_2D', 'compute_regsum_3D'] +__all__ = ['compute_regsum_2D', 'compute_regsum_3D', 'compute_regsum_levels'] -def compute_regsum_2D(var, basins, e1, e2): +def compute_regsum_2D(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. @@ -30,10 +30,8 @@ def compute_regsum_2D(var, basins, e1, e2): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2 : float32 - Masked array containing variable e2. + area : float32 + Masked array containing cell area. Returns ------- @@ -42,15 +40,15 @@ def compute_regsum_2D(var, basins, e1, e2): """ regsum = {} for basin, mask in basins.items(): - weighted_var = _weigh_var_2D(var, mask, e1, e2) + 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, e1, e2, e3, save3d): +def compute_regsum_3D(var, basins, volume): """Function that computes the regional sum for 3D vars in the cpu. - Computes the weights for each region and performs a weighted sum. + Computes the weights for each region and performs a weighted sum. Parameters ---------- @@ -58,10 +56,8 @@ def compute_regsum_3D_cpu(var, basins, e1, e2, e3, save3d): Masked array containing variable data. basins : float32 List containing basin names and masks. - e1 : float32 - Masked array containing variable e1. - e2 : float32 - Masked array containing variable e2. + volume : float32 + Masked array containing cell volume Returns ------- @@ -69,25 +65,45 @@ def compute_regsum_3D_cpu(var, basins, e1, e2, e3, save3d): List containing regional mean for variable var. """ regsum = {} - regsum_3D = {} for basin, mask in basins.items(): - weighted_var = _weigh_var_3D(var, mask, e1, e2, e3) + weighted_var = _weigh_var_3D(var, mask, volume) + regsum[basin] = np.sum(weighted_var, axis=(1, 2, 3)) + return regsum + + +def compute_regsum_levels(var, basins, volume): + """Function that computes the regional sum for 3D vars in the cpu. + + Computes the weights for each region and performs a weighted sum. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + volume : float32 + Masked array containing cell volume + Returns + ------- + regmean_total: float32 + List containing regional mean for variable var. + """ + regsum = {} + for basin, mask in basins.items(): + weighted_var = _weigh_var_3D(var, mask, volume) regsum[basin] = np.sum(weighted_var, axis=(2, 3)) - if save3d: - regsum_3D[basin] = np.sum(regsum[basin], axis=(1,)) - else: - regsum_3D = None - return regsum, regsum_3D + return regsum -@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _weigh_var_2D(var, mask, e1, e2): - weighted_var = var*mask*e1*e2 +@vectorize(['float32(float32, float32, float32)'], target='cpu') +def _weigh_var_2D(var, mask, area): + weighted_var = var*mask*area return weighted_var -@vectorize(['float32(float32, float32, float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32)'], target='cpu') -def _weigh_var_3D(var, mask, e1, e2, e3): - weighted_var = var*mask*e1*e2*e3 +def _weigh_var_3D(var, mask, volume): + weighted_var = var*mask*volume return weighted_var diff --git a/diagonals/siasie.py b/diagonals/siasie.py index 41b7b31..c4b6fb9 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -26,7 +26,7 @@ import diagonals __all__ = ['compute', 'load_data'] -def compute(gphit, e1t, e2t, sic_slices, basins): +def compute(gphit, area, sic_slices, basins): """Function that checks device and calls computing functions. Checks if the computations are going performed in the CPU or the GPU: @@ -35,10 +35,8 @@ def compute(gphit, e1t, e2t, sic_slices, basins): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Iris cubelist containing variable sic sliced over time. masks: float32 @@ -64,18 +62,18 @@ def compute(gphit, e1t, e2t, sic_slices, basins): if diagonals.CONFIG.use_gpu: device = 'GPU' extn, exts, arean, areas = _compute_sic_gpu( - gphit, e1t, e2t, sic_slices, basins + gphit, area, sic_slices, basins ) else: device = 'CPU' extn, exts, arean, areas = _compute_sic_cpu( - gphit, e1t, e2t, sic_slices, basins + gphit, area, sic_slices, basins ) return extn, exts, arean, areas, device -def _compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): +def _compute_sic_cpu(gphit, area, sic_slices, basins): """Function that computes sea ice are and extension on the CPU. Loops over time and basins and reduces over lat and lon: @@ -84,10 +82,8 @@ def _compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic_slices: float32 Iris cubelist containing variable sic sliced over time. basins: float32 @@ -121,24 +117,24 @@ def _compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): areas[time] = {} sic = sic_data.data for basin, tmask in basins.items(): - temp = _extn_cpu(gphit, e1t, e2t, tmask, sic) + temp = _extn_cpu(gphit, area, tmask, sic) extn[time][basin] = np.sum(temp, axis=(1, 2)) - temp = _exts_cpu(gphit, e1t, e2t, tmask, sic) + temp = _exts_cpu(gphit, area, tmask, sic) exts[time][basin] = np.sum(temp, axis=(1, 2)) - temp = _arean_cpu(gphit, e1t, e2t, tmask, sic) + temp = _arean_cpu(gphit, area, tmask, sic) arean[time][basin] = np.sum(temp, axis=(1, 2)) - temp = _areas_cpu(gphit, e1t, e2t, tmask, sic) + temp = _areas_cpu(gphit, area, tmask, sic) areas[time][basin] = np.sum(temp, axis=(1, 2)) - del gphit, e1t, e2t, sic_slices, temp + del gphit, area, sic_slices, temp return extn, exts, arean, areas -def _compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): +def _compute_sic_gpu(gphit, area, sic_slices, basins): """Function that computes sea ice are and extension on the GPU. Loops over time and basins and reduces over lat and lon: @@ -147,10 +143,8 @@ def _compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic_slices: float32 Iris cubelist containing variable sic sliced over time. basins: float32 @@ -177,12 +171,13 @@ def _compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): areas = {} gpu_gphit = cuda.to_device(gphit.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_temp = cuda.device_array(e1t.shape[1]*e1t.shape[2], dtype=np.float32) + gpu_area = cuda.to_device(area.astype(np.float32)) + gpu_temp = cuda.device_array(area.shape[1]*area.shape[2], dtype=np.float32) + del gphit, area + block = (128, 1) - grid_size = ((e1t.shape[2]) // block[0]) + 1 + grid_size = ((area.shape[2]) // block[0]) + 1 grid = (grid_size, e1t.shape[1]) for sic_data in sic_slices: @@ -196,24 +191,24 @@ def _compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): for basin, tmask in basins.items(): gpu_tmask = cuda.to_device(tmask.astype(np.float32)) - _extn_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) + _extn_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask, + gpu_sic, gpu_temp) extn[time][basin] = _sum_red_cuda(gpu_temp) - _exts_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) + _exts_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask, + gpu_sic, gpu_temp) exts[time][basin] = _sum_red_cuda(gpu_temp) - _arean_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) + _arean_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask, + gpu_sic, gpu_temp) arean[time][basin] = _sum_red_cuda(gpu_temp) - _areas_gpu[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) + _areas_gpu[grid, block](gpu_gphit, gpu_area, gpu_tmask, + gpu_sic, gpu_temp) areas[time][basin] = _sum_red_cuda(gpu_temp) - del gphit, e1t, e2t, sic_slices - del gpu_gphit, gpu_e1t, gpu_e2t, gpu_tmask, gpu_sic, gpu_temp + + del gpu_gphit, gpu_area, gpu_tmask, gpu_sic, gpu_temp return extn, exts, arean, areas @@ -260,9 +255,9 @@ def load_data(cube_gphit, cube_e1t, cube_e2t, cube_sic): return gphit, e1t, e2t, sic_slices -@vectorize(['float32(float32, float32, float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _arean_cpu(gphit, e1t, e2t, tmask, sic): +def _arean_cpu(gphit, area, tmask, sic): """Vectorized numba function executed on the cpu that computes the percentage of sea ice in the northern hemisphere for a given timestep and basin: @@ -271,10 +266,8 @@ def _arean_cpu(gphit, e1t, e2t, tmask, sic): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -286,12 +279,12 @@ def _arean_cpu(gphit, e1t, e2t, tmask, sic): Masked array to be summed over lat and lon to obtain sea ice area """ if gphit > 0: - temp = e1t * e2t * tmask * sic / 100 + temp = area * tmask * sic / 100 return temp @cuda.jit() -def _arean_gpu(gphit, e1t, e2t, tmask, sic, temp): +def _arean_gpu(gphit, area, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice in the northern hemisphere for a given timestep and basin: @@ -299,10 +292,8 @@ def _arean_gpu(gphit, e1t, e2t, tmask, sic, temp): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -316,18 +307,18 @@ def _arean_gpu(gphit, e1t, e2t, tmask, sic, temp): for the reduction kernel. """ i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[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] = (e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] * sic[j, i] + temp[x] = (area[0, j, i] * tmask[0, j, i] * sic[j, i] / 100) -@vectorize(['float32(float32, float32, float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _extn_cpu(gphit, e1t, e2t, tmask, sic): +def _extn_cpu(gphit, area, tmask, sic): """Vectorized numba function executed on the cpu that computes the sea ice extent per unit of area in the northern hemisphere for a given timestep and basin: @@ -336,10 +327,8 @@ def _extn_cpu(gphit, e1t, e2t, tmask, sic): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -352,12 +341,12 @@ def _extn_cpu(gphit, e1t, e2t, tmask, sic): value """ if(gphit > 0 and sic > 15): - temp = e1t * e2t * tmask + temp = area * tmask return temp @cuda.jit() -def _extn_gpu(gphit, e1t, e2t, tmask, sic, temp): +def _extn_gpu(gphit, area, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice extent value per unit of area in the northern hemisphere for a given timestep and basin: @@ -366,10 +355,8 @@ def _extn_gpu(gphit, e1t, e2t, tmask, sic, temp): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -383,17 +370,17 @@ def _extn_gpu(gphit, e1t, e2t, tmask, sic, temp): for the reduction kernel. """ i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): + x = i + j * sic.shape[1] + if(i >= sic.shape[1]): return temp[x] = 0.0 if(gphit[0, j, i] > 0 and sic[j, i] > 15): - temp[x] = e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] + temp[x] = area[0, j, i] * tmask[0, j, i] -@vectorize(['float32(float32, float32, float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _areas_cpu(gphit, e1t, e2t, tmask, sic): +def _areas_cpu(gphit, area, tmask, sic): """Vectorized numba function executed on the cpu that computes the percentage of sea ice in the southern hemisphere for a given timestep and basin: @@ -402,10 +389,8 @@ def _areas_cpu(gphit, e1t, e2t, tmask, sic): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -417,12 +402,12 @@ def _areas_cpu(gphit, e1t, e2t, tmask, sic): Masked array to be summed over lat and lon to obtain sea ice area """ if gphit < 0: - temp = e1t * e2t * tmask * sic / 100 + temp = area * tmask * sic / 100 return temp @cuda.jit() -def _areas_gpu(gphit, e1t, e2t, tmask, sic, temp): +def _areas_gpu(gphit, area, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice in the southern hemisphere for a given timestep and basin: @@ -430,10 +415,8 @@ def _areas_gpu(gphit, e1t, e2t, tmask, sic, temp): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -447,18 +430,17 @@ def _areas_gpu(gphit, e1t, e2t, tmask, sic, temp): for the reduction kernel. """ i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[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] = (e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] * sic[j, i] - / 100) + temp[x] = area[0, j, i] * tmask[0, j, i] * sic[j, i] / 100 -@vectorize(['float32(float32, float32, float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _exts_cpu(gphit, e1t, e2t, tmask, sic): +def _exts_cpu(gphit, area, tmask, sic): """Vectorized numba function executed on the cpu that computes the sea ice extent per unit of area in the southern hemisphere for a given timestep and basin: @@ -467,10 +449,8 @@ def _exts_cpu(gphit, e1t, e2t, tmask, sic): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -483,12 +463,12 @@ def _exts_cpu(gphit, e1t, e2t, tmask, sic): value """ if(gphit < 0 and sic > 15): - temp = e1t * e2t * tmask + temp = area * tmask return temp @cuda.jit() -def _exts_gpu(gphit, e1t, e2t, tmask, sic, temp): +def _exts_gpu(gphit, area, tmask, sic, temp): """Numba kernel executed on the gpu that computes the percentage of sea ice extent value per unit of area in the southern hemisphere for a given timestep and basin: @@ -497,10 +477,8 @@ def _exts_gpu(gphit, e1t, e2t, tmask, sic, temp): ---------- gphit : float32 Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. + area : float32 + Masked array containing cell area. sic: float32 Masked array containing variable sic for a given timestep. tmask: float32 @@ -514,12 +492,12 @@ def _exts_gpu(gphit, e1t, e2t, tmask, sic, temp): for the reduction kernel. """ i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): + x = i + j * sic.shape[1] + if(i >= sic.shape[1]): return temp[x] = 0.0 if(gphit[0, j, i] < 0 and sic[j, i] > 15): - temp[x] = e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] + temp[x] = area[0, j, i] * tmask[0, j, i] @vectorize(['float64(float32, float32)', 'float64(float64, float64)'], -- GitLab From 43ec00305a48343798337b76c09456c9f4494bc4 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 12:22:00 +0200 Subject: [PATCH 23/54] Remove unused function --- diagonals/examples/examples-oceanheatcontent.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py index 0d18414..9ede222 100644 --- a/diagonals/examples/examples-oceanheatcontent.py +++ b/diagonals/examples/examples-oceanheatcontent.py @@ -53,11 +53,6 @@ def main(): logger.info('Total ellapsed time on the %s: %s', device, ellapsed) -def load_mesh(): - depth = iris.load_cube(MESH_FILE, 'gdepw_0').data.astype(np.float32) - return depth - - def load_basins(): regions = iris.load(REGIONS_FILE) basins = {} -- GitLab From 58347e8db0ac1558645d7e8e4e23c98d7c1d9210 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 12:40:27 +0200 Subject: [PATCH 24/54] Fix style issues --- diagonals/examples/examples-regionmean.py | 4 +--- diagonals/moc.py | 2 +- diagonals/regmean.py | 8 ++++---- diagonals/regsum.py | 1 + diagonals/siasie.py | 5 +---- 5 files changed, 8 insertions(+), 12 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 43e05f5..7500fe5 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -22,6 +22,7 @@ REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" logger = logging.getLogger(__name__) + def main(): logging.basicConfig(level=logging.INFO) start = datetime.datetime.now() @@ -45,9 +46,6 @@ def main(): logger.info('Total ellapsed time on the %s: %s', device, ellapsed) - - - def load_masks(): basins = {} cubes = iris.load(REGIONS_FILE) diff --git a/diagonals/moc.py b/diagonals/moc.py index 28a884c..f5df58d 100644 --- a/diagonals/moc.py +++ b/diagonals/moc.py @@ -243,4 +243,4 @@ def _horizontal_integral(vo, area, moc): return for i in range(vo.shape[3]): temp += vo[t, lev, j, i] * area[0, lev, j, i] - moc[t, lev, j] = temp \ No newline at end of file + moc[t, lev, j] = temp diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 1451d21..681231b 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -374,14 +374,14 @@ def _compute_regmean_levels_gpu(var, basins, volume): for l in range(levs): reduced_weight = _sum_red_cuda(gpu_vector_weight[mask, l, :]) reduced_weighted_var = _sum_red_cuda(gpu_weighted_var[mask, t, - l, :]) + l, :]) regmean[m][t][l] = reduced_weighted_var / reduced_weight del gpu_weighted_var, gpu_vector_weight return regmean -@vectorize(['float32(float32, float32, float32)'], target='cpu') +@vectorize(['float32(float32, float32)'], target='cpu') def _compute_weights_2D(mask, area): """Function that computes the regional weights for 2D vars in the cpu. @@ -401,7 +401,7 @@ def _compute_weights_2D(mask, area): return weights -@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') +@vectorize(['float32(float32, float32)'], target='cpu') def _compute_weights_3D(mask, volume): """Function that computes the regional weights for 3D vars in the cpu. @@ -532,7 +532,7 @@ def _weigh_var_levels_cuda(var, volume, masks, weighted_var, for bas in range(basins): for lev in range(levels): vector_weight[bas, lev, x] = 0.0 - vector_weight[bas, lev, x] = volume[lev, j, i]* masks[bas, j, i] + vector_weight[bas, lev, x] = volume[lev, j, i] * masks[bas, j, i] weighted_var[bas, t, lev, x] = (var[t, lev, j, i] * vector_weight[bas, lev, x]) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index e080c7f..4aa813a 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -19,6 +19,7 @@ import diagonals __all__ = ['compute_regsum_2D', 'compute_regsum_3D', 'compute_regsum_levels'] + def compute_regsum_2D(var, basins, area): """Function that computes the regional sum for 2D vars in the cpu. diff --git a/diagonals/siasie.py b/diagonals/siasie.py index c4b6fb9..b2d6cb7 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -175,10 +175,9 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins): gpu_temp = cuda.device_array(area.shape[1]*area.shape[2], dtype=np.float32) del gphit, area - block = (128, 1) grid_size = ((area.shape[2]) // block[0]) + 1 - grid = (grid_size, e1t.shape[1]) + grid = (grid_size, area.shape[1]) for sic_data in sic_slices: time = sic_data.coord('time').points[0] @@ -207,9 +206,7 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins): gpu_sic, gpu_temp) areas[time][basin] = _sum_red_cuda(gpu_temp) - del gpu_gphit, gpu_area, gpu_tmask, gpu_sic, gpu_temp - return extn, exts, arean, areas -- GitLab From cbad532307d1760e81f6769f7c3e444508f4e350 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 13:22:56 +0200 Subject: [PATCH 25/54] Add tmask for 3d vars --- diagonals/regsum.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index 4aa813a..5ecf25d 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -46,7 +46,7 @@ def compute_regsum_2D(var, basins, area): return regsum -def compute_regsum_3D(var, basins, volume): +def compute_regsum_3D(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. @@ -67,12 +67,12 @@ def compute_regsum_3D(var, basins, volume): """ regsum = {} for basin, mask in basins.items(): - weighted_var = _weigh_var_3D(var, mask, volume) + weighted_var = _weigh_var_3D(var, mask, volume, tmask) regsum[basin] = np.sum(weighted_var, axis=(1, 2, 3)) return regsum -def compute_regsum_levels(var, basins, volume): +def compute_regsum_levels(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. @@ -92,19 +92,19 @@ def compute_regsum_levels(var, basins, volume): """ regsum = {} for basin, mask in basins.items(): - weighted_var = _weigh_var_3D(var, mask, volume) + 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): - weighted_var = var*mask*area + weighted_var = var * mask * area return weighted_var -@vectorize(['float32(float32, float32, float32)'], +@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def _weigh_var_3D(var, mask, volume): - weighted_var = var*mask*volume +def _weigh_var_3D(var, mask, volume, tmask): + weighted_var = var * mask * volume * tmask return weighted_var -- GitLab From cfc58542c40a1abc36fc906a097485b0d7744599 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 15:16:19 +0200 Subject: [PATCH 26/54] Adapted to user interface --- diagonals/examples/examples-regsum.py | 191 ++++---------------------- 1 file changed, 30 insertions(+), 161 deletions(-) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 9d9da18..c1e8ab0 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -1,17 +1,14 @@ -import os +import logging +import datetime import numpy as np + import iris import iris.cube import iris.analysis -import warnings -import numba -import datetime -from numba import jit, njit -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 + +import diagonals +import diagonals.regsum as regsum +from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" TOS_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ @@ -23,25 +20,31 @@ THETAO_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ "195001-195012.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" +logger = logging.getLogger(__name__) + def main(): - e1t, e2t, e3t_0, tmask = load_mesh() - basins = load_masks() - tos = load_tos() - thetao = load_thetao() - for i in range(1): - start = datetime.datetime.now() - # sum_cpu = compute_regsum_2D_cpu(tos, basins, e1t, e2t) - print('cpu {0}'.format(datetime.datetime.now()-start)) - start = datetime.datetime.now() - # sum_gpu = compute_regsum_2D_gpu(tos, basins, e1t, e2t) - print('gpu {0}'.format(datetime.datetime.now()-start)) - sum3D_cpu, sum_cpu = compute_regsum_3D_cpu(thetao, basins, e1t, e2t, - e3t_0, tmask, True) - a = 0 - # sum3D_gpu = compute_regsum_3D_gpu(thetao, basins, e1t, e2t, e3t_0) - # print(sum3D_cpu['Chukchi_Sea'][:]) - # print(sum3D_gpu[0]) + 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=True): + areacello = mesh.get_areacello() + tos = load_tos() + basin_name, basins = load_masks() + 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) + if diagonals.CONFIG.use_gpu: + device = 'GPU' + else: + device = 'CPU' + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) def load_mesh(): @@ -82,139 +85,5 @@ def load_thetao(): return thetao_data -def compute_regsum_2D_cpu(var, basins, e1t, e2t): - regsum = {} - for basin, mask in basins.items(): - weighted_var = compute_weighted_var_2D(var, mask, e1t, e2t) - regsum[basin] = np.sum(weighted_var, axis=(1, 2)) - return regsum - - -def compute_regsum_3D_cpu(var, basins, e1t, e2t, e3t, tmask, save3d): - regsum = {} - regsum_3D = {} - for basin, mask in basins.items(): - weighted_var = compute_weighted_var_3D(var, mask, e1t, e2t, e3t, tmask) - regsum[basin] = np.sum(weighted_var, axis=(2, 3)) - if save3d: - regsum_3D[basin] = np.sum(regsum[basin], axis=(1,)) - else: - regsum_3D = None - return regsum, regsum_3D - - -def compute_regsum_2D_gpu(var, basins, e1t, e2t): - times = var.shape[0] - lats = var.shape[1] - lons = var.shape[2] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, lats*lons), - dtype=np.float32) - - variable_weight_cuda[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_masks, - gpu_weighted_var) - - regsum = {} - for m in range(mask): - regsum[m] = {} - for t in range(times): - regsum[m][t] = sum_red_cuda(gpu_weighted_var[m, t, :]) - - return regsum - - -def compute_regsum_3D_gpu(var, basins, e1t, e2t, e3t): - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_e3t = cuda.to_device(e3t.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), - dtype=np.float32) - - variable_weight_cuda_3D[grid, block](gpu_var, gpu_e1t, gpu_e2t, gpu_e3t, - gpu_masks, gpu_weighted_var) - - regsum = {} - for m in range(mask): - regsum[m] = {} - for t in range(times): - regsum[m][t] = sum_red_cuda(gpu_weighted_var[m, t, :]) - - return regsum - -@vectorize(['float64(float32, float32, float32, float32, float32, int8)'], - target='cpu') -def compute_weighted_var_3D(var, mask, e1t, e2t, e3t, tmask): - weighted_var = var*mask*e1t*e2t*e3t*tmask - return weighted_var - - -@vectorize(['float32(float32, float32, float32, float32)'], target='cpu') -def compute_weighted_var_2D(var, mask, e1t, e2t): - weighted_var = var*mask*e1t*e2t - return weighted_var - - -@cuda.jit -def variable_weight_cuda(var, e1t, e2t, masks, weighted_var): - i, j, t = cuda.grid(3) - basins = masks.shape[0] - x = i+j*var.shape[2] - if(i >= var.shape[2]): - return - for bas in range(basins): - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = (var[t, j, i] * e1t[j, i] * e2t[j, i] * - masks[bas, j, i]) - - -@cuda.jit -def variable_weight_cuda_3D(var, e1t, e2t, e3t, masks, weighted_var): - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - if(i >= var.shape[3]): - return - for bas in range(basins): - for lev in range(levels): - x = i+var.shape[3]*(j+var.shape[2]*lev) - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = (var[t, j, i] * e1t[j, i] * e2t[j, i] * - e3t[j, i] * masks[bas, j, i]) - - -@cuda.reduce -def sum_red_cuda(x, y): - return x+y - - if __name__ == '__main__': main() -- GitLab From ab736228a452e04c79fc89b4897d2ebdf7eb55b0 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 15:47:20 +0200 Subject: [PATCH 27/54] Add documentation --- diagonals/regsum.py | 289 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 288 insertions(+), 1 deletion(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index 5ecf25d..af6c46a 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -21,6 +21,66 @@ __all__ = ['compute_regsum_2D', 'compute_regsum_3D', 'compute_regsum_levels'] 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: + + 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: + regsum = _compute_regsum_2D_gpu( + var, basins, area + ) + else: + regsum = _compute_regsum_2D_cpu( + var, basins, area + ) + return regsum + + +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: + + 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: + regsum = _compute_regsum_3D_gpu( + var, basins, volume, tmask + ) + else: + regsum = _compute_regsum_3D_cpu( + var, basins, volume, tmask + ) + return regsum + + +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. @@ -46,7 +106,7 @@ def compute_regsum_2D(var, basins, area): return regsum -def compute_regsum_3D(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. @@ -97,8 +157,129 @@ def compute_regsum_levels(var, basins, volume, tmask): return regsum +def _compute_regsum_2D_gpu(var, basins, area): + """Function that computes the regional sum for 2D vars in the GPU. + + Computes the weights for each region and performs a weighted sum. + + 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_total: float32 + List containing regional mean for variable var. + """ + times = var.shape[0] + lats = var.shape[1] + lons = var.shape[2] + mask = len(basins) + + masks = np.empty((mask, lats, lons)) + for i, basin in enumerate(basins): + masks[i, :, :] = basins[basin] + + block = (128, 1, 1) + grid_size = (lons // block[0]) + 1 + grid = (grid_size, lats, times) + + gpu_var = cuda.to_device(var.astype(np.float32)) + gpu_area = cuda.to_device(area.astype(np.float32)) + gpu_masks = cuda.to_device(masks) + gpu_weighted_var = cuda.device_array((mask, times, lats*lons), + dtype=np.float32) + + _variable_weight_cuda[grid, block](gpu_var, gpu_area, gpu_masks, + gpu_weighted_var) + + regsum = {} + for m in range(mask): + regsum[m] = {} + for t in range(times): + regsum[m][t] = _sum_red_cuda(gpu_weighted_var[m, t, :]) + + return regsum + + +def _compute_regsum_3D_gpu(var, basins, volume, tmask): + """Function that computes the regional sum for 3D vars in the GPU. + + Computes the weights for each region and performs a weighted sum. + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + volume : float32 + Masked array containing cell volume. + tmask: float32 + Masked array containing land-sea mask for any grid-point type + + Returns + ------- + regmean_total: float32 + List containing regional mean for variable var. + """ + times = var.shape[0] + levs = var.shape[1] + lats = var.shape[2] + lons = var.shape[3] + mask = len(basins) + + masks = np.empty((mask, lats, lons)) + for i, basin in enumerate(basins): + masks[i, :, :] = basins[basin] + + block = (128, 1, 1) + grid_size = (lons // block[0]) + 1 + grid = (grid_size, lats, times) + + gpu_var = cuda.to_device(var.astype(np.float32)) + gpu_volume = cuda.to_device(volume.astype(np.float32)) + gpu_tmask = cuda.to_device(tmask.astype(np.float32)) + gpu_masks = cuda.to_device(masks) + gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), + dtype=np.float32) + + _variable_weight_cuda_3D[grid, block](gpu_var, gpu_volume, gpu_tmask, + gpu_masks, gpu_weighted_var) + + regsum = {} + for m in range(mask): + regsum[m] = {} + for t in range(times): + regsum[m][t] = _sum_red_cuda(gpu_weighted_var[m, t, :]) + + return regsum + + + @vectorize(['float32(float32, float32, float32)'], target='cpu') def _weigh_var_2D(var, mask, area): + """Function that weights a 2D variable for each region. + + Parameters + ---------- + var: float32 + Masked array containing variable to sum. + mask : float32 + Mask array containing a region mask. + area: float32 + Masked array containing cell area. + + Returns + ------- + weights: float32 + Masked array containing the weights for a given region. + """ weighted_var = var * mask * area return weighted_var @@ -106,5 +287,111 @@ def _weigh_var_2D(var, mask, area): @vectorize(['float32(float32, float32, float32, float32)'], target='cpu') def _weigh_var_3D(var, mask, volume, tmask): + """Function that weights a 3D variable for each region. + + Parameters + ---------- + var: float32 + Masked array containing variable to sum. + mask : float32 + Mask array containing a region mask. + volume: float32 + Masked array containing cell volume. + tmask: float32 + Masked array containing land-sea mask for any grid-point type + + Returns + ------- + weights: float32 + Masked array containing the weights for a given region. + """ weighted_var = var * mask * volume * tmask return weighted_var + + +@cuda.jit +def _variable_weight_cuda(var, area, masks, weighted_var): + """Numba kernel executed on the gpu that weights a 2D variable for each + region. + + Parameters + ---------- + var: float32 + Masked array containing variable to sum. + area: float32 + Masked array containing cell area. + masks : float32 + Masked array containing region masks. + weighted_var: float32 + Empty array to store the results. + + Returns + ------- + weighted_var: float32 + Array containing the weighted variable for a given region. + """ + i, j, t = cuda.grid(3) + basins = masks.shape[0] + x = i+j*var.shape[2] + if(i >= var.shape[2]): + return + for bas in range(basins): + weighted_var[bas, t, x] = 0.0 + weighted_var[bas, t, x] = var[t, j, i] * area[j, i] * masks[bas, j, i] + + +@cuda.jit +def _variable_weight_cuda_3D(var, volume, tmask, masks, weighted_var): + """Numba kernel executed on the gpu that computes the weights and + weights a 3D variable for each region. + + Parameters + ---------- + var: float32 + Masked array containing variable to average. + volume: float32 + Masked array containing cell volume. + tmask: float32 + Masked array containing land-sea mask for any grid-point type + masks : float32 + Masked array containing region masks. + weighted_var: float32 + Empty array to store the results. + + Returns + ------- + weighted_var: float32 + Array containing the weighted variable for a given region. + """ + i, j, t = cuda.grid(3) + basins = masks.shape[0] + levels = var.shape[1] + if(i >= var.shape[3]): + return + for bas in range(basins): + for lev in range(levels): + x = i+var.shape[3]*(j+var.shape[2]*lev) + weighted_var[bas, t, x] = 0.0 + weighted_var[bas, t, x] = (var[t, j, i] * volume[lev, j, i] * + tmask[lev, j, i] * masks[bas, j, i]) + + +@cuda.reduce +def _sum_red_cuda(x, y): + """Numba kernel executed on the gpu that performs a sum reduction: + + Parameters + ---------- + x : float32 + Numpy array. + y : float32 + Numpy array. + + Returns + ------- + x+y : float64 + Reduction operation. + """ + return x+y + + -- GitLab From bd1ef10a94197fbe9fd39b0eb731dd7fc58454ab Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 15:48:55 +0200 Subject: [PATCH 28/54] Correct documentation --- diagonals/regmean.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 681231b..3acf154 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -501,7 +501,7 @@ def _weigh_3D_var_cuda(var, volume, masks, weighted_var, vector_weight): def _weigh_var_levels_cuda(var, volume, masks, weighted_var, vector_weight): """Numba kernel executed on the gpu that computes the weights and - weights a 3D variable for each region. + weights a 3D variable for each region at each level. Parameters ---------- -- GitLab From 8bb1713edd8bcd6f69e12d0bbb5c385f40cc7b55 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 15:49:07 +0200 Subject: [PATCH 29/54] Adapt to user interface --- diagonals/examples/examples-regsum.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index c1e8ab0..4357eee 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -47,18 +47,6 @@ def main(): logger.info('Total ellapsed time on the %s: %s', device, ellapsed) -def load_mesh(): - e1t = iris.load_cube(MESH_FILE, 'e1t').data.astype(np.float32) - e1t = np.squeeze(e1t) - e2t = iris.load_cube(MESH_FILE, 'e2t').data.astype(np.float32) - e2t = np.squeeze(e2t) - e3t = iris.load_cube(MESH_FILE, 'e3t_0').data.astype(np.float32) - e3t = np.squeeze(e3t) - tmask = iris.load_cube(MESH_FILE, 'tmask').data.astype(np.int8) - tmask = np.squeeze(tmask) - return e1t, e2t, e3t, tmask - - def load_masks(): basins = {} cubes = iris.load(REGIONS_FILE) -- GitLab From f0aecb6beccbc9c18ef04b51558aa58b351f3ca3 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 15:52:56 +0200 Subject: [PATCH 30/54] Fix style issues --- diagonals/regsum.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index af6c46a..b72f4ed 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -196,7 +196,7 @@ def _compute_regsum_2D_gpu(var, basins, area): dtype=np.float32) _variable_weight_cuda[grid, block](gpu_var, gpu_area, gpu_masks, - gpu_weighted_var) + gpu_weighted_var) regsum = {} for m in range(mask): @@ -250,7 +250,7 @@ def _compute_regsum_3D_gpu(var, basins, volume, tmask): dtype=np.float32) _variable_weight_cuda_3D[grid, block](gpu_var, gpu_volume, gpu_tmask, - gpu_masks, gpu_weighted_var) + gpu_masks, gpu_weighted_var) regsum = {} for m in range(mask): @@ -261,7 +261,6 @@ def _compute_regsum_3D_gpu(var, basins, volume, tmask): return regsum - @vectorize(['float32(float32, float32, float32)'], target='cpu') def _weigh_var_2D(var, mask, area): """Function that weights a 2D variable for each region. @@ -393,5 +392,3 @@ def _sum_red_cuda(x, y): Reduction operation. """ return x+y - - -- GitLab From 17c8a03130d3e90286bcb2c553177b6c0c3df015 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 16:02:38 +0200 Subject: [PATCH 31/54] Return basin name --- diagonals/examples/examples-regionmean.py | 4 +++- diagonals/examples/examples-regsum.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 7500fe5..e1d26d0 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -48,14 +48,16 @@ def main(): def load_masks(): basins = {} + basin_name = [] cubes = iris.load(REGIONS_FILE) for cube in cubes: name = cube.name() if name in ('nav_lat', 'nav_lon', 'Caspian_Sea'): continue + basin_name.append(name) cube = cube.extract(iris.Constraint(z=1)) basins[name] = cube.data.astype(np.float32) - return basins + return basin_name, basins def load_tos(): diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 4357eee..721208f 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -49,16 +49,16 @@ def main(): def load_masks(): basins = {} + basin_name = [] cubes = iris.load(REGIONS_FILE) for cube in cubes: name = cube.name() if name in ('nav_lat', 'nav_lon', 'Caspian_Sea'): continue - if name not in ('Global_Ocean'): - continue + basin_name.append(name) cube = cube.extract(iris.Constraint(z=1)) basins[name] = cube.data.astype(np.float32) - return basins + return basin_name, basins def load_tos(): -- GitLab From 4b16d8750c6660a5c31b33daccfb78c9af601671 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 16:03:21 +0200 Subject: [PATCH 32/54] Do not change missing values for 0 --- diagonals/examples/examples-regsum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 721208f..02043af 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -69,7 +69,7 @@ def load_tos(): def load_thetao(): thetao = iris.load_cube(THETAO_FILE, 'sea_water_potential_temperature') - thetao_data = np.ma.filled(thetao.data, 0.0).astype(np.float32) + thetao_data = thetao.data.astype(np.float32) return thetao_data -- GitLab From 22d816d31f880e3396e79ccddb7e5875d2a93caa Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 16:52:03 +0200 Subject: [PATCH 33/54] Add loading for gphi --- diagonals/mesh_helpers/nemo.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py index ce338b5..38b1df0 100644 --- a/diagonals/mesh_helpers/nemo.py +++ b/diagonals/mesh_helpers/nemo.py @@ -51,6 +51,11 @@ class Nemo(): cell_point = self.default_cell_point return self.get_mesh_var(cell_point.lower() + 'mask', dtype) + def get_grid_latitude(self, cell_point=None, dtype=np.float32): + if cell_point is None: + cell_point = self.default_cell_point + return self.get_mesh_var('gphi' + cell_point.lower(), dtype) + def get_mesh_var(self, var, dtype): return iris.load_cube(self.mesh_file, var).data.astype(dtype) -- GitLab From 751c92c0a2b7249f16445f983bf61c808a166fed Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 16:52:19 +0200 Subject: [PATCH 34/54] Adapt to user interface --- diagonals/examples/examples-siarea.py | 247 +++----------------------- diagonals/siasie.py | 49 +---- 2 files changed, 29 insertions(+), 267 deletions(-) diff --git a/diagonals/examples/examples-siarea.py b/diagonals/examples/examples-siarea.py index f0bbd05..32f6e7c 100644 --- a/diagonals/examples/examples-siarea.py +++ b/diagonals/examples/examples-siarea.py @@ -1,9 +1,6 @@ -import os -import numpy as np -import glob -import time -import operator +import logging import datetime +import numpy as np import iris import iris.cube @@ -11,17 +8,11 @@ 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 jit -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 +import diagonals +import diagonals.siasie as siasie +from diagonals.mesh_helpers.nemo import Nemo + MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O25L75.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O25L75.nc" @@ -29,74 +20,31 @@ SIC_FILE = "/esarchive/exp/ecearth/a0qp/cmorfiles/CMIP/EC-Earth-Consortium"\ "/EC-Earth3-HR/historical/r1i1p1f1/SImon/siconc/gn/v20180406"\ "/siconc_SImon_EC-Earth3-HR_historical_r1i1p1f1_gn_195001-195012.nc" +logger = logging.getLogger(__name__) + def main(): + logging.basicConfig(level=logging.INFO) start = datetime.datetime.now() - cube_gphit = load_cube(MESH_FILE, 'gphit') - cube_e1t = load_cube(MESH_FILE, 'e1t') - cube_e2t = load_cube(MESH_FILE, 'e2t') - cube_sic = load_cube(SIC_FILE, 'sea_ice_area_fraction') - gphit, e1t, e2t, sic_slices = load_data(cube_gphit, cube_e1t, cube_e2t, - cube_sic) - masks = load_masks(REGIONS_FILE) - extn, exts, arean, areas, device = check_device_and_compute(gphit, e1t, - e2t, - sic_slices, - masks) - save_cubes(extn, exts, arean, areas, masks, device) - seconds_total = datetime.datetime.now() - start - print('Total ellapsed time on the {1}: {0}'.format(seconds_total, device)) - - -def check_device_and_compute(gphit, e1t, e2t, sic_slices, basins): - check_gpu = numba.cuda.is_available() - if check_gpu is False: - device = 'CPU' - extn, exts, arean, areas = compute_sic_cpu(gphit, e1t, e2t, sic_slices, - basins) - else: - use_gpu = True - if use_gpu is True: + logger.info('Starting at %s', start) + mesh = Nemo(MESH_FILE, REGIONS_FILE) + with diagonals.CONFIG.context(use_gpu=True): + basin_name, basins = load_masks(REGIONS_FILE) + areacello = mesh.get_areacello() + gphit = mesh.get_grid_latitude() + cube_sic = load_cube(SIC_FILE, 'sea_ice_area_fraction') + sic_slices = load_data(cube_sic) + masks = load_masks(REGIONS_FILE) + extn, exts, arean, areas = siasie.compute(gphit, areacello, + sic_slices, masks) + del areacello, gphit, cube_sic, sic_slices + if diagonals.CONFIG.use_gpu: device = 'GPU' - extn, exts, arean, areas = compute_sic_gpu(gphit, e1t, e2t, - sic_slices, basins) else: device = 'CPU' - extn, exts, arean, areas = compute_sic_cpu(gphit, e1t, e2t, - sic_slices, basins) - - return extn, exts, arean, areas, device - - -def compute_sic_cpu(gphit, e1t, e2t, sic_slices, basins): - extn = {} - exts = {} - arean = {} - areas = {} - - for sic_data in sic_slices: - time = sic_data.coord('time').points[0] - extn[time] = {} - exts[time] = {} - arean[time] = {} - areas[time] = {} - sic = sic_data.data - for basin, tmask in basins.items(): - temp = extn_numba(gphit, e1t, e2t, tmask, sic) - extn[time][basin] = np.sum(temp, axis=(1, 2)) - - temp = exts_numba(gphit, e1t, e2t, tmask, sic) - exts[time][basin] = np.sum(temp, axis=(1, 2)) - - temp = arean_numba(gphit, e1t, e2t, tmask, sic) - arean[time][basin] = np.sum(temp, axis=(1, 2)) - - temp = areas_numba(gphit, e1t, e2t, tmask, sic) - areas[time][basin] = np.sum(temp, axis=(1, 2)) - - del gphit, e1t, e2t, sic_slices, temp - - return extn, exts, arean, areas + save_cubes(extn, exts, arean, areas, masks, device) + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) def save_cubes(extn, exts, arean, areas, basins, device): @@ -136,64 +84,12 @@ def save_cubes(extn, exts, arean, areas, basins, device): 'numba_outputs/areas_{0}.nc'.format(device), zlib=True) -def compute_sic_gpu(gphit, e1t, e2t, sic_slices, basins): - extn = {} - exts = {} - arean = {} - areas = {} - - gpu_gphit = cuda.to_device(gphit.astype(np.float32)) - gpu_e1t = cuda.to_device(e1t.astype(np.float32)) - gpu_e2t = cuda.to_device(e2t.astype(np.float32)) - gpu_temp = cuda.device_array(e1t.shape[1]*e1t.shape[2], dtype=np.float32) - - block = (128, 1) - grid_size = ((e1t.shape[2]) // block[0]) + 1 - grid = (grid_size, e1t.shape[1]) - - for sic_data in sic_slices: - time = sic_data.coord('time').points[0] - extn[time] = {} - exts[time] = {} - arean[time] = {} - areas[time] = {} - sic = sic_data.data.astype(np.float32) - gpu_sic = cuda.to_device(sic.astype(np.float32)) - for basin, tmask in basins.items(): - gpu_tmask = cuda.to_device(tmask.astype(np.float32)) - - extn_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - extn[time][basin] = sum_red_cuda(gpu_temp) - - exts_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - exts[time][basin] = sum_red_cuda(gpu_temp) - - arean_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - arean[time][basin] = sum_red_cuda(gpu_temp) - - areas_numba_cuda[grid, block](gpu_gphit, gpu_e1t, gpu_e2t, - gpu_tmask, gpu_sic, gpu_temp) - areas[time][basin] = sum_red_cuda(gpu_temp) - - del gphit, e1t, e2t, sic_slices - del gpu_gphit, gpu_e1t, gpu_e2t, gpu_tmask, gpu_sic, gpu_temp - - return extn, exts, arean, areas - - -def load_data(cube_gphit, cube_e1t, cube_e2t, cube_sic): - gphit = cube_gphit.data.astype(np.float32) - e1t = cube_e1t.data.astype(np.float32) - e2t = cube_e2t.data.astype(np.float32) +def load_data(cube_sic): sic_slices = [] for sic_data in cube_sic.slices_over('time'): sic_data.data = np.ma.filled(sic_data.data, 0.0).astype(np.float32) sic_slices.append(sic_data) - - return gphit, e1t, e2t, sic_slices + return sic_slices def load_cube(filepath, variable): @@ -219,94 +115,5 @@ def load_masks(filepath): return basins -@vectorize(['float32(float32, float32, float32, float32, float32)'], - target='cpu') -def arean_numba(gphit, e1t, e2t, tmask, sic): - if gphit > 0: - temp = e1t * e2t * tmask * sic / 100 - return temp - - -@cuda.jit() -def arean_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): - i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): - return - temp[x] = 0.0 - if(gphit[0, j, i] > 0): - temp[x] = (e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] * sic[j, i] - / 100) - - -@vectorize(['float32(float32, float32, float32, float32, float32)'], - target='cpu') -def extn_numba(gphit, e1t, e2t, tmask, sic): - if(gphit > 0 and sic > 15): - temp = e1t * e2t * tmask - return temp - - -@cuda.jit() -def extn_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): - i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): - return - temp[x] = 0.0 - if(gphit[0, j, i] > 0 and sic[j, i] > 15): - temp[x] = e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] - - -@vectorize(['float32(float32, float32, float32, float32, float32)'], - target='cpu') -def areas_numba(gphit, e1t, e2t, tmask, sic): - if gphit < 0: - temp = e1t * e2t * tmask * sic / 100 - return temp - - -@cuda.jit() -def areas_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): - i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): - return - temp[x] = 0.0 - if(gphit[0, j, i] < 0): - temp[x] = (e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] * sic[j, i] - / 100) - - -@vectorize(['float32(float32, float32, float32, float32, float32)'], - target='cpu') -def exts_numba(gphit, e1t, e2t, tmask, sic): - if(gphit < 0 and sic > 15): - temp = e1t * e2t * tmask - return temp - - -@cuda.jit() -def exts_numba_cuda(gphit, e1t, e2t, tmask, sic, temp): - i, j = cuda.grid(2) - x = i + j * e1t.shape[2] - if(i >= e1t.shape[2]): - return - temp[x] = 0.0 - if(gphit[0, j, i] < 0 and sic[j, i] > 15): - temp[x] = e1t[0, j, i] * e2t[0, j, i] * tmask[0, j, i] - - -@vectorize(['float64(float32, float32)', 'float64(float64, float64)'], - target='cpu') -def sum_red(x, y): - return x+y - - -@cuda.reduce -def sum_red_cuda(x, y): - return x+y - - if __name__ == '__main__': main() diff --git a/diagonals/siasie.py b/diagonals/siasie.py index b2d6cb7..21e5607 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -1,9 +1,5 @@ import os import numpy as np -import glob -import time -import operator -import datetime import iris import iris.cube @@ -15,6 +11,7 @@ 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 @@ -23,7 +20,7 @@ from numba.cuda.cudadrv import driver import diagonals -__all__ = ['compute', 'load_data'] +__all__ = ['compute'] def compute(gphit, area, sic_slices, basins): @@ -210,48 +207,6 @@ def _compute_sic_gpu(gphit, area, sic_slices, basins): return extn, exts, arean, areas -def load_data(cube_gphit, cube_e1t, cube_e2t, cube_sic): - """Function that loads data from Iris cubes. - - Loads the data from Iris cubes as float32 for cube_gphit, cube_e1t and - cube_e2t. For cube_sic, the cube gets sliced over time: - - Parameters - ---------- - cube_gphit : Cube - Iris cube containing variable gphit. - cube_e1t : Cube - Iris cube containing variable e1t. - cube_e2t: Cube - Iris cube containing variable e2t. - cube_sic: Cube - Iris cube containing variable sic (sea_ice_area_fraction). - - Returns - ------- - gphit : float32 - Masked array containing variable gphit. - e1t : float32 - Masked array containing variable e1t. - e2t: float32 - Masked array containing variable e2t. - sic_slices: float32 - Iris cubelist containing variable sic sliced over time. - basins: float32 - List containing basin names and masks. - - """ - gphit = cube_gphit.data.astype(np.float32) - e1t = cube_e1t.data.astype(np.float32) - e2t = cube_e2t.data.astype(np.float32) - sic_slices = [] - for sic_data in cube_sic.slices_over('time'): - sic_data.data = np.ma.filled(sic_data.data, 0.0).astype(np.float32) - sic_slices.append(sic_data) - - return gphit, e1t, e2t, sic_slices - - @vectorize(['float32(float32, float32, float32, float32)'], target='cpu') def _arean_cpu(gphit, area, tmask, sic): -- GitLab From 16de0f5e10dc063c069b64262cf9ef6218d52c62 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 17:09:51 +0200 Subject: [PATCH 35/54] Adapt to user interface --- diagonals/examples/examples-moc.py | 142 ++++++----------------------- 1 file changed, 28 insertions(+), 114 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index 00573a0..198a185 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -1,80 +1,46 @@ -import os +import logging +import datetime import numpy as np -import math -import six + import iris import iris.cube -from iris.cube import CubeList import iris.analysis import iris.util import iris.coords -import warnings -import numba -import datetime -from numba import vectorize, guvectorize -from numba import cuda -from numba import int32, int64, float32, float64 -from numba.cuda.cudadrv import driver + +import diagonals +import diagonals.moc as moc +from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O25L75.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O25L75.nc" VO_FILE = "/esarchive/exp/ecearth/a0qp/cmorfiles/CMIP/EC-Earth-Consortium/" \ "EC-Earth3-HR/historical/r1i1p1f1/Omon/vo/gn/v20190411/" \ "vo_Omon_EC-Earth3-HR_historical_r1i1p1f1_gn_201301-201312.nc" -RES = '/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium' \ - '/EC-Earth3-LR/historical/r1i1p1f1/Omon/vsftmyz/gn/v20180711' \ - '/vsftmyz_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.nc' - - -def main(): - e3v, e1v = load_mesh() - basin_name, basins = load_masks() - area = {} - for basin in basins: - area[basin] = compute_area_basin(e1v, e3v, basins[basin]) - del e3v, e1v - vo = load_vo() - for i in range(1): - start = datetime.datetime.now() - vsftmyz = compute_vsftmyz_cpu(vo, area) - print('np {0}'.format(datetime.datetime.now()-start)) - start = datetime.datetime.now() - moc = compute_vsftmyz_gpu(vo, area) - print('gp {0}'.format(datetime.datetime.now()-start)) - save_data(moc) - save_data_cpu(vsftmyz) - -def compute_vsftmyz_cpu(vo, area): - vsftmyz = [] - for basin in area: - vsftmyz_basin = np.sum(multiply_vo_basin(vo, area[basin]), axis=3) - compute_moc(vsftmyz_basin, vsftmyz_basin) - vsftmyz.append(vsftmyz_basin) - return vsftmyz +logger = logging.getLogger(__name__) -def compute_vsftmyz_gpu(vo, area): - times = vo.shape[0] - levels = vo.shape[1] - lats = vo.shape[2] - lons = vo.shape[3] - - block = (128, 1, 1) - grid_size = (lats // block[0]) + 1 - grid3D = (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) - - vsftmyz = [] - for basin in area: - gpu_area = cuda.to_device(area[basin].astype(np.float32)) - horizontal_integral[grid3D, block](gpu_vo, gpu_area, gpu_moc) - compute_moc_gpu[grid2D, block](gpu_moc) - vsftmyz.append(gpu_moc.copy_to_host()) - - return vsftmyz +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=True): + e1v = mesh.get_i_length(cell_point='V') + e3v = mesh.get_k_length(cell_point='V') + basin_name, basins = load_masks() + vo = load_vo() + moc_index = moc.compute(basins, e1v, e3v, vo) + del e1v, e3v, vo + if diagonals.CONFIG.use_gpu: + device = 'GPU' + else: + device = 'CPU' + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) + + # save_data(moc) def load_vo(): @@ -83,12 +49,6 @@ def load_vo(): return vo_data -def load_mesh(): - e3v = iris.load_cube(MESH_FILE, 'e3v_0').data.astype(np.float32) - e1v = iris.load_cube(MESH_FILE, 'e1v').data.astype(np.float32) - return e3v, e1v - - def load_masks(): basins = {} global_mask = iris.load_cube(MESH_FILE, 'vmask') @@ -115,51 +75,5 @@ def save_data(vsftmyz): zlib=True) -def save_data_cpu(vsftmyz): - moc_cube = [] - for basin in range(len(vsftmyz)): - moc_cube.append(iris.cube.Cube(vsftmyz[basin])) - iris.save(moc_cube, '/esarchive/scratch/sloosvel/numba_outputs/vsftmyz.nc', - zlib=True) - - -@vectorize(['float32(float32, float32, float32)'], target='cpu') -def compute_area_basin(e1v, e3v, basin): - return - e1v * e3v * basin / 1e6 - - -@vectorize(['float32(float32, float32)'], target='cpu') -def multiply_vo_basin(vo, basin): - return vo*basin - - -@guvectorize([(float32[:, :, :], float32[:, :, :])], '(t, l, j)->(t, l, j)', - target='cpu') -def compute_moc(moc, out): - for lev in reversed(range(moc.shape[1]-1)): - moc[:, lev, :] = moc[:, lev, :] + moc[:, lev+1, :] - - -@cuda.jit() -def compute_moc_gpu(moc): - j, t = cuda.grid(2) - if(j >= moc.shape[2]): - return - for lev in range(moc.shape[1]-2, -1, -1): - moc[t, lev, j] = moc[t, lev, j] + moc[t, lev+1, j] - - -@cuda.jit() -def horizontal_integral(vo, area, moc): - j, lev, t = cuda.grid(3) - moc[t, lev, j] = 0.0 - temp = 0.0 - if(j >= vo.shape[2]): - return - for i in range(vo.shape[3]): - temp += vo[t, lev, j, i] * area[0, lev, j, i] - moc[t, lev, j] = temp - - if __name__ == '__main__': main() -- GitLab From 9fa6cc8cea80090ff385e0afc6d0397ea0191c06 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 17:15:16 +0200 Subject: [PATCH 36/54] Delete unused areamoc --- diagonals/examples/examples-areamoc.py | 78 -------------------------- 1 file changed, 78 deletions(-) delete mode 100644 diagonals/examples/examples-areamoc.py diff --git a/diagonals/examples/examples-areamoc.py b/diagonals/examples/examples-areamoc.py deleted file mode 100644 index ff54767..0000000 --- a/diagonals/examples/examples-areamoc.py +++ /dev/null @@ -1,78 +0,0 @@ -import os -import numpy as np -import math -import iris -import iris.cube -import iris.analysis -import iris.util -from iris.coords import DimCoord -import warnings -import numba -import datetime -from numba import jit, njit -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 - -VSFTMYZ_FILE = '/esarchive/exp/ecearth/a16l/cmorfiles/CMIP'\ - '/EC-Earth-Consortium/EC-Earth3-LR/historical/r1i1p1f1/Omon'\ - '/vsftmyz/gn/v20180711/vsftmyz_Omon_EC-Earth3-LR_historical_'\ - 'r1i1p1f1_gn_196101-196112.nc' - - -def main(): - one_layer = ((0, 300),) - min_lat = 40 - max_lat = 50 - min_depth = 1000 - max_depth = 2000 - vsftmyz = load_vsftmyz(min_lat, max_lat, min_depth, max_depth) - vsftmyz = np.squeeze(vsftmyz) - levels = vsftmyz.shape[1] - lats = vsftmyz.shape[2] - basins = vsftmyz.shape[3] - - -def load_vsftmyz(min_lat, max_lat, min_depth, max_depth): - long_name = 'Ocean meridional overturning volume streamfunction' - vsftmyz = iris.load_cube(VSFTMYZ_FILE, long_name) - vsftmyz = add_coord_names(vsftmyz) - vsftmyz = latitude_constraint(min_lat, max_lat, vsftmyz) - vsftmyz = depth_constraint(min_depth, max_depth, vsftmyz) - vsftmyz_data = np.ma.filled(vsftmyz.data, 0.0).astype(np.float32) - return vsftmyz_data - - -def add_coord_names(cube): - dims = len(cube.shape) - cube.add_dim_coord( - DimCoord(np.arange(cube.shape[dims - 1]), var_name='basin'), - dims - 1 - ) - cube.add_dim_coord( - DimCoord(np.arange(cube.shape[dims - 2]), var_name='latitude'), - dims - 2 - ) - cube.add_dim_coord( - DimCoord(np.arange(cube.shape[dims - 3]), var_name='longitude'), - dims - 3 - ) - return cube - - -def latitude_constraint(min_lat, max_lat, cube): - constraint = iris.Constraint(latitude=lambda c: min_lat <= c <= max_lat) - cube = cube.extract(constraint) - return cube - - -def depth_constraint(min_depth, max_depth, cube): - constraint = iris.Constraint(depth=lambda c: min_depth <= c <= max_depth) - cube = cube.extract(constraint) - return cube - - -if __name__ == '__main__': - main() -- GitLab From a0ef986f18ce808503416254e8ef398a5edf7855 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Mon, 13 May 2019 17:57:10 +0200 Subject: [PATCH 37/54] Add save functions --- diagonals/examples/examples-regionmean.py | 28 +++++++++++++++++++++++ diagonals/examples/examples-regsum.py | 28 +++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index e1d26d0..73ae5cc 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -42,6 +42,7 @@ def main(): device = 'GPU' else: device = 'CPU' + save_cubes(tosmean, thetaomean, thetaomean3D, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) @@ -72,5 +73,32 @@ def load_thetao(): return thetao_data +def save_cubes(tosmean, thetaomean, thetaomean3D, basins, device): + cubes_tosmean = iris.cube.CubeList() + cubes_thetaomean = iris.cube.CubeList() + cubes_thetaomean3D = iris.cube.CubeList() + for basin in basins.items(): + cube_tosmean = iris.cube.Cube(tosmean[basin]) + cube_thetaomean = iris.cube.Cube(thetaomean[basin]) + cube_thetaomean3D = iris.cube.Cube(thetaomean3D[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')) + + cubes_tosmean.append(cube_tosmean) + cubes_thetaomean.append(cube_thetaomean) + cubes_thetaomean3D.append(cube_thetaomean3D) + + 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/' + 'numba_outputs/thetaomean3D_{0}.nc'.format(device), zlib=True) + + if __name__ == '__main__': main() diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 02043af..cdc3be2 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -43,6 +43,7 @@ def main(): device = 'GPU' else: device = 'CPU' + save_cubes(tossum, thetaosum, thetaosum3D, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) @@ -73,5 +74,32 @@ def load_thetao(): return thetao_data +def save_cubes(tossum, thetaosum, thetaosum3D, basins, device): + cubes_tossum = iris.cube.CubeList() + cubes_thetaosum = iris.cube.CubeList() + cubes_thetaosum3D = iris.cube.CubeList() + for basin in basins.items(): + cube_tosmean = iris.cube.Cube(tossum[basin]) + cube_thetaosum = iris.cube.Cube(thetaosum[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')) + + cubes_tossum.append(cube_tossum) + cubes_thetaosum.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/' + 'numba_outputs/thetaosum3D_{0}.nc'.format(device), zlib=True) + + if __name__ == '__main__': main() -- GitLab From 8f4b8f68981cb9fb0db17404caacd269b8f5dc71 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 09:43:29 +0200 Subject: [PATCH 38/54] Fix variable bad name --- diagonals/examples/examples-regsum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index cdc3be2..4251d62 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -79,7 +79,7 @@ def save_cubes(tossum, thetaosum, thetaosum3D, basins, device): cubes_thetaosum = iris.cube.CubeList() cubes_thetaosum3D = iris.cube.CubeList() for basin in basins.items(): - cube_tosmean = iris.cube.Cube(tossum[basin]) + cube_tossum = iris.cube.Cube(tossum[basin]) cube_thetaosum = iris.cube.Cube(thetaosum[basin]) cube_thetaosum3D = iris.cube.Cube(thetaosum3D[basin]) -- GitLab From d1fa87696f4e5b22bd7c298e81d8b6c9c1ae2e18 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 10:14:03 +0200 Subject: [PATCH 39/54] Fix style issues --- diagonals/examples/examples-moc.py | 18 +++++++++----- diagonals/examples/examples-regionmean.py | 28 ++++++++++----------- diagonals/examples/examples-regsum.py | 30 ++++++++++------------- diagonals/moc.py | 16 ++++++------ 4 files changed, 46 insertions(+), 46 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index 198a185..ab4d3c4 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -37,6 +37,7 @@ def main(): device = 'GPU' else: device = 'CPU' + save_cubes(moc, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) @@ -67,12 +68,17 @@ def load_masks(): return basin_name, basins -def save_data(vsftmyz): - moc_cube = [] - for basin in range(len(vsftmyz)): - moc_cube.append(iris.cube.Cube(vsftmyz[basin])) - iris.save(moc_cube, '/esarchive/scratch/sloosvel/numba_outputs/moc.nc', - zlib=True) +def save_cubes(moc, basins, device): + cubes_moc = iris.cube.CubeList() + cubes_moc = iris.cube.CubeList() + cubes_moc = iris.cube.CubeList() + for basin in basins.items(): + cube_moc = iris.cube.Cube(moc[basin]) + cube_moc.add_aux_coord(iris.coords.AuxCoord(basin, 'region')) + cubes_moc.append(cube_moc) + + iris.save(cubes_moc.merge_cube(), '/esarchive/scratch/sloosvel/' + 'numba_outputs/moc_{0}.nc'.format(device), zlib=True) if __name__ == '__main__': diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 73ae5cc..67c1119 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -32,7 +32,7 @@ def main(): with diagonals.CONFIG.context(use_gpu=True): areacello = mesh.get_areacello() tos = load_tos() - basin_name, basins = load_masks() + basins = load_masks() tosmean = regmean.compute_regmean_2D(tos, basins, areacello) volcello = mesh.get_volcello() thetao = load_thetao() @@ -58,7 +58,7 @@ def load_masks(): basin_name.append(name) cube = cube.extract(iris.Constraint(z=1)) basins[name] = cube.data.astype(np.float32) - return basin_name, basins + return basins def load_tos(): @@ -78,19 +78,17 @@ def save_cubes(tosmean, thetaomean, thetaomean3D, basins, device): cubes_thetaomean = iris.cube.CubeList() cubes_thetaomean3D = iris.cube.CubeList() for basin in basins.items(): - cube_tosmean = iris.cube.Cube(tosmean[basin]) - cube_thetaomean = iris.cube.Cube(thetaomean[basin]) - cube_thetaomean3D = iris.cube.Cube(thetaomean3D[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')) - - cubes_tosmean.append(cube_tosmean) - cubes_thetaomean.append(cube_thetaomean) - cubes_thetaomean3D.append(cube_thetaomean3D) + cube_tosmean = iris.cube.Cube(tosmean[basin]) + cube_thetaomean = iris.cube.Cube(thetaomean[basin]) + cube_thetaomean3D = iris.cube.Cube(thetaomean3D[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')) + + cubes_tosmean.append(cube_tosmean) + cubes_thetaomean.append(cube_thetaomean) + cubes_thetaomean3D.append(cube_thetaomean3D) iris.save(cubes_tosmean.merge_cube(), '/esarchive/scratch/sloosvel/' 'numba_outputs/tosmean_{0}.nc'.format(device), zlib=True) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 4251d62..7a34ef8 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -31,7 +31,7 @@ def main(): with diagonals.CONFIG.context(use_gpu=True): areacello = mesh.get_areacello() tos = load_tos() - basin_name, basins = load_masks() + basins = load_masks() tossum = regsum.compute_regsum_2D(tos, basins, areacello) volcello = mesh.get_volcello() tmask = mesh.get_landsea_mask() @@ -50,16 +50,14 @@ def main(): def load_masks(): basins = {} - basin_name = [] cubes = iris.load(REGIONS_FILE) for cube in cubes: name = cube.name() if name in ('nav_lat', 'nav_lon', 'Caspian_Sea'): continue - basin_name.append(name) cube = cube.extract(iris.Constraint(z=1)) basins[name] = cube.data.astype(np.float32) - return basin_name, basins + return basins def load_tos(): @@ -79,19 +77,17 @@ def save_cubes(tossum, thetaosum, thetaosum3D, basins, device): cubes_thetaosum = 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_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')) - - cubes_tossum.append(cube_tossum) - cubes_thetaosum.append(cube_thetaosum) - cubes_thetaosum3D.append(cube_thetaosum) + cube_tossum = iris.cube.Cube(tossum[basin]) + cube_thetaosum = iris.cube.Cube(thetaosum[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')) + + cubes_tossum.append(cube_tossum) + cubes_thetaosum.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) diff --git a/diagonals/moc.py b/diagonals/moc.py index f5df58d..926ca98 100644 --- a/diagonals/moc.py +++ b/diagonals/moc.py @@ -77,11 +77,11 @@ def _compute_moc_cpu(vo, area): moc: list List of masked arrays containing the moc index for each basin. """ - moc = [] - for basin in area: - moc_basin = np.sum(_multiply_vo_basin(vo, area[basin]), axis=3) + moc = {} + for basin, mask in area.items(): + moc_basin = np.sum(_multiply_vo_basin(vo, mask), axis=3) _vertical_cumsum(moc_basin, moc_basin) - moc.append(moc_basin) + moc[basin] = moc_basin return moc @@ -114,12 +114,12 @@ def _compute_moc_gpu(vo, area): gpu_moc = cuda.device_array((times, levels, lats), dtype=np.float32) grid2D = (grid_size, times) - moc = [] - for basin in area: - gpu_area = cuda.to_device(area[basin].astype(np.float32)) + 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) - moc.append(gpu_moc.copy_to_host()) + moc[basin] = gpu_moc.copy_to_host() del gpu_area, gpu_moc, gpu_vo -- GitLab From 5051a179dfc7ab3d994704346e9c1a8295e04915 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 12:45:44 +0200 Subject: [PATCH 40/54] Delete unused variable --- diagonals/examples/examples-siarea.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/examples/examples-siarea.py b/diagonals/examples/examples-siarea.py index 32f6e7c..6f2c5ca 100644 --- a/diagonals/examples/examples-siarea.py +++ b/diagonals/examples/examples-siarea.py @@ -29,7 +29,7 @@ def main(): logger.info('Starting at %s', start) mesh = Nemo(MESH_FILE, REGIONS_FILE) with diagonals.CONFIG.context(use_gpu=True): - basin_name, basins = load_masks(REGIONS_FILE) + basins = load_masks(REGIONS_FILE) areacello = mesh.get_areacello() gphit = mesh.get_grid_latitude() cube_sic = load_cube(SIC_FILE, 'sea_ice_area_fraction') -- GitLab From ecde73a5b3d65ace3f94d5df3ce00cc6923030e6 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 13:11:48 +0200 Subject: [PATCH 41/54] Remove device from compute function --- diagonals/siasie.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/diagonals/siasie.py b/diagonals/siasie.py index 21e5607..2a5c49b 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -57,17 +57,15 @@ def compute(gphit, area, sic_slices, basins): Device where the computations are being performed (CPU or GPU) """ if diagonals.CONFIG.use_gpu: - device = 'GPU' extn, exts, arean, areas = _compute_sic_gpu( gphit, area, sic_slices, basins ) else: - device = 'CPU' extn, exts, arean, areas = _compute_sic_cpu( gphit, area, sic_slices, basins ) - return extn, exts, arean, areas, device + return extn, exts, arean, areas def _compute_sic_cpu(gphit, area, sic_slices, basins): -- GitLab From ebe08ed1e3fb1ad939ef1ffad96adb2009264caa Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 15:26:35 +0200 Subject: [PATCH 42/54] Fix big in save functions --- diagonals/examples/examples-moc.py | 2 +- diagonals/examples/examples-siarea.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index ab4d3c4..02963f5 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -72,7 +72,7 @@ def save_cubes(moc, basins, device): cubes_moc = iris.cube.CubeList() cubes_moc = iris.cube.CubeList() cubes_moc = iris.cube.CubeList() - for basin in basins.items(): + for basin in basins.keys(): cube_moc = iris.cube.Cube(moc[basin]) cube_moc.add_aux_coord(iris.coords.AuxCoord(basin, 'region')) cubes_moc.append(cube_moc) diff --git a/diagonals/examples/examples-siarea.py b/diagonals/examples/examples-siarea.py index 6f2c5ca..389c68c 100644 --- a/diagonals/examples/examples-siarea.py +++ b/diagonals/examples/examples-siarea.py @@ -53,7 +53,7 @@ def save_cubes(extn, exts, arean, areas, basins, device): cubes_arean = iris.cube.CubeList() cubes_areas = iris.cube.CubeList() for time in extn.keys(): - for basin in basins.items(): + for basin in basins.keys(): cube_extn = iris.cube.Cube(extn[time][basin] / 1e12) cube_exts = iris.cube.Cube(exts[time][basin] / 1e12) cube_arean = iris.cube.Cube(arean[time][basin] / 1e12) -- GitLab From 44a623f4e520fe6429b373409cb2b98da5d00d7a Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 14 May 2019 18:05:20 +0200 Subject: [PATCH 43/54] Delete unused var --- diagonals/examples/examples-moc.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index 02963f5..b6487fd 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -29,7 +29,7 @@ def main(): with diagonals.CONFIG.context(use_gpu=True): e1v = mesh.get_i_length(cell_point='V') e3v = mesh.get_k_length(cell_point='V') - basin_name, basins = load_masks() + basins = load_masks() vo = load_vo() moc_index = moc.compute(basins, e1v, e3v, vo) del e1v, e3v, vo @@ -41,8 +41,6 @@ def main(): ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) - # save_data(moc) - def load_vo(): vo = iris.load_cube(VO_FILE, 'sea_water_y_velocity') @@ -65,7 +63,7 @@ def load_masks(): continue cube = cube.extract(iris.Constraint(z=1)) basins[name] = cube.data.astype(np.float32) - return basin_name, basins + return basins def save_cubes(moc, basins, device): -- GitLab From c24343c5ab9a19f06093134d8f005765ad130286 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 15 May 2019 11:53:10 +0200 Subject: [PATCH 44/54] Change input file, original one got deleted --- diagonals/examples/examples-moc.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index b6487fd..d03ec5c 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -12,11 +12,11 @@ import diagonals import diagonals.moc as moc from diagonals.mesh_helpers.nemo import Nemo -MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O25L75.nc" -REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O25L75.nc" -VO_FILE = "/esarchive/exp/ecearth/a0qp/cmorfiles/CMIP/EC-Earth-Consortium/" \ - "EC-Earth3-HR/historical/r1i1p1f1/Omon/vo/gn/v20190411/" \ - "vo_Omon_EC-Earth3-HR_historical_r1i1p1f1_gn_201301-201312.nc" +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" +VO_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium/" \ + "EC-Earth3-LR/historical/r1i1p1f1/Omon/vo/gn/v20180711/" \ + "vo_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_201301-201312.nc" logger = logging.getLogger(__name__) @@ -37,7 +37,7 @@ def main(): device = 'GPU' else: device = 'CPU' - save_cubes(moc, basins, device) + save_cubes(moc_index, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) -- GitLab From b96baa4c7912d7b09dc48e387f34a6d3a253d1f6 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 15 May 2019 15:11:44 +0200 Subject: [PATCH 45/54] Compare merged and unmerged cubelists --- diagonals/examples/examples-moc.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index d03ec5c..e9433fe 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -16,7 +16,7 @@ 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" VO_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium/" \ "EC-Earth3-LR/historical/r1i1p1f1/Omon/vo/gn/v20180711/" \ - "vo_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_201301-201312.nc" + "vo_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.nc" logger = logging.getLogger(__name__) @@ -57,7 +57,7 @@ def load_masks(): global_mask.data[..., -1] = 0.0 basins['Global_Ocean'] = global_mask.data.astype(np.float32) cubes = iris.load(REGIONS_FILE) - for cube in cubes: + for cube in cubes[0:2]: name = cube.name() if name in ('nav_lat', 'nav_lon', 'Caspian_Sea', 'Global_Ocean'): continue @@ -67,16 +67,15 @@ def load_masks(): def save_cubes(moc, basins, device): - cubes_moc = iris.cube.CubeList() - cubes_moc = iris.cube.CubeList() cubes_moc = iris.cube.CubeList() for basin in basins.keys(): cube_moc = iris.cube.Cube(moc[basin]) cube_moc.add_aux_coord(iris.coords.AuxCoord(basin, 'region')) cubes_moc.append(cube_moc) - + iris.save(cubes_moc, '/esarchive/scratch/sloosvel/' + 'numba_outputs/moc_{0}_notmerged.nc'.format(device), zlib=True) iris.save(cubes_moc.merge_cube(), '/esarchive/scratch/sloosvel/' - 'numba_outputs/moc_{0}.nc'.format(device), zlib=True) + 'numba_outputs/moc_{0}_merged.nc'.format(device), zlib=True) if __name__ == '__main__': -- GitLab From e5df571d7afda5a548605156818dc1be9e0ad87d Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 15 May 2019 15:18:10 +0200 Subject: [PATCH 46/54] Change filename --- diagonals/examples/examples-oceanheatcontent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/examples/examples-oceanheatcontent.py b/diagonals/examples/examples-oceanheatcontent.py index 9ede222..867d818 100644 --- a/diagonals/examples/examples-oceanheatcontent.py +++ b/diagonals/examples/examples-oceanheatcontent.py @@ -93,7 +93,7 @@ def save_data(layers, basins, ohc_2D, ohc1D, device): long_name='{0}'.format(basin))) iris.save(ohc_1D, '/esarchive/scratch/sloosvel/numba_outputs' - '/ohc_1D_{1}_{0[0]}_{0[1]}_numba_vec.nc' + '/ohc_1D_{1}_{0[0]}_{0[1]}.nc' .format(layer, device), zlib=True) iris.save(ohc_cube, '/esarchive/scratch/sloosvel/numba_outputs/ohc_{0}.nc' -- GitLab From c757aef31b9a4a53ccf1e375d2f064bc602930cd Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 15 May 2019 17:24:37 +0200 Subject: [PATCH 47/54] Fix loading of e3 as k length --- diagonals/mesh_helpers/nemo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py index 38b1df0..63fd4b7 100644 --- a/diagonals/mesh_helpers/nemo.py +++ b/diagonals/mesh_helpers/nemo.py @@ -23,7 +23,7 @@ class Nemo(): cell_point = self.default_cell_point e1 = self.get_i_length(cell_point, dtype) e2 = self.get_j_length(cell_point, dtype) - e3 = self.get_j_length(cell_point, dtype) + e3 = self.get_k_length(cell_point, dtype) return _get_volume(e1, e2, e3) def get_i_length(self, cell_point=None, dtype=np.float32): -- GitLab From 74d3cff3a6b6bb245f9513a7ce529f40e38f6cb5 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 15 May 2019 18:15:21 +0200 Subject: [PATCH 48/54] Fix in save routine --- diagonals/examples/examples-regionmean.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 67c1119..e092ad6 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -11,13 +11,13 @@ import diagonals.regmean as regmean from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" -TOS_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ - "/EC-Earth3-LR/historical/r1i1p1f1/Omon/tos/gn/v20180501"\ - "/tos_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_195001-195012.nc" -THETAO_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ - "/EC-Earth3-LR/historical/r1i1p1f1/Omon/thetao/gn/v20180501/"\ +TOS_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium"\ + "/EC-Earth3-LR/historical/r1i1p1f1/Omon/tos/gn/v20180711"\ + "/tos_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.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_"\ - "195001-195012.nc" + "199001-199012.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" logger = logging.getLogger(__name__) @@ -68,7 +68,7 @@ def load_tos(): def load_thetao(): - thetao = iris.load_cube(THETAO_FILE, 'sea_water_potential_temperature') + thetao = iris.load_cube(THETAO_FILE, 'sea_water_conservative_temperature') thetao_data = thetao.data.astype(np.float32) return thetao_data @@ -77,7 +77,7 @@ def save_cubes(tosmean, thetaomean, thetaomean3D, basins, device): cubes_tosmean = iris.cube.CubeList() cubes_thetaomean = iris.cube.CubeList() cubes_thetaomean3D = iris.cube.CubeList() - for basin in basins.items(): + 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]) -- GitLab From 896365b4ed823e7eda01c961cfdcf6c5994625d4 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 10:09:30 +0200 Subject: [PATCH 49/54] Change regmean levs gpu --- diagonals/regmean.py | 39 +++++++++++++++------------------------ 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 3acf154..968095a 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -342,11 +342,6 @@ def _compute_regmean_levels_gpu(var, basins, volume): levs = var.shape[1] lats = var.shape[2] lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] block = (128, 1, 1) grid_size = (lons // block[0]) + 1 @@ -354,28 +349,26 @@ def _compute_regmean_levels_gpu(var, basins, volume): gpu_var = cuda.to_device(var.astype(np.float32)) gpu_volume = cuda.to_device(volume.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, levs, lats*lons), + gpu_weighted_var = cuda.device_array((times, levs, lats*lons), dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, levs, lats*lons), + gpu_vector_weight = cuda.device_array((levs, lats*lons), dtype=np.float32) - _weigh_var_levels_cuda[grid, block](gpu_var, gpu_volume, - gpu_masks, gpu_weighted_var, - gpu_vector_weight) - del gpu_var, gpu_volume, gpu_masks - - regmean = {} - for m in range(mask): - regmean[m] = {} + regmean_total = {} + for m, basin in enumerate(basins): + regmean = np.empty((times, levs)) + gpu_masks = cuda.to_device(basins[basin]) + _weigh_var_levels_cuda[grid, block](gpu_var, gpu_volume, + gpu_masks, gpu_weighted_var, + gpu_vector_weight) for t in range(times): - regmean[m][t] = {} for l in range(levs): reduced_weight = _sum_red_cuda(gpu_vector_weight[mask, l, :]) reduced_weighted_var = _sum_red_cuda(gpu_weighted_var[mask, t, l, :]) - regmean[m][t][l] = reduced_weighted_var / reduced_weight + regmean[t][l] = reduced_weighted_var / reduced_weight + regmean_total[basin] = regmean del gpu_weighted_var, gpu_vector_weight return regmean @@ -529,12 +522,10 @@ def _weigh_var_levels_cuda(var, volume, masks, weighted_var, x = i+j*var.shape[3] if(i >= var.shape[3]): return - for bas in range(basins): - for lev in range(levels): - vector_weight[bas, lev, x] = 0.0 - vector_weight[bas, lev, x] = volume[lev, j, i] * masks[bas, j, i] - weighted_var[bas, t, lev, x] = (var[t, lev, j, i] * - vector_weight[bas, lev, x]) + for lev in range(levels): + vector_weight[lev, x] = 0.0 + vector_weight[lev, x] = volume[lev, j, i] * masks[0, j, i] + weighted_var[t, lev, x] = (var[t, lev, j, i] * vector_weight[lev, x]) @cuda.reduce -- GitLab From c752856e7d33adb30b7c43dd33d11638369d8443 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 10:59:50 +0200 Subject: [PATCH 50/54] Delete save for unmerged cube --- diagonals/examples/examples-moc.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index e9433fe..7058718 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -72,8 +72,6 @@ def save_cubes(moc, basins, device): cube_moc = iris.cube.Cube(moc[basin]) cube_moc.add_aux_coord(iris.coords.AuxCoord(basin, 'region')) cubes_moc.append(cube_moc) - iris.save(cubes_moc, '/esarchive/scratch/sloosvel/' - 'numba_outputs/moc_{0}_notmerged.nc'.format(device), zlib=True) iris.save(cubes_moc.merge_cube(), '/esarchive/scratch/sloosvel/' 'numba_outputs/moc_{0}_merged.nc'.format(device), zlib=True) -- GitLab From 12d77d3a7cd8b142043beec9ca192d9c1d2015da Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 11:00:07 +0200 Subject: [PATCH 51/54] Delete GPU routines, not worth it --- diagonals/regmean.py | 325 ++----------------------------------------- 1 file changed, 8 insertions(+), 317 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index 968095a..c97d31c 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -41,13 +41,9 @@ def compute_regmean_2D(var, basins, area): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - regmean = _compute_regmean_2D_gpu( - var, basins, area - ) - else: - regmean = _compute_regmean_2D_cpu( - var, basins, area - ) + logger.warning('GPU routines not implemented for regmean diagnostic. + 'Using CPU instead') + regmean = _compute_regmean_2D_cpu(var, basins, area) return regmean @@ -71,13 +67,9 @@ def compute_regmean_3D(var, basins, volume): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - regmean = _compute_regmean_3D_gpu( - var, basins, volume - ) - else: - regmean = _compute_regmean_3D_cpu( - var, basins, volume - ) + logger.warning('GPU routines not implemented for regmean diagnostic. + 'Using CPU instead') + regmean = _compute_regmean_3D_cpu(var, basins, volume) return regmean @@ -101,9 +93,8 @@ def compute_regmean_levels(var, basins, volume): List containing regional mean for variable var at each level. """ if diagonals.CONFIG.use_gpu: - regmean = _compute_regmean_levels_gpu( - var, basins, volume - ) + logger.warning('GPU routines not implemented for regmean diagnostic. + 'Using CPU instead') else: regmean = _compute_regmean_levels_cpu( var, basins, volume @@ -142,60 +133,6 @@ def _compute_regmean_2D_cpu(var, basins, area): return regmean_total -def _compute_regmean_2D_gpu(var, basins, area): - """Function that computes the regional mean for 2D vars in the gpu. - - Computes the weights for each region and performs a weighted average. - - 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_total: float32 - List containing regional mean for variable var. - """ - times = var.shape[0] - lats = var.shape[1] - lons = var.shape[2] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - return masks - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_area = cuda.to_device(area.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, lats*lons), dtype=np.float32) - - _weigh_2D_var_cuda[grid, block](gpu_var, gpu_area, gpu_masks, - gpu_weighted_var, gpu_vector_weight) - - regmean = {} - for m in range(mask): - reduced_weight = _sum_red_cuda(gpu_vector_weight[m, :]) - regmean[m] = {} - for t in range(times): - reduced_var = _sum_red_cuda(gpu_weighted_var[m, t, :]) - regmean[m][t] = reduced_var / reduced_weight - - return regmean - - def _compute_regmean_3D_cpu(var, basins, volume): """Function that computes the regional mean for 3D vars in the cpu. @@ -227,63 +164,6 @@ def _compute_regmean_3D_cpu(var, basins, volume): return regmean_total -def _compute_regmean_3D_gpu(var, basins, volume): - """Function that computes the regional mean for 3D vars in the gpu. - - Computes the weights for each region and performs a weighted average. - - Parameters - ---------- - var : float32 - Masked array containing variable data. - basins : float32 - List containing basin names and masks. - volume : float32 - Masked array containing cell volume. - - Returns - ------- - regmean: float32 - List containing the regional mean for variable var. - """ - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_volume = cuda.to_device(volume.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - - gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((mask, levs*lats*lons), - dtype=np.float32) - - _weigh_3D_var_cuda[grid, block](gpu_var, gpu_volume, gpu_masks, - gpu_weighted_var, gpu_vector_weight) - del gpu_var, gpu_volume, gpu_masks - - regmean = {} - for m in range(mask): - reduced_weight = _sum_red_cuda(gpu_vector_weight[m, :]) - regmean[m] = {} - for t in range(times): - reduced_var = _sum_red_cuda(gpu_weighted_var[m, t, :]) - regmean[m][t] = reduced_var / reduced_weight - - return regmean - - def _compute_regmean_levels_cpu(var, basins, volume): """Function that computes the regional mean at every depth level for 3D vars in the cpu. @@ -318,62 +198,6 @@ def _compute_regmean_levels_cpu(var, basins, volume): return regmean_total -def _compute_regmean_levels_gpu(var, basins, volume): - """Function that computes the regional mean at every depth level - for 3D vars in the cpu. - - Computes the weights for each region and performs a weighted average. - - Parameters - ---------- - var : float32 - Masked array containing variable data. - basins : float32 - List containing basin names and masks. - volume : float32 - Masked array containing cell volume. - - Returns - ------- - regmean: float32 - List containing regional mean at every depth level for variable var. - """ - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_volume = cuda.to_device(volume.astype(np.float32)) - - gpu_weighted_var = cuda.device_array((times, levs, lats*lons), - dtype=np.float32) - gpu_vector_weight = cuda.device_array((levs, lats*lons), - dtype=np.float32) - - regmean_total = {} - for m, basin in enumerate(basins): - regmean = np.empty((times, levs)) - gpu_masks = cuda.to_device(basins[basin]) - _weigh_var_levels_cuda[grid, block](gpu_var, gpu_volume, - gpu_masks, gpu_weighted_var, - gpu_vector_weight) - for t in range(times): - for l in range(levs): - reduced_weight = _sum_red_cuda(gpu_vector_weight[mask, l, :]) - reduced_weighted_var = _sum_red_cuda(gpu_weighted_var[mask, t, - l, :]) - regmean[t][l] = reduced_weighted_var / reduced_weight - regmean_total[basin] = regmean - del gpu_weighted_var, gpu_vector_weight - - return regmean - - @vectorize(['float32(float32, float32)'], target='cpu') def _compute_weights_2D(mask, area): """Function that computes the regional weights for 2D vars in the cpu. @@ -412,136 +236,3 @@ def _compute_weights_3D(mask, volume): """ weights = mask*volume return weights - - -@cuda.jit -def _weigh_2D_var_cuda(var, area, masks, weighted_var, vector_weight): - """Numba kernel executed on the gpu that computes the weights and - weights a 2D variable for each region - - Parameters - ---------- - var: float32 - Masked array containing variable to average. - area: float32 - Masked array containing cell area. - masks : float32 - Masked array containing region masks. - weighted_var: float32 - Empty array to store the results. - vector_var: float32 - Empty array to store the results. - - Returns - ------- - weighted_var: float32 - Array containing the weighted variable for a given region. - vector_weight: float32 - Array containing the weights for a given region. - """ - i, j, t = cuda.grid(3) - basins = masks.shape[0] - x = i+j*var.shape[2] - if(i >= var.shape[2]): - return - for bas in range(basins): - vector_weight[bas, x] = 0.0 - weighted_var[bas, t, x] = 0.0 - vector_weight[bas, x] = area[j, i]*masks[bas, j, i] - weighted_var[bas, t, x] = var[t, j, i] * vector_weight[bas, x] - - -@cuda.jit -def _weigh_3D_var_cuda(var, volume, masks, weighted_var, vector_weight): - """Numba kernel executed on the gpu that computes the weights and - weights a 3D variable for each region. - - Parameters - ---------- - var: float32 - Masked array containing variable to average. - volume: float32 - Masked array containing cell volume. - masks : float32 - Masked array containing region masks. - weighted_var: float32 - Empty array to store the results. - vector_weight: float32 - Empty array to store the results. - - Returns - ------- - weighted_var: float32 - Array containing the weighted variable for a given region. - vector_weight: float32 - Array containing the weights for a given region. - """ - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - if(i >= var.shape[3]): - return - for bas in range(basins): - for lev in range(levels): - x = i+var.shape[3]*(j+var.shape[2]*lev) - vector_weight[bas, x] = 0.0 - vector_weight[bas, x] = volume[lev, j, i] * masks[bas, j, i] - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = var[t, lev, j, i] * vector_weight[bas, x] - - -@cuda.jit -def _weigh_var_levels_cuda(var, volume, masks, weighted_var, - vector_weight): - """Numba kernel executed on the gpu that computes the weights and - weights a 3D variable for each region at each level. - - Parameters - ---------- - var: float32 - Masked array containing variable to average. - volume: float32 - Masked array containing cell volume. - masks : float32 - Masked array containing region masks. - weighted_var: float32 - Empty array to store the results. - vector_weight: float32 - Empty array to store the results. - - Returns - ------- - weighted_var: float32 - Array containing the weighted variable for a given region. - vector_weight: float32 - Array containing the weights for a given region. - """ - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - x = i+j*var.shape[3] - if(i >= var.shape[3]): - return - for lev in range(levels): - vector_weight[lev, x] = 0.0 - vector_weight[lev, x] = volume[lev, j, i] * masks[0, j, i] - weighted_var[t, lev, x] = (var[t, lev, j, i] * vector_weight[lev, x]) - - -@cuda.reduce -def _sum_red_cuda(x, y): - """Numba kernel executed on the gpu that performs a sum reduction: - - Parameters - ---------- - x : float32 - Numpy array. - y : float32 - Numpy array. - - Returns - ------- - x+y : float64 - Reduction operation. - """ - return x+y -- GitLab From 4df77bf2b8d7418aff27f38a756fcd4b5370986d Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 11:13:45 +0200 Subject: [PATCH 52/54] Remove GPU routines from regsum and regmean --- diagonals/examples/examples-regionmean.py | 5 +- diagonals/examples/examples-regsum.py | 5 +- diagonals/regsum.py | 210 +--------------------- 3 files changed, 8 insertions(+), 212 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index e092ad6..3983906 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -38,10 +38,7 @@ def main(): thetao = load_thetao() thetaomean = regmean.compute_regmean_levels(thetao, basins, volcello) thetaomean3D = regmean.compute_regmean_3D(thetao, basins, volcello) - if diagonals.CONFIG.use_gpu: - device = 'GPU' - else: - device = 'CPU' + device = 'CPU' save_cubes(tosmean, thetaomean, thetaomean3D, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 7a34ef8..646a2e5 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -39,10 +39,7 @@ def main(): thetaosum = regsum.compute_regsum_levels(thetao, basins, volcello, tmask) thetaosum3D = regsum.compute_regsum_3D(thetao, basins, volcello, tmask) - if diagonals.CONFIG.use_gpu: - device = 'GPU' - else: - device = 'CPU' + device = 'CPU' save_cubes(tossum, thetaosum, thetaosum3D, basins, device) ellapsed = datetime.datetime.now() - start logger.info('Total ellapsed time on the %s: %s', device, ellapsed) diff --git a/diagonals/regsum.py b/diagonals/regsum.py index b72f4ed..d37a9a3 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -40,13 +40,9 @@ def compute_regsum_2D(var, basins, area): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - regsum = _compute_regsum_2D_gpu( - var, basins, area - ) - else: - regsum = _compute_regsum_2D_cpu( - var, basins, area - ) + logger.warning('GPU routines not implemented for regmean diagnostic. + 'Using CPU instead') + regsum = _compute_regsum_2D_cpu(var, basins, area) return regsum @@ -70,13 +66,9 @@ def compute_regsum_3D(var, basins, volume, tmask): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - regsum = _compute_regsum_3D_gpu( - var, basins, volume, tmask - ) - else: - regsum = _compute_regsum_3D_cpu( - var, basins, volume, tmask - ) + logger.warning('GPU routines not implemented for regmean diagnostic. + 'Using CPU instead') + regsum = _compute_regsum_3D_cpu(var, basins, volume, tmask) return regsum @@ -157,110 +149,6 @@ def compute_regsum_levels(var, basins, volume, tmask): return regsum -def _compute_regsum_2D_gpu(var, basins, area): - """Function that computes the regional sum for 2D vars in the GPU. - - Computes the weights for each region and performs a weighted sum. - - 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_total: float32 - List containing regional mean for variable var. - """ - times = var.shape[0] - lats = var.shape[1] - lons = var.shape[2] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_area = cuda.to_device(area.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, lats*lons), - dtype=np.float32) - - _variable_weight_cuda[grid, block](gpu_var, gpu_area, gpu_masks, - gpu_weighted_var) - - regsum = {} - for m in range(mask): - regsum[m] = {} - for t in range(times): - regsum[m][t] = _sum_red_cuda(gpu_weighted_var[m, t, :]) - - return regsum - - -def _compute_regsum_3D_gpu(var, basins, volume, tmask): - """Function that computes the regional sum for 3D vars in the GPU. - - Computes the weights for each region and performs a weighted sum. - - Parameters - ---------- - var : float32 - Masked array containing variable data. - basins : float32 - List containing basin names and masks. - volume : float32 - Masked array containing cell volume. - tmask: float32 - Masked array containing land-sea mask for any grid-point type - - Returns - ------- - regmean_total: float32 - List containing regional mean for variable var. - """ - times = var.shape[0] - levs = var.shape[1] - lats = var.shape[2] - lons = var.shape[3] - mask = len(basins) - - masks = np.empty((mask, lats, lons)) - for i, basin in enumerate(basins): - masks[i, :, :] = basins[basin] - - block = (128, 1, 1) - grid_size = (lons // block[0]) + 1 - grid = (grid_size, lats, times) - - gpu_var = cuda.to_device(var.astype(np.float32)) - gpu_volume = cuda.to_device(volume.astype(np.float32)) - gpu_tmask = cuda.to_device(tmask.astype(np.float32)) - gpu_masks = cuda.to_device(masks) - gpu_weighted_var = cuda.device_array((mask, times, levs*lats*lons), - dtype=np.float32) - - _variable_weight_cuda_3D[grid, block](gpu_var, gpu_volume, gpu_tmask, - gpu_masks, gpu_weighted_var) - - regsum = {} - for m in range(mask): - regsum[m] = {} - for t in range(times): - regsum[m][t] = _sum_red_cuda(gpu_weighted_var[m, t, :]) - - return regsum - - @vectorize(['float32(float32, float32, float32)'], target='cpu') def _weigh_var_2D(var, mask, area): """Function that weights a 2D variable for each region. @@ -306,89 +194,3 @@ def _weigh_var_3D(var, mask, volume, tmask): """ weighted_var = var * mask * volume * tmask return weighted_var - - -@cuda.jit -def _variable_weight_cuda(var, area, masks, weighted_var): - """Numba kernel executed on the gpu that weights a 2D variable for each - region. - - Parameters - ---------- - var: float32 - Masked array containing variable to sum. - area: float32 - Masked array containing cell area. - masks : float32 - Masked array containing region masks. - weighted_var: float32 - Empty array to store the results. - - Returns - ------- - weighted_var: float32 - Array containing the weighted variable for a given region. - """ - i, j, t = cuda.grid(3) - basins = masks.shape[0] - x = i+j*var.shape[2] - if(i >= var.shape[2]): - return - for bas in range(basins): - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = var[t, j, i] * area[j, i] * masks[bas, j, i] - - -@cuda.jit -def _variable_weight_cuda_3D(var, volume, tmask, masks, weighted_var): - """Numba kernel executed on the gpu that computes the weights and - weights a 3D variable for each region. - - Parameters - ---------- - var: float32 - Masked array containing variable to average. - volume: float32 - Masked array containing cell volume. - tmask: float32 - Masked array containing land-sea mask for any grid-point type - masks : float32 - Masked array containing region masks. - weighted_var: float32 - Empty array to store the results. - - Returns - ------- - weighted_var: float32 - Array containing the weighted variable for a given region. - """ - i, j, t = cuda.grid(3) - basins = masks.shape[0] - levels = var.shape[1] - if(i >= var.shape[3]): - return - for bas in range(basins): - for lev in range(levels): - x = i+var.shape[3]*(j+var.shape[2]*lev) - weighted_var[bas, t, x] = 0.0 - weighted_var[bas, t, x] = (var[t, j, i] * volume[lev, j, i] * - tmask[lev, j, i] * masks[bas, j, i]) - - -@cuda.reduce -def _sum_red_cuda(x, y): - """Numba kernel executed on the gpu that performs a sum reduction: - - Parameters - ---------- - x : float32 - Numpy array. - y : float32 - Numpy array. - - Returns - ------- - x+y : float64 - Reduction operation. - """ - return x+y -- GitLab From ab5aa728ecc9518870c078e9fe83a6baadcfb75f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 11:18:19 +0200 Subject: [PATCH 53/54] Change input files, old ones got deleted --- diagonals/examples/examples-regionmean.py | 2 +- diagonals/examples/examples-regsum.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/diagonals/examples/examples-regionmean.py b/diagonals/examples/examples-regionmean.py index 3983906..4f18cb1 100644 --- a/diagonals/examples/examples-regionmean.py +++ b/diagonals/examples/examples-regionmean.py @@ -17,7 +17,7 @@ TOS_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium"\ 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_"\ - "199001-199012.nc" + "196101-196112.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" logger = logging.getLogger(__name__) diff --git a/diagonals/examples/examples-regsum.py b/diagonals/examples/examples-regsum.py index 646a2e5..6a49975 100644 --- a/diagonals/examples/examples-regsum.py +++ b/diagonals/examples/examples-regsum.py @@ -11,13 +11,13 @@ import diagonals.regsum as regsum from diagonals.mesh_helpers.nemo import Nemo MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" -TOS_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ - "/EC-Earth3-LR/historical/r1i1p1f1/Omon/tos/gn/v20180501"\ - "/tos_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_195001-195012.nc" -THETAO_FILE = "/esarchive/exp/ecearth/a13c/cmorfiles/CMIP/EC-Earth-Consortium"\ - "/EC-Earth3-LR/historical/r1i1p1f1/Omon/thetao/gn/v20180501/"\ +TOS_FILE = "/esarchive/exp/ecearth/a16l/cmorfiles/CMIP/EC-Earth-Consortium"\ + "/EC-Earth3-LR/historical/r1i1p1f1/Omon/tos/gn/v20180711"\ + "/tos_Omon_EC-Earth3-LR_historical_r1i1p1f1_gn_196101-196112.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_"\ - "195001-195012.nc" + "196101-196112.nc" REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" logger = logging.getLogger(__name__) -- GitLab From 5ed4ca870ee399de9176452841f2509e7ef946e3 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 16 May 2019 11:37:23 +0200 Subject: [PATCH 54/54] Fix style issues --- diagonals/regmean.py | 6 +++--- diagonals/regsum.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/diagonals/regmean.py b/diagonals/regmean.py index c97d31c..5e6eb9a 100644 --- a/diagonals/regmean.py +++ b/diagonals/regmean.py @@ -41,7 +41,7 @@ def compute_regmean_2D(var, basins, area): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - logger.warning('GPU routines not implemented for regmean diagnostic. + logger.warning('GPU routines not implemented for regmean diagnostic.' 'Using CPU instead') regmean = _compute_regmean_2D_cpu(var, basins, area) return regmean @@ -67,7 +67,7 @@ def compute_regmean_3D(var, basins, volume): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - logger.warning('GPU routines not implemented for regmean diagnostic. + logger.warning('GPU routines not implemented for regmean diagnostic.' 'Using CPU instead') regmean = _compute_regmean_3D_cpu(var, basins, volume) return regmean @@ -93,7 +93,7 @@ def compute_regmean_levels(var, basins, volume): List containing regional mean for variable var at each level. """ if diagonals.CONFIG.use_gpu: - logger.warning('GPU routines not implemented for regmean diagnostic. + logger.warning('GPU routines not implemented for regmean diagnostic.' 'Using CPU instead') else: regmean = _compute_regmean_levels_cpu( diff --git a/diagonals/regsum.py b/diagonals/regsum.py index d37a9a3..4051c5f 100644 --- a/diagonals/regsum.py +++ b/diagonals/regsum.py @@ -40,7 +40,7 @@ def compute_regsum_2D(var, basins, area): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - logger.warning('GPU routines not implemented for regmean diagnostic. + logger.warning('GPU routines not implemented for regmean diagnostic.' 'Using CPU instead') regsum = _compute_regsum_2D_cpu(var, basins, area) return regsum @@ -66,7 +66,7 @@ def compute_regsum_3D(var, basins, volume, tmask): List containing regional mean for variable var """ if diagonals.CONFIG.use_gpu: - logger.warning('GPU routines not implemented for regmean diagnostic. + logger.warning('GPU routines not implemented for regmean diagnostic.' 'Using CPU instead') regsum = _compute_regsum_3D_cpu(var, basins, volume, tmask) return regsum -- GitLab