diff --git a/diagonals/examples/examples-zonalmean.py b/diagonals/examples/examples-zonalmean.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5b3a713fea964008d145a3a2cd44a918a9ffe0 --- /dev/null +++ b/diagonals/examples/examples-zonalmean.py @@ -0,0 +1,76 @@ +import logging +import datetime +import numpy as np + +import iris +import iris.cube +import iris.analysis +import iris.coords +import iris.coord_categorisation +import iris.analysis + +import diagonals +import diagonals.zonmean as zonmean +from diagonals.mesh_helpers.nemo import Nemo + +MESH_FILE = "/esarchive/autosubmit/con_files/mesh_mask_nemo.Ec3.2_O1L75.nc" +REGIONS_FILE = "/esarchive/autosubmit/con_files/mask.regions.Ec3.2_O1L75.nc" +THETAO_FILE = "/esarchive/exp/ecearth/a1tr/cmorfiles/CMIP/EC-Earth-Consortium"\ + "/EC-Earth3/historical/r24i1p1f1/Omon/thetao/gn/v20190312/"\ + "thetao_Omon_EC-Earth3_historical_r24i1p1f1_gn_"\ + "185001-185012.nc" + +logger = logging.getLogger(__name__) + + +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): + lats = mesh.get_mesh_var('nav_lat', dtype=np.float32) + areacello = mesh.get_areacello() + var = load_thetao() + basins = load_masks() + area_basin = zonmean.get_basin_area(areacello, basins) + zonalmean = zonmean.compute_zonmean(var, lats, area_basin) + if diagonals.CONFIG.use_gpu: + device = 'GPU' + else: + device = 'CPU' + save_data(zonalmean, basins, device) + ellapsed = datetime.datetime.now() - start + logger.info('Total ellapsed time on the %s: %s', device, ellapsed) + + +def load_thetao(): + thetao = iris.load_cube(THETAO_FILE, 'sea_water_potential_temperature') + thetao_data = thetao.data.astype(np.float32) + return thetao_data + + +def load_masks(): + basins = {} + cubes = iris.load(REGIONS_FILE) + for cube in cubes: + name = cube.name() + if name in ('nav_lat', 'nav_lon', 'Caspian_Sea'): + continue + cube = cube.extract(iris.Constraint(z=1)) + basins[name] = cube.data.astype(np.float32) + return basins + + +def save_data(zonalmean, basins, device): + cubes_zonmean = iris.cube.CubeList() + for basin in basins.keys(): + cube_zonmean = iris.cube.Cube(zonalmean[basin]) + cube_zonmean.add_aux_coord(iris.coords.AuxCoord(basin, 'region')) + cubes_zonmean.append(cube_zonmean) + iris.save(cubes_zonmean.merge_cube(), '/esarchive/scratch/sloosvel/' + 'numba_outputs/zonmean_{0}.nc'.format(device), zlib=True) + + +if __name__ == '__main__': + main() diff --git a/diagonals/siasie.py b/diagonals/siasie.py index 95f383228ecef4a19183d03665ca868e73225e8f..b4280d36f9fe6f8c311d07a89c72881cc4113ed9 100644 --- a/diagonals/siasie.py +++ b/diagonals/siasie.py @@ -137,10 +137,10 @@ def _compute_sic_cpu(gphit, area, sic_slices, basins, sithick): if sithick: temp = _voln_cpu(gphit, area, tmask, sic, sit) - voln[b][time] = np.sum(temp, axis=(1,2)) + voln[b][time] = np.sum(temp, axis=(1, 2)) temp = _vols_cpu(gphit, area, tmask, sic, sit) - vols[b][time] = np.sum(temp, axis=(1,2)) + vols[b][time] = np.sum(temp, axis=(1, 2)) b += 1 diff --git a/diagonals/zonmean.py b/diagonals/zonmean.py new file mode 100644 index 0000000000000000000000000000000000000000..ccffc24d212f4e0e4cfa9d9131383ae69fd9491a --- /dev/null +++ b/diagonals/zonmean.py @@ -0,0 +1,145 @@ +import os + +import numpy as np + +import iris +import iris.cube +import iris.analysis + +import warnings +import datetime + +import numba +from numba import vectorize +from numba import cuda +from numba import njit +from numba import int32, int64, float32, float64 +from numba.cuda.cudadrv import driver + +import diagonals + +__all__ = ['compute_zonmean', 'get_basin_area'] + + +def compute_zonmean(var, lats, area): + """Function that checks device and calls computing functions. + + Checks if the computations are going performed in the CPU or the GPU: + + Parameters + ---------- + var : float32 + Masked array containing variable data. + basins : float32 + List containing basin names and masks. + area : float32 + Masked array containing cell area. + + Returns + ------- + regmean: float32 + List containing regional mean for variable var + """ + if diagonals.CONFIG.use_gpu: + logger.warning('GPU routines not implemented yet' + 'for zonmean diagnostic due to a compiler bug.' + 'Using CPU instead until bug has been fixed') + + zonmean = _compute_zonal_mean_cpu(var, lats, area) + return zonmean + + +def _compute_zonal_mean_cpu(var, lats, area): + times = var.shape[0] + levs = var.shape[1] + value = {} + for basin, mask in area.items(): + weight = np.squeeze(area[basin]) + value[basin] = np.empty((times, levs, 180)) + for t in range(times): + for lev in range(levs): + value[basin][t][lev][:] = _zonal_mean_cpu(var[t, lev, :, :], + weight, + lats) + return value + + +def _compute_zonal_mean_gpu(var, latitudes, area): + times = var.shape[0] + levs = var.shape[1] + lats = var.shape[2] + lons = var.shape[3] + value = {} + + block = (128, 1, 1) + grid_size = (lons // block[0]) + 1 + grid = (grid_size, lats) + gpu_var = cuda.to_device(var.astype(np.float32)) + gpu_lats = cuda.to_device(latitudes.astype(np.float32)) + gpu_total = cuda.device_array(180, dtype=np.float32) + gpu_weights = cuda.device_array(180, dtype=np.float32) + for basin, mask in area.items(): + value[basin] = np.empty((times, levs, 180)) + gpu_area = cuda.to_device(np.squeeze(area[basin]).astype(np.float32)) + for t in range(times): + for lev in range(levs): + _zonal_mean_gpu[block, grid](gpu_var[t, lev, :, :], gpu_area, + gpu_lats, gpu_total, gpu_weights) + value[basin][t][lev][:] = (gpu_total.copy_to_host() / + gpu_weights.copy_to_host()) + del gpu_area, gpu_value, gpu_var + + return value + + +@numba.njit() +def _zonal_mean_cpu(variable, weight, latitude): + total = np.zeros(180, np.float32) + weights = np.zeros(180, np.float32) + for i in range(variable.shape[0]): + for j in range(variable.shape[1]): + if weight[i, j] == 0: + continue + bin_value = int(round(latitude[i, j]) + 90) + weights[bin_value] += weight[i, j] + total[bin_value] += variable[i, j] * weight[i, j] + return total / weights + + +@cuda.jit() +def _zonal_mean_gpu(variable, weight, latitude, total, weights): + i, j = cuda.grid(2) + if(i >= total.shape[0]): + return + if(weight[i, j] != 0.0): + + bin_value = int(round(latitude[i, j]) + 90) + weights[bin_value] += weight[i, j] + total[bin_value] += variable[i, j] * weight[i, j] + + +def get_basin_area(areacello, basins): + basin_areas = {} + for basin in basins: + basin_areas[basin] = _compute_basin_area(areacello, basins[basin]) + return basin_areas + + +@vectorize(['float32(float32, float32)'], target='cpu') +def _compute_basin_area(areacello, basin): + """Vectorized numba function executed on the cpu that computes the area + for each basin: + + Parameters + ---------- + areacello : float32 + Masked array containing areacello. + basin: float32 + Masked array containing the mask for a given basin. + + Returns + ------- + areacello*basin : float32 + Masked array containing the area for a given basin. + """ + return areacello*basin