diff --git a/bin/earthdiags b/bin/earthdiags index 200086ee283ea3c8b0eb15435d3a2d347dbc238d..99a4a3777ff981c412040f8a7092e53ccf111cae 100644 --- a/bin/earthdiags +++ b/bin/earthdiags @@ -5,10 +5,6 @@ import os import sys -scriptdir = os.path.abspath(os.path.dirname(sys.argv[0])) -assert sys.path[0] == scriptdir -sys.path[0] = os.path.normpath(os.path.join(scriptdir, os.pardir)) - # noinspection PyUnresolvedReferences,PyPep8 from earthdiagnostics.earthdiags import EarthDiags diff --git a/earthdiagnostics/ocean/zonalmean.py b/earthdiagnostics/ocean/zonalmean.py index 5e9177907d36fcf050f8a36e8693e8a5ac335a27..b7b9908e970ed95c25c71071a23b36dd292df3ba 100644 --- a/earthdiagnostics/ocean/zonalmean.py +++ b/earthdiagnostics/ocean/zonalmean.py @@ -11,13 +11,18 @@ from iris.cube import Cube, CubeList import numpy as np import numba +import netCDF4 + from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticIntOption, DiagnosticDomainOption, \ - DiagnosticBoolOption, DiagnosticBasinOption, DiagnosticVariableOption + DiagnosticBoolOption, DiagnosticBasinListOption, DiagnosticVariableOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile +import diagonals.zonmean as zonmean +from diagonals.mesh_helpers.nemo import Nemo + class ZonalMean(Diagnostic): """ @@ -40,14 +45,14 @@ class ZonalMean(Diagnostic): alias = 'zonmean' "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, basin, grid_point): + def __init__(self, data_manager, startdate, member, chunk, domain, variable, basins, grid_point): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.domain = domain self.variable = variable - self.basin = basin + self.basins = basins self.grid_point = grid_point self.declared = {} @@ -83,14 +88,14 @@ class ZonalMean(Diagnostic): DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), DiagnosticOption('grid_point', 'T'), - DiagnosticBasinOption('basin', Basins().Global), + DiagnosticBasinListOption('basins', Basins().Global), ) options = cls.process_options(options, options_available) job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job = ZonalMean(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], options['basin'], + options['domain'], options['variable'], options['basins'], options['grid_point'].lower()) job_list.append(job) @@ -106,76 +111,24 @@ class ZonalMean(Diagnostic): """Declare data to be generated by the diagnostic""" self.zonal_mean = self.declare_chunk( ModelingRealms.ocean, self.variable + 'zonal', - self.startdate, self.member, self.chunk, region=self.basin + self.startdate, self.member, self.chunk, region=self.basins ) def compute(self): """Run the diagnostic""" self._fix_file_metadata() + mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') + areacello = mesh.get_areacello(cell_point=self.grid_point) + lats = mesh.get_mesh_var('latitude', dtype=np.float32) data = self._load_data() - self._meand_3d_variable(data) - - def _meand_3d_variable(self, data): - e1 = self._try_load_cube(1) - e2 = self._try_load_cube(2) - mask = np.squeeze(Utils.get_mask(self.basin, True)) - mask = e1.data * e2.data * mask - if len(mask.shape) == 2: - data.add_aux_coord( - iris.coords.AuxCoord(mask.data, long_name='mask'), - data.coord_dims('latitude') - ) - else: - data.add_aux_coord( - iris.coords.AuxCoord(mask.data, long_name='mask'), - data.coord_dims('depth') + data.coord_dims('latitude') - ) - - @numba.njit() - def get_zonal_mean(variable, weight, latitude): - total = np.zeros(180, np.float64) - weights = np.zeros(180, np.float64) - for i in range(variable.shape[0]): - for j in range(variable.shape[1]): - if weight[i, j] == 0: - continue - bin_value = int(round(latitude[i, j]) + 90) - weights[bin_value] += weight[i, j] - total[bin_value] += variable[i, j] * weight[i, j] - return total / weights - - mean = iris.cube.CubeList() - lat_coord = None - for map_slice in data.slices_over('time'): - # Force data loading - map_slice.data - surface_cubes = iris.cube.CubeList() - for surface_slice in map_slice.slices_over('depth'): - value = get_zonal_mean( - surface_slice.data, - surface_slice.coord('mask').points, - surface_slice.coord('latitude').points, - ) - cube = Cube(value) - cube.add_aux_coord(surface_slice.coord('depth')) - if lat_coord is None: - lat_coord = surface_slice.coord('latitude') - lat_coord = lat_coord.copy( - np.arange(-90, 90, dtype=np.float32) - ) - lat_coord = iris.coords.DimCoord.from_coord(lat_coord) - cube.add_dim_coord(lat_coord, 0) - surface_cubes.append(cube) - time_cube = surface_cubes.merge_cube() - time_cube.add_aux_coord(map_slice.coord('time')) - mean.append(time_cube) - cube = mean.merge_cube() - cube.var_name = 'result' - cube.units = data.units - cube.attributes = data.attributes - temp = TempFile.get() - iris.save(cube, temp) - self.zonal_mean.set_local_file(temp, rename_var='result', region=self.basin) + masks = {} + self.basins.sort() + for basin in self.basins: + masks[basin] = Utils.get_mask(basin) + area_basin = zonmean.get_basin_area(areacello, masks) + zonalmean = zonmean.compute_zonmean(data.data, lats, area_basin) + self._save_result(zonalmean, data) + def _try_load_cube(self, number): try: @@ -231,3 +184,40 @@ class ZonalMean(Diagnostic): var.coordinates = coordinates handler.close() return has_levels + + + def _save_result(self, result, data): + final_name = '{0}zonal'.format(self.variable) + temp = TempFile.get() + handler_source = Utils.open_cdf(self.variable_file.local_file) + handler_temp = Utils.open_cdf(temp, 'w') + Utils.copy_variable(handler_source, handler_temp, 'time', True, True) + handler_temp.createDimension('region', len(result)) + handler_temp.createDimension('region_length', 50) + handler_temp.createDimension('lev', data.shape[1]) + handler_temp.createDimension('lat', 180) + + var_lat = handler_temp.createVariable('lat', float, 'lat') + var_lat[...] = np.arange(-90, 90, dtype=np.float32) + var_lat.units = 'degrees_north' + var_lat.long_name = 'latitude' + var_lat.standard_name = 'latitude' + + var_level = handler_temp.createVariable('lev', float, 'lev') + var_level[...] = data.coord('depth').points + var_level.units = 'm' + var_level.axis = 'Z' + var_level.positive = 'down' + var_level.long_name = 'ocean depth coordinate' + var_level.standard_name = 'depth' + var_region = handler_temp.createVariable('region', 'S1', + ('region', 'region_length')) + var = handler_temp.createVariable('{0}zonal'.format(self.variable), + float, ('time', 'lev', 'lat','region',),) + var.units = '{0}'.format(data.units) + for i, basin in enumerate(result): + var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) + var[..., i] = result[basin] + handler_temp.close() + self.declared[final_name].set_local_file(temp, diagnostic=self) +