From 8c6bc53882aeb5df43c9e1fa905f79e918b188da Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 17 Jan 2020 10:04:54 +0100 Subject: [PATCH 01/14] Do not save basin if all values are 0 --- earthdiagnostics/ocean/siasiesiv.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/earthdiagnostics/ocean/siasiesiv.py b/earthdiagnostics/ocean/siasiesiv.py index 24cf3c0e..dd13da4e 100644 --- a/earthdiagnostics/ocean/siasiesiv.py +++ b/earthdiagnostics/ocean/siasiesiv.py @@ -210,7 +210,8 @@ class Siasiesiv(Diagnostic): else: var_res.units = 'm^2' for i, basin in enumerate(self.masks): - var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) - var_res[..., i] = res[i, ...] + if not np.all(res[i, ...]==0): + var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) + var_res[..., i] = res[i, ...] handler_temp.close() self.generated[var].set_local_file(temp, diagnostic=self) -- GitLab From 9f1d1d60c1e2b3f74b2af74f9d6cf2fbed53dfd4 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 17 Jan 2020 12:54:57 +0100 Subject: [PATCH 02/14] Fix cmor file processed counter --- earthdiagnostics/cmorizer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/earthdiagnostics/cmorizer.py b/earthdiagnostics/cmorizer.py index db461520..0a11acf4 100644 --- a/earthdiagnostics/cmorizer.py +++ b/earthdiagnostics/cmorizer.py @@ -110,9 +110,11 @@ class Cmorizer(object): self._unpack_tar_file(tarfile) self._cmorize_nc_files() Log.result('Oceanic file {0}/{1} finished'.format(count, len(tar_files))) + count += 1 except Exception as ex: Log.error('Could not CMORize oceanic file {0}: {1}', count, ex) result = False + count += 1 return result def _filter_files(self, file_list): -- GitLab From 38c2d04a29662b144d213284d91fab09b0ffd26e Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Mon, 20 Jan 2020 15:29:16 +0100 Subject: [PATCH 03/14] Add hash to psi --- earthdiagnostics/ocean/psi.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index b9b32f02..c9c06e69 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -45,6 +45,9 @@ class Psi(Diagnostic): def __str__(self): return 'PSI Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ -- GitLab From a9be405a50c7098ae36969fe60c4a661f974208e Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Tue, 21 Jan 2020 10:59:17 +0100 Subject: [PATCH 04/14] Add hashes --- earthdiagnostics/ocean/areamoc.py | 8 +++++--- earthdiagnostics/ocean/averagesection.py | 3 +++ earthdiagnostics/ocean/rotation.py | 3 +++ earthdiagnostics/ocean/siarea.py | 3 +++ earthdiagnostics/ocean/sivol2d.py | 3 +++ earthdiagnostics/ocean/sivolume.py | 3 +++ 6 files changed, 20 insertions(+), 3 deletions(-) diff --git a/earthdiagnostics/ocean/areamoc.py b/earthdiagnostics/ocean/areamoc.py index 42a7c331..cbc1ae83 100644 --- a/earthdiagnostics/ocean/areamoc.py +++ b/earthdiagnostics/ocean/areamoc.py @@ -62,9 +62,11 @@ class AreaMoc(Diagnostic): self.basin == other.basin and self.box == other.box def __str__(self): - return 'Area MOC Startdate: {0} Member: {1} Chunk: {2} Box: {3} Basin: {4}'.format(self.startdate, self.member, - self.chunk, self.box, - self.basin) + return 'Area MOC Startdate: {0} Member: {1} Chunk: {2} Box: {3} Basin: {4}'.format( + self.startdate, self.member, self.chunk, self.box, self.basin) + + def __hash__(self): + return hash(str(self)) @classmethod def generate_jobs(cls, diags, options): diff --git a/earthdiagnostics/ocean/averagesection.py b/earthdiagnostics/ocean/averagesection.py index b152a45d..74f10b94 100644 --- a/earthdiagnostics/ocean/averagesection.py +++ b/earthdiagnostics/ocean/averagesection.py @@ -59,6 +59,9 @@ class AverageSection(Diagnostic): return 'Average section Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Box: {0.box} ' \ 'Variable: {0.domain}:{0.variable} Grid: {0.grid}'.format(self) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ diff --git a/earthdiagnostics/ocean/rotation.py b/earthdiagnostics/ocean/rotation.py index e1bd15b2..22fdd405 100644 --- a/earthdiagnostics/ocean/rotation.py +++ b/earthdiagnostics/ocean/rotation.py @@ -56,6 +56,9 @@ class Rotation(Diagnostic): '{3}:{5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variableu, self.variablev) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ diff --git a/earthdiagnostics/ocean/siarea.py b/earthdiagnostics/ocean/siarea.py index 6ee702b5..9c867f54 100644 --- a/earthdiagnostics/ocean/siarea.py +++ b/earthdiagnostics/ocean/siarea.py @@ -62,6 +62,9 @@ class Siarea(Diagnostic): return 'Siarea Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ 'Basins: {1}'.format(self, ','.join(str(basin) for basin in self.masks.keys())) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ diff --git a/earthdiagnostics/ocean/sivol2d.py b/earthdiagnostics/ocean/sivol2d.py index 937b4c90..fe0e9dc7 100644 --- a/earthdiagnostics/ocean/sivol2d.py +++ b/earthdiagnostics/ocean/sivol2d.py @@ -58,6 +58,9 @@ class Sivol2d(Diagnostic): def __str__(self): return 'Sivol2d Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk}'.format(self) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ diff --git a/earthdiagnostics/ocean/sivolume.py b/earthdiagnostics/ocean/sivolume.py index 1ec4a4db..75c71e23 100644 --- a/earthdiagnostics/ocean/sivolume.py +++ b/earthdiagnostics/ocean/sivolume.py @@ -57,6 +57,9 @@ class Sivolume(Diagnostic): return 'Sivolume Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ 'Basins: {1} '.format(self, ','.join(str(basin) for basin in self.masks.keys())) + def __hash__(self): + return hash(str(self)) + @classmethod def generate_jobs(cls, diags, options): """ -- GitLab From d9857f6afc0b02d2494ee5a8c61a634518ac8a24 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 23 Jan 2020 16:29:24 +0100 Subject: [PATCH 05/14] Adapt psi to python --- earthdiagnostics/ocean/psi.py | 129 ++++++++++++++++++++++++++++++---- 1 file changed, 114 insertions(+), 15 deletions(-) diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index c9c06e69..d415d831 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -1,10 +1,27 @@ # coding=utf-8 """Compute the barotropic stream function""" -from earthdiagnostics import cdftools -from earthdiagnostics.diagnostic import Diagnostic +import os +import datetime + +import numpy as np + +import netCDF4 + +import iris +import iris.analysis +import iris.coords +import iris.util +from bscearth.utils.log import Log + +from earthdiagnostics.constants import Basins +from earthdiagnostics.diagnostic import Diagnostic, \ + DiagnosticBasinListOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile +import diagonals.psi as psi +from diagonals.mesh_helpers.nemo import Nemo + class Psi(Diagnostic): """ @@ -31,11 +48,12 @@ class Psi(Diagnostic): vsftbarot = 'vsftbarot' - def __init__(self, data_manager, startdate, member, chunk): + def __init__(self, data_manager, startdate, member, chunk, masks): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk + self.masks = masks def __eq__(self, other): if self._different_type(other): @@ -43,7 +61,11 @@ class Psi(Diagnostic): return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk def __str__(self): - return 'PSI Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk) + return 'PSI Startdate: {0} Member: {1} Chunk: {2} Basins: {3}'.format( + self.startdate, self.member, self.chunk, ','.join(str(basin) for basin in self.masks.keys())) + + def __hash__(self): + return hash(str(self)) def __hash__(self): return hash(str(self)) @@ -59,27 +81,104 @@ class Psi(Diagnostic): :type options: list[str] :return: """ - if len(options) > 1: - raise Exception('The PSI diagnostic has no options') + options_available = (DiagnosticBasinListOption('basins', + [Basins().Global]),) + options = cls.process_options(options, options_available) + + basins = options['basins'] + if not basins: + Log.error('Basins not recognized') + return () + + masks = {} + basins.sort() + for basin in basins: + masks[basin] = Utils.get_mask(basin) + job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Psi(diags.data_manager, startdate, member, chunk)) + job_list.append(Psi(diags.data_manager, startdate, member, chunk, + masks)) return job_list def request_data(self): """Request data required by the diagnostic""" - self.uo = self.request_chunk(ModelingRealms.ocean, 'uo', self.startdate, self.member, self.chunk) - self.vo = self.request_chunk(ModelingRealms.ocean, 'vo', self.startdate, self.member, self.chunk) + self.uo = self.request_chunk(ModelingRealms.ocean, 'uo', + self.startdate, self.member, self.chunk) + self.vo = self.request_chunk(ModelingRealms.ocean, 'vo', + self.startdate, self.member, self.chunk) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.psi = self.declare_chunk(ModelingRealms.ocean, Psi.vsftbarot, self.startdate, self.member, self.chunk) + self.psi = self.declare_chunk(ModelingRealms.ocean, Psi.vsftbarot, + self.startdate, self.member, self.chunk) def compute(self): """Run the diagnostic""" + self._fix_coordinates_attribute(self.uo.local_file, 'uo') + self._fix_coordinates_attribute(self.vo.local_file, 'vo') + uo_cube = iris.load_cube(self.uo.local_file) + vo_cube = iris.load_cube(self.vo.local_file) + + uo = np.ma.filled(uo_cube.data, 0.0).astype(np.float32) + vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) + + mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') + e2u = mesh.get_j_length(cell_point='U') + e3u = mesh.get_k_length(cell_point='U') + e1v = mesh.get_i_length(cell_point='V') + e3v = mesh.get_k_length(cell_point='V') + + results = psi.compute(self.masks, e2u, e1v, e3u, e3v, uo, vo) + + self.save(results) + + def save(self, result): temp = TempFile.get() - cdftools.run('cdfpsi', input_file=[self.uo.local_file, self.vo.local_file], output_file=temp, - options='-mean -mask') - Utils.rename_variable(temp, 'sobarstf', Psi.vsftbarot) - Utils.setminmax(temp, Psi.vsftbarot) - self.psi.set_local_file(temp) + handler_source = Utils.open_cdf(self.uo.local_file) + handler_temp = Utils.open_cdf(temp, 'w') + lat_name = next(alias for alias in ('lat', 'latitude') + if alias in handler_source.variables.keys()) + lon_name = next(alias for alias in ('lon', 'longitude') + if alias in handler_source.variables.keys()) + Utils.copy_variable(handler_source, handler_temp, 'time', True, True) + Utils.copy_variable(handler_source, handler_temp, 'i', False, True) + Utils.copy_variable(handler_source, handler_temp, 'j', False, True) + + Utils.copy_variable(handler_source, handler_temp, lat_name, True, True) + Utils.copy_variable(handler_source, handler_temp, lon_name, True, True) + handler_temp.createDimension('region', len(result)) + handler_temp.createDimension('region_length', 50) + var_region = handler_temp.createVariable('region', 'S1', + ('region', 'region_length')) + var = handler_temp.createVariable('vsftbarot', float, + ('time', 'j', 'i', 'region')) + var.units = 'm3/s' + var.coordinates = ' '.join((lat_name, lon_name)) + var.missing_value = 1e20 + var.fill_value = 1e20 + var.valid_min = -300e6 + var.valid_max = 300e6 + var.long_name = 'Barotropic_Stream_Function' + + for i, basin in enumerate(result): + var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) + result[basin].mask = self.masks[basin]<1 + var[..., i] = result[basin] + handler_temp.close() + self.psi.set_local_file(temp, diagnostic=self) + + def _fix_coordinates_attribute(self, filepath, var_name): + add_coordinates = { + 'time_centered', 'leadtime' + } + handler = Utils.open_cdf(filepath) + coordinates = handler.variables[var_name].coordinates.split() + handler.variables[var_name].coordinates = \ + ' '.join(set(coordinates) | add_coordinates) + handler.close() + + + + + -- GitLab From 777c5d2c7d067539f62c0528ce252e5a69492fa0 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 24 Jan 2020 12:34:06 +0100 Subject: [PATCH 06/14] Improved cmorization --- earthdiagnostics/cmorizer.py | 38 ++++++++++++++++++++++++++------- earthdiagnostics/cmormanager.py | 5 ++++- earthdiagnostics/config.py | 2 +- earthdiagnostics/utils.py | 5 ++++- 4 files changed, 39 insertions(+), 11 deletions(-) diff --git a/earthdiagnostics/cmorizer.py b/earthdiagnostics/cmorizer.py index 0a11acf4..7fd1016e 100644 --- a/earthdiagnostics/cmorizer.py +++ b/earthdiagnostics/cmorizer.py @@ -13,6 +13,7 @@ from bscearth.utils.date import parse_date, chunk_end_date, previous_day, date2s from bscearth.utils.log import Log from cdo import CDOException import iris +from iris.coords import DimCoord import iris.coord_categorisation import iris.analysis import iris.util @@ -66,9 +67,11 @@ class Cmorizer(object): self.lon_name = self.config.data_convention.lon_name self.lat_name = self.config.data_convention.lat_name - self.alt_coord_names = {'time_counter': 'time', 'time_counter_bnds': 'time_bnds', - 'time_counter_bounds': 'time_bnds', - 'tbnds': 'bnds', 'nav_lat': self.lat_name, 'nav_lon': self.lon_name, 'x': 'i', 'y': 'j'} + self.alt_coord_names = { + 'time_counter': 'time', 'time_counter_bnds': 'time_bnds', + 'time_counter_bounds': 'time_bnds', + 'tbnds': 'bnds', 'nav_lat': self.lat_name, 'nav_lon': self.lon_name, + 'x': 'i', 'y': 'j'} @property def path_icm(self): @@ -121,10 +124,9 @@ class Cmorizer(object): if not self.cmor.filter_files: return file_list filtered = list() - filters = self.cmor.filter_files.split(' ') for file_path in file_list: filename = os.path.basename(file_path) - if any(f in filename for f in filters): + if any(f in filename for f in self.cmor.filter_files): filtered.append(file_path) else: self._remove(file_path) @@ -168,7 +170,7 @@ class Cmorizer(object): def _unpack_tar_file(self, tarfile): self._clean_cmor_scratch() os.makedirs(self.cmor_scratch) - Utils.untar((tarfile,), self.cmor_scratch) + Utils.untar((tarfile,), self.cmor_scratch, self.config.cmor.filter_files) if os.path.isdir(os.path.join(self.cmor_scratch, 'backup')): for filepath in glob.glob(os.path.join(self.cmor_scratch, 'backup', '*')): Log.debug('Moving file {0}', filepath) @@ -516,12 +518,32 @@ class Cmorizer(object): cube = iris.load_cube(file_path, iris.Constraint(cube_func=lambda c: c.var_name == variable)) for lev_original, lev_target in six.iteritems(lev_dimensions): try: - cube.coord(lev_original).var_name = lev_target + cube.coord(var_name=lev_original).var_name = lev_target except iris.exceptions.CoordinateNotFoundError: pass + try: + cube.coord('j'), + except iris.exceptions.CoordinateNotFoundError: + cube.add_dim_coord( + DimCoord( + range(cube.shape[-2]), + var_name='j', + long_name='Cell index along second dimension', + units=1.0, + ), len(cube.shape) - 2 + ) + try: + cube.coord('i'), + except iris.exceptions.CoordinateNotFoundError: + cube.add_dim_coord( + DimCoord( + range(cube.shape[-1]), + var_name='i', + long_name='Cell index along first dimension', + units=1.0, + ), len(cube.shape) - 1) iris.save(cube, temp, zlib=True) - Utils.rename_variables(temp, {'dim1': 'j', 'dim2': 'i'}, must_exist=False, rename_dimension=True) def _set_coordinates_attribute(self, file_path, var_cmor, variable): handler = Utils.open_cdf(file_path) diff --git a/earthdiagnostics/cmormanager.py b/earthdiagnostics/cmormanager.py index 91ffb067..1ae7570a 100644 --- a/earthdiagnostics/cmormanager.py +++ b/earthdiagnostics/cmormanager.py @@ -359,7 +359,10 @@ class CMORManager(DataManager): filepaths = self._get_transferred_cmor_data_filepaths(startdate, member, chunk, 'tar') if len(filepaths) > 0: Log.info('Unpacking cmorized data for {0} {1} {2}...', startdate, member, chunk) - Utils.untar(filepaths, os.path.join(self.cmor_path, 'cmorfiles')) + Utils.untar( + filepaths, + os.path.join(self.cmor_path, 'cmorfiles') + ) self._correct_paths(startdate) self.convention.create_links(startdate, member) return True diff --git a/earthdiagnostics/config.py b/earthdiagnostics/config.py index bd1e45af..cc628298 100644 --- a/earthdiagnostics/config.py +++ b/earthdiagnostics/config.py @@ -229,7 +229,7 @@ class CMORConfig(object): self.force = parser.get_bool_option('CMOR', 'FORCE', False) self.force_untar = parser.get_bool_option('CMOR', 'FORCE_UNTAR', False) self.skip_prepare = parser.get_bool_option('CMOR', 'SKIP_PREPARE', False) - self.filter_files = parser.get_option('CMOR', 'FILTER_FILES', '') + self.filter_files = parser.get_list_option('CMOR', 'FILTER_FILES', []) self.ocean = parser.get_bool_option('CMOR', 'OCEAN_FILES', True) self.atmosphere = parser.get_bool_option('CMOR', 'ATMOSPHERE_FILES', True) self.use_grib = parser.get_bool_option('CMOR', 'USE_GRIB', True) diff --git a/earthdiagnostics/utils.py b/earthdiagnostics/utils.py index cc81b2bd..990b14d0 100644 --- a/earthdiagnostics/utils.py +++ b/earthdiagnostics/utils.py @@ -837,7 +837,7 @@ class Utils(object): var_handler.units = new_units @staticmethod - def untar(files, destiny_path): + def untar(files, destiny_path, filter_files=None): """ Untar files to a given destiny @@ -855,6 +855,9 @@ class Utils(object): if os.path.isdir(os.path.join(destiny_path, file_compressed.name)): continue else: + if filter_files: + if not any(f in file_compressed.name for f in filter_files): + continue if os.path.exists(os.path.join(destiny_path, file_compressed.name)): os.remove(os.path.join( destiny_path, file_compressed.name)) -- GitLab From 0674c892dabf9d0489241f7670bda8d76f7c5165 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 24 Jan 2020 15:15:08 +0100 Subject: [PATCH 07/14] Fix cmorizer for atmos and tests --- earthdiagnostics/cmorizer.py | 43 +++++++++++----------- earthdiagnostics/datafile.py | 7 ++-- earthdiagnostics/ocean/heatcontentlayer.py | 10 ++--- earthdiagnostics/utils.py | 5 +++ test/integration/test_cmorizer.py | 8 ++-- test/unit/ocean/test_heatcontentlayer.py | 2 +- test/unit/ocean/test_region_mean.py | 16 ++++---- test/unit/test_config.py | 2 +- 8 files changed, 50 insertions(+), 43 deletions(-) diff --git a/earthdiagnostics/cmorizer.py b/earthdiagnostics/cmorizer.py index 7fd1016e..6905c34a 100644 --- a/earthdiagnostics/cmorizer.py +++ b/earthdiagnostics/cmorizer.py @@ -521,27 +521,28 @@ class Cmorizer(object): cube.coord(var_name=lev_original).var_name = lev_target except iris.exceptions.CoordinateNotFoundError: pass - try: - cube.coord('j'), - except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord( - DimCoord( - range(cube.shape[-2]), - var_name='j', - long_name='Cell index along second dimension', - units=1.0, - ), len(cube.shape) - 2 - ) - try: - cube.coord('i'), - except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord( - DimCoord( - range(cube.shape[-1]), - var_name='i', - long_name='Cell index along first dimension', - units=1.0, - ), len(cube.shape) - 1) + if cube.coord('latitude').ndim > 1: + try: + cube.coord('j'), + except iris.exceptions.CoordinateNotFoundError: + cube.add_dim_coord( + DimCoord( + range(cube.shape[-2]), + var_name='j', + long_name='Cell index along second dimension', + units=1.0, + ), len(cube.shape) - 2 + ) + try: + cube.coord('i'), + except iris.exceptions.CoordinateNotFoundError: + cube.add_dim_coord( + DimCoord( + range(cube.shape[-1]), + var_name='i', + long_name='Cell index along first dimension', + units=1.0, + ), len(cube.shape) - 1) iris.save(cube, temp, zlib=True) diff --git a/earthdiagnostics/datafile.py b/earthdiagnostics/datafile.py index 4d02204c..98052c69 100644 --- a/earthdiagnostics/datafile.py +++ b/earthdiagnostics/datafile.py @@ -662,9 +662,10 @@ class NetCDFFile(DataFile): if self.data_convention == 'meteofrance': Log.debug('Converting variable names from meteofrance convention') - alt_coord_names = {'time_counter': 'time', 'time_counter_bounds': 'time_bnds', - 'tbnds': 'bnds', 'nav_lat': 'lat', 'nav_lon': 'lon', 'x': 'i', - 'y': 'j'} + alt_coord_names = { + 'time_counter': 'time', 'time_counter_bounds': 'time_bnds', + 'tbnds': 'bnds', 'nav_lat': 'lat', 'nav_lon': 'lon', + 'x': 'i', 'y': 'j'} Utils.rename_variables(self.local_file, alt_coord_names, must_exist=False) Log.info('File {0} ready!', path) self.local_status = LocalStatus.READY diff --git a/earthdiagnostics/ocean/heatcontentlayer.py b/earthdiagnostics/ocean/heatcontentlayer.py index 1f64bcfc..f6b51ea2 100644 --- a/earthdiagnostics/ocean/heatcontentlayer.py +++ b/earthdiagnostics/ocean/heatcontentlayer.py @@ -109,7 +109,7 @@ class HeatContentLayer(Diagnostic): depth = mesh.get_depth(cell_point='W') weight = ohc.get_weights(layers, mask, e3t, depth) max_level, min_level = cls._get_used_levels(weight) - weight[0] = weight[0][:, min_level:max_level,:, :] + weight[0] = weight[0][:, min_level:max_level, :, :] del mask, depth, e3t @@ -184,9 +184,10 @@ class HeatContentLayer(Diagnostic): handler = Utils.open_cdf(thetao_file) Utils.convert_units(handler.variables['thetao'], 'K') - heatc_sl, heatc_sl1D = ohc.compute(self.layers, self.weight, - handler.variables['thetao'][:,self.min_level:self.max_level,:,:], - self.areas) + heatc_sl, heatc_sl1D = ohc.compute( + self.layers, self.weight, + handler.variables['thetao'][:, self.min_level:self.max_level, :, :], + self.areas) handler.sync() handler.renameVariable('thetao', 'heatc_sl') @@ -229,4 +230,3 @@ class HeatContentLayer(Diagnostic): Utils.setminmax(results, 'heatc') self.heatc.set_local_file(results) self.heatcsum.set_local_file(results1D) - diff --git a/earthdiagnostics/utils.py b/earthdiagnostics/utils.py index 990b14d0..d9a69c48 100644 --- a/earthdiagnostics/utils.py +++ b/earthdiagnostics/utils.py @@ -849,6 +849,7 @@ class Utils(object): """ for filepath in files: Log.debug('Unpacking {0}', filepath) + extracted = False tar = tarfile.open(filepath) for file_compressed in tar.getmembers(): if file_compressed.isdir(): @@ -861,10 +862,14 @@ class Utils(object): if os.path.exists(os.path.join(destiny_path, file_compressed.name)): os.remove(os.path.join( destiny_path, file_compressed.name)) + extracted = True tar.extract(file_compressed, destiny_path) Log.debug('File {0} extracted', os.path.basename(file_compressed.name)) + tar.close() + if filter_files and not extracted: + ValueError('No files extracted. Check filters') @staticmethod def unzip(files, force=False): diff --git a/test/integration/test_cmorizer.py b/test/integration/test_cmorizer.py index 5b15145a..a9eb3130 100644 --- a/test/integration/test_cmorizer.py +++ b/test/integration/test_cmorizer.py @@ -76,7 +76,7 @@ class TestCmorizer(TestCase): self.data_manager.config.cmor.ocean = True self.data_manager.config.cmor.atmosphere = True self.data_manager.config.cmor.use_grib = True - self.data_manager.config.cmor.filter_files = '' + self.data_manager.config.cmor.filter_files = [] self.data_manager.config.cmor.associated_experiment = 'associated_experiment' self.data_manager.config.cmor.initialization_method = 'initialization_method' self.data_manager.config.cmor.initialization_description = 'initialization_description' @@ -280,14 +280,14 @@ class TestCmorizer(TestCase): def test_ocean_cmorization_with_filter(self): """Test ocean cmorization filtering files""" self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self.data_manager.config.cmor.filter_files = 'expid' + self.data_manager.config.cmor.filter_files = ['expid'] self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) def test_ocean_cmorization_with_bad_filter(self): """Test ocean cmorization fails if a bad filter is added""" self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self.data_manager.config.cmor.filter_files = 'badfilter' - self._test_ocean_cmor(warnings=True, check_vars={'var1': False, 'var2': False}) + self.data_manager.config.cmor.filter_files = ['badfilter'] + self._test_ocean_cmor(error==True, check_vars={'var1': False, 'var2': False}) def test_ocean_cmorization_gzip(self): """Test ocean cmorization if tars are also zipped""" diff --git a/test/unit/ocean/test_heatcontentlayer.py b/test/unit/ocean/test_heatcontentlayer.py index f4a7947a..c46d6575 100644 --- a/test/unit/ocean/test_heatcontentlayer.py +++ b/test/unit/ocean/test_heatcontentlayer.py @@ -21,5 +21,5 @@ class TestHeatContentLayer(TestCase): self.box.max_depth = 100 def test_str(self): - diag = HeatContentLayer(self.data_manager, '20000101', 1, 1, self.box, self.weight, 0, 10, None, Mock()) + diag = HeatContentLayer(self.data_manager, '20000101', 1, 1, self.box, self.weight, 0, 10, None, Mock(), 0, 0) self.assertEqual(str(diag), 'Heat content layer Startdate: 20000101 Member: 1 Chunk: 1 Box: 0-100m') diff --git a/test/unit/ocean/test_region_mean.py b/test/unit/ocean/test_region_mean.py index d9cf5d04..eba2e84b 100644 --- a/test/unit/ocean/test_region_mean.py +++ b/test/unit/ocean/test_region_mean.py @@ -39,7 +39,7 @@ class TestRegionMean(TestCase): box.max_depth = -1 jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', 'var']) + self.diags, ['diagnostic', 'ocean', ['var']]) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', box, True, False, Basins().Global, 't', Frequencies.monthly)) @@ -47,7 +47,7 @@ class TestRegionMean(TestCase): box, True, False, Basins().Global, 't', Frequencies.monthly)) jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', 'var', '', 'U']) + self.diags, ['diagnostic', 'ocean', ['var'], '', 'U']) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', box, True, False, Basins().Global, 'u', Frequencies.monthly)) @@ -55,7 +55,7 @@ class TestRegionMean(TestCase): box, True, False, Basins().Global, 'u', Frequencies.monthly)) jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', 'var', 'global', 'U']) + self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U']) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', box, True, False, Basins().Global, 'u', Frequencies.monthly)) @@ -67,7 +67,7 @@ class TestRegionMean(TestCase): box.max_depth = 10.0 jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', 'var', 'global', 'U', '1', '10']) + self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10']) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', box, True, False, Basins().Global, 'u', Frequencies.monthly)) @@ -76,7 +76,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', 'var', 'global', 'U', + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', '', '', '', '', 'false'] ) self.assertEqual(len(jobs), 2) @@ -87,7 +87,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', 'var', 'global', 'U', + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', '', '', '', '', 'false', 'True'] ) self.assertEqual(len(jobs), 2) @@ -98,7 +98,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', 'var', 'global', 'U', '1', + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', '', '', '', '', 'false', 'True', 'grid'] ) self.assertEqual(len(jobs), 2) @@ -113,7 +113,7 @@ class TestRegionMean(TestCase): with self.assertRaises(DiagnosticOptionError): RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', 'var', 'global', 'U', '1', '10', '', '', '', '', 'false', + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', '', '', '', '', 'false', 'True', 'grid', 'extra'] ) diff --git a/test/unit/test_config.py b/test/unit/test_config.py index b18d7ebb..b7ead20d 100644 --- a/test/unit/test_config.py +++ b/test/unit/test_config.py @@ -128,7 +128,7 @@ class TestCMORConfig(TestCase): self.assertEqual(config.version, '') self.assertEqual(config.physics_version, '1') self.assertEqual(config.physics_description, 'to be filled') - self.assertEqual(config.filter_files, '') + self.assertEqual(config.filter_files, []) self.assertEqual(config.default_atmos_grid, 'gr') self.assertEqual(config.default_ocean_grid, 'gn') self.assertEqual(config.min_cmorized_vars, 10) -- GitLab From 3711c757bf2fd3d8ca31a1cee714bda1d0d4dd74 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 24 Jan 2020 16:09:37 +0100 Subject: [PATCH 08/14] Regions are updated now on the file --- VERSION | 2 +- earthdiagnostics/datafile.py | 200 ++++++++--------------- earthdiagnostics/ocean/psi.py | 21 +-- earthdiagnostics/ocean/regionmean.py | 2 +- earthdiagnostics/ocean/siasiesiv.py | 2 +- src/mixdiags | 1 - test/integration/test_cmorizer.py | 2 +- test/unit/ocean/test_heatcontentlayer.py | 13 +- test/unit/ocean/test_psi.py | 25 ++- test/unit/ocean/test_region_mean.py | 22 ++- 10 files changed, 121 insertions(+), 169 deletions(-) delete mode 160000 src/mixdiags diff --git a/VERSION b/VERSION index 944880fa..15a27998 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -3.2.0 +3.3.0 diff --git a/earthdiagnostics/datafile.py b/earthdiagnostics/datafile.py index 98052c69..f9640240 100644 --- a/earthdiagnostics/datafile.py +++ b/earthdiagnostics/datafile.py @@ -2,16 +2,14 @@ """Module for classes to manage storage manipulation""" import csv import os -import shutil from datetime import datetime -import six import iris +import iris.cube import iris.coords import iris.exceptions -import numpy as np +import iris.experimental.equalise_cubes from bscearth.utils.log import Log -import netCDF4 from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -92,7 +90,8 @@ class DataFile(Publisher): if self.local_status != LocalStatus.READY or self.suscribers or self.upload_required() or \ self.storage_status == StorageStatus.UPLOADING: return - Log.debug('File {0} no longer needed. Deleting from scratch...'.format(self.remote_file)) + Log.debug('File {0} no longer needed. Deleting from scratch...'.format( + self.remote_file)) os.remove(self.local_file) Log.debug('File {0} deleted from scratch'.format(self.remote_file)) self.local_file = None @@ -252,17 +251,16 @@ class DataFile(Publisher): else: original_name = self.var if self.final_name != original_name: - Utils.rename_variable(self.local_file, original_name, self.final_name) + Utils.rename_variable( + self.local_file, original_name, self.final_name) self._rename_coordinate_variables() self._correct_metadata() self._prepare_region() self.add_diagnostic_history() Utils.convert2netcdf4(self.local_file) - if self.region is not None: - self.upload() def upload(self): - """Send a loal file to the storage""" + """Send a local file to the storage""" self.storage_status = StorageStatus.UPLOADING remote_file = self.remote_file try: @@ -278,7 +276,7 @@ class DataFile(Publisher): # self.create_link() self.storage_status = StorageStatus.READY - def set_local_file(self, local_file, diagnostic=None, rename_var='', region=None): + def set_local_file(self, local_file, diagnostic=None, rename_var=''): """ Set the local file generated by EarthDiagnostics @@ -289,19 +287,15 @@ class DataFile(Publisher): local_file: str diagnostic: Diagnostic or None rename_var: str - region: Basin or None Returns ------- None """ + self.storage_status = StorageStatus.PENDING if diagnostic in self._modifiers: self._modifiers.remove(diagnostic) - if region is not None: - self.region = region - else: - self.region = None self.local_file = local_file self.prepare_to_upload(rename_var) self.local_status = LocalStatus.READY @@ -313,9 +307,15 @@ class DataFile(Publisher): def _correct_metadata(self): handler = Utils.open_cdf(self.local_file) var_handler = handler.variables[self.final_name] - coords = set.intersection({'time', 'lev', self.lat_name, self.lon_name, 'leadtime', 'region', 'time_centered'}, - set(handler.variables.keys())) - var_handler.coordinates = Utils.convert_to_ascii_if_possible(' '.join(coords)) + coords = set.intersection( + { + 'time', 'lev', self.lat_name, self.lon_name, + 'leadtime', 'region', 'time_centered' + }, + set(handler.variables.keys()) + ) + var_handler.coordinates = Utils.convert_to_ascii_if_possible( + ' '.join(coords)) if 'time_centered' in handler.variables: if hasattr(handler.variables['time_centered'], 'standard_name'): del handler.variables['time_centered'].standard_name @@ -400,99 +400,46 @@ class DataFile(Publisher): var_handler[:] = var_handler[:] * factor + offset if 'valid_min' in var_handler.ncattrs(): - var_handler.valid_min = float(var_handler.valid_min) * factor + offset + var_handler.valid_min = float( + var_handler.valid_min) * factor + offset if 'valid_max' in var_handler.ncattrs(): - var_handler.valid_max = float(var_handler.valid_max) * factor + offset + var_handler.valid_max = float( + var_handler.valid_max) * factor + offset def _prepare_region(self): - if not self.region: + if not self.check_is_in_storage(update_status=False): return - self._add_region_dimension_to_var() - if os.path.exists(self.remote_file): - self._update_var_with_region_data() - self._correct_metadata() - Utils.nco().ncks(input=self.local_file, output=self.local_file, options=['--fix_rec_dmn region']) - handler = Utils.open_cdf(self.local_file) - regions = handler.variables['region'][...].tolist() - if len(regions) > 1: - ordered_regions = sorted(regions) - new_indexes = [regions.index(region) for region in ordered_regions] - - for var in handler.variables.values(): - if 'region' not in var.dimensions: - continue - index_region = var.dimensions.index('region') - var_values = var[...] - var_ordered = np.take(var_values, new_indexes, index_region) - var[...] = var_ordered - handler.close() - - def _update_var_with_region_data(self): - temp = TempFile.get() - shutil.copyfile(self.remote_file, temp) - handler = Utils.open_cdf(temp) - var_handler = handler.variables[self.final_name] - var_type = var_handler.dtype - handler.close() - self._fix_values_metadata(var_type, temp) - - Utils.nco().ncks(input=temp, output=temp, options=['--mk_rec_dmn region']) - cubes = iris.load(self.local_file) - for cube in cubes: - if self.final_name == cube.var_name: - value = cube - break - for index_region, region in enumerate(value.coord('region').points): - handler = Utils.open_cdf(temp) - region_slice = value.data[...] - original_regions = handler.variables['region'][...] - str_regions = [] - for x in range(original_regions.shape[0]): - str_regions.append(''.join( - [str(value) for value in original_regions[x, ...] - if value != '-'] - )) - Log.debug(str(str_regions)) - var = handler.variables[self.final_name] - if region in str_regions: - region_index = str_regions.index(region) - else: - region_index = original_regions.shape[0] - handler.variables['region'][region_index, ...] = netCDF4.stringtoarr(region, 50) - indices = list() - for dim in var.dimensions: - if dim == 'region': - indices.append(region_index) - else: - indices.append(slice(None)) - var[indices] = region_slice - handler.close() - - # handler.close() - Utils.move_file(temp, self.local_file) - - def _add_region_dimension_to_var(self): - handler = Utils.open_cdf(self.local_file) - if 'region' in handler.variables: - handler.close() + cube = iris.load_cube(self.local_file) + try: + cube.coord('region') + except iris.exceptions.CoordinateNotFoundError: return - handler.createDimension('region') - handler.createDimension('region_length', 50) - var_region = handler.createVariable('region', 'S1', ('region', 'region_length')) - var_region[0, ...] = netCDF4.stringtoarr(self.region.name, 50) - original_var = handler.variables[self.final_name] - new_var = handler.createVariable( - 'new_var', - original_var.datatype, - original_var.dimensions + ('region',), + try: + old_cube = iris.load_cube(self.remote_file) + except Exception: + old_cube = iris.load_cube( + self.remote_file.replace('/cmorfiles/', '/diags/') + ) + new_data = {} + for region_slice in cube.slices_over('region'): + Log.debug(region_slice.coord('region').points[0]) + new_data[region_slice.coord('region').points[0]] = region_slice + for region_slice in old_cube.slices_over('region'): + region = region_slice.coord('region').points[0] + Log.debug(region) + if region not in new_data: + new_data[region] = region_slice + cube_list = iris.cube.CubeList( + [new_data[region] for region in sorted(new_data)] ) - new_var.setncatts({k: original_var.getncattr(k) for k in original_var.ncattrs()}) - new_var.coordinates = new_var.coordinates + ' region' - value = original_var[:] - new_var[..., 0] = value - handler.close() - Utils.nco().ncks(input=self.local_file, output=self.local_file, options=('-x -v {0}'.format(self.final_name),)) - Utils.rename_variable(self.local_file, 'new_var', self.final_name) + if len(cube_list) == 1: + return + iris.experimental.equalise_cubes.equalise_attributes(cube_list) + final_cube = cube_list.merge_cube() + temp = TempFile.get() + iris.save(final_cube, temp, zlib=True) + Utils.move_file(temp, self.local_file) + self._correct_metadata() def _rename_coordinate_variables(self): variables = dict() @@ -518,7 +465,8 @@ class DataFile(Publisher): def add_cmorization_history(self): """Add the history line corresponding to the cmorization to the local file""" from earthdiagnostics.earthdiags import EarthDiags - history_line = 'CMORized with Earthdiagnostics version {0}'.format(EarthDiags.version) + history_line = 'CMORized with Earthdiagnostics version {0}'.format( + EarthDiags.version) self._add_history_line(history_line) def _add_history_line(self, history_line): @@ -564,7 +512,8 @@ class UnitConversion(object): for line in reader: if line[0] == 'original': continue - cls.add_conversion(UnitConversion(line[0], line[1], line[2], line[3])) + cls.add_conversion(UnitConversion( + line[0], line[1], line[2], line[3])) @classmethod def add_conversion(cls, conversion): @@ -574,7 +523,8 @@ class UnitConversion(object): :param conversion: conversion to add :type conversion: UnitConversion """ - cls._dict_conversions[(conversion.source, conversion.destiny)] = conversion + cls._dict_conversions[( + conversion.source, conversion.destiny)] = conversion @classmethod def get_conversion_factor_offset(cls, input_units, output_units): @@ -677,36 +627,17 @@ class NetCDFFile(DataFile): Log.error('File {0} not available: {1}', path, ex) self.local_status = LocalStatus.FAILED return - Log.error('File {0} not available: {1}', self.remote_file, 'FileNotFound') + Log.error('File {0} not available: {1}', + self.remote_file, 'FileNotFound') self.local_status = LocalStatus.FAILED - def check_is_in_storage(self): + def check_is_in_storage(self, update_status=True): for path in (self.remote_file, self.remote_file.replace('/cmorfiles/', '/diags/')): if os.path.isfile(path): - if self.region: - try: - cubes = iris.load(path) - self._check_regions(cubes) - except iris.exceptions.TranslationError as ex: - # If the check goes wrong, we must execute everything - os.remove(path) - except Exception as ex: - Log.debug('Exception when checking file {0}: {1}', path, ex) - else: + if update_status: self.storage_status = StorageStatus.READY - return - - def _check_regions(self, cubes): - for cube in cubes: - try: - if isinstance(self.region, six.string_types): - regions = {self.region.name} - else: - regions = {basin.name for basin in self.region} - if regions.issubset(set(cube.coord('region').points)): - self.storage_status = StorageStatus.READY - except iris.exceptions.CoordinateNotFoundError: - pass + return True + return False def create_link(self): """Create a link from the original data in the _ folder""" @@ -714,7 +645,8 @@ class NetCDFFile(DataFile): self.data_convention.create_link(self.domain, self.remote_file, self.frequency, self.final_name, self.grid, True, self.var_type) except (ValueError, Exception) as ex: - Log.error('Can not create link to {1}: {0}'.format(ex, self.remote_file)) + Log.error('Can not create link to {1}: {0}'.format( + ex, self.remote_file)) def _get_size(self): try: diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index d415d831..e13a20be 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -67,9 +67,6 @@ class Psi(Diagnostic): def __hash__(self): return hash(str(self)) - def __hash__(self): - return hash(str(self)) - @classmethod def generate_jobs(cls, diags, options): """ @@ -81,8 +78,7 @@ class Psi(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption('basins', - [Basins().Global]),) + options_available = (DiagnosticBasinListOption('basins', 'global'),) options = cls.process_options(options, options_available) basins = options['basins'] @@ -120,8 +116,8 @@ class Psi(Diagnostic): uo_cube = iris.load_cube(self.uo.local_file) vo_cube = iris.load_cube(self.vo.local_file) - uo = np.ma.filled(uo_cube.data, 0.0).astype(np.float32) - vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) + uo = np.ma.filled(uo_cube.data, 0.0).astype(np.float32) + vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') e2u = mesh.get_j_length(cell_point='U') @@ -151,8 +147,8 @@ class Psi(Diagnostic): handler_temp.createDimension('region_length', 50) var_region = handler_temp.createVariable('region', 'S1', ('region', 'region_length')) - var = handler_temp.createVariable('vsftbarot', float, - ('time', 'j', 'i', 'region')) + var = handler_temp.createVariable( + 'vsftbarot', float, ('time', 'j', 'i', 'region')) var.units = 'm3/s' var.coordinates = ' '.join((lat_name, lon_name)) var.missing_value = 1e20 @@ -163,7 +159,7 @@ class Psi(Diagnostic): for i, basin in enumerate(result): var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) - result[basin].mask = self.masks[basin]<1 + result[basin].mask = self.masks[basin] < 1 var[..., i] = result[basin] handler_temp.close() self.psi.set_local_file(temp, diagnostic=self) @@ -177,8 +173,3 @@ class Psi(Diagnostic): handler.variables[var_name].coordinates = \ ' '.join(set(coordinates) | add_coordinates) handler.close() - - - - - diff --git a/earthdiagnostics/ocean/regionmean.py b/earthdiagnostics/ocean/regionmean.py index 50da4315..b05c1e18 100644 --- a/earthdiagnostics/ocean/regionmean.py +++ b/earthdiagnostics/ocean/regionmean.py @@ -185,7 +185,7 @@ class RegionMean(Diagnostic): ) e3 = e3.extract(depth_constraint) data = data.extract(depth_constraint) - volcello = areacello*e3.data.astype(np.float32) + volcello = areacello * e3.data.astype(np.float32) mean = regmean.compute_regmean_3D(data.data, masks, volcello) self._save_result_2D('mean', mean, data) if self.save3d: diff --git a/earthdiagnostics/ocean/siasiesiv.py b/earthdiagnostics/ocean/siasiesiv.py index dd13da4e..0b81ff6c 100644 --- a/earthdiagnostics/ocean/siasiesiv.py +++ b/earthdiagnostics/ocean/siasiesiv.py @@ -210,7 +210,7 @@ class Siasiesiv(Diagnostic): else: var_res.units = 'm^2' for i, basin in enumerate(self.masks): - if not np.all(res[i, ...]==0): + if not np.all(res[i, ...] == 0): var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) var_res[..., i] = res[i, ...] handler_temp.close() diff --git a/src/mixdiags b/src/mixdiags deleted file mode 160000 index 19997970..00000000 --- a/src/mixdiags +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 199979700e38d3918a82bd2052855d46375e48ab diff --git a/test/integration/test_cmorizer.py b/test/integration/test_cmorizer.py index a9eb3130..b6b35e4a 100644 --- a/test/integration/test_cmorizer.py +++ b/test/integration/test_cmorizer.py @@ -287,7 +287,7 @@ class TestCmorizer(TestCase): """Test ocean cmorization fails if a bad filter is added""" self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') self.data_manager.config.cmor.filter_files = ['badfilter'] - self._test_ocean_cmor(error==True, check_vars={'var1': False, 'var2': False}) + self._test_ocean_cmor(success=False, error=True, check_vars={'var1': False, 'var2': False}) def test_ocean_cmorization_gzip(self): """Test ocean cmorization if tars are also zipped""" diff --git a/test/unit/ocean/test_heatcontentlayer.py b/test/unit/ocean/test_heatcontentlayer.py index c46d6575..6000a732 100644 --- a/test/unit/ocean/test_heatcontentlayer.py +++ b/test/unit/ocean/test_heatcontentlayer.py @@ -12,7 +12,9 @@ class TestHeatContentLayer(TestCase): self.diags = Mock() self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ('20010101', 0, 0), ('20010101', 0, 1) + ) self.weight = Mock() @@ -21,5 +23,10 @@ class TestHeatContentLayer(TestCase): self.box.max_depth = 100 def test_str(self): - diag = HeatContentLayer(self.data_manager, '20000101', 1, 1, self.box, self.weight, 0, 10, None, Mock(), 0, 0) - self.assertEqual(str(diag), 'Heat content layer Startdate: 20000101 Member: 1 Chunk: 1 Box: 0-100m') + diag = HeatContentLayer( + self.data_manager, '20000101', 1, 1, self.box, self.weight, 0, 10, None, Mock(), 0, 0 + ) + self.assertEqual( + str(diag), + 'Heat content layer Startdate: 20000101 Member: 1 Chunk: 1 Box: 0-100m' + ) diff --git a/test/unit/ocean/test_psi.py b/test/unit/ocean/test_psi.py index 1bdf597c..e80ffae4 100644 --- a/test/unit/ocean/test_psi.py +++ b/test/unit/ocean/test_psi.py @@ -1,7 +1,10 @@ # coding=utf-8 from unittest import TestCase +from mock import Mock, patch + +from earthdiagnostics.constants import Basins from earthdiagnostics.ocean.psi import Psi -from mock import Mock +from earthdiagnostics.utils import Utils class TestPsi(TestCase): @@ -10,16 +13,26 @@ class TestPsi(TestCase): self.data_manager = Mock() self.diags = Mock() self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.psi = Psi(self.data_manager, '20000101', 1, 1) + @staticmethod + def fake_get(basin): + return None + + @patch.object(Utils, 'get_mask', fake_get) def test_generate_jobs(self): jobs = Psi.generate_jobs(self.diags, ['diagnostic']) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Psi(self.data_manager, '20010101', 0, 0)) - self.assertEqual(jobs[1], Psi(self.data_manager, '20010101', 0, 1)) + self.assertEqual(jobs[0], Psi(self.data_manager, '20010101', 0, 0, {Basins().Global: None})) + self.assertEqual(jobs[1], Psi(self.data_manager, '20010101', 0, 1, {Basins().Global: None})) + + jobs = Psi.generate_jobs(self.diags, ['diagnostic', 'atl']) + self.assertEqual(len(jobs), 2) + self.assertEqual(jobs[0], Psi(self.data_manager, '20010101', 0, 0, {Basins().Atlantic: None})) + self.assertEqual(jobs[1], Psi(self.data_manager, '20010101', 0, 1, {Basins().Atlantic: None})) with self.assertRaises(Exception): - Psi.generate_jobs(self.diags, ['diagnostic', 'badoption']) + Psi.generate_jobs(self.diags, ['diagnostic', 'atl', 'badoption']) def test_str(self): - self.assertEqual(str(self.psi), 'PSI Startdate: 20000101 Member: 1 Chunk: 1') + psi = Psi(self.data_manager, '20000101', 1, 1, {Basins().Global: None}) + self.assertEqual(str(psi), 'PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global') diff --git a/test/unit/ocean/test_region_mean.py b/test/unit/ocean/test_region_mean.py index eba2e84b..fb554abf 100644 --- a/test/unit/ocean/test_region_mean.py +++ b/test/unit/ocean/test_region_mean.py @@ -33,7 +33,6 @@ class TestRegionMean(TestCase): @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) @patch.object(TempFile, 'get', fake_get) def test_generate_jobs(self): - box = Box() box.min_depth = -1 box.max_depth = -1 @@ -77,7 +76,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U', - '1', '10', '', '', '', '', 'false'] + '1', '10', 'false'] ) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', @@ -88,7 +87,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U', - '1', '10', '', '', '', '', 'false', 'True'] + '1', '10', 'false', 'True'] ) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', @@ -99,7 +98,7 @@ class TestRegionMean(TestCase): jobs = RegionMean.generate_jobs( self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', - '10', '', '', '', '', 'false', 'True', 'grid'] + '10', 'false', 'True', 'grid'] ) self.assertEqual(len(jobs), 2) self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', @@ -107,14 +106,25 @@ class TestRegionMean(TestCase): self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', box, False, Basins().Global, True, 'grid', Frequencies.monthly)) + jobs = RegionMean.generate_jobs( + self.diags, + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', + '10', 'false', 'True', 'grid', 'day'] + ) + self.assertEqual(len(jobs), 2) + self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', + box, False, Basins().Global, True, 'grid', Frequencies.daily)) + self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', + box, False, Basins().Global, True, 'grid', Frequencies.daily)) + with self.assertRaises(DiagnosticOptionError): RegionMean.generate_jobs(self.diags, ['diagnostic']) with self.assertRaises(DiagnosticOptionError): RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', '', '', '', '', 'false', - 'True', 'grid', 'extra'] + ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', 'false', + 'True', 'grid', 'day', 'extra'] ) def test_str(self): -- GitLab From 4de7da5792bd8f0b2303c3ce5645c7372415ff98 Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Mon, 27 Jan 2020 12:43:36 +0100 Subject: [PATCH 09/14] Update CI to clone master diagonals each time --- .gitlab-ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index e6fe6598..684f1213 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -24,6 +24,9 @@ test: - git submodule update --init --recursive - conda env update -f environment.yml -n earthdiagnostics3 python=3.7 - source activate earthdiagnostics3 + - rm -r ../diagonals + - git clone https://earth.bsc.es/gitlab/es/diagonals.git ../diagonals + - pip install ../diagonals - pip install .[develop] - python setup.py test -- GitLab From 3c11d3db4afd8b8ccb21442a8d025d6bcf36cc4e Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Mon, 27 Jan 2020 15:19:14 +0100 Subject: [PATCH 10/14] Fix bugs in datafile --- earthdiagnostics/datafile.py | 54 +++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 25 deletions(-) diff --git a/earthdiagnostics/datafile.py b/earthdiagnostics/datafile.py index f9640240..accc2c49 100644 --- a/earthdiagnostics/datafile.py +++ b/earthdiagnostics/datafile.py @@ -65,8 +65,6 @@ class DataFile(Publisher): self.job_added = False self._modifiers = [] self._size = None - self.lon_name = None - self.lat_name = None def __str__(self): return 'Data file for {0}'.format(self.remote_file) @@ -239,13 +237,6 @@ class DataFile(Publisher): This includes renaming the variable if necessary, updating the metadata and adding the history and managing the possibility of multiple regions """ - if self.data_convention in ('primavera', 'cmip6'): - self.lon_name = 'longitude' - self.lat_name = 'latitude' - else: - self.lon_name = 'lon' - self.lat_name = 'lat' - if rename_var: original_name = rename_var else: @@ -309,7 +300,8 @@ class DataFile(Publisher): var_handler = handler.variables[self.final_name] coords = set.intersection( { - 'time', 'lev', self.lat_name, self.lon_name, + 'time', 'lev', + self.data_convention.lat_name, self.data_convention.lon_name, 'leadtime', 'region', 'time_centered' }, set(handler.variables.keys()) @@ -365,16 +357,18 @@ class DataFile(Publisher): ) def _fix_coordinate_variables_metadata(self, handler): + lat_name = self.data_convention.lat_name + lon_name = self.data_convention.lon_name if 'lev' in handler.variables: handler.variables['lev'].short_name = 'lev' if self.domain == ModelingRealms.ocean: handler.variables['lev'].standard_name = 'depth' - if self.lon_name in handler.variables: - handler.variables[self.lon_name].short_name = self.lon_name - handler.variables[self.lon_name].standard_name = 'longitude' - if self.lat_name in handler.variables: - handler.variables[self.lat_name].short_name = self.lat_name - handler.variables[self.lat_name].standard_name = 'latitude' + if lon_name in handler.variables: + handler.variables[lon_name].short_name = lon_name + handler.variables[lon_name].standard_name = 'longitude' + if lat_name in handler.variables: + handler.variables[lat_name].short_name = lat_name + handler.variables[lat_name].standard_name = 'latitude' def _fix_units(self, var_handler): if 'units' not in var_handler.ncattrs(): @@ -417,9 +411,13 @@ class DataFile(Publisher): try: old_cube = iris.load_cube(self.remote_file) except Exception: - old_cube = iris.load_cube( - self.remote_file.replace('/cmorfiles/', '/diags/') - ) + try: + old_cube = iris.load_cube( + self.remote_file.replace('/cmorfiles/', '/diags/') + ) + except Exception: + # Bad data, overwrite + return new_data = {} for region_slice in cube.slices_over('region'): Log.debug(region_slice.coord('region').points[0]) @@ -438,6 +436,12 @@ class DataFile(Publisher): final_cube = cube_list.merge_cube() temp = TempFile.get() iris.save(final_cube, temp, zlib=True) + if '-' in final_cube.var_name: + Utils.rename_variable( + temp, + final_cube.var_name.replace('-', '_'), + final_cube.var_name, must_exist=False + ) Utils.move_file(temp, self.local_file) self._correct_metadata() @@ -445,12 +449,12 @@ class DataFile(Publisher): variables = dict() variables['x'] = 'i' variables['y'] = 'j' - variables['nav_lat_grid_V'] = self.lat_name - variables['nav_lon_grid_V'] = self.lon_name - variables['nav_lat_grid_U'] = self.lat_name - variables['nav_lon_grid_U'] = self.lon_name - variables['nav_lat_grid_T'] = self.lat_name - variables['nav_lon_grid_T'] = self.lon_name + variables['nav_lat_grid_V'] = self.data_convention.lat_name + variables['nav_lon_grid_V'] = self.data_convention.lon_name + variables['nav_lat_grid_U'] = self.data_convention.lat_name + variables['nav_lon_grid_U'] = self.data_convention.lon_name + variables['nav_lat_grid_T'] = self.data_convention.lat_name + variables['nav_lon_grid_T'] = self.data_convention.lon_name Utils.rename_variables(self.local_file, variables, False) def add_diagnostic_history(self): -- GitLab From f257268eb812255932a72af8923c3740529eb39a Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Wed, 29 Jan 2020 11:32:22 +0100 Subject: [PATCH 11/14] Flake8 adoption --- .pylintrc | 2 +- doc/source/conf.py | 212 ++--- earthdiagnostics/__init__.py | 2 +- earthdiagnostics/box.py | 69 +- earthdiagnostics/cdftools.py | 44 +- earthdiagnostics/cmorizer.py | 789 +++++++++++++----- earthdiagnostics/cmormanager.py | 449 +++++++--- earthdiagnostics/config.py | 420 +++++++--- earthdiagnostics/constants.py | 187 +++-- earthdiagnostics/data_convention.py | 597 +++++++++---- earthdiagnostics/datafile.py | 257 +++--- earthdiagnostics/datamanager.py | 144 +++- earthdiagnostics/diagnostic.py | 268 ++++-- earthdiagnostics/earthdiags.py | 351 +++++--- earthdiagnostics/frequency.py | 118 ++- earthdiagnostics/general/attribute.py | 82 +- earthdiagnostics/general/fix_file.py | 67 +- earthdiagnostics/general/module.py | 103 ++- earthdiagnostics/general/relink.py | 73 +- earthdiagnostics/general/relinkall.py | 6 +- earthdiagnostics/general/rewrite.py | 4 +- earthdiagnostics/general/scale.py | 106 ++- earthdiagnostics/general/select_levels.py | 120 ++- .../general/simplify_dimensions.py | 149 ++-- earthdiagnostics/general/timemean.py | 233 ++++-- .../general/verticalmeanmetersiris.py | 85 +- earthdiagnostics/modelingrealm.py | 91 +- earthdiagnostics/obsreconmanager.py | 132 ++- earthdiagnostics/ocean/areamoc.py | 167 ++-- earthdiagnostics/ocean/averagesection.py | 121 ++- earthdiagnostics/ocean/convectionsites.py | 125 ++- earthdiagnostics/ocean/cutsection.py | 153 +++- earthdiagnostics/ocean/gyres.py | 185 ++-- earthdiagnostics/ocean/heatcontent.py | 191 +++-- earthdiagnostics/ocean/heatcontentlayer.py | 227 +++-- earthdiagnostics/ocean/interpolate.py | 201 +++-- earthdiagnostics/ocean/interpolatecdo.py | 299 +++++-- earthdiagnostics/ocean/mask_land.py | 113 ++- earthdiagnostics/ocean/maxmoc.py | 278 ++++-- .../ocean/mixedlayerheatcontent.py | 76 +- .../ocean/mixedlayersaltcontent.py | 70 +- earthdiagnostics/ocean/moc.py | 107 +-- earthdiagnostics/ocean/mxl.py | 63 +- earthdiagnostics/ocean/psi.py | 118 +-- earthdiagnostics/ocean/regionmean.py | 296 ++++--- earthdiagnostics/ocean/regionsum.py | 397 +++++---- earthdiagnostics/ocean/rotation.py | 196 +++-- earthdiagnostics/ocean/siarea.py | 172 ++-- earthdiagnostics/ocean/siasiesiv.py | 202 +++-- earthdiagnostics/ocean/sivol2d.py | 104 ++- earthdiagnostics/ocean/sivolume.py | 153 ++-- earthdiagnostics/ocean/verticalgradient.py | 102 ++- earthdiagnostics/ocean/verticalmean.py | 104 ++- earthdiagnostics/ocean/verticalmeanmeters.py | 130 ++- earthdiagnostics/ocean/zonalmean.py | 211 +++-- earthdiagnostics/publisher.py | 4 +- .../statistics/climatologicalpercentile.py | 139 ++- .../statistics/daysoverpercentile.py | 327 +++++--- earthdiagnostics/statistics/discretize.py | 223 +++-- .../statistics/monthlypercentile.py | 191 +++-- earthdiagnostics/threddsmanager.py | 313 +++++-- earthdiagnostics/utils.py | 407 ++++++--- earthdiagnostics/variable.py | 254 +++--- earthdiagnostics/work_manager.py | 174 ++-- setup.cfg | 31 + setup.py | 120 +-- test/integration/test_cmorizer.py | 581 +++++++++---- .../data_convention/test_data_convention.py | 295 +++++-- test/unit/data_convention/test_meteofrance.py | 241 ++++-- test/unit/data_convention/test_preface.py | 441 +++++++--- test/unit/data_convention/test_primavera.py | 584 +++++++++---- test/unit/data_convention/test_specs.py | 565 +++++++++---- test/unit/general/test_attribute.py | 124 ++- test/unit/general/test_dailymean.py | 118 ++- test/unit/general/test_module.py | 133 ++- test/unit/general/test_monthlymean.py | 119 ++- test/unit/general/test_relink.py | 135 ++- test/unit/general/test_relinkall.py | 22 +- test/unit/general/test_rewrite.py | 107 ++- test/unit/general/test_scale.py | 350 ++++++-- test/unit/general/test_select_levels.py | 210 +++-- test/unit/general/test_simplify_dimensions.py | 115 ++- .../general/test_verticalmeanmetersiris.py | 163 +++- test/unit/general/test_yearlymean.py | 113 ++- test/unit/ocean/test_areamoc.py | 91 +- test/unit/ocean/test_averagesection.py | 114 ++- test/unit/ocean/test_convectionsites.py | 37 +- test/unit/ocean/test_cutsection.py | 106 ++- test/unit/ocean/test_gyres.py | 34 +- test/unit/ocean/test_heatcontent.py | 74 +- test/unit/ocean/test_heatcontentlayer.py | 22 +- test/unit/ocean/test_interpolate.py | 277 +++++- test/unit/ocean/test_interpolatecdo.py | 348 ++++++-- test/unit/ocean/test_maskland.py | 147 +++- test/unit/ocean/test_maxmoc.py | 101 ++- test/unit/ocean/test_mixedlayerheatcontent.py | 29 +- test/unit/ocean/test_mixedlayersaltcontent.py | 29 +- test/unit/ocean/test_moc.py | 16 +- test/unit/ocean/test_mxl.py | 22 +- test/unit/ocean/test_psi.py | 47 +- test/unit/ocean/test_region_mean.py | 429 ++++++++-- test/unit/ocean/test_rotation.py | 108 ++- test/unit/ocean/test_siasiesiv.py | 26 +- test/unit/ocean/test_vertical_gradient.py | 71 +- test/unit/ocean/test_verticalmean.py | 73 +- test/unit/ocean/test_verticalmeanmeters.py | 141 +++- .../test_climatologicalpercentile.py | 52 +- .../statistics/test_daysoverpercentile.py | 66 +- test/unit/statistics/test_discretize.py | 147 +++- .../unit/statistics/test_monthlypercentile.py | 61 +- test/unit/test_box.py | 32 +- test/unit/test_cdftools.py | 124 ++- test/unit/test_cmormanager.py | 273 ++++-- test/unit/test_config.py | 580 ++++++++----- test/unit/test_constants.py | 48 +- test/unit/test_data_manager.py | 29 +- test/unit/test_datafile.py | 57 +- test/unit/test_diagnostic.py | 315 +++---- test/unit/test_earthdiags.py | 144 ++-- test/unit/test_frequency.py | 46 +- test/unit/test_lint.py | 15 +- test/unit/test_modelling_realm.py | 89 +- test/unit/test_obsreconmanager.py | 143 ++-- test/unit/test_publisher.py | 3 +- test/unit/test_utils.py | 75 +- test/unit/test_variable.py | 340 +++++--- test/unit/test_variable_type.py | 8 +- test/unit/test_workmanager.py | 116 ++- 128 files changed, 15121 insertions(+), 6304 deletions(-) create mode 100644 setup.cfg diff --git a/.pylintrc b/.pylintrc index db7741b9..4094c1f4 100644 --- a/.pylintrc +++ b/.pylintrc @@ -99,7 +99,7 @@ evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / stateme [FORMAT] # Maximum number of characters on a single line. -max-line-length=120 +max-line-length=79 # Regexp for a line that is allowed to be longer than the limit. ignore-long-lines=^\s*(# )??$ diff --git a/doc/source/conf.py b/doc/source/conf.py index 780291ea..a7356586 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -19,8 +19,8 @@ import sys # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. # sys.path.insert(0, os.path.abspath('.')) -sys.path.insert(0, os.path.abspath('../..')) -print(os.path.abspath('../..')) +sys.path.insert(0, os.path.abspath("../..")) +print(os.path.abspath("../..")) # -- General configuration ------------------------------------------------ @@ -31,51 +31,51 @@ print(os.path.abspath('../..')) # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.doctest', - 'sphinx.ext.intersphinx', - 'sphinx.ext.todo', - 'sphinx.ext.coverage', - 'sphinx.ext.pngmath', - 'sphinx.ext.ifconfig', - 'sphinx.ext.viewcode', - 'sphinx.ext.napoleon', + "sphinx.ext.autodoc", + "sphinx.ext.doctest", + "sphinx.ext.intersphinx", + "sphinx.ext.todo", + "sphinx.ext.coverage", + "sphinx.ext.pngmath", + "sphinx.ext.ifconfig", + "sphinx.ext.viewcode", + "sphinx.ext.napoleon", ] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix of source filenames. -source_suffix = '.rst' +source_suffix = ".rst" # The encoding of source files. -#source_encoding = 'utf-8-sig' +# source_encoding = 'utf-8-sig' # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'Earth Diagnostics' -copyright = u'2019, BSC-CNS Earth Sciences Department' +project = u"Earth Diagnostics" +copyright = u"2019, BSC-CNS Earth Sciences Department" # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents.source ~/vi # # The short X.Y version. -version = '3.2' +version = "3.2" # The full version, including alpha/beta/rc tags. -release = '3.2.0' +release = "3.2.0" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. -#language = None +# language = None # There are two options for replacing |today|: either, you set today to some # non-false value, then it is used: -#today = '' +# today = '' # Else, today_fmt is used as the format for a strftime call. -#today_fmt = '%B %d, %Y' +# today_fmt = '%B %d, %Y' # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. @@ -83,154 +83,157 @@ exclude_patterns = [] # The reST default role (used for this markup: `text`) to use for all # documents. -#default_role = None +# default_role = None # If true, '()' will be appended to :func: etc. cross-reference text. -#add_function_parentheses = True +# add_function_parentheses = True # If true, the current module name will be prepended to all description # unit titles (such as .. function::). -#add_module_names = True +# add_module_names = True # If true, sectionauthor and moduleauthor directives will be shown in the # output. They are ignored by default. show_authors = True # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # A list of ignored prefixes for module index sorting. -#modindex_common_prefix = [] +# modindex_common_prefix = [] # If true, keep warnings as "system message" paragraphs in the built documents. -#keep_warnings = False +# keep_warnings = False # -- Options for HTML output ---------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. -html_theme = 'default' +html_theme = "default" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the # documentation. -#html_theme_options = {} +# html_theme_options = {} # Add any paths that contain custom themes here, relative to this directory. -#html_theme_path = [] +# html_theme_path = [] # The name for this set of Sphinx documents. If None, it defaults to # " v documentation". -#html_title = None +# html_title = None # A shorter title for the navigation bar. Default is the same as html_title. -#html_short_title = None +# html_short_title = None # The name of an image file (relative to this directory) to place at the top # of the sidebar. -#html_logo = None +# html_logo = None # The name of an image file (within the static path) to use as favicon of the # docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 # pixels large. -#html_favicon = None +# html_favicon = None # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ['_static'] +html_static_path = ["_static"] # Add any extra paths that contain custom files (such as robots.txt or # .htaccess) here, relative to this directory. These files are copied # directly to the root of the documentation. -#html_extra_path = [] +# html_extra_path = [] # If not '', a 'Last updated on:' timestamp is inserted at every page bottom, # using the given strftime format. -#html_last_updated_fmt = '%b %d, %Y' +# html_last_updated_fmt = '%b %d, %Y' # If true, SmartyPants will be used to convert quotes and dashes to # typographically correct entities. -#html_use_smartypants = True +# html_use_smartypants = True # Custom sidebar templates, maps document names to template names. -#html_sidebars = {} +# html_sidebars = {} # Additional templates that should be rendered to pages, maps page names to # template names. -#html_additional_pages = {} +# html_additional_pages = {} # If false, no module index is generated. -#html_domain_indices = True +# html_domain_indices = True # If false, no index is generated. -#html_use_index = True +# html_use_index = True # If true, the index is split into individual pages for each letter. -#html_split_index = False +# html_split_index = False # If true, links to the reST sources are added to the pages. -#html_show_sourcelink = True +# html_show_sourcelink = True # If true, "Created using Sphinx" is shown in the HTML footer. Default is True. -#html_show_sphinx = True +# html_show_sphinx = True # If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. -#html_show_copyright = True +# html_show_copyright = True # If true, an OpenSearch description file will be output, and all pages will # contain a tag referring to it. The value of this option must be the # base URL from which the finished HTML is served. -#html_use_opensearch = '' +# html_use_opensearch = '' # This is the file name suffix for HTML files (e.g. ".xhtml"). -#html_file_suffix = None +# html_file_suffix = None # Output file base name for HTML help builder. -htmlhelp_basename = 'EarthDiagnosticsd' +htmlhelp_basename = "EarthDiagnosticsd" # -- Options for LaTeX output --------------------------------------------- latex_elements = { -# The paper size ('letterpaper' or 'a4paper'). -#'papersize': 'letterpaper', - -# The font size ('10pt', '11pt' or '12pt'). -#'pointsize': '10pt', - -# Additional stuff for the LaTeX preamble. -#'preamble': '', + # The paper size ('letterpaper' or 'a4paper'). + #'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + #'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + #'preamble': '', } # Grouping the document tree into LaTeX files. List of tuples # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - ('index', 'EarthDiagnostics.tex', u'Earth Diagnostics Documentation', - u'BSC-CNS Earth Sciences Department', 'manual'), + ( + "index", + "EarthDiagnostics.tex", + u"Earth Diagnostics Documentation", + u"BSC-CNS Earth Sciences Department", + "manual", + ), ] # The name of an image file (relative to this directory) to place at the top of # the title page. -#latex_logo = None +# latex_logo = None # For "manual" documents, if this is true, then toplevel headings are parts, # not chapters. -#latex_use_parts = False +# latex_use_parts = False # If true, show page references after internal links. -#latex_show_pagerefs = False +# latex_show_pagerefs = False # If true, show URL addresses after external links. -#latex_show_urls = False +# latex_show_urls = False # Documents to append as an appendix to all manuals. -#latex_appendices = [] +# latex_appendices = [] # If false, no module index is generated. -#latex_domain_indices = True +# latex_domain_indices = True # -- Options for manual page output --------------------------------------- @@ -238,12 +241,17 @@ latex_documents = [ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). man_pages = [ - ('index', 'earthdiagnostics', u'Earth Diagnostics Documentation', - [u'BSC-CNS Earth Sciences Department'], 1) + ( + "index", + "earthdiagnostics", + u"Earth Diagnostics Documentation", + [u"BSC-CNS Earth Sciences Department"], + 1, + ) ] # If true, show URL addresses after external links. -#man_show_urls = False +# man_show_urls = False # -- Options for Texinfo output ------------------------------------------- @@ -252,93 +260,99 @@ man_pages = [ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - ('index', 'EarthDiagnostics', u'Earth Diagnostics Documentation', - u'BSC-CNS Earth Sciences Department', 'EarthDiagnostics', 'One line description of project.', - 'Miscellaneous'), + ( + "index", + "EarthDiagnostics", + u"Earth Diagnostics Documentation", + u"BSC-CNS Earth Sciences Department", + "EarthDiagnostics", + "One line description of project.", + "Miscellaneous", + ), ] # Documents to append as an appendix to all manuals. -#texinfo_appendices = [] +# texinfo_appendices = [] # If false, no module index is generated. -#texinfo_domain_indices = True +# texinfo_domain_indices = True # How to display URL addresses: 'footnote', 'no', or 'inline'. -#texinfo_show_urls = 'footnote' +# texinfo_show_urls = 'footnote' # If true, do not generate a @detailmenu in the "Top" node's menu. -#texinfo_no_detailmenu = False +# texinfo_no_detailmenu = False # -- Options for Epub output ---------------------------------------------- # Bibliographic Dublin Core info. -epub_title = u'Earth Diagnostics' -epub_author = u'BSC-CNS Earth Sciences Department' -epub_publisher = u'BSC-CNS Earth Sciences Department' -epub_copyright = u'2016, BSC-CNS Earth Sciences Department' +epub_title = u"Earth Diagnostics" +epub_author = u"BSC-CNS Earth Sciences Department" +epub_publisher = u"BSC-CNS Earth Sciences Department" +epub_copyright = u"2016, BSC-CNS Earth Sciences Department" # The basename for the epub file. It defaults to the project name. -#epub_basename = u'Earth Diagnostics' +# epub_basename = u'Earth Diagnostics' # The HTML theme for the epub output. Since the default themes are not optimized # for small screen space, using the same theme for HTML and epub output is # usually not wise. This defaults to 'epub', a theme designed to save visual # space. -#epub_theme = 'epub' +# epub_theme = 'epub' # The language of the text. It defaults to the language option # or en if the language is not set. -#epub_language = '' +# epub_language = '' # The scheme of the identifier. Typical schemes are ISBN or URL. -#epub_scheme = '' +# epub_scheme = '' # The unique identifier of the text. This can be a ISBN number # or the project homepage. -#epub_identifier = '' +# epub_identifier = '' # A unique identification for the text. -#epub_uid = '' +# epub_uid = '' # A tuple containing the cover image and cover page html template filenames. -#epub_cover = () +# epub_cover = () # A sequence of (type, uri, title) tuples for the guide element of content.opf. -#epub_guide = () +# epub_guide = () # HTML files that should be inserted before the pages created by sphinx. # The format is a list of tuples containing the path and title. -#epub_pre_files = [] +# epub_pre_files = [] # HTML files shat should be inserted after the pages created by sphinx. # The format is a list of tuples containing the path and title. -#epub_post_files = [] +# epub_post_files = [] # A list of files that should not be packed into the epub file. -epub_exclude_files = ['search.html'] +epub_exclude_files = ["search.html"] # The depth of the table of contents in toc.ncx. -#epub_tocdepth = 3 +# epub_tocdepth = 3 # Allow duplicate toc entries. -#epub_tocdup = True +# epub_tocdup = True # Choose between 'default' and 'includehidden'. -#epub_tocscope = 'default' +# epub_tocscope = 'default' # Fix unsupported image types using the PIL. -#epub_fix_images = False +# epub_fix_images = False # Scale large images. -#epub_max_image_width = 0 +# epub_max_image_width = 0 # How to display URL addresses: 'footnote', 'no', or 'inline'. -#epub_show_urls = 'inline' +# epub_show_urls = 'inline' # If false, no index is generated. -#epub_use_index = True +# epub_use_index = True # Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = {'http://docs.python.org/': None} +intersphinx_mapping = {"http://docs.python.org/": None} diff --git a/earthdiagnostics/__init__.py b/earthdiagnostics/__init__.py index b6f6a1bd..5fa2c44a 100644 --- a/earthdiagnostics/__init__.py +++ b/earthdiagnostics/__init__.py @@ -5,4 +5,4 @@ import os from earthdiagnostics.cdftools import CDFTools cdftools = CDFTools() -DEVNULL = open(os.devnull, 'wb') +DEVNULL = open(os.devnull, "wb") diff --git a/earthdiagnostics/box.py b/earthdiagnostics/box.py index 4d68bb10..73bccb41 100644 --- a/earthdiagnostics/box.py +++ b/earthdiagnostics/box.py @@ -6,8 +6,8 @@ class Box(object): """ Represents a box in the 3D space. - Also allows easy conversion from the coordinate values to significant string - representations + Also allows easy conversion from the coordinate values to significant + string representations Parameters ---------- @@ -18,7 +18,7 @@ class Box(object): def __init__(self, depth_in_meters=False): self.depth_in_meters = depth_in_meters """ - If True, treats the depth as if it is given in meters. If False, as it is given in levels + If True, treats the depth as if it is meters. If False, as it is levels :rtype: bool """ self._max_lat = None @@ -37,9 +37,15 @@ class Box(object): """ def __eq__(self, other): - return self.depth_in_meters == other.depth_in_meters and self.max_lat == other.max_lat and \ - self.min_lat == other.min_lat and self.max_lon == other.max_lon and self.min_lon == other.min_lon and \ - self.max_depth == other.max_depth and self.min_depth == other.min_depth + return ( + self.depth_in_meters == other.depth_in_meters + and self.max_lat == other.max_lat + and self.min_lat == other.min_lat + and self.max_lon == other.max_lon + and self.min_lon == other.min_lon + and self.max_depth == other.max_depth + and self.min_depth == other.min_depth + ) def __str__(self): return self.get_lat_str() + self.get_lon_str() + self.get_depth_str() @@ -56,7 +62,9 @@ class Box(object): @max_lat.setter def max_lat(self, value): if value > 90 or value < -90: - raise ValueError('{0} is not a valid latitude. Must be between -90 and -90'.format(value)) + raise ValueError( + f"{value} is not a valid latitude. Must be between -90 and -90" + ) self._max_lat = value @property @@ -71,7 +79,9 @@ class Box(object): @min_lat.setter def min_lat(self, value): if value > 90 or value < -90: - raise ValueError('{0} is not a valid latitude. Must be between -90 and 90'.format(value)) + raise ValueError( + f"{value} is not a valid latitude. Must be between -90 and 90" + ) self._min_lat = value @property @@ -86,7 +96,9 @@ class Box(object): @max_lon.setter def max_lon(self, value): if value >= 360 or value <= -360: - raise ValueError('{0} is not a valid longitude. Must be between -360 and 360'.format(value)) + raise ValueError( + f"{value} is not a valid longitude. Must be in [-360, 360]" + ) self._max_lon = value @property @@ -101,7 +113,9 @@ class Box(object): @min_lon.setter def min_lon(self, value): if value >= 360 or value <= -360: - raise ValueError('{0} is not a valid longitude. Must be between -360 and 360'.format(value)) + raise ValueError( + "{value} is not a valid longitude. Must be in [-360, 360]" + ) self._min_lon = value def get_lat_str(self): @@ -114,19 +128,19 @@ class Box(object): :rtype: str """ if self.max_lat is None or self.min_lat is None: - return '' + return "" if self.min_lat < 0: - direction = 'S' + direction = "S" else: - direction = 'N' + direction = "N" string = str(abs(self.min_lat)) + direction if self.max_lat != self.min_lat: if self.max_lat < 0: - direction = 'S' + direction = "S" else: - direction = 'N' + direction = "N" string += str(abs(self.max_lat)) + direction return string @@ -141,19 +155,19 @@ class Box(object): :rtype: str """ if self.max_lon is None or self.min_lon is None: - return '' + return "" if self.min_lon < 0: - direction = 'W' + direction = "W" else: - direction = 'E' + direction = "E" string = str(abs(self.min_lon)) + direction if self.max_lon != self.min_lon: if self.max_lon < 0: - direction = 'W' + direction = "W" else: - direction = 'E' + direction = "E" string += str(abs(self.max_lon)) + direction return string @@ -162,21 +176,24 @@ class Box(object): Get a string representation of depth. For depth expressed in meters, it adds the character 'm' to the end - If min_depth is different from max_depth, it concatenates the two values + If min_depth is different from max_depth, it concatenates the two + values :return: string representation for depth :rtype: str """ if self.max_depth is None or self.min_depth is None: - return '' + return "" if self.depth_in_meters: - suffix = 'm' + suffix = "m" else: - suffix = '' + suffix = "" if self.min_depth != self.max_depth: - string = '{0:d}-{1:d}{2}'.format(int(abs(self.min_depth)), int(abs(self.max_depth)), suffix) + string = "{0:d}-{1:d}{2}".format( + int(abs(self.min_depth)), int(abs(self.max_depth)), suffix + ) else: - string = '{0:d}{1}'.format(int(abs(self.max_depth)), suffix) + string = "{0:d}{1}".format(int(abs(self.max_depth)), suffix) return string diff --git a/earthdiagnostics/cdftools.py b/earthdiagnostics/cdftools.py index 58bef080..33ff3e0a 100644 --- a/earthdiagnostics/cdftools.py +++ b/earthdiagnostics/cdftools.py @@ -16,11 +16,19 @@ class CDFTools(object): :type path: str """ - def __init__(self, path=''): + def __init__(self, path=""): self.path = path # noinspection PyShadowingBuiltins - def run(self, command, input_file, output_file=None, options=None, log_level=Log.INFO, input_option=None): + def run( + self, + command, + input_file, + output_file=None, + options=None, + log_level=Log.INFO, + input_option=None, + ): """ Run one of the CDFTools @@ -32,7 +40,8 @@ class CDFTools(object): :type options: str :param options: options for the tool. :type options: str | [str] | Tuple[str] | None - :param log_level: log level at which the output of the cdftool command will be added + :param log_level: log level at which the output of the cdftool + command will be added :type log_level: int :param input_option: option to add before input file :type input_option: str @@ -50,10 +59,12 @@ class CDFTools(object): line.append(str(option)) if output_file: if input_file == output_file: - raise ValueError('Input and output file can not be the same on CDFTools') - line.append('-o') + raise ValueError( + "Input and output file can not be the same on CDFTools" + ) + line.append("-o") line.append(output_file) - Log.debug('Executing {0}', ' '.join(line)) + Log.debug("Executing {0}", " ".join(line)) shell_output = Utils.execute_shell_command(line, log_level) @@ -64,7 +75,10 @@ class CDFTools(object): def _check_output_was_created(line, output): if output: if not os.path.isfile(output): - raise Exception('Error executing {0}\n Output file not created', ' '.join(line)) + raise Exception( + "Error executing {0}\n Output file not created", + " ".join(line), + ) # noinspection PyShadowingBuiltins @staticmethod @@ -73,12 +87,18 @@ class CDFTools(object): if isinstance(input_file, six.string_types): line.append(input_file) if not os.path.isfile(input_file): - raise ValueError('Error executing {0}\n Input file {1} file does not exist', command, input_file) + raise ValueError( + f"Error executing {command}\n" + f"Input file {input_file} file does not exist", + ) else: for element in input_file: line.append(element) if not os.path.isfile(element): - raise ValueError('Error executing {0}\n Input file {1} file does not exist', command, element) + raise ValueError( + "Error executing {command}\n" + f"Input file {element} file does not exist", + ) @staticmethod def _is_exe(fpath): @@ -94,4 +114,8 @@ class CDFTools(object): exe_file = os.path.join(path, command) if self._is_exe(exe_file): return - raise ValueError('Error executing {0}\n Command does not exist in {1}'.format(command, self.path)) + raise ValueError( + "Error executing {0}\n Command does not exist in {1}".format( + command, self.path + ) + ) diff --git a/earthdiagnostics/cmorizer.py b/earthdiagnostics/cmorizer.py index 6905c34a..ae200776 100644 --- a/earthdiagnostics/cmorizer.py +++ b/earthdiagnostics/cmorizer.py @@ -6,10 +6,16 @@ import shutil import uuid import traceback import time -from datetime import datetime, timedelta +from datetime import datetime import six -from bscearth.utils.date import parse_date, chunk_end_date, previous_day, date2str, add_months +from bscearth.utils.date import ( + parse_date, + chunk_end_date, + previous_day, + date2str, + add_months, +) from bscearth.utils.log import Log from cdo import CDOException import iris @@ -38,16 +44,54 @@ class Cmorizer(object): """ - NON_DATA_VARIABLES = ('lon', 'lat', 'longitude', 'latitude', 'plev', 'time', 'time_bnds', 'leadtime', 'lev', - 'lev_2', 'icethi', - 'deptht', 'depthu', 'depthw', 'depthv', 'time_centered', 'time_centered_bounds', - 'deptht_bounds', 'depthu_bounds', 'depthv_bounds', 'depthw_bounds', - 'deptht_bnds', 'depthu_bnds', 'depthv_bnds', 'depthw_bnds', - 'time_counter_bounds', 'ncatice', 'nav_lat_grid_V', 'nav_lat_grid_U', - 'nav_lat_grid_T', 'nav_lon_grid_V', 'nav_lon_grid_U', 'nav_lon_grid_T', - 'depth', 'depth_2', 'depth_3', 'depth_4', - 'depth_bnds', 'depth_2_bnds', 'depth_3_bnds', 'depth_4_bnds', - 'mlev', 'hyai', 'hybi', 'hyam', 'hybm') + NON_DATA_VARIABLES = ( + "lon", + "lat", + "longitude", + "latitude", + "plev", + "time", + "time_bnds", + "leadtime", + "lev", + "lev_2", + "icethi", + "deptht", + "depthu", + "depthw", + "depthv", + "time_centered", + "time_centered_bounds", + "deptht_bounds", + "depthu_bounds", + "depthv_bounds", + "depthw_bounds", + "deptht_bnds", + "depthu_bnds", + "depthv_bnds", + "depthw_bnds", + "time_counter_bounds", + "ncatice", + "nav_lat_grid_V", + "nav_lat_grid_U", + "nav_lat_grid_T", + "nav_lon_grid_V", + "nav_lon_grid_U", + "nav_lon_grid_T", + "depth", + "depth_2", + "depth_3", + "depth_4", + "depth_bnds", + "depth_2_bnds", + "depth_3_bnds", + "depth_4_bnds", + "mlev", + "hyai", + "hybi", + "hyam", + "hybm", + ) def __init__(self, data_manager, startdate, member): self.data_manager = data_manager @@ -59,43 +103,67 @@ class Cmorizer(object): self.cmor = self.config.cmor self.convetion = self.config.data_convention self.member_str = self.experiment.get_member_str(member) - self.original_files_path = os.path.join(self.config.data_dir, self.experiment.expid, 'original_files', - self.startdate, self.member_str, 'outputs') + self.original_files_path = os.path.join( + self.config.data_dir, + self.experiment.expid, + "original_files", + self.startdate, + self.member_str, + "outputs", + ) self.atmos_timestep = None - self.cmor_scratch = str(os.path.join(self.config.scratch_dir, 'CMOR', self.startdate, self.member_str)) + self.cmor_scratch = str( + os.path.join( + self.config.scratch_dir, + "CMOR", + self.startdate, + self.member_str, + ) + ) self.lon_name = self.config.data_convention.lon_name self.lat_name = self.config.data_convention.lat_name self.alt_coord_names = { - 'time_counter': 'time', 'time_counter_bnds': 'time_bnds', - 'time_counter_bounds': 'time_bnds', - 'tbnds': 'bnds', 'nav_lat': self.lat_name, 'nav_lon': self.lon_name, - 'x': 'i', 'y': 'j'} + "time_counter": "time", + "time_counter_bnds": "time_bnds", + "time_counter_bounds": "time_bnds", + "tbnds": "bnds", + "nav_lat": self.lat_name, + "nav_lon": self.lon_name, + "x": "i", + "y": "j", + } @property def path_icm(self): """Path to the ICM file""" - return os.path.join(self.config.scratch_dir, 'ICM') + return os.path.join(self.config.scratch_dir, "ICM") def cmorize_ocean(self): """Cmorize ocean files from MMO files""" if not self.cmor.ocean: - Log.info('Skipping ocean cmorization due to configuration') + Log.info("Skipping ocean cmorization due to configuration") return True - Log.info('\nCMORizing ocean\n') - return self._cmorize_ocean_files('MMO', 'PPO', 'diags') + Log.info("\nCMORizing ocean\n") + return self._cmorize_ocean_files("MMO", "PPO", "diags") def _cmorize_ocean_files(self, *args): tar_files = () for prefix in args: - tar_folder = os.path.join(self.original_files_path, '{0}*'.format(prefix)) + tar_folder = os.path.join( + self.original_files_path, "{0}*".format(prefix) + ) tar_files = glob.glob(tar_folder) tar_files.sort() if len(tar_files) > 0: break if not len(tar_files): - Log.error('No {1} files found in {0}'.format(self.original_files_path, args)) + Log.error( + "No {1} files found in {0}".format( + self.original_files_path, args + ) + ) return False count = 1 @@ -103,19 +171,37 @@ class Cmorizer(object): for tarfile in tar_files: if not self._cmorization_required( self._get_chunk(os.path.basename(tarfile)), - (ModelingRealms.ocean, ModelingRealms.seaIce, ModelingRealms.ocnBgchem) + ( + ModelingRealms.ocean, + ModelingRealms.seaIce, + ModelingRealms.ocnBgchem, + ), ): - Log.info('No need to unpack file {0}/{1}'.format(count, len(tar_files))) + Log.info( + "No need to unpack file {0}/{1}".format( + count, len(tar_files) + ) + ) count += 1 else: - Log.info('Unpacking oceanic file {0}/{1}'.format(count, len(tar_files))) + Log.info( + "Unpacking oceanic file {0}/{1}".format( + count, len(tar_files) + ) + ) try: self._unpack_tar_file(tarfile) self._cmorize_nc_files() - Log.result('Oceanic file {0}/{1} finished'.format(count, len(tar_files))) + Log.result( + "Oceanic file {0}/{1} finished".format( + count, len(tar_files) + ) + ) count += 1 except Exception as ex: - Log.error('Could not CMORize oceanic file {0}: {1}', count, ex) + Log.error( + "Could not CMORize oceanic file {0}: {1}", count, ex + ) result = False count += 1 return result @@ -130,29 +216,52 @@ class Cmorizer(object): filtered.append(file_path) else: self._remove(file_path) - if len(filtered) == 0: - Log.warning('Filters {0} do not match any of the files', filters) + if not filtered: + raise CMORException( + f"Filters {self.cmor.filter_files} do not match any file", + ) return filtered def _remove(self, file_path): os.remove(file_path) def _cmorize_nc_files(self): - nc_files = glob.glob(os.path.join(self.cmor_scratch, '*.nc')) + nc_files = glob.glob(os.path.join(self.cmor_scratch, "*.nc")) for filename in self._filter_files(nc_files): self._cmorize_nc_file(filename) self._clean_cmor_scratch() def _correct_fluxes(self): - fluxes_vars = [self.data_manager.variable_list.get_variable(cmor_var, True).short_name - for cmor_var in ('prc', 'prs', "prsn", "rss", "rls", "rsscs", "rsds", "rlds", "hfss", 'hfls')] - change_sign_vars = [self.data_manager.variable_list.get_variable(cmor_var, True).short_name - for cmor_var in ("hfss", 'hfls')] - total_seconds = (self.experiment.atmos_timestep * 3600) - for filename in glob.glob(os.path.join(self.cmor_scratch, '*.nc')): + fluxes_vars = [ + self.data_manager.variable_list.get_variable( + cmor_var, True + ).short_name + for cmor_var in ( + "prc", + "prs", + "prsn", + "rss", + "rls", + "rsscs", + "rsds", + "rlds", + "hfss", + "hfls", + ) + ] + change_sign_vars = [ + self.data_manager.variable_list.get_variable( + cmor_var, True + ).short_name + for cmor_var in ("hfss", "hfls") + ] + total_seconds = self.experiment.atmos_timestep * 3600 + for filename in glob.glob(os.path.join(self.cmor_scratch, "*.nc")): handler = Utils.open_cdf(filename) for varname in handler.variables.keys(): - cmor_var = self.data_manager.variable_list.get_variable(varname, True) + cmor_var = self.data_manager.variable_list.get_variable( + varname, True + ) if cmor_var is None or cmor_var.short_name not in fluxes_vars: continue @@ -164,24 +273,30 @@ class Cmorizer(object): var_handler = handler.variables[varname] var_handler[:] = sign * var_handler[:] / total_seconds - var_handler.units = '{0} {1}'.format(var_handler.units, 's-1') + var_handler.units = "{0} {1}".format(var_handler.units, "s-1") handler.close() def _unpack_tar_file(self, tarfile): self._clean_cmor_scratch() os.makedirs(self.cmor_scratch) - Utils.untar((tarfile,), self.cmor_scratch, self.config.cmor.filter_files) - if os.path.isdir(os.path.join(self.cmor_scratch, 'backup')): - for filepath in glob.glob(os.path.join(self.cmor_scratch, 'backup', '*')): - Log.debug('Moving file {0}', filepath) - shutil.move(filepath, filepath.replace('/backup/', '/')) - zip_files = glob.glob(os.path.join(self.cmor_scratch, '*.gz')) + Utils.untar( + (tarfile,), self.cmor_scratch, self.config.cmor.filter_files + ) + if os.path.isdir(os.path.join(self.cmor_scratch, "backup")): + for filepath in glob.glob( + os.path.join(self.cmor_scratch, "backup", "*") + ): + Log.debug("Moving file {0}", filepath) + shutil.move(filepath, filepath.replace("/backup/", "/")) + zip_files = glob.glob(os.path.join(self.cmor_scratch, "*.gz")) if zip_files: for zip_file in self._filter_files(zip_files): try: Utils.unzip(zip_file) except Utils.UnzipException as ex: - Log.error('File {0} could not be unzipped: {1}', tarfile, ex) + Log.error( + "File {0} could not be unzipped: {1}", tarfile, ex + ) def _clean_cmor_scratch(self): if os.path.exists(self.cmor_scratch): @@ -190,57 +305,87 @@ class Cmorizer(object): def _merge_mma_files(self, tarfile): temp = TempFile.get() - for grid in ['SH', 'GG']: - files = glob.glob(os.path.join(self.cmor_scratch, 'MMA_*_{}_*.nc'.format(grid))) + for grid in ["SH", "GG"]: + files = glob.glob( + os.path.join(self.cmor_scratch, "MMA_*_{}_*.nc".format(grid)) + ) if not files: continue merged = TempFile.get() - if grid == 'SH': + if grid == "SH": for filename in files: - Utils.cdo().sp2gpl(options='-O', input=filename, output=temp) + Utils.cdo().sp2gpl( + options="-O", input=filename, output=temp + ) shutil.move(temp, filename) Utils.cdo().mergetime(input=files, output=merged) for filename in files: self._remove(filename) - tar_startdate = os.path.basename(tarfile[0:-4]).split('_')[4].split('-') - filename = 'MMA{0}_1m_{1[0]}_{1[1]}.nc'.format(grid, tar_startdate) + tar_startdate = ( + os.path.basename(tarfile[0:-4]).split("_")[4].split("-") + ) + filename = "MMA{0}_1m_{1[0]}_{1[1]}.nc".format(grid, tar_startdate) shutil.move(merged, os.path.join(self.cmor_scratch, filename)) def cmorize_atmos(self): """Cmorize atmospheric data, from grib or MMA files""" if not self.cmor.atmosphere: - Log.info('Skipping atmosphere cmorization due to configuration') + Log.info("Skipping atmosphere cmorization due to configuration") return True - Log.info('\nCMORizing atmosphere\n') + Log.info("\nCMORizing atmosphere\n") if self.cmor.use_grib and self._gribfiles_available(): return self._cmorize_grib_files() else: return self._cmorize_mma_files() def _cmorize_mma_files(self): - tar_files = glob.glob(os.path.join(self.original_files_path, 'MMA*.tar')) + tar_files = glob.glob( + os.path.join(self.original_files_path, "MMA*.tar") + ) tar_files.sort() count = 1 if len(tar_files) == 0: - Log.error('MMA files not found in {0}'.format(self.original_files_path)) + Log.error( + "MMA files not found in {0}".format(self.original_files_path) + ) return False result = True for tarfile in tar_files: - if not self._cmorization_required(self._get_chunk(os.path.basename(tarfile)), (ModelingRealms.atmos,)): - Log.info('No need to unpack file {0}/{1}'.format(count, len(tar_files))) + if not self._cmorization_required( + self._get_chunk(os.path.basename(tarfile)), + (ModelingRealms.atmos,), + ): + Log.info( + "No need to unpack file {0}/{1}".format( + count, len(tar_files) + ) + ) count += 1 continue - Log.info('Unpacking atmospheric file {0}/{1}'.format(count, len(tar_files))) + Log.info( + "Unpacking atmospheric file {0}/{1}".format( + count, len(tar_files) + ) + ) try: self._unpack_tar_file(tarfile) self._merge_mma_files(tarfile) self._correct_fluxes() self._cmorize_nc_files() - Log.result('Atmospheric file {0}/{1} finished'.format(count, len(tar_files))) + Log.result( + "Atmospheric file {0}/{1} finished".format( + count, len(tar_files) + ) + ) except Exception as ex: - Log.error('Could not cmorize atmospheric file {0}: {1}\n {2}', count, ex, traceback.format_exc()) + Log.error( + "Could not cmorize atmospheric file {0}: {1}\n {2}", + count, + ex, + traceback.format_exc(), + ) result = False count += 1 @@ -250,106 +395,157 @@ class Cmorizer(object): chunk = 1 chunk_start = parse_date(self.startdate) result = True - while os.path.exists(self._get_original_grib_path(chunk_start, 'GG')) or \ - os.path.exists(self._get_original_grib_path(chunk_start, 'SH')): + while os.path.exists( + self._get_original_grib_path(chunk_start, "GG") + ) or os.path.exists(self._get_original_grib_path(chunk_start, "SH")): if self._cmorization_required(chunk, (ModelingRealms.atmos,)): - chunk_end = chunk_end_date(chunk_start, self.experiment.chunk_size, 'month', self.experiment.calendar) + chunk_end = chunk_end_date( + chunk_start, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) chunk_end = previous_day(chunk_end, self.experiment.calendar) - Log.info('CMORizing chunk {0}-{1}', date2str(chunk_start), date2str(chunk_end)) + Log.info( + "CMORizing chunk {0}-{1}", + date2str(chunk_start), + date2str(chunk_end), + ) try: - for grid in ('SH', 'GG'): - Log.info('Processing {0} variables', grid) + for grid in ("SH", "GG"): + Log.info("Processing {0} variables", grid) - first_grib = self._get_original_grib_path(chunk_start, grid) + first_grib = self._get_original_grib_path( + chunk_start, grid + ) if not os.path.exists(first_grib): continue var_list = Utils.cdo().showvar(input=first_grib)[0] - codes = {int(var.replace('var', '')) for var in var_list.split()} - if not codes.intersection(self.config.cmor.get_requested_codes()): - Log.info('No requested variables found in {0}. Skipping...', grid) + codes = { + int(var.replace("var", "")) + for var in var_list.split() + } + if not codes.intersection( + self.config.cmor.get_requested_codes() + ): + Log.info( + "No requested variables found in {0}. " + "Skipping...", + grid, + ) continue self._cmorize_grib_file(chunk_end, chunk_start, grid) except Exception as ex: - Log.error('Can not cmorize GRIB file for chunk {0}-{1}: {2}', - date2str(chunk_start), date2str(chunk_end), ex) + Log.error( + "Can not cmorize GRIB file for chunk {0}-{1}: {2}", + date2str(chunk_start), + date2str(chunk_end), + ex, + ) result = False - chunk_start = chunk_end_date(chunk_start, self.experiment.chunk_size, 'month', self.experiment.calendar) + chunk_start = chunk_end_date( + chunk_start, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) chunk += 1 return result def _cmorize_grib_file(self, chunk_end, chunk_start, grid): for month in range(0, self.experiment.chunk_size): - current_date = add_months(chunk_start, month, self.experiment.calendar) - original_gribfile = self._get_original_grib_path(current_date, grid) - Log.info('Processing month {1}', grid, date2str(current_date)) + current_date = add_months( + chunk_start, month, self.experiment.calendar + ) + original_gribfile = self._get_original_grib_path( + current_date, grid + ) + Log.info("Processing month {1}", grid, date2str(current_date)) gribfile = self._get_scratch_grib_path(current_date, grid) if not os.path.isfile(gribfile): - Log.info('Copying file...', grid, date2str(current_date)) + Log.info("Copying file...", grid, date2str(current_date)) Utils.copy_file(original_gribfile, gribfile) self._obtain_atmos_timestep(gribfile) full_file = self._get_monthly_grib(current_date, gribfile, grid) - if not self._unpack_grib(full_file, gribfile, grid, current_date.month): + if not self._unpack_grib( + full_file, gribfile, grid, current_date.month + ): self._remove(gribfile) return - next_gribfile = self._get_original_grib_path(add_months(current_date, 1, self.experiment.calendar), grid) + next_gribfile = self._get_original_grib_path( + add_months(current_date, 1, self.experiment.calendar), grid + ) if not os.path.exists(next_gribfile): self._remove(gribfile) self._ungrib_vars(gribfile, current_date.month) - for splited_file in glob.glob('{0}_*.128.nc'.format(gribfile)): + for splited_file in glob.glob("{0}_*.128.nc".format(gribfile)): self._remove(splited_file) - Log.result('Month {0}, {1} variables finished', date2str(current_date), grid) + Log.result( + "Month {0}, {1} variables finished", + date2str(current_date), + grid, + ) - self._merge_and_cmorize_atmos(chunk_start, chunk_end, grid, Frequencies.monthly) - self._merge_and_cmorize_atmos(chunk_start, chunk_end, grid, Frequencies.daily) - self._merge_and_cmorize_atmos(chunk_start, chunk_end, grid, - '{0}hr'.format(self.atmos_timestep)) + self._merge_and_cmorize_atmos( + chunk_start, chunk_end, grid, Frequencies.monthly + ) + self._merge_and_cmorize_atmos( + chunk_start, chunk_end, grid, Frequencies.daily + ) + self._merge_and_cmorize_atmos( + chunk_start, chunk_end, grid, "{0}hr".format(self.atmos_timestep) + ) def _unpack_grib(self, full_file, gribfile, grid, month): - Log.info('Unpacking... ') + Log.info("Unpacking... ") # remap on regular Gauss grid codes = self.cmor.get_requested_codes() if 228 in codes: codes.update((142, 143)) - codes_str = ','.join([str(code) for code in codes]) + codes_str = ",".join([str(code) for code in codes]) try: - if grid == 'SH': + if grid == "SH": Utils.cdo().splitparam( - input='-sp2gpl -selcode,{0} {1} '.format(codes_str, full_file), - output=gribfile + '_', - options='-f nc4 -t ecmwf' + input="-sp2gpl -selcode,{0} {1} ".format( + codes_str, full_file + ), + output=gribfile + "_", + options="-f nc4 -t ecmwf", ) else: Utils.cdo().splitparam( - input='-selcode,{0} {1}'.format(codes_str, full_file), - output=gribfile + '_', - options='-R -f nc4 -t ecmwf' + input="-selcode,{0} {1}".format(codes_str, full_file), + output=gribfile + "_", + options="-R -f nc4 -t ecmwf", ) # total precipitation (remove negative values) if 228 in codes: Utils.cdo().setcode( 228, - input='-chname,LSP,TP -setmisstoc,0 -setvrange,0,Inf ' - '-add {0}_142.128.nc {0}_143.128.nc'.format(gribfile), - output='{0}_228.128.nc'.format(gribfile), - options='-f nc4' + input="-chname,LSP,TP -setmisstoc,0 -setvrange,0,Inf " + "-add {0}_142.128.nc {0}_143.128.nc".format(gribfile), + output="{0}_228.128.nc".format(gribfile), + options="-f nc4", ) return True except CDOException: - Log.info('No requested codes found in {0} file'.format(grid)) + Log.info("No requested codes found in {0} file".format(grid)) return False finally: Utils.remove_file(self.path_icm) def _get_monthly_grib(self, current_date, gribfile, grid): - prev_gribfile = self._get_scratch_grib_path(add_months(current_date, -1, self.experiment.calendar), grid) + prev_gribfile = self._get_scratch_grib_path( + add_months(current_date, -1, self.experiment.calendar), grid + ) if os.path.exists(prev_gribfile): self._merge_grib_files(current_date, prev_gribfile, gribfile) @@ -359,31 +555,41 @@ class Cmorizer(object): return full_file def _get_scratch_grib_path(self, current_date, grid): - return os.path.join(self.config.scratch_dir, self._get_grib_filename(grid, current_date)) + return os.path.join( + self.config.scratch_dir, + self._get_grib_filename(grid, current_date), + ) def _obtain_atmos_timestep(self, gribfile): if self.atmos_timestep is None: self.atmos_timestep = self._get_atmos_timestep(gribfile) def _get_original_grib_path(self, current_date, grid): - return os.path.join(self.original_files_path, - self._get_grib_filename(grid, current_date)) + return os.path.join( + self.original_files_path, + self._get_grib_filename(grid, current_date), + ) def _get_grib_filename(self, grid, month): - return 'ICM{0}{1}+{2}.grb'.format(grid, self.experiment.expid, date2str(month)[:-2]) + return "ICM{0}{1}+{2}.grb".format( + grid, self.experiment.expid, date2str(month)[:-2] + ) def _get_atmos_timestep(self, gribfile): - Log.info('Getting timestep...') + Log.info("Getting timestep...") import cfgrib + grib = cfgrib.open_file(gribfile) dates = set() - valid_time = grib.variables['valid_time'] - for time in valid_time.data: - dates.add(cf_units.num2date( - time, - valid_time.attributes['units'], - valid_time.attributes['calendar'], - )) + valid_time = grib.variables["valid_time"] + for t in valid_time.data: + dates.add( + cf_units.num2date( + t, + valid_time.attributes["units"], + valid_time.attributes["calendar"], + ) + ) dates = list(dates) dates.sort() atmos_timestep = dates[1] - dates[0] @@ -392,7 +598,7 @@ class Cmorizer(object): return atmos_timestep def _cmorize_nc_file(self, filename): - Log.info('Processing file {0}', filename) + Log.info("Processing file {0}", filename) if not self._contains_requested_variables(filename): self._remove(filename) @@ -409,36 +615,38 @@ class Cmorizer(object): variables = handler.variables.keys() handler.close() - Log.info('Splitting file {0}', filename) + Log.info("Splitting file {0}", filename) for variable in variables: if variable in Cmorizer.NON_DATA_VARIABLES: continue try: - Log.debug('Checking variable {0}', variable) + Log.debug("Checking variable {0}", variable) self.extract_variable(filename, frequency, variable) except Exception as ex: - Log.error('Variable {0} can not be cmorized: {1}', variable, ex) - Log.result('File {0} cmorized!', filename) + Log.error( + "Variable {0} can not be cmorized: {1}", variable, ex + ) + Log.result("File {0} cmorized!", filename) self._remove(filename) @staticmethod def _remove_valid_limits(handler): for variable in handler.variables.keys(): var = handler.variables[variable] - if 'valid_min' in var.ncattrs(): + if "valid_min" in var.ncattrs(): del var.valid_min - if 'valid_max' in var.ncattrs(): + if "valid_max" in var.ncattrs(): del var.valid_max handler.sync() def _get_nc_file_frequency(self, filename): - file_parts = os.path.basename(filename).split('_') + file_parts = os.path.basename(filename).split("_") if self.experiment.expid in [file_parts[1], file_parts[2]]: - frequency = Frequency('m') + frequency = Frequency("m") elif self.experiment.expid == file_parts[0]: try: parse_date(file_parts[1]) - frequency = Frequency('m') + frequency = Frequency("m") except ValueError: frequency = Frequency(file_parts[1]) else: @@ -465,7 +673,9 @@ class Cmorizer(object): If the filename does not match any of the recognized patterns """ - alias, var_cmor = self.config.var_manager.get_variable_and_alias(variable) + alias, var_cmor = self.config.var_manager.get_variable_and_alias( + variable + ) if var_cmor is None: return @@ -473,7 +683,9 @@ class Cmorizer(object): return temp = TempFile.get() - lev_dimensions = self._set_coordinates_attribute(file_path, var_cmor, variable) + lev_dimensions = self._set_coordinates_attribute( + file_path, var_cmor, variable + ) self._rename_level_coords(file_path, lev_dimensions, temp, variable) if alias.basin is None: @@ -483,19 +695,31 @@ class Cmorizer(object): date_str = self._get_date_str(file_path) if date_str is None: - Log.error('Variable {0} can not be cmorized. Original filename does not match a recognized pattern', - var_cmor.short_name) - raise CMORException('Variable {0}:{1} can not be cmorized. Original filename does not match a recognized ' - 'pattern'.format(var_cmor.domain, var_cmor.short_name)) + Log.error( + f"Variable {var_cmor.short_name,} can not be cmorized. " + "Original filename does not match a recognized pattern", + ) + raise CMORException( + f"Variable {var_cmor.domain}:{var_cmor.short_name} can not " + "be cmorized. Original filename does not match a recognized " + "pattern" + ) netcdf_file = NetCDFFile() netcdf_file.data_manager = self.data_manager netcdf_file.local_file = temp - netcdf_file.remote_file = self.config.data_convention.get_file_path(self.startdate, self.member, - var_cmor.domain, var_cmor.short_name, - var_cmor, None, frequency, - grid=alias.grid, year=None, - date_str=date_str) + netcdf_file.remote_file = self.config.data_convention.get_file_path( + self.startdate, + self.member, + var_cmor.domain, + var_cmor.short_name, + var_cmor, + None, + frequency, + grid=alias.grid, + year=None, + date_str=date_str, + ) netcdf_file.data_convention = self.config.data_convention netcdf_file.region = region @@ -509,111 +733,138 @@ class Cmorizer(object): netcdf_file.upload() if region: - region_str = ' (Region {})'.format(region) + region_str = " (Region {})".format(region) else: - region_str = '' - Log.info('Variable {0.domain}:{0.short_name} processed{1}', var_cmor, region_str) + region_str = "" + Log.info( + "Variable {0.domain}:{0.short_name} processed{1}", + var_cmor, + region_str, + ) def _rename_level_coords(self, file_path, lev_dimensions, temp, variable): - cube = iris.load_cube(file_path, iris.Constraint(cube_func=lambda c: c.var_name == variable)) + cube = iris.load_cube( + file_path, + iris.Constraint(cube_func=lambda c: c.var_name == variable), + ) for lev_original, lev_target in six.iteritems(lev_dimensions): try: cube.coord(var_name=lev_original).var_name = lev_target except iris.exceptions.CoordinateNotFoundError: pass - if cube.coord('latitude').ndim > 1: + if cube.coord("latitude").ndim > 1: try: - cube.coord('j'), + cube.coord("j"), except iris.exceptions.CoordinateNotFoundError: cube.add_dim_coord( DimCoord( range(cube.shape[-2]), - var_name='j', - long_name='Cell index along second dimension', + var_name="j", + long_name="Cell index along second dimension", units=1.0, - ), len(cube.shape) - 2 + ), + len(cube.shape) - 2, ) try: - cube.coord('i'), + cube.coord("i"), except iris.exceptions.CoordinateNotFoundError: cube.add_dim_coord( DimCoord( range(cube.shape[-1]), - var_name='i', - long_name='Cell index along first dimension', + var_name="i", + long_name="Cell index along first dimension", units=1.0, - ), len(cube.shape) - 1) + ), + len(cube.shape) - 1, + ) iris.save(cube, temp, zlib=True) def _set_coordinates_attribute(self, file_path, var_cmor, variable): handler = Utils.open_cdf(file_path) - coords = [self.lon_name, self.lat_name, 'time'] - if 'leadtime' in handler.variables.keys(): - coords.append('leadtime') + coords = [self.lon_name, self.lat_name, "time"] + if "leadtime" in handler.variables.keys(): + coords.append("leadtime") lev_dimensions = self._get_lev_dimensions(var_cmor) for lev_dim in lev_dimensions.keys(): if lev_dim in handler.variables[variable].dimensions: coords.append(lev_dim) - handler.variables[variable].coordinates = ' '.join(set(coords)) + handler.variables[variable].coordinates = " ".join(set(coords)) handler.close() return lev_dimensions def _get_lev_dimensions(self, var_cmor): if var_cmor.domain == ModelingRealms.ocean: - lev_dimensions = {'deptht': 'lev', 'depthu': 'lev', 'depthw': 'lev', 'depthv': 'lev', - 'depth': 'lev'} + lev_dimensions = { + "deptht": "lev", + "depthu": "lev", + "depthw": "lev", + "depthv": "lev", + "depth": "lev", + } elif var_cmor.domain in [ModelingRealms.landIce, ModelingRealms.land]: - lev_dimensions = {'depth': 'sdepth', 'depth_2': 'sdepth', 'depth_3': 'sdepth', - 'depth_4': 'sdepth'} + lev_dimensions = { + "depth": "sdepth", + "depth_2": "sdepth", + "depth_3": "sdepth", + "depth_4": "sdepth", + } elif var_cmor.domain == ModelingRealms.atmos: - lev_dimensions = {'depth': 'plev'} + lev_dimensions = {"depth": "plev"} else: lev_dimensions = {} return lev_dimensions def _get_date_str(self, file_path): - file_parts = os.path.basename(file_path).split('_') - valid_starts = (self.experiment.expid, 'MMA', 'MMASH', 'MMAGG', 'MMO') - if file_parts[0] in valid_starts or file_parts[0].startswith('ORCA'): + file_parts = os.path.basename(file_path).split("_") + valid_starts = (self.experiment.expid, "MMA", "MMASH", "MMAGG", "MMO") + if file_parts[0] in valid_starts or file_parts[0].startswith("ORCA"): # Model output - if file_parts[-1].endswith('.tar'): - file_parts = file_parts[-1][0:-4].split('-') - return '{0}-{1}'.format(file_parts[0][0:6], file_parts[1][0:6]) + if file_parts[-1].endswith(".tar"): + file_parts = file_parts[-1][0:-4].split("-") + return "{0}-{1}".format(file_parts[0][0:6], file_parts[1][0:6]) else: - return '{0}-{1}'.format(file_parts[2][0:6], file_parts[3][0:6]) + return "{0}-{1}".format(file_parts[2][0:6], file_parts[3][0:6]) elif file_parts[1] == self.experiment.expid: # Files generated by the old version of the diagnostics - return '{0}-{1}'.format(file_parts[4][0:6], file_parts[5][0:6]) + return "{0}-{1}".format(file_parts[4][0:6], file_parts[5][0:6]) else: return None def _get_chunk(self, file_path): - chunk_start = parse_date(self._get_date_str(file_path).split('-')[0]) + chunk_start = parse_date(self._get_date_str(file_path).split("-")[0]) current_date = parse_date(self.startdate) chunk = 1 while current_date < chunk_start: - current_date = chunk_end_date(current_date, self.experiment.chunk_size, 'month', self.experiment.calendar) + current_date = chunk_end_date( + current_date, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) chunk += 1 if current_date != chunk_start: - raise Exception('File {0} start date is not a valid chunk start date'.format(file_path)) + raise Exception( + "File {0} start date is not a valid chunk start date".format( + file_path + ) + ) return chunk def _merge_grib_files(self, current_month, prev_gribfile, gribfile): - Log.info('Merging data from different files...') - temp = TempFile.get(suffix='.grb') - Utils.cdo().selmon(current_month.month, input=prev_gribfile, output=temp) - Utils.cdo().mergetime( - input=[temp, gribfile], - output=self.path_icm + Log.info("Merging data from different files...") + temp = TempFile.get(suffix=".grb") + Utils.cdo().selmon( + current_month.month, input=prev_gribfile, output=temp ) + Utils.cdo().mergetime(input=[temp, gribfile], output=self.path_icm) self._remove(prev_gribfile) self._remove(temp) def _ungrib_vars(self, gribfile, month): for var_code in self.cmor.get_requested_codes(): - file_path = '{0}_{1}.128.nc'.format(gribfile, var_code) + file_path = "{0}_{1}.128.nc".format(gribfile, var_code) if not os.path.exists(file_path): continue cube = iris.load_cube(file_path) @@ -622,7 +873,11 @@ class Cmorizer(object): cube = cube.extract(iris.Constraint(month_number=month)) cube = self._change_units(cube, var_code) - for frequency in (Frequencies.monthly, Frequencies.daily, Frequency('{0}hr'.format(self.atmos_timestep))): + for frequency in ( + Frequencies.monthly, + Frequencies.daily, + Frequency("{0}hr".format(self.atmos_timestep)), + ): if var_code not in self.cmor.get_variables(frequency): continue time_cube = self._get_time_average(cube, frequency, var_code) @@ -631,41 +886,72 @@ class Cmorizer(object): if levels: time_cube = time_cube.extract(level=levels) - if cube.var_name.endswith('_2'): + if cube.var_name.endswith("_2"): time_cube.var_name = cube.var_name[:-2] - out_file = '{0}_{1}_{2}.nc'.format(gribfile, var_code, frequency) - time_cube.remove_coord('month_number') - time_cube.remove_coord('day_of_month') + out_file = "{0}_{1}_{2}.nc".format( + gribfile, var_code, frequency + ) + time_cube.remove_coord("month_number") + time_cube.remove_coord("day_of_month") iris.save(time_cube, out_file, zlib=True) def _fix_time_coord(self, cube, var_code): - time = cube.coord('time') - target_units = 'days since 1950-01-01 00:00:00' - time.convert_units(cf_units.Unit(target_units, calendar=time.units.calendar)) + time = cube.coord("time") + target_units = "days since 1950-01-01 00:00:00" + time.convert_units( + cf_units.Unit(target_units, calendar=time.units.calendar) + ) time.units = target_units - if var_code in (144, 146, 147, 169, 175, 176, 177, 179, 180, 181, 182, 201, 202, 205, 212, 228): + if var_code in ( + 144, + 146, + 147, + 169, + 175, + 176, + 177, + 179, + 180, + 181, + 182, + 201, + 202, + 205, + 212, + 228, + ): time.points = time.points - (self.experiment.atmos_timestep / 24.0) - iris.coord_categorisation.add_day_of_month(cube, 'time') - iris.coord_categorisation.add_month_number(cube, 'time') + iris.coord_categorisation.add_day_of_month(cube, "time") + iris.coord_categorisation.add_month_number(cube, "time") return cube @staticmethod def _get_time_average(cube, frequency, var_code): if frequency == Frequencies.monthly: if var_code == 201: - cube = cube.aggregated_by(['month_number', 'day_of_month'], iris.analysis.MAX) + cube = cube.aggregated_by( + ["month_number", "day_of_month"], iris.analysis.MAX + ) elif var_code == 202: - cube = cube.aggregated_by(['month_number', 'day_of_month'], iris.analysis.MIN) - cube = cube.aggregated_by(['month_number'], iris.analysis.MEAN) + cube = cube.aggregated_by( + ["month_number", "day_of_month"], iris.analysis.MIN + ) + cube = cube.aggregated_by(["month_number"], iris.analysis.MEAN) elif frequency == Frequencies.daily: if var_code == 201: - cube = cube.aggregated_by(['month_number', 'day_of_month'], iris.analysis.MAX) + cube = cube.aggregated_by( + ["month_number", "day_of_month"], iris.analysis.MAX + ) elif var_code == 202: - cube = cube.aggregated_by(['month_number', 'day_of_month'], iris.analysis.MIN) + cube = cube.aggregated_by( + ["month_number", "day_of_month"], iris.analysis.MIN + ) else: - cube = cube.aggregated_by(['month_number', 'day_of_month'], iris.analysis.MEAN) + cube = cube.aggregated_by( + ["month_number", "day_of_month"], iris.analysis.MEAN + ) return cube def _change_units(self, cube, var_code): @@ -673,38 +959,54 @@ class Cmorizer(object): if var_code == 129: # geopotential cube = cube / 9.81 - cube.units = 'm' + cube.units = "m" elif var_code in (146, 147, 169, 175, 176, 177, 179, 212): # radiation cube = cube / (self.experiment.atmos_timestep * 3600) - cube.units = 'W m-2' + cube.units = "W m-2" elif var_code in (180, 181): # momentum flux cube = cube / (self.experiment.atmos_timestep * 3600) - cube.units = 'N m-2' + cube.units = "N m-2" elif var_code in (144, 182, 205, 228): # precipitation/evaporation/runoff cube = cube * 1000 / (self.experiment.atmos_timestep * 3600) - cube.units = 'kg m-2 s-1' + cube.units = "kg m-2 s-1" cube.var_name = var_name return cube - def _merge_and_cmorize_atmos(self, chunk_start, chunk_end, grid, frequency): - merged_file = 'MMA_{0}_{1}_{2}_{3}.nc'.format(frequency, date2str(chunk_start), date2str(chunk_end), grid) - files = glob.glob(os.path.join(self.config.scratch_dir, - '{0}_*_{1}.nc'.format(self._get_grib_filename(grid, chunk_start), frequency))) + def _merge_and_cmorize_atmos( + self, chunk_start, chunk_end, grid, frequency + ): + merged_file = "MMA_{0}_{1}_{2}_{3}.nc".format( + frequency, date2str(chunk_start), date2str(chunk_end), grid + ) + files = glob.glob( + os.path.join( + self.config.scratch_dir, + "{0}_*_{1}.nc".format( + self._get_grib_filename(grid, chunk_start), frequency + ), + ) + ) def _load_cube(cube, field, filename): - if 'history' in cube.attributes: - del cube.attributes['history'] + if "history" in cube.attributes: + del cube.attributes["history"] for first_file in files: var_files = [] current_month = chunk_start while current_month < chunk_end: - var_files.append(first_file.replace('+{0}.grb'.format(date2str(chunk_start)[:-2]), - '+{0}.grb'.format(date2str(current_month)[:-2]))) - current_month = add_months(current_month, 1, self.experiment.calendar) + var_files.append( + first_file.replace( + "+{0}.grb".format(date2str(chunk_start)[:-2]), + "+{0}.grb".format(date2str(current_month)[:-2]), + ) + ) + current_month = add_months( + current_month, 1, self.experiment.calendar + ) var_cubes = iris.load(var_files, callback=_load_cube) iris.util.unify_time_units(var_cubes) var_cube = var_cubes.concatenate_cube() @@ -714,30 +1016,45 @@ class Cmorizer(object): self._cmorize_nc_file(merged_file) def _update_time_variables(self, handler): - time_var = handler.variables['time'] - if hasattr(time_var, 'calendar'): + time_var = handler.variables["time"] + if hasattr(time_var, "calendar"): calendar = time_var.calendar else: - calendar = 'standard' + calendar = "standard" if "time_bnds" in handler.variables: time_var.bounds = "time_bnds" - handler.variables['time_bnds'].units = time_var.units - Utils.convert_units(handler.variables['time_bnds'], 'days since 1850-01-01 00:00:00', calendar, calendar) - Utils.convert_units(time_var, 'days since 1850-1-1 00:00:00', calendar) + handler.variables["time_bnds"].units = time_var.units + Utils.convert_units( + handler.variables["time_bnds"], + "days since 1850-01-01 00:00:00", + calendar, + calendar, + ) + Utils.convert_units(time_var, "days since 1850-1-1 00:00:00", calendar) self._set_leadtime_var(handler) def _set_leadtime_var(self, handler): - if 'leadtime' in handler.variables: - var = handler.variables['leadtime'] + if "leadtime" in handler.variables: + var = handler.variables["leadtime"] else: - var = handler.createVariable('leadtime', float, 'time') + var = handler.createVariable("leadtime", float, "time") var.units = "days" var.long_name = "Time elapsed since the start of the forecast" var.standard_name = "forecast_period" leadtime = Utils.get_datetime_from_netcdf(handler) startdate = parse_date(self.startdate) - leadtime = [datetime(time.year, time.month, time.day, time.hour, time.minute, time.second) - startdate - for time in leadtime] + leadtime = [ + datetime( + time.year, + time.month, + time.day, + time.hour, + time.minute, + time.second, + ) + - startdate + for time in leadtime + ] for lt, lead in enumerate(leadtime): var[lt] = lead.days @@ -745,15 +1062,22 @@ class Cmorizer(object): cmor = self.config.cmor experiment = self.config.experiment handler.associated_experiment = cmor.associated_experiment - handler.batch = '{0}{1}'.format(experiment.institute, datetime.now().strftime('%Y-%m-%d(T%H:%M:%SZ)')) - handler.contact = 'Pierre-Antoine Bretonniere, pierre-antoine.bretonniere@bsc.es , ' \ - 'Javier Vegas-Regidor, javier.vegas@bsc.es ' - handler.Conventions = 'CF-1.6' - handler.creation_date = datetime.now().strftime('%Y-%m-%d(T%H:%M:%SZ)') + handler.batch = "{0}{1}".format( + experiment.institute, + datetime.now().strftime("%Y-%m-%d(T%H:%M:%SZ)"), + ) + handler.contact = ( + "Pierre-Antoine Bretonniere, pierre-antoine.bretonniere@bsc.es , " + "Javier Vegas-Regidor, javier.vegas@bsc.es " + ) + handler.Conventions = "CF-1.6" + handler.creation_date = datetime.now().strftime("%Y-%m-%d(T%H:%M:%SZ)") handler.experiment_id = experiment.experiment_name startdate = parse_date(self.startdate) - handler.forecast_reference_time = '{0.year}-{0.month:02}-{0.day:02}' \ - '(T{0.hour:02}:{0.minute:02}:{0.second:02}Z)'.format(startdate) + handler.forecast_reference_time = ( + "{0.year}-{0.month:02}-{0.day:02}" + "(T{0.hour:02}:{0.minute:02}:{0.second:02}Z)".format(startdate) + ) handler.frequency = frequency.frequency handler.institute_id = experiment.institute handler.institution = experiment.institute @@ -766,13 +1090,16 @@ class Cmorizer(object): handler.project_id = self.config.data_convention.name.upper() handler.realization = str(self.member + 1) handler.source = cmor.source - handler.startdate = 'S{0}'.format(self.startdate) + handler.startdate = "S{0}".format(self.startdate) handler.tracking_id = str(uuid.uuid1()) - handler.title = "{0} model output prepared for {2} {1}".format(experiment.model, experiment.experiment_name, - self.config.data_convention.name.upper()) + handler.title = "{0} model output prepared for {2} {1}".format( + experiment.model, + experiment.experiment_name, + self.config.data_convention.name.upper(), + ) def _gribfiles_available(self): - grb_path = os.path.join(self.original_files_path, '*.grb') + grb_path = os.path.join(self.original_files_path, "*.grb") gribfiles = glob.glob(grb_path) return len(gribfiles) > 0 @@ -782,12 +1109,14 @@ class Cmorizer(object): if self.config.cmor.force: return True for domain in domains: - if self.data_manager.is_cmorized(self.startdate, self.member, chunk, domain): + if self.data_manager.is_cmorized( + self.startdate, self.member, chunk, domain + ): return False return True class CMORException(Exception): - """Exception to be launched when an error is encountered during cmorization""" + """Error during cmorization""" pass diff --git a/earthdiagnostics/cmormanager.py b/earthdiagnostics/cmormanager.py index 1ae7570a..9e9acaa3 100644 --- a/earthdiagnostics/cmormanager.py +++ b/earthdiagnostics/cmormanager.py @@ -31,7 +31,9 @@ class CMORManager(DataManager): super(CMORManager, self).__init__(config) self._dic_cmorized = dict() self.find_model_data() - self.cmor_path = os.path.join(self.config.data_dir, self.experiment.expid) + self.cmor_path = os.path.join( + self.config.data_dir, self.experiment.expid + ) self.convention = self.config.data_convention def find_model_data(self): @@ -45,30 +47,45 @@ class CMORManager(DataManager): Model has any '-' character removed and is passed to lower """ - data_folders = self.config.data_dir.split(':') + data_folders = self.config.data_dir.split(":") experiment_folder = self.experiment.model.lower() - if experiment_folder.startswith('ec-earth'): - experiment_folder = 'ecearth' + if experiment_folder.startswith("ec-earth"): + experiment_folder = "ecearth" self.config.data_dir = None for data_folder in data_folders: if os.path.isdir(os.path.join(data_folder, self.experiment.expid)): self.config.data_dir = data_folder break - test_folder = os.path.join(data_folder, self.experiment.model.lower().replace('-', '')) + test_folder = os.path.join( + data_folder, self.experiment.model.lower().replace("-", "") + ) if os.path.isdir(os.path.join(test_folder, self.experiment.expid)): self.config.data_dir = test_folder break - test_folder = os.path.join(data_folder, self.config.data_type, experiment_folder) + test_folder = os.path.join( + data_folder, self.config.data_type, experiment_folder + ) if os.path.isdir(os.path.join(test_folder, self.experiment.expid)): self.config.data_dir = test_folder break if not self.config.data_dir: - raise Exception('Can not find model data') + raise Exception("Can not find model data") # noinspection PyUnusedLocal - def file_exists(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - vartype=VariableType.MEAN, possible_versions=None): + def file_exists( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + possible_versions=None, + ): """ Check if a file exists in the storage @@ -91,20 +108,43 @@ class CMORManager(DataManager): """ cmor_var = self.variable_list.get_variable(var) - filepath = self.convention.get_file_path(startdate, member, domain, var, cmor_var, chunk, frequency, grid, - None, None) + filepath = self.convention.get_file_path( + startdate, + member, + domain, + var, + cmor_var, + chunk, + frequency, + grid, + None, + None, + ) if possible_versions is None: return os.path.isfile(filepath) else: for version in possible_versions: - if os.path.isfile(filepath.replace(self.config.cmor.version, version)): + if os.path.isfile( + filepath.replace(self.config.cmor.version, version) + ): return True return False - def request_chunk(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, vartype=None): + def request_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=None, + ): """ - Request a given file from the CMOR repository to the scratch folder and returns the path to the scratch's copy + Request a given file from the CMOR repository Parameters ---------- @@ -127,12 +167,33 @@ class CMORManager(DataManager): frequency = self.config.frequency cmor_var = self.variable_list.get_variable(var) var = self._get_final_var_name(box, var) - filepath = self.convention.get_file_path(startdate, member, domain, var, cmor_var, chunk, frequency, grid, - None, None) + filepath = self.convention.get_file_path( + startdate, + member, + domain, + var, + cmor_var, + chunk, + frequency, + grid, + None, + None, + ) return self._get_file_from_storage(filepath) - def request_year(self, diagnostic, domain, var, startdate, member, year, grid=None, box=None, frequency=None): + def request_year( + self, + diagnostic, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + frequency=None, + ): """ Request a given year for a variavle from a CMOR repository @@ -153,7 +214,9 @@ class CMORManager(DataManager): DataFile """ - job = MergeYear(self, domain, var, startdate, member, year, grid, box, frequency) + job = MergeYear( + self, domain, var, startdate, member, year, grid, box, frequency + ) job.request_data() job.declare_data_generated() if not job.year_file.job_added: @@ -161,8 +224,20 @@ class CMORManager(DataManager): job.year_file.job_added = True return job.year_file - def declare_chunk(self, domain, var, startdate, member, chunk, grid=None, region=None, box=None, frequency=None, - vartype=VariableType.MEAN, diagnostic=None): + def declare_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + region=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ Declare a variable chunk to be generated by a diagnostic @@ -193,15 +268,43 @@ class CMORManager(DataManager): var = cmor_var.short_name final_name = self._get_final_var_name(box, var) - filepath = self.convention.get_file_path(startdate, member, domain, final_name, cmor_var, chunk, - frequency, grid) - netcdf_file = self._declare_generated_file(filepath, domain, final_name, cmor_var, self.config.data_convention, - region, diagnostic, grid, vartype, original_name) + filepath = self.convention.get_file_path( + startdate, + member, + domain, + final_name, + cmor_var, + chunk, + frequency, + grid, + ) + netcdf_file = self._declare_generated_file( + filepath, + domain, + final_name, + cmor_var, + self.config.data_convention, + region, + diagnostic, + grid, + vartype, + original_name, + ) netcdf_file.frequency = frequency return netcdf_file - def declare_year(self, domain, var, startdate, member, year, grid=None, box=None, - vartype=VariableType.MEAN, diagnostic=None): + def declare_year( + self, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ Declare a variable year to be generated by a diagnostic @@ -228,15 +331,47 @@ class CMORManager(DataManager): var = cmor_var.short_name final_name = self._get_final_var_name(box, var) - filepath = self.convention.get_file_path(startdate, member, domain, final_name, cmor_var, None, - Frequencies.yearly, grid, year=year) - netcdf_file = self._declare_generated_file(filepath, domain, final_name, cmor_var, self.config.data_convention, - None, diagnostic, grid, vartype, original_name) + filepath = self.convention.get_file_path( + startdate, + member, + domain, + final_name, + cmor_var, + None, + Frequencies.yearly, + grid, + year=year, + ) + netcdf_file = self._declare_generated_file( + filepath, + domain, + final_name, + cmor_var, + self.config.data_convention, + None, + diagnostic, + grid, + vartype, + original_name, + ) netcdf_file.frequency = Frequencies.yearly return netcdf_file - def link_file(self, domain, var, cmor_var, startdate, member, chunk=None, grid=None, - frequency=None, year=None, date_str=None, move_old=False, vartype=VariableType.MEAN): + def link_file( + self, + domain, + var, + cmor_var, + startdate, + member, + chunk=None, + grid=None, + frequency=None, + year=None, + date_str=None, + move_old=False, + vartype=VariableType.MEAN, + ): """ Create the link of a given file from the CMOR repository. @@ -257,9 +392,21 @@ class CMORManager(DataManager): """ if frequency is None: frequency = self.config.frequency - filepath = self.convention.get_file_path(startdate, member, domain, var, cmor_var, chunk, frequency, - grid=grid, year=year, date_str=date_str) - self.convention.create_link(domain, filepath, frequency, var, grid, move_old, vartype) + filepath = self.convention.get_file_path( + startdate, + member, + domain, + var, + cmor_var, + chunk, + frequency, + grid=grid, + year=year, + date_str=date_str, + ) + self.convention.create_link( + domain, filepath, frequency, var, grid, move_old, vartype + ) def prepare(self): """ @@ -273,7 +420,7 @@ class CMORManager(DataManager): """ # Check if cmorized and convert if not - if self.config.data_convention.name == 'meteofrance': + if self.config.data_convention.name == "meteofrance": return if self.config.cmor.skip_prepare: return @@ -282,30 +429,36 @@ class CMORManager(DataManager): self._prepare_member(startdate, member) def _prepare_member(self, startdate, member): - Log.info('Checking data for startdate {0} member {1}', startdate, member) + Log.info( + "Checking data for startdate {0} member {1}", startdate, member + ) if not self.config.cmor.force: done = 0 for chunk in range(1, self.experiment.num_chunks + 1): if not self.config.cmor.chunk_cmorization_requested(chunk): - Log.debug('Skipping chunk {0}', chunk) + Log.debug("Skipping chunk {0}", chunk) done += 1 continue if not self.config.cmor.force_untar: - Log.info('Checking chunk {0}...', chunk) + Log.info("Checking chunk {0}...", chunk) skip = False - for domain in (ModelingRealms.atmos, ModelingRealms.ocean, ModelingRealms.seaIce): + for domain in ( + ModelingRealms.atmos, + ModelingRealms.ocean, + ModelingRealms.seaIce, + ): if self.is_cmorized(startdate, member, chunk, domain): - Log.debug('Chunk {0} ready', chunk) + Log.debug("Chunk {0} ready", chunk) skip = True break if skip: done += 1 continue if self._unpack_chunk(startdate, member, chunk): - Log.debug('Chunk {0} unpacked', chunk) + Log.debug("Chunk {0} unpacked", chunk) done += 1 if self.experiment.num_chunks == done: - Log.debug('Startdate {0} member {1} ready', startdate, member) + Log.debug("Startdate {0} member {1} ready", startdate, member) return self._cmorize_member(startdate, member) @@ -331,64 +484,103 @@ class CMORManager(DataManager): if identifier not in self._dic_cmorized: self._dic_cmorized[identifier] = {} if domain not in self._dic_cmorized[identifier]: - self._dic_cmorized[identifier][domain] = self.convention.is_cmorized(startdate, member, chunk, domain) + self._dic_cmorized[identifier][ + domain + ] = self.convention.is_cmorized(startdate, member, chunk, domain) return self._dic_cmorized[identifier][domain] def _cmorize_member(self, startdate, member): start_time = datetime.now() member_str = self.experiment.get_member_str(member) - Log.info('CMORizing startdate {0} member {1}. Starting at {2}', startdate, member_str, start_time) + Log.info( + "CMORizing startdate {0} member {1}. Starting at {2}", + startdate, + member_str, + start_time, + ) cmorizer = Cmorizer(self, startdate, member) result = cmorizer.cmorize_ocean() result &= cmorizer.cmorize_atmos() if result: - Log.result('CMORized startdate {0} member {1}! Elapsed time: {2}\n\n', startdate, member_str, - datetime.now() - start_time) + Log.result( + "CMORized startdate {0} member {1}! Elapsed time: {2}\n\n", + startdate, + member_str, + datetime.now() - start_time, + ) else: - raise(Exception( - 'Error appeared while cmorizing startdate {0}' - 'member {1}!'.format(startdate, member_str) - )) + raise ( + Exception( + "Error appeared while cmorizing startdate {0}" + "member {1}!".format(startdate, member_str) + ) + ) def _unpack_chunk(self, startdate, member, chunk): - filepaths = self._get_transferred_cmor_data_filepaths(startdate, member, chunk, 'tar.gz') + filepaths = self._get_transferred_cmor_data_filepaths( + startdate, member, chunk, "tar.gz" + ) if len(filepaths) > 0: - Log.info('Unzipping cmorized data for {0} {1} {2}...', startdate, member, chunk) + Log.info( + "Unzipping cmorized data for {0} {1} {2}...", + startdate, + member, + chunk, + ) Utils.unzip(filepaths, True) - filepaths = self._get_transferred_cmor_data_filepaths(startdate, member, chunk, 'tar') + filepaths = self._get_transferred_cmor_data_filepaths( + startdate, member, chunk, "tar" + ) if len(filepaths) > 0: - Log.info('Unpacking cmorized data for {0} {1} {2}...', startdate, member, chunk) - Utils.untar( - filepaths, - os.path.join(self.cmor_path, 'cmorfiles') + Log.info( + "Unpacking cmorized data for {0} {1} {2}...", + startdate, + member, + chunk, ) + Utils.untar(filepaths, os.path.join(self.cmor_path, "cmorfiles")) self._correct_paths(startdate) self.convention.create_links(startdate, member) return True return False - def _get_transferred_cmor_data_filepaths(self, startdate, member, chunk, extension): - tar_path = os.path.join(self.config.data_dir, self.experiment.expid, 'original_files', 'cmorfiles') - tar_original_files = os.path.join(self.config.data_dir, 'original_files', self.experiment.expid, - 'cmorfiles') + def _get_transferred_cmor_data_filepaths( + self, startdate, member, chunk, extension + ): + tar_path = os.path.join( + self.config.data_dir, + self.experiment.expid, + "original_files", + "cmorfiles", + ) + tar_original_files = os.path.join( + self.config.data_dir, + "original_files", + self.experiment.expid, + "cmorfiles", + ) filepaths = [] - for cmor_prefix in ('CMOR?', 'CMOR'): - file_name = '{5}_{0}_{1}_{2}_{3}-*.{4}'.format( - self.experiment.expid, startdate, + for cmor_prefix in ("CMOR?", "CMOR"): + file_name = "{5}_{0}_{1}_{2}_{3}-*.{4}".format( + self.experiment.expid, + startdate, self.experiment.get_member_str(member), self.experiment.get_chunk_start_str(startdate, chunk), - extension, cmor_prefix + extension, + cmor_prefix, ) filepaths += self._find_paths(tar_path, file_name) - filepaths += self._find_paths(tar_path, 'outputs', file_name) + filepaths += self._find_paths(tar_path, "outputs", file_name) filepaths += self._find_paths(tar_original_files, file_name) - filepaths += self._find_paths(tar_original_files, 'outputs', file_name) + filepaths += self._find_paths( + tar_original_files, "outputs", file_name + ) return filepaths def _find_paths(self, *args): pattern = os.path.join(*args) - Log.debug('Looking for pattern {0}', pattern) + Log.debug("Looking for pattern {0}", pattern) return glob.glob(pattern) def _correct_paths(self, startdate): @@ -398,35 +590,44 @@ class CMORManager(DataManager): def _fix_model_as_experiment_error(self, startdate): experiment_name = self.convention.experiment_name(startdate) if experiment_name != self.experiment.model: - bad_path = os.path.join(self.cmor_path, self.experiment.institute, self.experiment.model, - self.experiment.model) - Log.debug('Correcting double model appearance') + bad_path = os.path.join( + self.cmor_path, + self.experiment.institute, + self.experiment.model, + self.experiment.model, + ) + Log.debug("Correcting double model appearance") for (dirpath, _, filenames) in os.walk(bad_path, False): for filename in filenames: - if '_S{0}_'.format(startdate) in filename: + if "_S{0}_".format(startdate) in filename: continue filepath = os.path.join(dirpath, filename) good = filepath - good = good.replace('_{0}_output_'.format(self.experiment.model), - '_{0}_{1}_S{2}_'.format(self.experiment.model, - experiment_name, - startdate)) - - good = good.replace('/{0}/{0}'.format(self.experiment.model), - '/{0}/{1}'.format(self.experiment.model, - experiment_name)) + good = good.replace( + "_{0}_output_".format(self.experiment.model), + "_{0}_{1}_S{2}_".format( + self.experiment.model, experiment_name, startdate + ), + ) + + good = good.replace( + "/{0}/{0}".format(self.experiment.model), + "/{0}/{1}".format( + self.experiment.model, experiment_name + ), + ) Utils.move_file(filepath, good) if self.experiment.model != experiment_name: os.rmdir(dirpath) - Log.debug('Done') + Log.debug("Done") def _remove_extra_output_folder(self): - bad_path = os.path.join(self.cmor_path, 'output') + bad_path = os.path.join(self.cmor_path, "output") if os.path.exists(bad_path): - Log.debug('Moving CMOR files out of the output folder') + Log.debug("Moving CMOR files out of the output folder") Utils.move_tree(bad_path, self.cmor_path) - Log.debug('Done') + Log.debug("Done") class MergeYear(Diagnostic): @@ -449,13 +650,25 @@ class MergeYear(Diagnostic): @classmethod def generate_jobs(cls, diags, options): """ - Method to generate the required diagnostics from a section of the configuration file + Method to generate the required diagnostics from config - Required by the interface, does nothing as this diagnostic is not meant to be configured in the usal way + Required by the interface, does nothing as this diagnostic + is not meant to be configured in the usal way """ pass - def __init__(self, data_manager, domain, var, startdate, member, year, grid=None, box=None, frequency=None): + def __init__( + self, + data_manager, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + frequency=None, + ): super(MergeYear, self).__init__(data_manager) self.chunk_files = [] self.experiment = self.data_manager.experiment @@ -469,27 +682,59 @@ class MergeYear(Diagnostic): self.frequency = frequency def __str__(self): - return 'Merge year data Variable: {0.domain}:{0.var} Startdate: {0.startdate} Member: {0.member} ' \ - 'Year: {0.year} Grid: {0.grid} Box: {0.box} Frequency: {0.frequency}'.format(self) + return ( + "Merge year data Variable: {0.domain}:{0.var} " + "Startdate: {0.startdate} Member: {0.member} " + "Year: {0.year} Grid: {0.grid} Box: {0.box} " + "Frequency: {0.frequency}".format( + self + ) + ) def __eq__(self, other): - return self.domain == other.domain and self.var == other.var and self.startdate == other.startdate and \ - self.member == other.member and self.year == other.year and self.grid == other.grid and \ - self.box == other.box and self.frequency == other.frequency + return ( + self.domain == other.domain + and self.var == other.var + and self.startdate == other.startdate + and self.member == other.member + and self.year == other.year + and self.grid == other.grid + and self.box == other.box + and self.frequency == other.frequency + ) def __hash__(self): return hash(str(self)) def request_data(self): """Request all the data required by the diagnostic""" - for chunk in self.experiment.get_year_chunks(self.startdate, self.year): - self.chunk_files.append(self.request_chunk(self.domain, self.var, self.startdate, self.member, chunk, - grid=self.grid, box=self.box, frequency=self.frequency)) + for chunk in self.experiment.get_year_chunks( + self.startdate, self.year + ): + self.chunk_files.append( + self.request_chunk( + self.domain, + self.var, + self.startdate, + self.member, + chunk, + grid=self.grid, + box=self.box, + frequency=self.frequency, + ) + ) def declare_data_generated(self): """Declare all the data generated by the diagnostic""" - self.year_file = self.declare_year(self.domain, self.var, self.startdate, self.member, self.year, - grid=self.grid, box=self.box) + self.year_file = self.declare_year( + self.domain, + self.var, + self.startdate, + self.member, + self.year, + grid=self.grid, + box=self.box, + ) self.year_file.storage_status = StorageStatus.NO_STORE def compute(self): @@ -520,7 +765,11 @@ class MergeYear(Diagnostic): x += 1 if last_index is None: last_index = times.size - Utils.nco().ncks(input=data_file, output=temp2, options=['-d time,{0},{1}'.format(first_index, last_index - 1)]) + Utils.nco().ncks( + input=data_file, + output=temp2, + options=["-d time,{0},{1}".format(first_index, last_index - 1)], + ) return temp2 def _merge_chunk_files(self): @@ -529,7 +778,7 @@ class MergeYear(Diagnostic): Utils.copy_file(self.chunk_files[0].local_file, temp) return temp - Utils.nco().ncrcat(input=' '.join(self.chunk_files), output=temp) + Utils.nco().ncrcat(input=" ".join(self.chunk_files), output=temp) for chunk_file in self.chunk_files: os.remove(chunk_file) return temp diff --git a/earthdiagnostics/config.py b/earthdiagnostics/config.py index cc628298..d6480355 100644 --- a/earthdiagnostics/config.py +++ b/earthdiagnostics/config.py @@ -4,15 +4,28 @@ import os import six from bscearth.utils.config_parser import ConfigParser -from bscearth.utils.date import parse_date, chunk_start_date, chunk_end_date, date2str, add_years, add_months, add_days +from bscearth.utils.date import ( + parse_date, + chunk_start_date, + chunk_end_date, + date2str, + add_years, + add_months, + add_days, +) from bscearth.utils.log import Log import bscearth.utils.path from earthdiagnostics.frequency import Frequency, Frequencies from earthdiagnostics.modelingrealm import ModelingRealm from earthdiagnostics.variable import VariableManager -from earthdiagnostics.data_convention import SPECSConvention, CMIP6Convention, PrimaveraConvention, \ - MeteoFranceConvention, PrefaceConvention +from earthdiagnostics.data_convention import ( + SPECSConvention, + CMIP6Convention, + PrimaveraConvention, + MeteoFranceConvention, + PrefaceConvention, +) class ConfigException(Exception): @@ -65,7 +78,7 @@ class Config(object): self.parallel_uploads = None "Maximum number of simultaneous uploads" self.restore_meshes = None - "If True, forces the tool to copy all the mesh and mask files for the model, regardless of existence" + "If True, force the copy of all the mesh and mask files for the model" # Read experiment config self.experiment = ExperimentConfig() @@ -111,60 +124,102 @@ class Config(object): """ config_file_path = bscearth.utils.path.expand_path(path) if not os.path.isfile(config_file_path): - Log.critical('Configuration file {0} can not be found', config_file_path) - raise ValueError('Configuration file {0} can not be found'.format(config_file_path)) + Log.critical( + "Configuration file {0} can not be found", config_file_path + ) + raise ValueError( + "Configuration file {0} can not be found".format( + config_file_path + ) + ) parser = ConfigParser() parser.optionxform = str parser.read(config_file_path) # Read diags config - self.data_adaptor = parser.get_choice_option('DIAGNOSTICS', 'DATA_ADAPTOR', ('CMOR', 'THREDDS', 'OBSRECON'), - 'CMOR') - self.scratch_dir = parser.get_path_option('DIAGNOSTICS', 'SCRATCH_DIR') - self.use_ramdisk = parser.get_bool_option('DIAGNOSTICS', 'USE_RAMDISK', False) - self.auto_clean = parser.get_bool_option('DIAGNOSTICS', 'AUTO_CLEAN', True) + self.data_adaptor = parser.get_choice_option( + "DIAGNOSTICS", + "DATA_ADAPTOR", + ("CMOR", "THREDDS", "OBSRECON"), + "CMOR", + ) + self.scratch_dir = parser.get_path_option("DIAGNOSTICS", "SCRATCH_DIR") + self.use_ramdisk = parser.get_bool_option( + "DIAGNOSTICS", "USE_RAMDISK", False + ) + self.auto_clean = parser.get_bool_option( + "DIAGNOSTICS", "AUTO_CLEAN", True + ) if not self.auto_clean and self.use_ramdisk: - Log.warning('RAM disk scratch dir is always automatically cleaned.') + Log.warning( + "RAM disk scratch dir is always automatically cleaned." + ) self.auto_clean = True - self.scratch_masks = parser.get_path_option('DIAGNOSTICS', 'SCRATCH_MASKS', '/scratch/Earth/ocean_masks') - self.data_dir = parser.get_path_option('DIAGNOSTICS', 'DATA_DIR') - self.data_type = parser.get_choice_option('DIAGNOSTICS', 'DATA_TYPE', ('exp', 'obs', 'recon'), 'exp') - self.con_files = parser.get_path_option('DIAGNOSTICS', 'CON_FILES') - self.mesh_mask = parser.get_path_option('DIAGNOSTICS', 'MESH_MASK', '') - self.new_mask_glo = parser.get_path_option('DIAGNOSTICS', 'NEW_MASK_GLO', '') - self.mask_regions = parser.get_path_option('DIAGNOSTICS', 'MASK_REGIONS', '') - self.mask_regions_3d = parser.get_path_option('DIAGNOSTICS', 'MASK_REGIONS_3D', '') + self.scratch_masks = parser.get_path_option( + "DIAGNOSTICS", "SCRATCH_MASKS", "/scratch/Earth/ocean_masks" + ) + self.data_dir = parser.get_path_option("DIAGNOSTICS", "DATA_DIR") + self.data_type = parser.get_choice_option( + "DIAGNOSTICS", "DATA_TYPE", ("exp", "obs", "recon"), "exp" + ) + self.con_files = parser.get_path_option("DIAGNOSTICS", "CON_FILES") + self.mesh_mask = parser.get_path_option("DIAGNOSTICS", "MESH_MASK", "") + self.new_mask_glo = parser.get_path_option( + "DIAGNOSTICS", "NEW_MASK_GLO", "" + ) + self.mask_regions = parser.get_path_option( + "DIAGNOSTICS", "MASK_REGIONS", "" + ) + self.mask_regions_3d = parser.get_path_option( + "DIAGNOSTICS", "MASK_REGIONS_3D", "" + ) self._parse_dataconvention(parser) self.var_manager = VariableManager() self.var_manager.load_variables(self.data_convention.name) - self._diags = parser.get_option('DIAGNOSTICS', 'DIAGS') - self.skip_diags_done = parser.get_bool_option('DIAGNOSTICS', 'SKIP_DIAGS_DONE', True) - self.frequency = Frequency(parser.get_option('DIAGNOSTICS', 'FREQUENCY')) - - self.cdftools_path = parser.get_path_option('DIAGNOSTICS', 'CDFTOOLS_PATH', '') - self.max_cores = parser.get_int_option('DIAGNOSTICS', 'MAX_CORES', 0) - self.parallel_downloads = parser.get_int_option('DIAGNOSTICS', 'PARALLEL_DOWNLOADS', 1) - self.parallel_uploads = parser.get_int_option('DIAGNOSTICS', 'PARALLEL_UPLOADS', 1) - self.restore_meshes = parser.get_bool_option('DIAGNOSTICS', 'RESTORE_MESHES', False) + self._diags = parser.get_option("DIAGNOSTICS", "DIAGS") + self.skip_diags_done = parser.get_bool_option( + "DIAGNOSTICS", "SKIP_DIAGS_DONE", True + ) + self.frequency = Frequency( + parser.get_option("DIAGNOSTICS", "FREQUENCY") + ) + + self.cdftools_path = parser.get_path_option( + "DIAGNOSTICS", "CDFTOOLS_PATH", "" + ) + self.max_cores = parser.get_int_option("DIAGNOSTICS", "MAX_CORES", 0) + self.parallel_downloads = parser.get_int_option( + "DIAGNOSTICS", "PARALLEL_DOWNLOADS", 1 + ) + self.parallel_uploads = parser.get_int_option( + "DIAGNOSTICS", "PARALLEL_UPLOADS", 1 + ) + self.restore_meshes = parser.get_bool_option( + "DIAGNOSTICS", "RESTORE_MESHES", False + ) # Read experiment config self.experiment = ExperimentConfig() self.experiment.parse_ini(parser) # Read aliases self._aliases = dict() - if parser.has_section('ALIAS'): - for option in parser.options('ALIAS'): - self._aliases[option.lower()] = parser.get_list_option('ALIAS', option) - Log.debug('Preparing command list') + if parser.has_section("ALIAS"): + for option in parser.options("ALIAS"): + self._aliases[option.lower()] = parser.get_list_option( + "ALIAS", option + ) + Log.debug("Preparing command list") commands = self._diags.split() self._real_commands = list() self._apply_aliases(commands) - Log.debug('Command list ready ') + Log.debug("Command list ready ") - self.scratch_dir = os.path.join(self.scratch_dir, 'diags', self.experiment.expid) + self.scratch_dir = os.path.join( + self.scratch_dir, "diags", self.experiment.expid + ) self.cmor = CMORConfig(parser, self.var_manager) self.thredds = THREDDSConfig(parser) @@ -173,36 +228,47 @@ class Config(object): def _apply_aliases(self, commands): for command in commands: command = command.strip() - if command.startswith('#'): + if command.startswith("#"): break if command.lower() in self._aliases: added_commands = self._aliases[command.lower()] - Log.info('Changing alias {0} for {1}', command, ' '.join(added_commands)) + Log.info( + "Changing alias {0} for {1}", + command, + " ".join(added_commands), + ) for add_command in added_commands: self._real_commands.append(add_command) else: self._real_commands.append(command) def _parse_dataconvention(self, parser): - data_convention = parser.get_choice_option('DIAGNOSTICS', 'DATA_CONVENTION', - ('specs', 'primavera', 'cmip6', 'preface', 'meteofrance'), - 'specs', - ignore_case=True) - if data_convention == 'specs': + data_convention = parser.get_choice_option( + "DIAGNOSTICS", + "DATA_CONVENTION", + ("specs", "primavera", "cmip6", "preface", "meteofrance"), + "specs", + ignore_case=True, + ) + if data_convention == "specs": self.data_convention = SPECSConvention(data_convention, self) - elif data_convention == 'primavera': + elif data_convention == "primavera": self.data_convention = PrimaveraConvention(data_convention, self) - elif data_convention == 'cmip6': + elif data_convention == "cmip6": self.data_convention = CMIP6Convention(data_convention, self) - elif data_convention == 'preface': + elif data_convention == "preface": self.data_convention = PrefaceConvention(data_convention, self) - elif data_convention == 'meteofrance': + elif data_convention == "meteofrance": self.data_convention = MeteoFranceConvention(data_convention, self) - self.scratch_masks = self.data_convention.get_scratch_masks(self.scratch_masks) - namelist_file = os.path.join(os.path.dirname(__file__), - 'CDFTOOLS_{0}.namlist'.format(self.data_convention.name)) - Log.debug('Setting namelist {0}', namelist_file) - os.environ['NAM_CDF_NAMES'] = namelist_file + self.scratch_masks = self.data_convention.get_scratch_masks( + self.scratch_masks + ) + namelist_file = os.path.join( + os.path.dirname(__file__), + "CDFTOOLS_{0}.namlist".format(self.data_convention.name), + ) + Log.debug("Setting namelist {0}", namelist_file) + os.environ["NAM_CDF_NAMES"] = namelist_file def get_commands(self): """ @@ -226,58 +292,106 @@ class CMORConfig(object): """ def __init__(self, parser, var_manager): - self.force = parser.get_bool_option('CMOR', 'FORCE', False) - self.force_untar = parser.get_bool_option('CMOR', 'FORCE_UNTAR', False) - self.skip_prepare = parser.get_bool_option('CMOR', 'SKIP_PREPARE', False) - self.filter_files = parser.get_list_option('CMOR', 'FILTER_FILES', []) - self.ocean = parser.get_bool_option('CMOR', 'OCEAN_FILES', True) - self.atmosphere = parser.get_bool_option('CMOR', 'ATMOSPHERE_FILES', True) - self.use_grib = parser.get_bool_option('CMOR', 'USE_GRIB', True) - self._chunks = parser.get_int_list_option('CMOR', 'CHUNKS') - self.associated_experiment = parser.get_option('CMOR', 'ASSOCIATED_EXPERIMENT', 'to be filled') - self.associated_model = parser.get_option('CMOR', 'ASSOCIATED_MODEL', 'to be filled') - self.initialization_description = parser.get_option('CMOR', 'INITIALIZATION_DESCRIPTION', 'to be filled') - self.initialization_method = parser.get_option('CMOR', 'INITIALIZATION_METHOD', '1') - self.initialization_number = parser.get_int_option('CMOR', 'INITIALIZATION_NUMBER', 1) - self.physics_description = parser.get_option('CMOR', 'PHYSICS_DESCRIPTION', 'to be filled') - self.physics_version = parser.get_option('CMOR', 'PHYSICS_VERSION', '1') - self.source = parser.get_option('CMOR', 'SOURCE', 'to be filled') - self.version = parser.get_option('CMOR', 'VERSION', '') - self.default_ocean_grid = parser.get_option('CMOR', 'DEFAULT_OCEAN_GRID', 'gn') - self.default_atmos_grid = parser.get_option('CMOR', 'DEFAULT_ATMOS_GRID', 'gr') - self.activity = parser.get_option('CMOR', 'ACTIVITY', 'CMIP') - self.min_cmorized_vars = parser.get_int_option('CMOR', 'MIN_CMORIZED_VARS', 10) - self.append_startdate = parser.get_bool_option('CMOR', 'APPEND_STARTDATE', False) - self.append_startdate_year_only = parser.get_bool_option('CMOR', 'APPEND_STARTDATE_YEAR_ONLY', False) - - vars_string = parser.get_option('CMOR', 'VARIABLE_LIST', '') + self.force = parser.get_bool_option("CMOR", "FORCE", False) + self.force_untar = parser.get_bool_option("CMOR", "FORCE_UNTAR", False) + self.skip_prepare = parser.get_bool_option( + "CMOR", "SKIP_PREPARE", False + ) + self.filter_files = parser.get_list_option("CMOR", "FILTER_FILES", []) + self.ocean = parser.get_bool_option("CMOR", "OCEAN_FILES", True) + self.atmosphere = parser.get_bool_option( + "CMOR", "ATMOSPHERE_FILES", True + ) + self.use_grib = parser.get_bool_option("CMOR", "USE_GRIB", True) + self._chunks = parser.get_int_list_option("CMOR", "CHUNKS") + self.associated_experiment = parser.get_option( + "CMOR", "ASSOCIATED_EXPERIMENT", "to be filled" + ) + self.associated_model = parser.get_option( + "CMOR", "ASSOCIATED_MODEL", "to be filled" + ) + self.initialization_description = parser.get_option( + "CMOR", "INITIALIZATION_DESCRIPTION", "to be filled" + ) + self.initialization_method = parser.get_option( + "CMOR", "INITIALIZATION_METHOD", "1" + ) + self.initialization_number = parser.get_int_option( + "CMOR", "INITIALIZATION_NUMBER", 1 + ) + self.physics_description = parser.get_option( + "CMOR", "PHYSICS_DESCRIPTION", "to be filled" + ) + self.physics_version = parser.get_option( + "CMOR", "PHYSICS_VERSION", "1" + ) + self.source = parser.get_option("CMOR", "SOURCE", "to be filled") + self.version = parser.get_option("CMOR", "VERSION", "") + self.default_ocean_grid = parser.get_option( + "CMOR", "DEFAULT_OCEAN_GRID", "gn" + ) + self.default_atmos_grid = parser.get_option( + "CMOR", "DEFAULT_ATMOS_GRID", "gr" + ) + self.activity = parser.get_option("CMOR", "ACTIVITY", "CMIP") + self.min_cmorized_vars = parser.get_int_option( + "CMOR", "MIN_CMORIZED_VARS", 10 + ) + self.append_startdate = parser.get_bool_option( + "CMOR", "APPEND_STARTDATE", False + ) + self.append_startdate_year_only = parser.get_bool_option( + "CMOR", "APPEND_STARTDATE_YEAR_ONLY", False + ) + + vars_string = parser.get_option("CMOR", "VARIABLE_LIST", "") self.var_manager = var_manager if vars_string: self._variable_list = list() - for domain_var in vars_string.split(' '): + for domain_var in vars_string.split(" "): domain_var = domain_var.strip() - if domain_var.startswith('#'): + if domain_var.startswith("#"): break if not domain_var: continue - splitted = domain_var.split(':') - cmor_var = self.var_manager.get_variable(splitted[1], silent=True) + splitted = domain_var.split(":") + cmor_var = self.var_manager.get_variable( + splitted[1], silent=True + ) if not cmor_var: - Log.warning('Variable {0} not recognized. It will not be cmorized', domain_var) + Log.warning( + "Variable {0} not recognized. It will not be cmorized", + domain_var, + ) continue if ModelingRealm(splitted[0]) != cmor_var.domain: - Log.warning('Domain {0} for variable {1} is not correct: is {2}', splitted[0], splitted[1], - cmor_var.domain) + Log.warning( + "Domain {0} for variable {1} is not correct: is {2}", + splitted[0], + splitted[1], + cmor_var.domain, + ) continue - self._variable_list.append('{0.domain}:{0.short_name}'.format(cmor_var)) + self._variable_list.append( + "{0.domain}:{0.short_name}".format(cmor_var) + ) if len(self._variable_list) == 0: - raise ConfigException('Variable list value is specified, but no variables were found') + raise ConfigException( + "Variable list value is specified, but no variables were " + "found" + ) else: self._variable_list = None - self._var_hourly = CMORConfig._parse_variables(parser.get_option('CMOR', 'ATMOS_HOURLY_VARS', '')) - self._var_daily = CMORConfig._parse_variables(parser.get_option('CMOR', 'ATMOS_DAILY_VARS', '')) - self._var_monthly = CMORConfig._parse_variables(parser.get_option('CMOR', 'ATMOS_MONTHLY_VARS', '')) + self._var_hourly = CMORConfig._parse_variables( + parser.get_option("CMOR", "ATMOS_HOURLY_VARS", "") + ) + self._var_daily = CMORConfig._parse_variables( + parser.get_option("CMOR", "ATMOS_DAILY_VARS", "") + ) + self._var_monthly = CMORConfig._parse_variables( + parser.get_option("CMOR", "ATMOS_MONTHLY_VARS", "") + ) def cmorize(self, var_cmor): """ @@ -292,7 +406,10 @@ class CMORConfig(object): return True if not var_cmor: return False - return '{0}:{1}'.format(var_cmor.domain, var_cmor.short_name) in self._variable_list + return ( + "{0}:{1}".format(var_cmor.domain, var_cmor.short_name) + in self._variable_list + ) def any_required(self, variables): """ @@ -336,20 +453,22 @@ class CMORConfig(object): def _parse_variables(raw_string): variables = dict() if raw_string: - splitted = raw_string.split(',') + splitted = raw_string.split(",") for var_section in splitted: - splitted_var = var_section.split(':') + splitted_var = var_section.split(":") if len(splitted_var) == 1: levels = None else: - levels = ','.join(map(str, CMORConfig._parse_levels(splitted_var[1:]))) + levels = ",".join( + map(str, CMORConfig._parse_levels(splitted_var[1:])) + ) variables[int(splitted_var[0])] = levels return variables @staticmethod def _parse_levels(levels_splitted): if len(levels_splitted) == 1: - return map(int, levels_splitted[0].split('-')) + return map(int, levels_splitted[0].split("-")) start = int(levels_splitted[0]) end = int(levels_splitted[1]) if len(levels_splitted) == 3: @@ -382,7 +501,7 @@ class CMORConfig(object): return self._var_daily elif frequency == Frequencies.monthly: return self._var_monthly - raise ValueError('Frequency not recognized: {0}'.format(frequency)) + raise ValueError("Frequency not recognized: {0}".format(frequency)) def get_requested_codes(self): """ @@ -393,7 +512,11 @@ class CMORConfig(object): set of int """ - return set(list(self._var_hourly.keys()) + list(self._var_daily.keys()) + list(self._var_monthly.keys())) + return set( + list(self._var_hourly.keys()) + + list(self._var_daily.keys()) + + list(self._var_monthly.keys()) + ) def get_levels(self, frequency, variable): """ @@ -423,7 +546,7 @@ class THREDDSConfig(object): """ def __init__(self, parser): - self.server_url = parser.get_option('THREDDS', 'SERVER_URL', '') + self.server_url = parser.get_option("THREDDS", "SERVER_URL", "") class ExperimentConfig(object): @@ -441,37 +564,54 @@ class ExperimentConfig(object): parser: ConfigParser """ - self.institute = parser.get_option('EXPERIMENT', 'INSTITUTE') - self.expid = parser.get_option('EXPERIMENT', 'EXPID') - self.experiment_name = parser.get_option('EXPERIMENT', 'NAME', self.expid) - self.members = parser.get_list_option('EXPERIMENT', 'MEMBERS') - self.member_digits = parser.get_int_option('EXPERIMENT', 'MEMBER_DIGITS', 1) - self.member_prefix = parser.get_option('EXPERIMENT', 'MEMBER_PREFIX', 'fc') - self.member_count_start = parser.get_int_option('EXPERIMENT', 'MEMBER_COUNT_START', 0) + self.institute = parser.get_option("EXPERIMENT", "INSTITUTE") + self.expid = parser.get_option("EXPERIMENT", "EXPID") + self.experiment_name = parser.get_option( + "EXPERIMENT", "NAME", self.expid + ) + self.members = parser.get_list_option("EXPERIMENT", "MEMBERS") + self.member_digits = parser.get_int_option( + "EXPERIMENT", "MEMBER_DIGITS", 1 + ) + self.member_prefix = parser.get_option( + "EXPERIMENT", "MEMBER_PREFIX", "fc" + ) + self.member_count_start = parser.get_int_option( + "EXPERIMENT", "MEMBER_COUNT_START", 0 + ) self._parse_members() - self.calendar = parser.get_option('EXPERIMENT', 'CALENDAR', 'standard') + self.calendar = parser.get_option("EXPERIMENT", "CALENDAR", "standard") self._parse_startdates(parser) - self.chunk_size = parser.get_int_option('EXPERIMENT', 'CHUNK_SIZE') - self.num_chunks = parser.get_int_option('EXPERIMENT', 'CHUNKS') - self.chunk_list = parser.get_int_list_option('EXPERIMENT', 'CHUNK_LIST', []) - - self.model = parser.get_option('EXPERIMENT', 'MODEL') - self.model_version = parser.get_option('EXPERIMENT', 'MODEL_VERSION', '') - self.atmos_grid = parser.get_option('EXPERIMENT', 'ATMOS_GRID', '') - self.atmos_timestep = parser.get_int_option('EXPERIMENT', 'ATMOS_TIMESTEP', 6) - self.ocean_timestep = parser.get_int_option('EXPERIMENT', 'OCEAN_TIMESTEP', 6) + self.chunk_size = parser.get_int_option("EXPERIMENT", "CHUNK_SIZE") + self.num_chunks = parser.get_int_option("EXPERIMENT", "CHUNKS") + self.chunk_list = parser.get_int_list_option( + "EXPERIMENT", "CHUNK_LIST", [] + ) + + self.model = parser.get_option("EXPERIMENT", "MODEL") + self.model_version = parser.get_option( + "EXPERIMENT", "MODEL_VERSION", "" + ) + self.atmos_grid = parser.get_option("EXPERIMENT", "ATMOS_GRID", "") + self.atmos_timestep = parser.get_int_option( + "EXPERIMENT", "ATMOS_TIMESTEP", 6 + ) + self.ocean_timestep = parser.get_int_option( + "EXPERIMENT", "OCEAN_TIMESTEP", 6 + ) def _parse_startdates(self, parser): - startdates = parser.get_list_option('EXPERIMENT', 'STARTDATES') + startdates = parser.get_list_option("EXPERIMENT", "STARTDATES") import exrex + self.startdates = [] for startdate_pattern in startdates: startdate_pattern = startdate_pattern.strip() if not startdate_pattern: continue - if startdate_pattern[0] == '{' and startdate_pattern[-1] == '}': + if startdate_pattern[0] == "{" and startdate_pattern[-1] == "}": self._read_startdates(startdate_pattern[1:-1]) else: for startdate in exrex.generate(startdate_pattern): @@ -481,8 +621,8 @@ class ExperimentConfig(object): def _parse_members(self): members = [] for mem in self.members: - if '-' in mem: - start, end = mem.split('-') + if "-" in mem: + start, end = mem.split("-") if start.startswith(self.member_prefix): start = start[len(self.member_prefix):] if end.startswith(self.member_prefix): @@ -496,7 +636,7 @@ class ExperimentConfig(object): self.members = members def _read_startdates(self, pattern): - pattern = pattern.split(',') + pattern = pattern.split(",") start = parse_date(pattern[0].strip()) end = parse_date(pattern[1].strip()) interval = pattern[2].strip() @@ -507,16 +647,20 @@ class ExperimentConfig(object): interval = interval[-1].upper() while start <= end: self.startdates.append(date2str(start)) - if interval == 'Y': + if interval == "Y": start = add_years(start, factor) - elif interval == 'M': + elif interval == "M": start = add_months(start, factor, cal=self.calendar) - elif interval == 'W': + elif interval == "W": start = add_days(start, factor * 7, cal=self.calendar) - elif interval == 'D': + elif interval == "D": start = add_days(start, factor, cal=self.calendar) else: - raise ConfigException('Interval {0} not supported in STARTDATES definition: {1}', interval, pattern) + raise ConfigException( + "Interval {0} not supported in STARTDATES definition: {1}", + interval, + pattern, + ) def get_chunk_list(self): """ @@ -566,8 +710,13 @@ class ExperimentConfig(object): chunk_start = self.get_chunk_start(date, chunk) if chunk_start.year > year: break - elif chunk_start.year == year or chunk_end_date(chunk_start, self.chunk_size, 'month', - self.calendar).year == year: + elif ( + chunk_start.year == year + or chunk_end_date( + chunk_start, self.chunk_size, "month", self.calendar + ).year + == year + ): chunks.append(chunk) return chunks @@ -593,7 +742,9 @@ class ExperimentConfig(object): # noinspection PyTypeChecker if isinstance(startdate, six.string_types): startdate = parse_date(startdate) - return chunk_start_date(startdate, chunk, self.chunk_size, 'month', self.calendar) + return chunk_start_date( + startdate, chunk, self.chunk_size, "month", self.calendar + ) def get_chunk_start_str(self, startdate, chunk): """ @@ -633,7 +784,12 @@ class ExperimentConfig(object): get_chunk_end_str """ - return chunk_end_date(self.get_chunk_start(startdate, chunk), self.chunk_size, 'month', self.calendar) + return chunk_end_date( + self.get_chunk_start(startdate, chunk), + self.chunk_size, + "month", + self.calendar, + ) def get_chunk_end_str(self, startdate, chunk): """ @@ -692,7 +848,9 @@ class ExperimentConfig(object): :rtype: str """ - return '{0}{1}'.format(self.member_prefix, str(member).zfill(self.member_digits)) + return "{0}{1}".format( + self.member_prefix, str(member).zfill(self.member_digits) + ) class ReportConfig(object): @@ -706,5 +864,7 @@ class ReportConfig(object): """ def __init__(self, parser): - self.maximum_priority = parser.get_int_option('REPORT', 'MAXIMUM_PRIORITY', 10) - self.path = parser.get_path_option('REPORT', 'PATH', '') + self.maximum_priority = parser.get_int_option( + "REPORT", "MAXIMUM_PRIORITY", 10 + ) + self.path = parser.get_path_option("REPORT", "PATH", "") diff --git a/earthdiagnostics/constants.py b/earthdiagnostics/constants.py index bb8378e0..ac883796 100644 --- a/earthdiagnostics/constants.py +++ b/earthdiagnostics/constants.py @@ -68,86 +68,90 @@ class Basins(object): def __initialize(self): self.aliases = { - 'Global': ('Global', 'glob'), - 'Atlantic_Ocean': ('atl', 'atlantic'), - 'North_Atlantic_Ocean': ('natl', 'north_atlantic'), - 'Tropical_Atlantic_Ocean': ('tatl', 'tropical_atlantic'), - - 'Pacific_Ocean': ('pac', 'pacific'), - 'North_Pacific_Ocean': ('npac', 'north_pacific'), - 'Tropical_Pacific_Ocean': ('tpac', 'tropical_pacific'), - 'Indo_Pacific_Ocean': ('indpac', 'indo_pacific'), - - 'Indian_Ocean': ('ind', 'indian'), - 'Tropical_Indian_Ocean': ('tind', 'tropical_indian'), - - 'Antarctic_Ocean': ('anta', 'antarctiv'), - 'Antarctic_Atlantic_Sector': ('antaatl', 'antarctic_atlantic'), - 'Antarctic_Indian_Sector': ('antaind', 'antarctic_indian'), - - 'Arctic_Ocean': ('arct', 'arctic'), - 'Arctic_Ocean_North_Atlantic': ('arctnatl', 'artic_north_atlantic'), - 'Arctic_Marginal_Seas': ('arctmarg', 'arctic_marginal'), - - 'Baffin': ('Baffin',), - 'Baffin_Bay': ('BaffBay',), - 'Baltic_Sea': ('Baltic',), - 'BarKara': ('BarKara',), - 'Barents_Sea': ('Barents',), - 'Beaufort_Chukchi_Sea': ('BeaufortChukchi',), - 'Beaufort_Sea': ('Beaufort',), - 'Bellingshausen_Sea': ('Bellingshausen',), - 'Bering': ('Bering',), - 'Bering_Strait': ('BeringStr',), - 'CanArch': ('CanArch',), - 'Canadian_Waters': ('Canadian',), - 'Caspian_Sea': ('Caspian',), - 'Central_Arctic': ('CArct',), - 'Chukchi_Sea': ('Chukchi',), - 'East_Siberian_Sea': ('ESiberian',), - 'Eastern_Central_Arctic': ('ECArct',), - 'Fram_Strait': ('Fram',), - 'Greenland_Sea': ('Greenland',), - 'Grnland': ('Grnland',), - 'Hudson': ('Hudson',), - 'Icelandic_Sea': ('Iceland',), - 'Irminger_Sea': ('Irminger',), - 'Kara_Gate_Strait': ('KaraGate',), - 'Kara_Sea': ('Kara',), - 'Labrador_Sea': ('Labrador',), - 'Laptev_East_Siberian_Chukchi_Seas': ('LaptevESiberianChukchi',), - 'Laptev_East_Siberian_Seas': ('LaptevESiberian',), - 'Laptev_Sea': ('Laptev',), - 'Lincoln_Sea': ('Lincoln',), - 'Mediterranean_Sea': ('Medit',), - 'Nares_Strait': ('Nares',), - 'Nordic_Barents_Seas': ('NordicBarents',), - 'Nordic_Seas': ('Nordic',), - 'NorthWest_Passage': ('NWPass',), - 'North_Atlantic-Arctic': ('North_Atlantic-Arctic',), - 'North_Hemisphere_Ocean': ('NHem',), - 'Norwegian_Sea': ('Norwegian',), - 'Okhotsk': ('Okhotsk',), - 'OpenOcean': ('OpenOcean',), - 'Ross_Sea': ('Ross',), - 'Serreze_Arctic': ('SerArc',), - 'Southern_Hemisphere': ('SHem',), - 'StLawr': ('StLawr',), - 'Subpolar_Gyre': ('Subpolar_Gyre',), - 'TotalArc': ('TotalArc',), - 'Vilkitsky_Strait': ('Vilkitsky_Strait',), - 'Weddell_Sea': ('Weddell_Sea',), - 'Western_Central_Arctic': ('Western_Central_Arctic',), + "Global": ("Global", "glob"), + "Atlantic_Ocean": ("atl", "atlantic"), + "North_Atlantic_Ocean": ("natl", "north_atlantic"), + "Tropical_Atlantic_Ocean": ("tatl", "tropical_atlantic"), + "Pacific_Ocean": ("pac", "pacific"), + "North_Pacific_Ocean": ("npac", "north_pacific"), + "Tropical_Pacific_Ocean": ("tpac", "tropical_pacific"), + "Indo_Pacific_Ocean": ("indpac", "indo_pacific"), + "Indian_Ocean": ("ind", "indian"), + "Tropical_Indian_Ocean": ("tind", "tropical_indian"), + "Antarctic_Ocean": ("anta", "antarctiv"), + "Antarctic_Atlantic_Sector": ("antaatl", "antarctic_atlantic"), + "Antarctic_Indian_Sector": ("antaind", "antarctic_indian"), + "Arctic_Ocean": ("arct", "arctic"), + "Arctic_Ocean_North_Atlantic": ( + "arctnatl", + "artic_north_atlantic", + ), + "Arctic_Marginal_Seas": ("arctmarg", "arctic_marginal"), + "Baffin": ("Baffin",), + "Baffin_Bay": ("BaffBay",), + "Baltic_Sea": ("Baltic",), + "BarKara": ("BarKara",), + "Barents_Sea": ("Barents",), + "Beaufort_Chukchi_Sea": ("BeaufortChukchi",), + "Beaufort_Sea": ("Beaufort",), + "Bellingshausen_Sea": ("Bellingshausen",), + "Bering": ("Bering",), + "Bering_Strait": ("BeringStr",), + "CanArch": ("CanArch",), + "Canadian_Waters": ("Canadian",), + "Caspian_Sea": ("Caspian",), + "Central_Arctic": ("CArct",), + "Chukchi_Sea": ("Chukchi",), + "East_Siberian_Sea": ("ESiberian",), + "Eastern_Central_Arctic": ("ECArct",), + "Fram_Strait": ("Fram",), + "Greenland_Sea": ("Greenland",), + "Grnland": ("Grnland",), + "Hudson": ("Hudson",), + "Icelandic_Sea": ("Iceland",), + "Irminger_Sea": ("Irminger",), + "Kara_Gate_Strait": ("KaraGate",), + "Kara_Sea": ("Kara",), + "Labrador_Sea": ("Labrador",), + "Laptev_East_Siberian_Chukchi_Seas": ("LaptevESiberianChukchi",), + "Laptev_East_Siberian_Seas": ("LaptevESiberian",), + "Laptev_Sea": ("Laptev",), + "Lincoln_Sea": ("Lincoln",), + "Mediterranean_Sea": ("Medit",), + "Nares_Strait": ("Nares",), + "Nordic_Barents_Seas": ("NordicBarents",), + "Nordic_Seas": ("Nordic",), + "NorthWest_Passage": ("NWPass",), + "North_Atlantic-Arctic": ("North_Atlantic-Arctic",), + "North_Hemisphere_Ocean": ("NHem",), + "Norwegian_Sea": ("Norwegian",), + "Okhotsk": ("Okhotsk",), + "OpenOcean": ("OpenOcean",), + "Ross_Sea": ("Ross",), + "Serreze_Arctic": ("SerArc",), + "Southern_Hemisphere": ("SHem",), + "StLawr": ("StLawr",), + "Subpolar_Gyre": ("Subpolar_Gyre",), + "TotalArc": ("TotalArc",), + "Vilkitsky_Strait": ("Vilkitsky_Strait",), + "Weddell_Sea": ("Weddell_Sea",), + "Western_Central_Arctic": ("Western_Central_Arctic",), } - self.Global = Basin('Global') - self.Atlantic = Basin('Atlantic_Ocean') - self.Pacific = Basin('Pacific_Ocean') - self.IndoPacific = Basin('Indo_Pacific_Ocean') - self.Indian = Basin('Indian_Ocean') + self.Global = Basin("Global") + self.Atlantic = Basin("Atlantic_Ocean") + self.Pacific = Basin("Pacific_Ocean") + self.IndoPacific = Basin("Indo_Pacific_Ocean") + self.Indian = Basin("Indian_Ocean") self._known_aliases = {} - self._add_alias('glob', self.Global) - for basin in (self.Global, self.Atlantic, self.Pacific, self.IndoPacific, self.Indian): + self._add_alias("glob", self.Global) + for basin in ( + self.Global, + self.Atlantic, + self.Pacific, + self.IndoPacific, + self.Indian, + ): for alias in self.aliases[basin.name]: self._add_alias(alias, basin) @@ -161,7 +165,7 @@ class Basins(object): :type handler: netCDF4.Dataset """ - for basin in cube.coord('region').points: + for basin in cube.coord("region").points: basin_object = Basin(basin) setattr(self, basin, basin_object) self._add_alias(basin, basin_object) @@ -180,7 +184,8 @@ class Basins(object): If the parameter basin is a Basin instance, directly returns the same instance. This bahaviour is intended to facilitate the development of - methods that can either accept a nameor a Basin instance to characterize the basin. + methods that can either accept a nameor a Basin instance to + characterize the basin. :param basin: basin name or basin instance :type basin: str | Basin @@ -199,29 +204,29 @@ class Basins(object): class Models(object): """Predefined models""" - ECEARTH_2_3_O1L42 = 'Ec2.3_O1L42' + ECEARTH_2_3_O1L42 = "Ec2.3_O1L42" """ EC-Earth 2.3 ORCA1 L42""" - ECEARTH_3_0_O1L46 = 'Ec3.0_O1L46' + ECEARTH_3_0_O1L46 = "Ec3.0_O1L46" """ EC-Earth 3 ORCA1 L46 """ - ECEARTH_3_0_O25L46 = 'Ec3.0_O25L46' + ECEARTH_3_0_O25L46 = "Ec3.0_O25L46" """ EC-Earth 3 ORCA0.25 L46 """ - ECEARTH_3_0_O25L75 = 'Ec3.0_O25L75' + ECEARTH_3_0_O25L75 = "Ec3.0_O25L75" """ EC-Earth 3 ORCA0.25 L75 """ - ECEARTH_3_1_O25L75 = 'Ec3.1_O25L75' + ECEARTH_3_1_O25L75 = "Ec3.1_O25L75" """ EC-Earth 3.1 ORCA0.25 L75 """ - ECEARTH_3_2_O1L75 = 'Ec3.2_O1L75' + ECEARTH_3_2_O1L75 = "Ec3.2_O1L75" """ EC-Earth 3.2 ORCA1 L75 """ - ECEARTH_3_2_O25L75 = 'Ec3.2_O25L75' + ECEARTH_3_2_O25L75 = "Ec3.2_O25L75" """ EC-Earth 3.2 ORCA0.25 L75 """ - NEMO_3_2_O1L42 = 'N3.2_O1L42' + NEMO_3_2_O1L42 = "N3.2_O1L42" """ NEMO 3.2 ORCA1 L42 """ - NEMO_3_3_O1L46 = 'N3.3_O1L46' + NEMO_3_3_O1L46 = "N3.3_O1L46" """ NEMO 3.3 ORCA1 L46 """ - NEMO_3_6_O1L46 = 'N3.6_O1L75' + NEMO_3_6_O1L46 = "N3.6_O1L75" """ NEMO 3.6 ORCA1 L75 """ - NEMOVAR_O1L42 = 'nemovar_O1L42' + NEMOVAR_O1L42 = "nemovar_O1L42" """ NEMOVAR ORCA1 L42 """ - GLORYS2_V1_O25L75 = 'glorys2v1_O25L75' + GLORYS2_V1_O25L75 = "glorys2v1_O25L75" """ GLORYS2v1 ORCA0.25 L75 """ diff --git a/earthdiagnostics/data_convention.py b/earthdiagnostics/data_convention.py index 1c54d35c..fc1f76bf 100644 --- a/earthdiagnostics/data_convention.py +++ b/earthdiagnostics/data_convention.py @@ -1,10 +1,16 @@ -"""Module to manage the different data conventions supported by EarthDiagnostics""" +"""Module to manage the data conventions supported by EarthDiagnostics""" import os import shutil import re import threading -from bscearth.utils.date import parse_date, chunk_start_date, chunk_end_date, previous_day, add_hours +from bscearth.utils.date import ( + parse_date, + chunk_start_date, + chunk_end_date, + previous_day, + add_hours, +) from bscearth.utils.log import Log from earthdiagnostics.frequency import Frequency, Frequencies @@ -19,9 +25,9 @@ class DataConvention(object): def __init__(self, name, config): self.config = config self.name = name - self.lat_name = 'lat' - self.lon_name = 'lon' - self.time_separator = '-' + self.lat_name = "lat" + self.lon_name = "lon" + self.time_separator = "-" self.lock = threading.Lock() self._checked_vars = list() @@ -40,8 +46,19 @@ class DataConvention(object): """ return scratch_masks - def get_file_path(self, startdate, member, domain, var, cmor_var, chunk, frequency, - grid=None, year=None, date_str=None): + def get_file_path( + self, + startdate, + member, + domain, + var, + cmor_var, + chunk, + frequency, + grid=None, + year=None, + date_str=None, + ): """ Return the path to a concrete file @@ -65,20 +82,45 @@ class DataConvention(object): Raises ------ ValueError - If you provide two or more parameters from chunk, year or date_str or none at all + If you provide two or more parameters from chunk, year or + date_str or none at all """ if frequency is None: frequency = self.config.frequency frequency = Frequency.parse(frequency) - folder_path = self.get_cmor_folder_path(startdate, member, domain, var, frequency, grid, cmor_var) - file_name = self.get_file_name(startdate, member, domain, var, cmor_var, frequency, - chunk, year, date_str, grid) + folder_path = self.get_cmor_folder_path( + startdate, member, domain, var, frequency, grid, cmor_var + ) + file_name = self.get_file_name( + startdate, + member, + domain, + var, + cmor_var, + frequency, + chunk, + year, + date_str, + grid, + ) filepath = os.path.join(folder_path, file_name) return filepath - def get_file_name(self, startdate, member, domain, var, cmor_var, frequency, chunk, year, date_str, grid, ): + def get_file_name( + self, + startdate, + member, + domain, + var, + cmor_var, + frequency, + chunk, + year, + date_str, + grid, + ): """ Get filename for a given configuration @@ -107,7 +149,9 @@ class DataConvention(object): """ raise NotImplementedError - def get_cmor_folder_path(self, startdate, member, domain, var, frequency, grid, cmor_var): + def get_cmor_folder_path( + self, startdate, member, domain, var, frequency, grid, cmor_var + ): """ Get the folder path following current CMOR convention @@ -146,9 +190,15 @@ class DataConvention(object): str """ - return os.path.join(self.config.data_dir, self.config.experiment.expid, 'cmorfiles', self.config.cmor.activity, - self.config.experiment.institute, self.config.experiment.model, - self.experiment_name(startdate)) + return os.path.join( + self.config.data_dir, + self.config.experiment.expid, + "cmorfiles", + self.config.cmor.activity, + self.config.experiment.institute, + self.config.experiment.model, + self.experiment_name(startdate), + ) def experiment_name(self, startdate): """ @@ -164,7 +214,9 @@ class DataConvention(object): """ if self.config.cmor.append_startdate: - return '{}S{}'.format(self.config.experiment.experiment_name, startdate) + return "{}S{}".format( + self.config.experiment.experiment_name, startdate + ) else: return self.config.experiment.experiment_name @@ -202,15 +254,17 @@ class DataConvention(object): member_str = self.get_member_str(member) else: member_str = None - Log.info('Creating links for CMOR files ({0})', startdate) + Log.info("Creating links for CMOR files ({0})", startdate) path = self.get_startdate_path(startdate) self._link_startdate(path, member_str) - Log.debug('Links ready') + Log.debug("Links ready") def _link_startdate(self, path, member_str): raise NotImplementedError - def create_link(self, domain, filepath, frequency, var, grid, move_old, vartype): + def create_link( + self, domain, filepath, frequency, var, grid, move_old, vartype + ): """ Create file link @@ -228,24 +282,43 @@ class DataConvention(object): freq_str = frequency.folder_name(vartype) if not grid: - grid = 'original' - variable_folder = domain.get_varfolder(var, self.config.experiment.ocean_timestep, - self.config.experiment.atmos_timestep) - vargrid_folder = domain.get_varfolder(var, self.config.experiment.ocean_timestep, - self.config.experiment.atmos_timestep, grid=grid) + grid = "original" + variable_folder = domain.get_varfolder( + var, + self.config.experiment.ocean_timestep, + self.config.experiment.atmos_timestep, + ) + vargrid_folder = domain.get_varfolder( + var, + self.config.experiment.ocean_timestep, + self.config.experiment.atmos_timestep, + grid=grid, + ) self.lock.acquire() try: expid = self.config.experiment.expid - if grid == 'original': - link_path = os.path.join(self.config.data_dir, expid, freq_str, variable_folder) + if grid == "original": + link_path = os.path.join( + self.config.data_dir, expid, freq_str, variable_folder + ) Utils.create_folder_tree(link_path) else: - link_path = os.path.join(self.config.data_dir, expid, freq_str, vargrid_folder) + link_path = os.path.join( + self.config.data_dir, expid, freq_str, vargrid_folder + ) Utils.create_folder_tree(link_path) - default_path = os.path.join(self.config.data_dir, expid, freq_str, variable_folder) - original_path = os.path.join(self.config.data_dir, expid, freq_str, - vargrid_folder.replace('-{0}_f'.format(grid), '-original_f')) + default_path = os.path.join( + self.config.data_dir, expid, freq_str, variable_folder + ) + original_path = os.path.join( + self.config.data_dir, + expid, + freq_str, + vargrid_folder.replace( + "-{0}_f".format(grid), "-original_f" + ), + ) if os.path.islink(default_path): os.remove(default_path) @@ -255,15 +328,19 @@ class DataConvention(object): if move_old and link_path not in self._checked_vars: self._checked_vars.append(link_path) - old_path = os.path.join(self.config.data_dir, expid, freq_str, - 'old_{0}'.format(os.path.basename(link_path))) - regex = re.compile(var + '_[0-9]{6,8}[.]nc') + old_path = os.path.join( + self.config.data_dir, + expid, + freq_str, + "old_{0}".format(os.path.basename(link_path)), + ) + regex = re.compile(var + "_[0-9]{6,8}[.]nc") for filename in os.listdir(link_path): if regex.match(filename): Utils.create_folder_tree(old_path) Utils.move_file( os.path.join(link_path, filename), - os.path.join(old_path, filename) + os.path.join(old_path, filename), ) link_path = os.path.join(link_path, os.path.basename(filepath)) @@ -273,8 +350,12 @@ class DataConvention(object): except OSError: pass if not os.path.exists(filepath): - raise ValueError('Original file {0} does not exists'.format(filepath)) - relative_path = os.path.relpath(filepath, os.path.dirname(link_path)) + raise ValueError( + "Original file {0} does not exists".format(filepath) + ) + relative_path = os.path.relpath( + filepath, os.path.dirname(link_path) + ) try: os.symlink(relative_path, link_path) except OSError: @@ -286,34 +367,60 @@ class DataConvention(object): def _get_time_component(self, chunk, date_str, frequency, startdate, year): if len([x for x in (chunk, date_str, year) if x is not None]) > 1: - raise ValueError('Only one of the parameters chunk, year or date_str may be provided') + raise ValueError( + "Only one of the parameters chunk, year or date_str may be " + "provided" + ) if chunk is not None: - time_bound = self._get_chunk_time_bounds(startdate, chunk, frequency) + time_bound = self._get_chunk_time_bounds( + startdate, chunk, frequency + ) elif year: if frequency != Frequencies.yearly: - raise ValueError('Year may be provided instead of chunk only if frequency is "yr"') + raise ValueError( + 'Year may be provided instead of chunk only if ' + 'frequency is "yr"' + ) time_bound = str(year) elif date_str: time_bound = date_str else: - raise ValueError('Time info not provided') + raise ValueError("Time info not provided") return time_bound def _get_chunk_time_bounds(self, startdate, chunk, frequency): start = parse_date(startdate) - chunk_start = chunk_start_date(start, chunk, self.config.experiment.chunk_size, 'month', - self.config.experiment.calendar) - chunk_end = chunk_end_date(chunk_start, self.config.experiment.chunk_size, 'month', - self.config.experiment.calendar) + chunk_start = chunk_start_date( + start, + chunk, + self.config.experiment.chunk_size, + "month", + self.config.experiment.calendar, + ) + chunk_end = chunk_end_date( + chunk_start, + self.config.experiment.chunk_size, + "month", + self.config.experiment.calendar, + ) chunk_end = previous_day(chunk_end, self.config.experiment.calendar) - time_bound = "{0:04}{1:02}{4}{2:04}{3:02}".format(chunk_start.year, chunk_start.month, chunk_end.year, - chunk_end.month, self.time_separator) + time_bound = "{0:04}{1:02}{4}{2:04}{3:02}".format( + chunk_start.year, + chunk_start.month, + chunk_end.year, + chunk_end.month, + self.time_separator, + ) return time_bound - def _check_var_presence(self, folder, current_count, startdate, member, domain, chunk, freq): + def _check_var_presence( + self, folder, current_count, startdate, member, domain, chunk, freq + ): for var in os.listdir(folder): cmor_var = self.config.var_manager.get_variable(var, True) - var_path = self.get_file_path(startdate, member, domain, var, cmor_var, chunk, frequency=freq) + var_path = self.get_file_path( + startdate, member, domain, var, cmor_var, chunk, frequency=freq + ) if os.path.isfile(var_path): current_count += 1 if current_count >= self.config.cmor.min_cmorized_vars: @@ -347,7 +454,19 @@ class DataConvention(object): class Cmor2Convention(DataConvention): """Base class for CMOR2-based conventions""" - def get_file_name(self, startdate, member, domain, var, cmor_var, frequency, chunk, year, date_str, grid, ): + def get_file_name( + self, + startdate, + member, + domain, + var, + cmor_var, + frequency, + chunk, + year, + date_str, + grid, + ): """ Get filename for a given configuration @@ -370,21 +489,40 @@ class Cmor2Convention(DataConvention): """ if cmor_var is None: - cmor_table = domain.get_table(frequency, self.config.data_convention) + cmor_table = domain.get_table( + frequency, self.config.data_convention + ) else: - cmor_table = cmor_var.get_table(frequency, self.config.data_convention) - - time_bound = self._get_time_component(chunk, date_str, frequency, startdate, year) - time_bound = '_{0}.nc'.format(time_bound) - - file_name = '{0}_{1}_{2}_{3}_S{4}_{5}{6}'.format(var, cmor_table.name, self.config.experiment.model, - self.experiment_name(startdate), startdate, - self.get_member_str(member), time_bound) + cmor_table = cmor_var.get_table( + frequency, self.config.data_convention + ) + + time_bound = self._get_time_component( + chunk, date_str, frequency, startdate, year + ) + time_bound = "_{0}.nc".format(time_bound) + + file_name = "{0}_{1}_{2}_{3}_S{4}_{5}{6}".format( + var, + cmor_table.name, + self.config.experiment.model, + self.experiment_name(startdate), + startdate, + self.get_member_str(member), + time_bound, + ) return file_name - def get_cmor_folder_path(self, startdate, member, domain, var, frequency, grid, cmor_var): - folder_path = os.path.join(self.get_startdate_path(startdate), str(frequency), domain.name, var) - if grid and grid != 'original': + def get_cmor_folder_path( + self, startdate, member, domain, var, frequency, grid, cmor_var + ): + folder_path = os.path.join( + self.get_startdate_path(startdate), + str(frequency), + domain.name, + var, + ) + if grid and grid != "original": folder_path = os.path.join(folder_path, grid) folder_path = os.path.join(folder_path, self.get_member_str(member)) if self.config.cmor.version: @@ -409,29 +547,51 @@ class Cmor2Convention(DataConvention): If not implemented by derived classes """ - template = 'r{0}i{1}p1' - return template.format(member + 1 - self.config.experiment.member_count_start, - self.config.cmor.initialization_number) + template = "r{0}i{1}p1" + return template.format( + member + 1 - self.config.experiment.member_count_start, + self.config.cmor.initialization_number, + ) def _link_startdate(self, path, member_str): for freq in os.listdir(path): - Log.debug('Creating links for frequency {0}', freq) + Log.debug("Creating links for frequency {0}", freq) frequency = Frequency.parse(freq) for domain in os.listdir(os.path.join(path, freq)): - Log.debug('Creating links for domain {0}', domain) + Log.debug("Creating links for domain {0}", domain) for var in os.listdir(os.path.join(path, freq, domain)): - for member in os.listdir(os.path.join(path, freq, domain, var)): + for member in os.listdir( + os.path.join(path, freq, domain, var) + ): if member_str is not None and member_str != member: continue - for name in os.listdir(os.path.join(path, freq, domain, var, member)): - filepath = os.path.join(path, freq, domain, var, member, name) + for name in os.listdir( + os.path.join(path, freq, domain, var, member) + ): + filepath = os.path.join( + path, freq, domain, var, member, name + ) if os.path.isfile(filepath): - self.create_link(ModelingRealms.parse(domain), filepath, frequency, var, "", False, - vartype=VariableType.MEAN) + self.create_link( + ModelingRealms.parse(domain), + filepath, + frequency, + var, + "", + False, + vartype=VariableType.MEAN, + ) else: for filename in os.listdir(filepath): - self.create_link(ModelingRealms.parse(domain), os.path.join(filepath, filename), - frequency, var, "", False, vartype=VariableType.MEAN) + self.create_link( + ModelingRealms.parse(domain), + os.path.join(filepath, filename), + frequency, + var, + "", + False, + vartype=VariableType.MEAN, + ) def is_cmorized(self, startdate, member, chunk, domain): """ @@ -456,7 +616,9 @@ class Cmor2Convention(DataConvention): for freq in os.listdir(startdate_path): domain_path = os.path.join(startdate_path, freq, domain.name) if os.path.isdir(domain_path): - count = self._check_var_presence(domain_path, count, startdate, member, domain, chunk, freq) + count = self._check_var_presence( + domain_path, count, startdate, member, domain, chunk, freq + ) if count >= self.config.cmor.min_cmorized_vars: return True return False @@ -478,9 +640,15 @@ class SPECSConvention(Cmor2Convention): str """ - return os.path.join(self.config.data_dir, self.config.experiment.expid, 'cmorfiles', - self.config.experiment.institute, - self.config.experiment.model, self.experiment_name(startdate), 'S' + startdate) + return os.path.join( + self.config.data_dir, + self.config.experiment.expid, + "cmorfiles", + self.config.experiment.institute, + self.config.experiment.model, + self.experiment_name(startdate), + "S" + startdate, + ) class PrefaceConvention(Cmor2Convention): @@ -496,7 +664,7 @@ class PrefaceConvention(Cmor2Convention): def __init__(self, name, config): super(PrefaceConvention, self).__init__(name, config) - self.time_separator = '_' + self.time_separator = "_" def get_startdate_path(self, startdate): """ @@ -511,9 +679,14 @@ class PrefaceConvention(Cmor2Convention): str """ - return os.path.join(self.config.data_dir, self.config.experiment.expid, 'cmorfiles', - self.config.experiment.institute, - self.experiment_name(startdate), 'S' + startdate) + return os.path.join( + self.config.data_dir, + self.config.experiment.expid, + "cmorfiles", + self.config.experiment.institute, + self.experiment_name(startdate), + "S" + startdate, + ) class Cmor3Convention(DataConvention): @@ -529,8 +702,8 @@ class Cmor3Convention(DataConvention): def __init__(self, name, config): super(Cmor3Convention, self).__init__(name, config) - self.lat_name = 'latitude' - self.lon_name = 'longitude' + self.lat_name = "latitude" + self.lon_name = "longitude" def get_scratch_masks(self, scratch_masks): """ @@ -549,7 +722,19 @@ class Cmor3Convention(DataConvention): """ return os.path.join(scratch_masks, self.name) - def get_file_name(self, startdate, member, domain, var, cmor_var, frequency, chunk, year, date_str, grid, ): + def get_file_name( + self, + startdate, + member, + domain, + var, + cmor_var, + frequency, + chunk, + year, + date_str, + grid, + ): """ Get filename for a given configuration @@ -572,102 +757,176 @@ class Cmor3Convention(DataConvention): """ if cmor_var is None: - cmor_table = domain.get_table(frequency, self.config.data_convention) + cmor_table = domain.get_table( + frequency, self.config.data_convention + ) else: - cmor_table = cmor_var.get_table(frequency, self.config.data_convention) + cmor_table = cmor_var.get_table( + frequency, self.config.data_convention + ) - time_bound = self._get_time_component(chunk, date_str, frequency, startdate, year) - time_bound = '_{0}.nc'.format(time_bound) + time_bound = self._get_time_component( + chunk, date_str, frequency, startdate, year + ) + time_bound = "_{0}.nc".format(time_bound) if not grid: - if domain in [ModelingRealms.ocnBgchem, ModelingRealms.seaIce, ModelingRealms.ocean]: + if domain in [ + ModelingRealms.ocnBgchem, + ModelingRealms.seaIce, + ModelingRealms.ocean, + ]: grid = self.config.cmor.default_ocean_grid else: grid = self.config.cmor.default_atmos_grid if self.config.cmor.append_startdate: if self.config.cmor.append_startdate_year_only: startdate = startdate[0:4] - subexp_id = 's{}-'.format(startdate) + subexp_id = "s{}-".format(startdate) else: - subexp_id = '' - - file_name = '{var_name}_{table}_{model}_{exp_id}_{subexp_id}{member}_{grid}{time}'.format( - var_name=var, - table=cmor_table.name, - model=self.config.experiment.model, - exp_id=self.experiment_name(startdate), - subexp_id=subexp_id, - member=self.get_member_str(member), - grid=grid, - time=time_bound) + subexp_id = "" + + file_name = ( + f"{var}_{cmor_table.name}_{self.config.experiment.model}" + f"_{self.experiment_name(startdate)}_{subexp_id}" + f"{self.get_member_str(member)}_{grid}{time_bound}" + ) return file_name def _get_chunk_time_bounds(self, startdate, chunk, frequency): start = parse_date(startdate) - chunk_start = chunk_start_date(start, chunk, self.config.experiment.chunk_size, 'month', - self.config.experiment.calendar) - chunk_end = chunk_end_date(chunk_start, self.config.experiment.chunk_size, 'month', - self.config.experiment.calendar) + chunk_start = chunk_start_date( + start, + chunk, + self.config.experiment.chunk_size, + "month", + self.config.experiment.calendar, + ) + chunk_end = chunk_end_date( + chunk_start, + self.config.experiment.chunk_size, + "month", + self.config.experiment.calendar, + ) if frequency == Frequencies.monthly: - chunk_end = previous_day(chunk_end, self.config.experiment.calendar) - time_bound = "{0:04}{1:02}{4}{2:04}{3:02}".format(chunk_start.year, chunk_start.month, chunk_end.year, - chunk_end.month, self.time_separator) + chunk_end = previous_day( + chunk_end, self.config.experiment.calendar + ) + time_bound = "{0:04}{1:02}{4}{2:04}{3:02}".format( + chunk_start.year, + chunk_start.month, + chunk_end.year, + chunk_end.month, + self.time_separator, + ) elif frequency == Frequencies.daily: - chunk_end = previous_day(chunk_end, self.config.experiment.calendar) - time_bound = "{0.year:04}{0.month:02}{0.day:02}{2}" \ - "{1.year:04}{1.month:02}{1.day:02}".format(chunk_start, chunk_end, self.time_separator) - elif frequency.frequency.endswith('hr'): - chunk_end = add_hours(chunk_end, -int(frequency.frequency[:-2]), self.config.experiment.calendar) - time_bound = "{0.year:04}{0.month:02}{0.day:02}{0.hour:02}{0.minute:02}{2}" \ - "{1.year:04}{1.month:02}{1.day:02}{1.hour:02}{1.minute:02}".format(chunk_start, - chunk_end, - self.time_separator) + chunk_end = previous_day( + chunk_end, self.config.experiment.calendar + ) + time_bound = ( + "{0.year:04}{0.month:02}{0.day:02}{2}" + "{1.year:04}{1.month:02}{1.day:02}".format( + chunk_start, chunk_end, self.time_separator + ) + ) + elif frequency.frequency.endswith("hr"): + chunk_end = add_hours( + chunk_end, + -int(frequency.frequency[:-2]), + self.config.experiment.calendar, + ) + time_bound = ( + "{0.year:04}{0.month:02}{0.day:02}{0.hour:02}{0.minute:02}{2}" + "{1.year:04}{1.month:02}{1.day:02}{1.hour:02}" + "{1.minute:02}".format( + chunk_start, chunk_end, self.time_separator + ) + ) return time_bound - def get_cmor_folder_path(self, startdate, member, domain, var, frequency, grid, cmor_var): + def get_cmor_folder_path( + self, startdate, member, domain, var, frequency, grid, cmor_var + ): if not self.config.cmor.version: - raise ValueError('CMOR version is mandatory for PRIMAVERA and CMIP6') + raise ValueError( + "CMOR version is mandatory for PRIMAVERA and CMIP6" + ) if not grid: - if domain in [ModelingRealms.ocnBgchem, ModelingRealms.seaIce, ModelingRealms.ocean]: + if domain in [ + ModelingRealms.ocnBgchem, + ModelingRealms.seaIce, + ModelingRealms.ocean, + ]: grid = self.config.cmor.default_ocean_grid else: grid = self.config.cmor.default_atmos_grid if cmor_var is None: - table_name = domain.get_table(frequency, self.config.data_convention).name + table_name = domain.get_table( + frequency, self.config.data_convention + ).name else: - table_name = cmor_var.get_table(frequency, self.config.data_convention).name - folder_path = os.path.join(self.get_startdate_path(startdate), self.get_member_str(member), - table_name, var, - grid, self.config.cmor.version) + table_name = cmor_var.get_table( + frequency, self.config.data_convention + ).name + folder_path = os.path.join( + self.get_startdate_path(startdate), + self.get_member_str(member), + table_name, + var, + grid, + self.config.cmor.version, + ) return folder_path def _link_startdate(self, path, member_str): for member in os.listdir(path): for table in os.listdir(os.path.join(path, member)): frequency = self.config.var_manager.tables[table].frequency - Log.debug('Creating links for table {0}', table) + Log.debug("Creating links for table {0}", table) for var in os.listdir(os.path.join(path, member, table)): - cmor_var = self.config.var_manager.get_variable(var, silent=True) + cmor_var = self.config.var_manager.get_variable( + var, silent=True + ) domain = None if cmor_var is not None: domain = cmor_var.domain if domain is None: domain = self.config.var_manager.tables[table].domain - for grid in os.listdir(os.path.join(path, member, table, var)): + for grid in os.listdir( + os.path.join(path, member, table, var) + ): if member_str is not None and member_str != member: continue - for name in os.listdir(os.path.join(path, member, table, var, grid)): - filepath = os.path.join(path, member, table, var, grid, name) + for name in os.listdir( + os.path.join(path, member, table, var, grid) + ): + filepath = os.path.join( + path, member, table, var, grid, name + ) if os.path.isfile(filepath): - self.create_link(domain, filepath, frequency, var, "", False, - vartype=VariableType.MEAN) + self.create_link( + domain, + filepath, + frequency, + var, + "", + False, + vartype=VariableType.MEAN, + ) else: for filename in os.listdir(filepath): cmorfile = os.path.join(filepath, filename) - self.create_link(domain, cmorfile, frequency, var, "", - False, vartype=VariableType.MEAN) + self.create_link( + domain, + cmorfile, + frequency, + var, + "", + False, + vartype=VariableType.MEAN, + ) def experiment_name(self, startdate): """ @@ -699,9 +958,11 @@ class Cmor3Convention(DataConvention): str """ - template = 'r{0}i{1}p1f1' - return template.format(member + 1 - self.config.experiment.member_count_start, - self.config.cmor.initialization_number) + template = "r{0}i{1}p1f1" + return template.format( + member + 1 - self.config.experiment.member_count_start, + self.config.cmor.initialization_number, + ) def is_cmorized(self, startdate, member, chunk, domain): """ @@ -731,7 +992,9 @@ class Cmor3Convention(DataConvention): table_dir = os.path.join(member_path, table.name) if not os.path.isdir(table_dir): return False - count = self._check_var_presence(table_dir, count, startdate, member, domain, chunk, freq) + count = self._check_var_presence( + table_dir, count, startdate, member, domain, chunk, freq + ) if count >= self.config.cmor.min_cmorized_vars: return True return False @@ -752,7 +1015,19 @@ class PrimaveraConvention(Cmor3Convention): class MeteoFranceConvention(DataConvention): """Class managing MeteoFrance file conventions""" - def get_file_name(self, startdate, member, domain, var, cmor_var, frequency, chunk, year, date_str, grid,): + def get_file_name( + self, + startdate, + member, + domain, + var, + cmor_var, + frequency, + chunk, + year, + date_str, + grid, + ): """ Get filename for a given configuration @@ -775,19 +1050,30 @@ class MeteoFranceConvention(DataConvention): """ if year is not None: - raise ValueError('Year not supported with MeteoFrance convention') + raise ValueError("Year not supported with MeteoFrance convention") if date_str is not None: - raise ValueError('Date_str not supported with MeteoFrance convention') + raise ValueError( + "Date_str not supported with MeteoFrance convention" + ) if chunk is None: - raise ValueError('Chunk must be provided in MeteoFrance convention') + raise ValueError( + "Chunk must be provided in MeteoFrance convention" + ) time_bound = self._get_chunk_time_bounds(startdate, chunk, frequency) - file_name = '{0}_{1}_{2}_{3}.nc'.format(var, frequency, time_bound, self.get_member_str(member)) + file_name = "{0}_{1}_{2}_{3}.nc".format( + var, frequency, time_bound, self.get_member_str(member) + ) return file_name - def get_cmor_folder_path(self, startdate, member, domain, var, frequency, grid, cmor_var): - folder_path = os.path.join(self.config.data_dir, self.experiment_name(startdate), - 'H{0}'.format(chr(64 + int(startdate[4:6]))), - startdate[0:4]) + def get_cmor_folder_path( + self, startdate, member, domain, var, frequency, grid, cmor_var + ): + folder_path = os.path.join( + self.config.data_dir, + self.experiment_name(startdate), + "H{0}".format(chr(64 + int(startdate[4:6]))), + startdate[0:4], + ) return folder_path def get_member_str(self, member): @@ -803,16 +1089,23 @@ class MeteoFranceConvention(DataConvention): str """ - return '{0:02d}'.format(member) + return "{0:02d}".format(member) def _get_chunk_time_bounds(self, startdate, chunk, frequency): start = parse_date(startdate) - chunk_start = chunk_start_date(start, chunk, self.config.experiment.chunk_size, 'month', - self.config.experiment.calendar) + chunk_start = chunk_start_date( + start, + chunk, + self.config.experiment.chunk_size, + "month", + self.config.experiment.calendar, + ) time_bound = "{0:04}{1:02}".format(chunk_start.year, chunk_start.month) return time_bound - def create_link(self, domain, filepath, frequency, var, grid, move_old, vartype): + def create_link( + self, domain, filepath, frequency, var, grid, move_old, vartype + ): """ Create file link diff --git a/earthdiagnostics/datafile.py b/earthdiagnostics/datafile.py index accc2c49..e32c5cdf 100644 --- a/earthdiagnostics/datafile.py +++ b/earthdiagnostics/datafile.py @@ -67,7 +67,7 @@ class DataFile(Publisher): self._size = None def __str__(self): - return 'Data file for {0}'.format(self.remote_file) + return "Data file for {0}".format(self.remote_file) @property def size(self): @@ -85,13 +85,20 @@ class DataFile(Publisher): def clean_local(self): """Check if a local file is still needed and remove it if not""" - if self.local_status != LocalStatus.READY or self.suscribers or self.upload_required() or \ - self.storage_status == StorageStatus.UPLOADING: + if ( + self.local_status != LocalStatus.READY + or self.suscribers + or self.upload_required() + or self.storage_status == StorageStatus.UPLOADING + ): return - Log.debug('File {0} no longer needed. Deleting from scratch...'.format( - self.remote_file)) + Log.debug( + "File {0} no longer needed. Deleting from scratch...".format( + self.remote_file + ) + ) os.remove(self.local_file) - Log.debug('File {0} deleted from scratch'.format(self.remote_file)) + Log.debug("File {0} deleted from scratch".format(self.remote_file)) self.local_file = None self.local_status = LocalStatus.PENDING @@ -104,7 +111,10 @@ class DataFile(Publisher): bool """ - return self.local_status == LocalStatus.READY and self.storage_status == StorageStatus.PENDING + return ( + self.local_status == LocalStatus.READY + and self.storage_status == StorageStatus.PENDING + ) def download_required(self): """ @@ -129,8 +139,9 @@ class DataFile(Publisher): """ Register a diagnostic as a modifier of this data - A modifier diagnostic is a diagnostic that read this data and changes it in any way. - The diagnostic must be a modifier even if it only affects the metadata + A modifier diagnostic is a diagnostic that read this data and changes + it in any way. The diagnostic must be a modifier even if it only + affects the metadata Parameters ---------- @@ -154,7 +165,8 @@ class DataFile(Publisher): """ Check if the data is ready to run for a given diagnostics - To be ready to run, the datafile should be in the local storage and no modifiers can be pending. + To be ready to run, the datafile should be in the local storage + and no modifiers can be pending. Parameters ---------- @@ -209,7 +221,7 @@ class DataFile(Publisher): @classmethod def to_storage(cls, remote_file, data_convention): - """Create a new datafile object for a file that is going to be generated and stored""" + """Creates new datafile object for a file to be generated and stored""" new_object = cls() new_object.remote_file = remote_file new_object.storage_status = StorageStatus.PENDING @@ -228,13 +240,14 @@ class DataFile(Publisher): If the derived classes do not override this """ - raise NotImplementedError('Class must implement the download method') + raise NotImplementedError("Class must implement the download method") def prepare_to_upload(self, rename_var): """ Prepare a local file to be uploaded - This includes renaming the variable if necessary, updating the metadata and adding the history and + This includes renaming the variable if necessary, + updating the metadata and adding the history and managing the possibility of multiple regions """ if rename_var: @@ -243,7 +256,8 @@ class DataFile(Publisher): original_name = self.var if self.final_name != original_name: Utils.rename_variable( - self.local_file, original_name, self.final_name) + self.local_file, original_name, self.final_name + ) self._rename_coordinate_variables() self._correct_metadata() self._prepare_region() @@ -255,19 +269,19 @@ class DataFile(Publisher): self.storage_status = StorageStatus.UPLOADING remote_file = self.remote_file try: - if '/cmorfiles/' in remote_file: - remote_file = remote_file.replace('/cmorfiles/', '/diags/') + if "/cmorfiles/" in remote_file: + remote_file = remote_file.replace("/cmorfiles/", "/diags/") Utils.copy_file(self.local_file, remote_file, save_hash=True) except (OSError, Exception) as ex: - Log.error('File {0} can not be uploaded: {1}', remote_file, ex) + Log.error("File {0} can not be uploaded: {1}", remote_file, ex) self.storage_status = StorageStatus.FAILED return - Log.info('File {0} uploaded!', self.remote_file) + Log.info("File {0} uploaded!", self.remote_file) # self.create_link() self.storage_status = StorageStatus.READY - def set_local_file(self, local_file, diagnostic=None, rename_var=''): + def set_local_file(self, local_file, diagnostic=None, rename_var=""): """ Set the local file generated by EarthDiagnostics @@ -292,7 +306,7 @@ class DataFile(Publisher): self.local_status = LocalStatus.READY def create_link(self): - """Create a link from the original data in the _ folder""" + """Create a link from the original in _ folder""" pass def _correct_metadata(self): @@ -304,13 +318,14 @@ class DataFile(Publisher): self.data_convention.lat_name, self.data_convention.lon_name, 'leadtime', 'region', 'time_centered' }, - set(handler.variables.keys()) + set(handler.variables.keys()), ) var_handler.coordinates = Utils.convert_to_ascii_if_possible( - ' '.join(coords)) - if 'time_centered' in handler.variables: - if hasattr(handler.variables['time_centered'], 'standard_name'): - del handler.variables['time_centered'].standard_name + " ".join(coords) + ) + if "time_centered" in handler.variables: + if hasattr(handler.variables["time_centered"], "standard_name"): + del handler.variables["time_centered"].standard_name if not self.cmor_var: handler.close() return @@ -318,7 +333,7 @@ class DataFile(Publisher): self._fix_variable_name(var_handler) handler.modeling_realm = self.cmor_var.domain.name table = self.cmor_var.get_table(self.frequency, self.data_convention) - handler.table_id = 'Table {0} ({1})'.format(table.name, table.date) + handler.table_id = "Table {0} ({1})".format(table.name, table.date) if self.cmor_var.units: self._fix_units(var_handler) handler.sync() @@ -334,26 +349,29 @@ class DataFile(Publisher): def _fix_values_metadata(self, var_type, file_path=None): if file_path is None: file_path = self.local_file - valid_min = '' - valid_max = '' + valid_min = "" + valid_max = "" if self.cmor_var is not None: if self.cmor_var.valid_min: - valid_min = '-a valid_min,{0},o,{1},"{2}" '.format(self.final_name, var_type.char, - self.cmor_var.valid_min) + valid_min = '-a valid_min,{0},o,{1},"{2}" '.format( + self.final_name, var_type.char, self.cmor_var.valid_min + ) if self.cmor_var.valid_max: - valid_max = '-a valid_max,{0},o,{1},"{2}" '.format(self.final_name, var_type.char, - self.cmor_var.valid_max) + valid_max = '-a valid_max,{0},o,{1},"{2}" '.format( + self.final_name, var_type.char, self.cmor_var.valid_max + ) Utils.nco().ncatted( - input=file_path, output=file_path, + input=file_path, + output=file_path, options=( '-O -a _FillValue,{0},o,{1},"1.e20" ' '-a missingValue,{0},o,{1},"1.e20" {2}{3}'.format( - self.final_name, var_type.char, - valid_min, valid_max), - ) + self.final_name, var_type.char, valid_min, valid_max + ), + ), ) def _fix_coordinate_variables_metadata(self, handler): @@ -371,14 +389,14 @@ class DataFile(Publisher): handler.variables[lat_name].standard_name = 'latitude' def _fix_units(self, var_handler): - if 'units' not in var_handler.ncattrs(): + if "units" not in var_handler.ncattrs(): return - if var_handler.units == '-': - var_handler.units = '1.0' - if var_handler.units == 'PSU': - var_handler.units = 'psu' - if var_handler.units == 'C' and self.cmor_var.units == 'K': - var_handler.units = 'deg_C' + if var_handler.units == "-": + var_handler.units = "1.0" + if var_handler.units == "PSU": + var_handler.units = "psu" + if var_handler.units == "C" and self.cmor_var.units == "K": + var_handler.units = "deg_C" if self.cmor_var.units != var_handler.units: self._convert_units(var_handler) var_handler.units = self.cmor_var.units @@ -387,25 +405,33 @@ class DataFile(Publisher): try: Utils.convert_units(var_handler, self.cmor_var.units) except ValueError as ex: - Log.warning('Can not convert {3} from {0} to {1}: {2}', var_handler.units, self.cmor_var.units, ex, - self.cmor_var.short_name) - factor, offset = UnitConversion.get_conversion_factor_offset(var_handler.units, - self.cmor_var.units) + Log.warning( + "Can not convert {3} from {0} to {1}: {2}", + var_handler.units, + self.cmor_var.units, + ex, + self.cmor_var.short_name, + ) + factor, offset = UnitConversion.get_conversion_factor_offset( + var_handler.units, self.cmor_var.units + ) var_handler[:] = var_handler[:] * factor + offset - if 'valid_min' in var_handler.ncattrs(): - var_handler.valid_min = float( - var_handler.valid_min) * factor + offset - if 'valid_max' in var_handler.ncattrs(): - var_handler.valid_max = float( - var_handler.valid_max) * factor + offset + if "valid_min" in var_handler.ncattrs(): + var_handler.valid_min = ( + float(var_handler.valid_min) * factor + offset + ) + if "valid_max" in var_handler.ncattrs(): + var_handler.valid_max = ( + float(var_handler.valid_max) * factor + offset + ) def _prepare_region(self): if not self.check_is_in_storage(update_status=False): return cube = iris.load_cube(self.local_file) try: - cube.coord('region') + cube.coord("region") except iris.exceptions.CoordinateNotFoundError: return try: @@ -419,11 +445,11 @@ class DataFile(Publisher): # Bad data, overwrite return new_data = {} - for region_slice in cube.slices_over('region'): - Log.debug(region_slice.coord('region').points[0]) - new_data[region_slice.coord('region').points[0]] = region_slice - for region_slice in old_cube.slices_over('region'): - region = region_slice.coord('region').points[0] + for region_slice in cube.slices_over("region"): + Log.debug(region_slice.coord("region").points[0]) + new_data[region_slice.coord("region").points[0]] = region_slice + for region_slice in old_cube.slices_over("region"): + region = region_slice.coord("region").points[0] Log.debug(region) if region not in new_data: new_data[region] = region_slice @@ -458,24 +484,29 @@ class DataFile(Publisher): Utils.rename_variables(self.local_file, variables, False) def add_diagnostic_history(self): - """Add the history line corresponding to the diagnostic to the local file""" + """Add diagnostic history line to local file""" if not self.diagnostic: return from earthdiagnostics.earthdiags import EarthDiags - history_line = 'Diagnostic {1} calculated with EarthDiagnostics version {0}'.format(EarthDiags.version, - self.diagnostic) + + history_line = ( + f"Diagnostic {self.diagnostic} calculated with EarthDiagnostics " + f"version {EarthDiags.version}" + ) self._add_history_line(history_line) def add_cmorization_history(self): - """Add the history line corresponding to the cmorization to the local file""" + """Add cmorization history line to local file""" from earthdiagnostics.earthdiags import EarthDiags - history_line = 'CMORized with Earthdiagnostics version {0}'.format( - EarthDiags.version) + + history_line = "CMORized with Earthdiagnostics version {0}".format( + EarthDiags.version + ) self._add_history_line(history_line) def _add_history_line(self, history_line): - utc_datetime = 'UTC ' + datetime.utcnow().isoformat() - history_line = '{0}: {1};'.format(utc_datetime, history_line) + utc_datetime = "UTC " + datetime.utcnow().isoformat() + history_line = "{0}: {1};".format(utc_datetime, history_line) handler = Utils.open_cdf(self.local_file) try: @@ -511,13 +542,19 @@ class UnitConversion(object): def load_conversions(cls): """Load conversions from the configuration file""" cls._dict_conversions = dict() - with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'conversions.csv'), 'r') as csvfile: - reader = csv.reader(csvfile, dialect='excel') + with open( + os.path.join( + os.path.dirname(os.path.realpath(__file__)), "conversions.csv" + ), + "r", + ) as csvfile: + reader = csv.reader(csvfile, dialect="excel") for line in reader: - if line[0] == 'original': + if line[0] == "original": continue - cls.add_conversion(UnitConversion( - line[0], line[1], line[2], line[3])) + cls.add_conversion( + UnitConversion(line[0], line[1], line[2], line[3]) + ) @classmethod def add_conversion(cls, conversion): @@ -527,8 +564,9 @@ class UnitConversion(object): :param conversion: conversion to add :type conversion: UnitConversion """ - cls._dict_conversions[( - conversion.source, conversion.destiny)] = conversion + cls._dict_conversions[ + (conversion.source, conversion.destiny) + ] = conversion @classmethod def get_conversion_factor_offset(cls, input_units, output_units): @@ -550,8 +588,8 @@ class UnitConversion(object): scale_unit = 1 unit = units[0] else: - if '^' in units[0]: - values = units[0].split('^') + if "^" in units[0]: + values = units[0].split("^") scale_unit = pow(int(values[0]), int(values[1])) else: scale_unit = float(units[0]) @@ -562,8 +600,8 @@ class UnitConversion(object): scale_new_unit = 1 new_unit = units[0] else: - if '^' in units[0]: - values = units[0].split('^') + if "^" in units[0]: + values = units[0].split("^") scale_new_unit = pow(int(values[0]), int(values[1])) else: scale_new_unit = float(units[0]) @@ -600,43 +638,58 @@ class NetCDFFile(DataFile): def download(self): """Get data from remote storage to the local one""" - for path in (self.remote_file.replace('/cmorfiles/', '/diags/'), self.remote_file): + for path in ( + self.remote_file.replace("/cmorfiles/", "/diags/"), + self.remote_file, + ): if os.path.isfile(path): try: self.local_status = LocalStatus.DOWNLOADING - Log.debug('Downloading file {0}...', path) + Log.debug("Downloading file {0}...", path) if not self.local_file: self.local_file = TempFile.get() - # Utils.get_file_hash(self.remote_file, use_stored=True, save=True) try: Utils.copy_file(path, self.local_file, retrials=1) except Utils.CopyException: - # Utils.get_file_hash(self.remote_file, use_stored=False, save=True) Utils.copy_file(path, self.local_file, retrials=2) - if self.data_convention == 'meteofrance': - Log.debug('Converting variable names from meteofrance convention') + if self.data_convention == "meteofrance": + Log.debug( + "Converting variable names from " + "meteofrance convention" + ) alt_coord_names = { - 'time_counter': 'time', 'time_counter_bounds': 'time_bnds', - 'tbnds': 'bnds', 'nav_lat': 'lat', 'nav_lon': 'lon', - 'x': 'i', 'y': 'j'} - Utils.rename_variables(self.local_file, alt_coord_names, must_exist=False) - Log.info('File {0} ready!', path) + "time_counter": "time", + "time_counter_bounds": "time_bnds", + "tbnds": "bnds", + "nav_lat": "lat", + "nav_lon": "lon", + "x": "i", + "y": "j", + } + Utils.rename_variables( + self.local_file, alt_coord_names, must_exist=False + ) + Log.info("File {0} ready!", path) self.local_status = LocalStatus.READY return except Exception as ex: if os.path.isfile(self.local_file): os.remove(self.local_file) - Log.error('File {0} not available: {1}', path, ex) + Log.error("File {0} not available: {1}", path, ex) self.local_status = LocalStatus.FAILED return - Log.error('File {0} not available: {1}', - self.remote_file, 'FileNotFound') + Log.error( + "File {0} not available: {1}", self.remote_file, "FileNotFound" + ) self.local_status = LocalStatus.FAILED def check_is_in_storage(self, update_status=True): - for path in (self.remote_file, self.remote_file.replace('/cmorfiles/', '/diags/')): + for path in ( + self.remote_file, + self.remote_file.replace("/cmorfiles/", "/diags/"), + ): if os.path.isfile(path): if update_status: self.storage_status = StorageStatus.READY @@ -644,13 +697,21 @@ class NetCDFFile(DataFile): return False def create_link(self): - """Create a link from the original data in the _ folder""" + """Create a link from original data to _ folder""" try: - self.data_convention.create_link(self.domain, self.remote_file, self.frequency, self.final_name, - self.grid, True, self.var_type) + self.data_convention.create_link( + self.domain, + self.remote_file, + self.frequency, + self.final_name, + self.grid, + True, + self.var_type, + ) except (ValueError, Exception) as ex: - Log.error('Can not create link to {1}: {0}'.format( - ex, self.remote_file)) + Log.error( + "Can not create link to {1}: {0}".format(ex, self.remote_file) + ) def _get_size(self): try: diff --git a/earthdiagnostics/datamanager.py b/earthdiagnostics/datamanager.py index 4077c24e..c938154e 100644 --- a/earthdiagnostics/datamanager.py +++ b/earthdiagnostics/datamanager.py @@ -1,7 +1,12 @@ # coding: utf-8 """Base data manager for Earth diagnostics""" -from earthdiagnostics.datafile import NetCDFFile as NCfile, StorageStatus, LocalStatus, UnitConversion +from earthdiagnostics.datafile import ( + NetCDFFile as NCfile, + StorageStatus, + LocalStatus, + UnitConversion, +) from earthdiagnostics.variable import VariableType @@ -23,15 +28,30 @@ class DataManager(object): def _get_file_from_storage(self, filepath): if filepath not in self.requested_files: - self.requested_files[filepath] = NCfile.from_storage(filepath, self.config.data_convention) + self.requested_files[filepath] = NCfile.from_storage( + filepath, self.config.data_convention + ) file_object = self.requested_files[filepath] file_object.local_status = LocalStatus.PENDING return self.requested_files[filepath] - def _declare_generated_file(self, remote_file, domain, final_var, cmor_var, data_convention, - region, diagnostic, grid, var_type, original_var): + def _declare_generated_file( + self, + remote_file, + domain, + final_var, + cmor_var, + data_convention, + region, + diagnostic, + grid, + var_type, + original_var, + ): if remote_file not in self.requested_files: - self.requested_files[remote_file] = NCfile.to_storage(remote_file, data_convention) + self.requested_files[remote_file] = NCfile.to_storage( + remote_file, data_convention + ) file_object = self.requested_files[remote_file] file_object.diagnostic = diagnostic file_object.var_type = var_type @@ -51,15 +71,29 @@ class DataManager(object): var += box.get_lon_str() + box.get_lat_str() + box.get_depth_str() return var - def link_file(self, domain, var, cmor_var, startdate, member, chunk=None, grid=None, - frequency=None, year=None, date_str=None, move_old=False, vartype=VariableType.MEAN): + def link_file( + self, + domain, + var, + cmor_var, + startdate, + member, + chunk=None, + grid=None, + frequency=None, + year=None, + date_str=None, + move_old=False, + vartype=VariableType.MEAN, + ): """ Create the link of a given file from the CMOR repository. :param cmor_var: :param move_old: :param date_str: - :param year: if frequency is yearly, this parameter is used to give the corresponding year + :param year: if frequency is yearly, this parameter is used to give + the corresponding year :type year: int :param domain: CMOR domain :type domain: Domain @@ -73,7 +107,8 @@ class DataManager(object): :type chunk: int :param grid: file's grid (only needed if it is not the original) :type grid: str - :param frequency: file's frequency (only needed if it is different from the default) + :param frequency: file's frequency (only needed if it is different + from the default) :type frequency: str :param vartype: Variable type (mean, statistic) :type vartype: VariableType @@ -86,9 +121,20 @@ class DataManager(object): """Prepare the data to be used by Earth Diagnostics""" pass - def request_chunk(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, vartype=None): + def request_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=None, + ): """ - Request a given file from the CMOR repository to the scratch folder and returns the path to the scratch's copy + Request a given file from the CMOR repository Parameters ---------- @@ -112,9 +158,20 @@ class DataManager(object): If not implemented by derived classes """ - raise NotImplementedError('Class must override request_chunk method') - - def request_year(self, diagnostic, domain, var, startdate, member, year, grid=None, box=None, frequency=None): + raise NotImplementedError("Class must override request_chunk method") + + def request_year( + self, + diagnostic, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + frequency=None, + ): """ Request a given year for a variavle from a CMOR repository @@ -140,10 +197,22 @@ class DataManager(object): If not implemented by derived classes """ - raise NotImplementedError('Class must override request_year method') - - def declare_chunk(self, domain, var, startdate, member, chunk, grid=None, region=None, box=None, frequency=None, - vartype=VariableType.MEAN, diagnostic=None): + raise NotImplementedError("Class must override request_year method") + + def declare_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + region=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ Declare a variable chunk to be generated by a diagnostic @@ -171,10 +240,20 @@ class DataManager(object): If not implemented by derived classes """ - raise NotImplementedError('Class must override declare_chunk method') - - def declare_year(self, domain, var, startdate, member, year, grid=None, box=None, - vartype=VariableType.MEAN, diagnostic=None): + raise NotImplementedError("Class must override declare_chunk method") + + def declare_year( + self, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ Declare a variable year to be generated by a diagnostic @@ -200,10 +279,21 @@ class DataManager(object): If not implemented by derived classes """ - raise NotImplementedError('Class must override declare_year method') - - def file_exists(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - vartype=VariableType.MEAN, possible_versions=None): + raise NotImplementedError("Class must override declare_year method") + + def file_exists( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + possible_versions=None, + ): """ Check if a file exists in the storage @@ -230,4 +320,4 @@ class DataManager(object): bool """ - raise NotImplementedError('Class must override file_exists method') + raise NotImplementedError("Class must override file_exists method") diff --git a/earthdiagnostics/diagnostic.py b/earthdiagnostics/diagnostic.py index 7874afc8..2d461cfd 100644 --- a/earthdiagnostics/diagnostic.py +++ b/earthdiagnostics/diagnostic.py @@ -1,5 +1,7 @@ # coding=utf-8 -"""This module contains the Diagnostic base class and all the classes for parsing the options passed to them""" +""" +This module contains the Diagnostic base class and all option classes +""" import datetime from bscearth.utils.log import Log @@ -26,14 +28,16 @@ class Diagnostic(Publisher): """ Base class for the diagnostics. - Provides a common interface for them and also has a mechanism that allows diagnostic retrieval by name. + Provides a common interface for them and also has a mechanism that allows + diagnostic retrieval by name. - :param data_manager: data manager that will be used to store and retrieve the necessary data + :param data_manager: data manager that will be used to store and + retrieve the necessary data :type data_manager: DataManager """ alias = None - """ Alias to call the diagnostic. Must be overridden at the derived clases""" + """ Alias to call the diagnostic. Must be overridden""" _diag_list = dict() def __init__(self, data_manager): @@ -58,7 +62,8 @@ class Diagnostic(Publisher): """ Check if a diagnostic is different than other - Implementation is just the negation of the equal, that should be implemented by the derived classes + Implementation is just the negation of the equal, that should be + implemented by the derived classes Parameters ---------- @@ -79,7 +84,8 @@ class Diagnostic(Publisher): """ Check if a diagnostic calculation can be skipped - Looks if the data to be generated is already there and is not going to be modified + Looks if the data to be generated is already there and is not going + to be modified Returns ------- @@ -93,7 +99,11 @@ class Diagnostic(Publisher): if file_generated.storage_status != StorageStatus.READY: return False if file_generated.has_modifiers(): - Log.warning('Can not skip diagnostics run when data is going to be modified: {0}'.format(self)) + Log.warning( + "Can not skip diagnostics run when data is going to be " + "modified: {0}", + self + ) return False return True @@ -115,7 +125,10 @@ class Diagnostic(Publisher): if self.status == DiagnosticStatus.RUNNING: for generated_file in self._generated_files: generated_file.local_status = LocalStatus.COMPUTING - if self.status in (DiagnosticStatus.FAILED, DiagnosticStatus.COMPLETED): + if self.status in ( + DiagnosticStatus.FAILED, + DiagnosticStatus.COMPLETED, + ): self._unsuscribe_requests() self.dispatch(self, old_status) @@ -132,9 +145,17 @@ class Diagnostic(Publisher): """ if not issubclass(diagnostic_class, Diagnostic): - raise ValueError('Class {0} must be derived from Diagnostic'.format(diagnostic_class)) + raise ValueError( + "Class {0} must be derived from Diagnostic".format( + diagnostic_class + ) + ) if diagnostic_class.alias is None: - raise ValueError('Diagnostic class {0} must have defined an alias'.format(diagnostic_class)) + raise ValueError( + "Diagnostic class {0} must have defined an alias".format( + diagnostic_class + ) + ) Diagnostic._diag_list[diagnostic_class.alias] = diagnostic_class # noinspection PyProtectedMember @@ -178,10 +199,23 @@ class Diagnostic(Publisher): Must be implemented by derived classes """ - raise NotImplementedError("Class must override declare_data_generated method") - - def declare_chunk(self, domain, var, startdate, member, chunk, grid=None, region=None, box=None, frequency=None, - vartype=VariableType.MEAN): + raise NotImplementedError( + "Class must override declare_data_generated method" + ) + + def declare_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + region=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + ): """ Declare a chunk that is going to be generated by the diagnostic @@ -205,15 +239,35 @@ class Diagnostic(Publisher): """ if isinstance(region, Basin): region = region.name - generated_chunk = self.data_manager.declare_chunk(domain, var, startdate, member, chunk, grid, region, box, - diagnostic=self, vartype=vartype, frequency=frequency) + generated_chunk = self.data_manager.declare_chunk( + domain, + var, + startdate, + member, + chunk, + grid, + region, + box, + diagnostic=self, + vartype=vartype, + frequency=frequency, + ) # if region is not None: # generated_chunk.add_modifier(self) self._generated_files.append(generated_chunk) return generated_chunk - def declare_year(self, domain, var, startdate, member, year, grid=None, box=None, - vartype=VariableType.MEAN): + def declare_year( + self, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + vartype=VariableType.MEAN, + ): """ Declare a year that is going to be generated by the diagnostic @@ -233,15 +287,24 @@ class Diagnostic(Publisher): DataFile """ - generated_year = self.data_manager.declare_year(domain, var, startdate, member, year, grid, box, - diagnostic=self, vartype=vartype) + generated_year = self.data_manager.declare_year( + domain, + var, + startdate, + member, + year, + grid, + box, + diagnostic=self, + vartype=vartype, + ) self._generated_files.append(generated_year) return generated_year @classmethod def generate_jobs(cls, diags, options): """ - Generate the instances of the diagnostics that will be run by the manager + Generate diagnostics instance to be run by the manager Must be implemented by derived classes. @@ -255,7 +318,9 @@ class Diagnostic(Publisher): list of Diagnostic """ - raise NotImplementedError("Class must override generate_jobs class method") + raise NotImplementedError( + "Class must override generate_jobs class method" + ) @classmethod def process_options(cls, options, options_available): @@ -281,14 +346,18 @@ class Diagnostic(Publisher): processed = dict() options = options[1:] if len(options) > len(options_available): - raise DiagnosticOptionError('You have specified more options than available for diagnostic ' - '{0}'.format(cls.alias)) + raise DiagnosticOptionError( + "You have specified more options than available for diagnostic" + f" {cls.alias}" + ) for x, option_definition in enumerate(options_available): if len(options) <= x: - option_value = '' + option_value = "" else: option_value = options[x] - processed[option_definition.name] = option_definition.parse(option_value) + processed[option_definition.name] = option_definition.parse( + option_value + ) return processed def __str__(self): @@ -297,7 +366,7 @@ class Diagnostic(Publisher): Must be implemented by derived classesgit """ - return 'Developer must override base class __str__ method' + return "Developer must override base class __str__ method" def add_subjob(self, subjob): """ @@ -316,8 +385,19 @@ class Diagnostic(Publisher): def _subjob_status_changed(self, job, status): self.check_is_ready() - def request_chunk(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - to_modify=False, vartype=VariableType.MEAN): + def request_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + to_modify=False, + vartype=VariableType.MEAN, + ): """ Request one chunk of data required by the diagnostic @@ -332,8 +412,9 @@ class Diagnostic(Publisher): box: Box or None frequency: Frequency or str or None to_modify: bool - Flag that must be active if the diagnostic is going to generate a modified version of this data. In this - case this data must not be declared as an output of the diagnostic + Flag that must be active if the diagnostic is going to generate a + modified version of this data. In this case this data must not be + declared as an output of the diagnostic vartype: VariableType Returns @@ -347,14 +428,35 @@ class Diagnostic(Publisher): declare_year """ - request = self.data_manager.request_chunk(domain, var, startdate, member, chunk, grid, box, frequency, vartype) + request = self.data_manager.request_chunk( + domain, + var, + startdate, + member, + chunk, + grid, + box, + frequency, + vartype, + ) if to_modify: request.add_modifier(self) self._requests.append(request) request.subscribe(self, self._updated_request) return request - def request_year(self, domain, var, startdate, member, year, grid=None, box=None, frequency=None, to_modify=False): + def request_year( + self, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + frequency=None, + to_modify=False, + ): """ Request one year of data that is required for the diagnostic @@ -381,7 +483,9 @@ class Diagnostic(Publisher): declare_year """ - request = self.data_manager.request_year(self, domain, var, startdate, member, year, grid, box, frequency) + request = self.data_manager.request_year( + self, domain, var, startdate, member, year, grid, box, frequency + ) if to_modify: request.add_modifier(self) self._requests.append(request) @@ -392,7 +496,9 @@ class Diagnostic(Publisher): if self.status != DiagnosticStatus.WAITING: return if request.local_status == LocalStatus.FAILED: - self.message = 'Required file {0} is not available'.format(request.remote_file) + self.message = "Required file {0} is not available".format( + request.remote_file + ) self.status = DiagnosticStatus.FAILED return @@ -400,9 +506,15 @@ class Diagnostic(Publisher): self.check_is_ready() def check_is_ready(self): - """Check if a diagnostic is ready to run and change its status accordingly""" - if all([request.ready_to_run(self) for request in self._requests]) and\ - all([subjob.status == DiagnosticStatus.COMPLETED for subjob in self.subjobs]): + """Check if diagnostic is ready to run and change status accordingly""" + if all( + [request.ready_to_run(self) for request in self._requests] + ) and all( + [ + subjob.status == DiagnosticStatus.COMPLETED + for subjob in self.subjobs + ] + ): self.status = DiagnosticStatus.READY def _unsuscribe_requests(self): @@ -429,9 +541,14 @@ class Diagnostic(Publisher): int """ - return len([request for request in self._requests - if request.storage_status != StorageStatus.READY or - request.local_status != LocalStatus.READY]) + return len( + [ + request + for request in self._requests + if request.storage_status != StorageStatus.READY + or request.local_status != LocalStatus.READY + ] + ) def _different_type(self, other): return type(self) is not type(other) @@ -448,7 +565,8 @@ class DiagnosticOption(object): ---------- name: str default_value: object, optional - If None, the option is required and an exception will be thrown at parse time if the value is empty + If None, the option is required and an exception will be + thrown at parse time if the value is empty """ self.name = name @@ -478,9 +596,11 @@ class DiagnosticOption(object): return option_value def _check_default(self, option_value): - if option_value == '': + if option_value == "": if self.default_value is None: - raise DiagnosticOptionError('Option {0} is not optional'.format(self.name)) + raise DiagnosticOptionError( + "Option {0} is not optional".format(self.name) + ) else: return self.default_value return option_value @@ -519,7 +639,9 @@ class DiagnosticIntOption(DiagnosticOption): """ - def __init__(self, name, default_value=None, min_limit=None, max_limit=None): + def __init__( + self, name, default_value=None, min_limit=None, max_limit=None + ): super(DiagnosticIntOption, self).__init__(name, default_value) self.min_limit = min_limit self.max_limit = max_limit @@ -547,9 +669,17 @@ class DiagnosticIntOption(DiagnosticOption): def _check_limits(self, value): if self.min_limit is not None and value < self.min_limit: - raise DiagnosticOptionError('Value {0} is lower than minimum ({1})'.format(value, self.min_limit)) + raise DiagnosticOptionError( + "Value {0} is lower than minimum ({1})".format( + value, self.min_limit + ) + ) if self.max_limit is not None and value > self.max_limit: - raise DiagnosticOptionError('Value {0} is higher than maximum ({1})'.format(value, self.max_limit)) + raise DiagnosticOptionError( + "Value {0} is higher than maximum ({1})".format( + value, self.max_limit + ) + ) class DiagnosticListIntOption(DiagnosticIntOption): @@ -567,7 +697,9 @@ class DiagnosticListIntOption(DiagnosticIntOption): """ - def __init__(self, name, default_value=None, min_limit=None, max_limit=None): + def __init__( + self, name, default_value=None, min_limit=None, max_limit=None + ): super(DiagnosticListIntOption, self).__init__(name, default_value) self.min_limit = min_limit """ Lower limit """ @@ -594,7 +726,7 @@ class DiagnosticListIntOption(DiagnosticIntOption): option_value = self._check_default(option_value) if isinstance(option_value, tuple) or isinstance(option_value, list): return option_value - values = [int(i) for i in option_value.split('-')] + values = [int(i) for i in option_value.split("-")] for value in values: self._check_limits(value) @@ -613,7 +745,9 @@ class DiagnosticListFrequenciesOption(DiagnosticOption): """ def __init__(self, name, default_value=None): - super(DiagnosticListFrequenciesOption, self).__init__(name, default_value) + super(DiagnosticListFrequenciesOption, self).__init__( + name, default_value + ) def parse(self, option_value): """ @@ -627,7 +761,7 @@ class DiagnosticListFrequenciesOption(DiagnosticOption): option_value = self._check_default(option_value) if isinstance(option_value, (tuple, list)): return option_value - values = [Frequency(i) for i in option_value.split('-')] + values = [Frequency(i) for i in option_value.split("-")] return values @@ -643,7 +777,7 @@ class DiagnosticVariableOption(DiagnosticOption): """ - def __init__(self, var_manager, name='variable', default_value=None): + def __init__(self, var_manager, name="variable", default_value=None): super(DiagnosticVariableOption, self).__init__(name, default_value) self.var_manager = var_manager @@ -690,7 +824,7 @@ class DiagnosticVariableListOption(DiagnosticOption): """ option_value = self._check_default(option_value) var_names = [] - for value in option_value.split(':'): + for value in option_value.split(":"): real_name = self.var_manager.get_variable(value, False) if real_name is None: var_names.append(value) @@ -710,7 +844,7 @@ class DiagnosticDomainOption(DiagnosticOption): """ - def __init__(self, name='domain', default_value=None): + def __init__(self, name="domain", default_value=None): super(DiagnosticDomainOption, self).__init__(name, default_value) def parse(self, option_value): @@ -735,7 +869,7 @@ class DiagnosticFrequencyOption(DiagnosticOption): """ - def __init__(self, name='frequency', default_value=None): + def __init__(self, name="frequency", default_value=None): super(DiagnosticFrequencyOption, self).__init__(name, default_value) def parse(self, option_value): @@ -772,10 +906,12 @@ class DiagnosticBasinListOption(DiagnosticOption): """ option_value = self._check_default(option_value) basins = [] - for value in option_value.split(':'): + for value in option_value.split(":"): basin = Basins().parse(value) if basin is None: - raise DiagnosticOptionError('Basin {0} not recognized'.format(value)) + raise DiagnosticOptionError( + "Basin {0} not recognized".format(value) + ) basins.append(basin) return basins @@ -799,7 +935,9 @@ class DiagnosticBasinOption(DiagnosticOption): value = self._check_default(option_value) basin = Basins().parse(value) if basin is None: - raise DiagnosticOptionError('Basin {0} not recognized'.format(value)) + raise DiagnosticOptionError( + "Basin {0} not recognized".format(value) + ) return basin @@ -823,7 +961,11 @@ class DiagnosticComplexStrOption(DiagnosticOption): ------- str """ - return self._check_default(option_value).replace('&;', ',').replace('&.', ' ') + return ( + self._check_default(option_value) + .replace("&;", ",") + .replace("&.", " ") + ) class DiagnosticBoolOption(DiagnosticOption): @@ -844,7 +986,7 @@ class DiagnosticBoolOption(DiagnosticOption): option_value = self._check_default(option_value) if isinstance(option_value, bool): return option_value - return option_value.lower() in ('true', 't', 'yes') + return option_value.lower() in ("true", "t", "yes") class DiagnosticChoiceOption(DiagnosticOption): @@ -893,8 +1035,10 @@ class DiagnosticChoiceOption(DiagnosticOption): else: if option_value in self.choices: return option_value - raise DiagnosticOptionError('Value {1} in option {0} is not a valid choice. ' - 'Options are {2}'.format(self.name, option_value, self.choices)) + raise DiagnosticOptionError( + "Value {1} in option {0} is not a valid choice. " + "Options are {2}".format(self.name, option_value, self.choices) + ) class DiagnosticOptionError(Exception): diff --git a/earthdiagnostics/earthdiags.py b/earthdiagnostics/earthdiags.py index d6ad4c72..dfcba88a 100755 --- a/earthdiagnostics/earthdiags.py +++ b/earthdiagnostics/earthdiags.py @@ -34,14 +34,15 @@ class EarthDiags(object): config_file: str """ - # Get the version number from the relevant file. If not, from autosubmit package + # Get the version number from the relevant file. + # If not, from earthdiagnostics package scriptdir = os.path.abspath(os.path.dirname(__file__)) - if not os.path.exists(os.path.join(scriptdir, 'VERSION')): + if not os.path.exists(os.path.join(scriptdir, "VERSION")): scriptdir = os.path.join(scriptdir, os.path.pardir) - version_path = os.path.join(scriptdir, 'VERSION') - readme_path = os.path.join(scriptdir, 'README') - changes_path = os.path.join(scriptdir, 'CHANGELOG') + version_path = os.path.join(scriptdir, "VERSION") + readme_path = os.path.join(scriptdir, "README") + changes_path = os.path.join(scriptdir, "CHANGELOG") if os.path.isfile(version_path): with open(version_path) as f: version = f.read().strip() @@ -68,15 +69,20 @@ class EarthDiags(object): """ Log.info( - 'Initialising Earth Diagnostics Version {0}', EarthDiags.version) + "Initialising Earth Diagnostics Version {0}", EarthDiags.version + ) self.config.parse(config_file) - os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE' + os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" TempFile.scratch_folder = self.config.scratch_dir cdftools.path = self.config.cdftools_path self._create_dic_variables() - Log.debug('Diags ready') - Log.info('Running diags for experiment {0}, startdates {1}, members {2}', self.config.experiment.expid, - self.config.experiment.startdates, self.config.experiment.members) + Log.debug("Diags ready") + Log.info( + "Running diags for experiment {0}, startdates {1}, members {2}", + self.config.experiment.expid, + self.config.experiment.startdates, + self.config.experiment.members, + ) @staticmethod def parse_args(args): @@ -87,28 +93,70 @@ class EarthDiags(object): """ # try: parser = argparse.ArgumentParser( - description='Main executable for Earth Diagnostics.') - parser.add_argument('-v', '--version', action='version', version=EarthDiags.version, - help="returns Earth Diagnostics's version number and exit") - parser.add_argument('--doc', action='store_true', - help="opens documentation and exits") - parser.add_argument('--clean', action='store_true', - help="clean the scratch folder and exits") - parser.add_argument('--report', action='store_true', - help="generates a report about the available files") - parser.add_argument('-lf', '--logfile', choices=('EVERYTHING', 'DEBUG', 'INFO', 'RESULT', 'USER_WARNING', - 'WARNING', 'ERROR', 'CRITICAL', 'NO_LOG'), - default='DEBUG', type=str, - help="sets file's log level.") - parser.add_argument('-lc', '--logconsole', choices=('EVERYTHING', 'DEBUG', 'INFO', 'RESULT', 'USER_WARNING', - 'WARNING', 'ERROR', 'CRITICAL', 'NO_LOG'), - default='INFO', type=str, - help="sets console's log level") - - parser.add_argument('-log', '--logfilepath', default=None, type=str) - - parser.add_argument('-f', '--configfile', - default='diags.conf', type=str) + description="Main executable for Earth Diagnostics." + ) + parser.add_argument( + "-v", + "--version", + action="version", + version=EarthDiags.version, + help="returns Earth Diagnostics's version number and exit", + ) + parser.add_argument( + "--doc", action="store_true", help="opens documentation and exits" + ) + parser.add_argument( + "--clean", + action="store_true", + help="clean the scratch folder and exits", + ) + parser.add_argument( + "--report", + action="store_true", + help="generates a report about the available files", + ) + parser.add_argument( + "-lf", + "--logfile", + choices=( + "EVERYTHING", + "DEBUG", + "INFO", + "RESULT", + "USER_WARNING", + "WARNING", + "ERROR", + "CRITICAL", + "NO_LOG", + ), + default="DEBUG", + type=str, + help="sets file's log level.", + ) + parser.add_argument( + "-lc", + "--logconsole", + choices=( + "EVERYTHING", + "DEBUG", + "INFO", + "RESULT", + "USER_WARNING", + "WARNING", + "ERROR", + "CRITICAL", + "NO_LOG", + ), + default="INFO", + type=str, + help="sets console's log level", + ) + + parser.add_argument("-log", "--logfilepath", default=None, type=str) + + parser.add_argument( + "-f", "--configfile", default="diags.conf", type=str + ) args = parser.parse_args(args) if args.doc: @@ -145,23 +193,24 @@ class EarthDiags(object): bool: True if successful """ - Log.info('Opening documentation...') + Log.info("Opening documentation...") doc_path = os.path.join( - 'http://earthdiagnostics.readthedocs.io/en/latest') - Utils.execute_shell_command(('xdg-open', doc_path)) - Log.result('Documentation opened!') + "http://earthdiagnostics.readthedocs.io/en/latest" + ) + Utils.execute_shell_command(("xdg-open", doc_path)) + Log.result("Documentation opened!") return True def _create_dic_variables(self): self.dic_variables = dict() - self.dic_variables['x'] = 'i' - self.dic_variables['y'] = 'j' - self.dic_variables['z'] = 'lev' - self.dic_variables['nav_lon'] = self.config.data_convention.lon_name - self.dic_variables['nav_lat'] = self.config.data_convention.lat_name - self.dic_variables['nav_lev'] = 'lev' - self.dic_variables['time_counter'] = 'time' - self.dic_variables['t'] = 'time' + self.dic_variables["x"] = "i" + self.dic_variables["y"] = "j" + self.dic_variables["z"] = "lev" + self.dic_variables["nav_lon"] = self.config.data_convention.lon_name + self.dic_variables["nav_lat"] = self.config.data_convention.lat_name + self.dic_variables["nav_lev"] = "lev" + self.dic_variables["time_counter"] = "time" + self.dic_variables["t"] = "time" def run(self): """ @@ -174,18 +223,18 @@ class EarthDiags(object): """ start = datetime.now() self.had_errors = False - Log.debug('Using netCDF version {0}', netCDF4.getlibversion()) + Log.debug("Using netCDF version {0}", netCDF4.getlibversion()) self._prepare_scratch_dir() self._prepare_mesh_files() self._initialize_basins() - Log.info('Time to prepare: {}', datetime.now() - start) + Log.info("Time to prepare: {}", datetime.now() - start) self._prepare_data_manager() # Run diagnostics - Log.info('Running diagnostics') + Log.info("Running diagnostics") work_manager = WorkManager(self.config, self.data_manager) work_manager.prepare_job_list() @@ -197,7 +246,7 @@ class EarthDiags(object): return result def _initialize_basins(self): - self._read_basins_from_file('basins.nc') + self._read_basins_from_file("basins.nc") @staticmethod def _read_basins_from_file(filename): @@ -208,20 +257,21 @@ class EarthDiags(object): def _prepare_scratch_dir(self): if self.config.use_ramdisk: self._remove_scratch_dir() - tempfile.mkdtemp(dir='/dev/shm') - os.symlink(tempfile.mkdtemp(dir='/dev/shm'), - self.config.scratch_dir) + tempfile.mkdtemp(dir="/dev/shm") + os.symlink( + tempfile.mkdtemp(dir="/dev/shm"), self.config.scratch_dir + ) else: if not os.path.exists(self.config.scratch_dir): os.makedirs(self.config.scratch_dir) os.chdir(self.config.scratch_dir) def _prepare_data_manager(self): - if self.config.data_adaptor == 'CMOR': + if self.config.data_adaptor == "CMOR": self.data_manager = CMORManager(self.config) - elif self.config.data_adaptor == 'THREDDS': + elif self.config.data_adaptor == "THREDDS": self.data_manager = THREDDSManager(self.config) - elif self.config.data_adaptor == 'OBSRECON': + elif self.config.data_adaptor == "OBSRECON": self.data_manager = ObsReconManager(self.config) self.data_manager.prepare() @@ -234,9 +284,9 @@ class EarthDiags(object): bool """ - Log.info('Removing scratch folder...') + Log.info("Removing scratch folder...") self._remove_scratch_dir() - Log.result('Scratch folder removed') + Log.result("Scratch folder removed") return True def _remove_scratch_dir(self): @@ -244,13 +294,15 @@ class EarthDiags(object): # time.sleep(4) # shutil.rmtree(os.path.realpath(self.config.scratch_dir)) Utils.execute_shell_command( - 'rm -r {0}'.format(os.path.realpath(self.config.scratch_dir))) + "rm -r {0}".format(os.path.realpath(self.config.scratch_dir)) + ) os.remove(self.config.scratch_dir) elif os.path.isdir(self.config.scratch_dir): # time.sleep(4) # shutil.rmtree(self.config.scratch_dir) Utils.execute_shell_command( - 'rm -r {0}'.format(self.config.scratch_dir)) + "rm -r {0}".format(self.config.scratch_dir) + ) def report(self): """ @@ -260,7 +312,7 @@ class EarthDiags(object): ------- bool """ - Log.info('Looking for existing vars...') + Log.info("Looking for existing vars...") self._prepare_data_manager() base_folder = self.config.report.path if not base_folder: @@ -270,13 +322,17 @@ class EarthDiags(object): for member in self.config.experiment.members: results = self._get_variable_report(startdate, member) - report_path = os.path.join(base_folder, - '{0}_{1}.report'.format(startdate, - self.config.experiment.get_member_str(member))) + report_path = os.path.join( + base_folder, + "{0}_{1}.report".format( + startdate, + self.config.experiment.get_member_str(member), + ), + ) self._create_report(report_path, results) - Log.result('Report finished') + Log.result("Report finished") return True def _get_variable_report(self, startdate, member): @@ -286,16 +342,29 @@ class EarthDiags(object): if var.domain is None: continue for table, priority in var.tables: - if priority is None or priority > self.config.report.maximum_priority: + if ( + priority is None + or priority > self.config.report.maximum_priority + ): continue - if not self.data_manager.file_exists(var.domain, var.short_name, startdate, member, 1, - frequency=table.frequency): + if not self.data_manager.file_exists( + var.domain, + var.short_name, + startdate, + member, + 1, + frequency=table.frequency, + ): results.append((var, table, priority)) Log.debug( - 'Variable {0.short_name} not found in {1.name}', var, table) + "Variable {0.short_name} not found in {1.name}", + var, + table, + ) else: Log.result( - 'Variable {0.short_name} found in {1.name}', var, table) + "Variable {0.short_name} found in {1.name}", var, table + ) return results @@ -303,43 +372,52 @@ class EarthDiags(object): def _create_report(report_path, results): tables = set([result[1].name for result in results]) for table in tables: - file_handler = open('{0}.{1}'.format(report_path, table), 'w') + file_handler = open("{0}.{1}".format(report_path, table), "w") table_results = [ - result for result in results if result[1].name == table] + result for result in results if result[1].name == table + ] - file_handler.write('\nTable {0}\n'.format(table)) - file_handler.write('===================================\n') + file_handler.write("\nTable {0}\n".format(table)) + file_handler.write("===================================\n") priorities = set([result[2] for result in table_results]) priorities = sorted(priorities) for priority in priorities: - priority_results = [result[0] - for result in table_results if result[2] == priority] + priority_results = [ + result[0] + for result in table_results + if result[2] == priority + ] priority_results = sorted( - priority_results, key=lambda v: v.short_name) + priority_results, key=lambda v: v.short_name + ) file_handler.write( - '\nMissing variables with priority {0}:\n'.format(priority)) - file_handler.write('--------------------------------------\n') + "\nMissing variables with priority {0}:\n".format(priority) + ) + file_handler.write("--------------------------------------\n") for var in priority_results: - file_handler.write('{0:12}: {1}\n'.format( - var.short_name, var.standard_name)) + file_handler.write( + "{0:12}: {1}\n".format( + var.short_name, var.standard_name + ) + ) file_handler.flush() file_handler.close() def _prepare_mesh_files(self): model_version = self.config.experiment.model_version if not model_version: - Log.info('No model version defined. Skipping mesh files copy!') + Log.info("No model version defined. Skipping mesh files copy!") return - Log.info('Copying mesh files') + Log.info("Copying mesh files") con_files = self.config.con_files model_version = self.config.experiment.model_version restore_meshes = self.config.restore_meshes - mesh_mask = 'mesh_mask_nemo.{0}.nc'.format(model_version) - new_mask_glo = 'new_maskglo.{0}.nc'.format(model_version) - basins = 'basins.{0}.nc'.format(model_version) + mesh_mask = "mesh_mask_nemo.{0}.nc".format(model_version) + new_mask_glo = "new_maskglo.{0}.nc".format(model_version) + basins = "basins.{0}.nc".format(model_version) if self.config.mesh_mask: mesh_mask_path = self.config.mesh_mask @@ -358,74 +436,97 @@ class EarthDiags(object): if self.config.scratch_masks: self._prepare_mesh_using_scratch( - basins, basins_path, - mesh_mask, mesh_mask_path, - new_mask_glo, new_mask_glo_path, - restore_meshes + basins, + basins_path, + mesh_mask, + mesh_mask_path, + new_mask_glo, + new_mask_glo_path, + restore_meshes, ) else: - self._copy_file(mesh_mask_path, 'mesh_hgr.nc', restore_meshes) - self._link_file('mesh_hgr.nc', 'mesh_zgr.nc') - self._link_file('mesh_hgr.nc', 'mask.nc') - self._copy_file(new_mask_glo_path, 'new_maskglo.nc', - restore_meshes) - self._copy_file(basins_path, 'basins.nc', restore_meshes) - - Log.result('Mesh files ready!') - - def _prepare_mesh_using_scratch(self, basins, basins_path, - mesh_mask, mesh_mask_path, - new_mask_glo, new_mask_glo_path, - restore_meshes): + self._copy_file(mesh_mask_path, "mesh_hgr.nc", restore_meshes) + self._link_file("mesh_hgr.nc", "mesh_zgr.nc") + self._link_file("mesh_hgr.nc", "mask.nc") + self._copy_file( + new_mask_glo_path, "new_maskglo.nc", restore_meshes + ) + self._copy_file(basins_path, "basins.nc", restore_meshes) + + Log.result("Mesh files ready!") + + def _prepare_mesh_using_scratch( + self, + basins, + basins_path, + mesh_mask, + mesh_mask_path, + new_mask_glo, + new_mask_glo_path, + restore_meshes, + ): Utils.create_folder_tree(self.config.scratch_masks) Utils.give_group_write_permissions(self.config.scratch_masks) mesh_mask_scratch_path = os.path.join( - self.config.scratch_masks, mesh_mask) - if self._copy_file(mesh_mask_path, mesh_mask_scratch_path, - restore_meshes): + self.config.scratch_masks, mesh_mask + ) + if self._copy_file( + mesh_mask_path, mesh_mask_scratch_path, restore_meshes + ): Utils.give_group_write_permissions(mesh_mask_scratch_path) - self._link_file(mesh_mask_scratch_path, 'mesh_hgr.nc') - self._link_file(mesh_mask_scratch_path, 'mesh_zgr.nc') - self._link_file(mesh_mask_scratch_path, 'mask.nc') + self._link_file(mesh_mask_scratch_path, "mesh_hgr.nc") + self._link_file(mesh_mask_scratch_path, "mesh_zgr.nc") + self._link_file(mesh_mask_scratch_path, "mask.nc") new_maskglo_scratch_path = os.path.join( - self.config.scratch_masks, new_mask_glo) - if self._copy_file(new_mask_glo_path, - new_maskglo_scratch_path, restore_meshes): + self.config.scratch_masks, new_mask_glo + ) + if self._copy_file( + new_mask_glo_path, new_maskglo_scratch_path, restore_meshes + ): Utils.give_group_write_permissions(new_maskglo_scratch_path) - self._link_file(new_maskglo_scratch_path, 'new_maskglo.nc') + self._link_file(new_maskglo_scratch_path, "new_maskglo.nc") - basins_scratch_path = os.path.join( - self.config.scratch_masks, basins) + basins_scratch_path = os.path.join(self.config.scratch_masks, basins) if self._copy_file(basins_path, basins_scratch_path, restore_meshes): Utils.give_group_write_permissions(basins_path) - self._link_file(basins_scratch_path, 'basins.nc') + self._link_file(basins_scratch_path, "basins.nc") def _copy_file(self, source, destiny, force): if not os.path.exists(source): - Log.user_warning('File {0} is not available for {1}', - destiny, self.config.experiment.model_version) - Log.debug('Looking for it in {0}', source) + Log.user_warning( + "File {0} is not available for {1}", + destiny, + self.config.experiment.model_version, + ) + Log.debug("Looking for it in {0}", source) return False if not force and os.path.exists(destiny): # Small size differences can be due to the renaming of variables reference_size = os.stat(source).st_size delta_size = abs(reference_size - os.stat(destiny).st_size) - if delta_size < 2048 or delta_size / reference_size < 1 / 1000 or True: - Log.info('File {0} already exists', destiny) + if ( + delta_size < 2048 + or delta_size / reference_size < 1 / 1000 + or True + ): + Log.info("File {0} already exists", destiny) return True - Log.info('Copying file {0}', destiny) + Log.info("Copying file {0}", destiny) shutil.copyfile(source, destiny) - Log.info('File {0} ready', destiny) + Log.info("File {0} ready", destiny) Utils.rename_variables(destiny, self.dic_variables, False) return True def _link_file(self, source, destiny): if not os.path.exists(source): - Log.user_warning('File {0} is not available for {1}', - destiny, self.config.experiment.model_version) + Log.user_warning( + "File {0} is not available for {1}", + destiny, + self.config.experiment.model_version, + ) return if os.path.lexists(destiny): @@ -433,11 +534,11 @@ class EarthDiags(object): os.remove(destiny) except OSError as ex: if ex.errno == 13: # Permission denied - Log.info('Link already created') + Log.info("Link already created") return os.symlink(source, destiny) - Log.info('File {0} ready', destiny) + Log.info("File {0} ready", destiny) def main(): diff --git a/earthdiagnostics/frequency.py b/earthdiagnostics/frequency.py index 45391649..9134352a 100644 --- a/earthdiagnostics/frequency.py +++ b/earthdiagnostics/frequency.py @@ -6,34 +6,73 @@ class Frequency(object): """Time frequency""" _recognized = { - 'f': 'fx', 'fx': 'fx', 'fixed': 'fx', - 'c': 'clim', 'clim': 'clim', 'climatology': 'clim', 'monclim': 'clim', '1hrclimmon': 'clim', - 'monc': 'clim', - 'dec': 'dec', 'decadal': 'dec', - 'y': 'year', 'yr': 'year', 'year': 'year', 'yearly': 'year', - 'm': 'mon', '1m': 'mon', 'mon': 'mon', 'monthly': 'mon', 'mm': 'mon', - 'w': 'week', '1w': 'week', 'week': 'week', 'weekly': 'week', - 'd': 'day', '1d': 'day', 'daily': 'day', 'day': 'day', - '15': '15hr', '15h': '15hr', '15hr': '15hr', '15_hourly': '15hr', '15hourly': '15hr', - '15 hourly': '15hr', - '6': '6hr', '6h': '6hr', '6hr': '6hr', '6_hourly': '6hr', '6hourly': '6hr', '6 hourly': '6hr', - '3': '3hr', '3h': '3hr', '3hr': '3hr', '3_hourly': '3hr', '3hourly': '3hr', '3 hourly': '3hr', - '1': '1hr', 'hr': '1hr', '1h': '1hr', 'hourly': '1hr', '1hr': '1hr', '1 hourly': '1hr', - '450mn': '450mn', - 'subhr': 'subhr' + "f": "fx", + "fx": "fx", + "fixed": "fx", + "c": "clim", + "clim": "clim", + "climatology": "clim", + "monclim": "clim", + "1hrclimmon": "clim", + "monc": "clim", + "dec": "dec", + "decadal": "dec", + "y": "year", + "yr": "year", + "year": "year", + "yearly": "year", + "m": "mon", + "1m": "mon", + "mon": "mon", + "monthly": "mon", + "mm": "mon", + "w": "week", + "1w": "week", + "week": "week", + "weekly": "week", + "d": "day", + "1d": "day", + "daily": "day", + "day": "day", + "15": "15hr", + "15h": "15hr", + "15hr": "15hr", + "15_hourly": "15hr", + "15hourly": "15hr", + "15 hourly": "15hr", + "6": "6hr", + "6h": "6hr", + "6hr": "6hr", + "6_hourly": "6hr", + "6hourly": "6hr", + "6 hourly": "6hr", + "3": "3hr", + "3h": "3hr", + "3hr": "3hr", + "3_hourly": "3hr", + "3hourly": "3hr", + "3 hourly": "3hr", + "1": "1hr", + "hr": "1hr", + "1h": "1hr", + "hourly": "1hr", + "1hr": "1hr", + "1 hourly": "1hr", + "450mn": "450mn", + "subhr": "subhr", } def __init__(self, freq): freq = freq.lower() - if freq.endswith('cm'): + if freq.endswith("cm"): freq = freq[:-2] - self.point = freq.endswith('pt') + self.point = freq.endswith("pt") if self.point: freq = freq[:-2] try: self.frequency = Frequency._recognized[freq] except KeyError: - raise ValueError('Frequency {0} not supported'.format(freq)) + raise ValueError("Frequency {0} not supported".format(freq)) def __eq__(self, other): return self.frequency == other.frequency @@ -57,25 +96,26 @@ class Frequency(object): str """ from earthdiagnostics.variable import VariableType + if self == Frequencies.climatology: - return 'clim' + return "clim" elif self == Frequencies.fixed: - return 'fx' + return "fx" elif self == Frequencies.decadal: - freq_str = 'decadal' + freq_str = "decadal" elif self == Frequencies.yearly: - freq_str = 'yearly' + freq_str = "yearly" elif self == Frequencies.monthly: - freq_str = 'monthly' + freq_str = "monthly" elif self == Frequencies.weekly: - freq_str = 'weekly' + freq_str = "weekly" elif self == Frequencies.daily: - freq_str = 'daily' - elif self.frequency.endswith('hr'): - freq_str = self.frequency[:-2] + 'hourly' + freq_str = "daily" + elif self.frequency.endswith("hr"): + freq_str = self.frequency[:-2] + "hourly" self.point = True if not self.point or vartype != VariableType.MEAN: - freq_str = '{0}_{1}'.format(freq_str, VariableType.to_str(vartype)) + freq_str = "{0}_{1}".format(freq_str, VariableType.to_str(vartype)) return freq_str @staticmethod @@ -101,14 +141,14 @@ class Frequency(object): class Frequencies(object): """Enumeration of supported frequencies""" - fixed = Frequency('fx') - climatology = Frequency('clim') - decadal = Frequency('dec') - yearly = Frequency('year') - monthly = Frequency('mon') - weekly = Frequency('week') - daily = Frequency('day') - six_hourly = Frequency('6hr') - three_hourly = Frequency('3hr') - hourly = Frequency('hr') - subhourly = Frequency('subhr') + fixed = Frequency("fx") + climatology = Frequency("clim") + decadal = Frequency("dec") + yearly = Frequency("year") + monthly = Frequency("mon") + weekly = Frequency("week") + daily = Frequency("day") + six_hourly = Frequency("6hr") + three_hourly = Frequency("3hr") + hourly = Frequency("hr") + subhourly = Frequency("subhr") diff --git a/earthdiagnostics/general/attribute.py b/earthdiagnostics/general/attribute.py index 8609ba66..0be1f292 100644 --- a/earthdiagnostics/general/attribute.py +++ b/earthdiagnostics/general/attribute.py @@ -2,8 +2,12 @@ """Set attributtes in netCDF files""" from earthdiagnostics.general.fix_file import FixFile from earthdiagnostics.utils import Utils -from earthdiagnostics.diagnostic import DiagnosticDomainOption, \ - DiagnosticVariableOption, DiagnosticOption, DiagnosticComplexStrOption +from earthdiagnostics.diagnostic import ( + DiagnosticDomainOption, + DiagnosticVariableOption, + DiagnosticOption, + DiagnosticComplexStrOption, +) class Attribute(FixFile): @@ -30,35 +34,57 @@ class Attribute(FixFile): :type domain: ModelingRealm """ - alias = 'att' + alias = "att" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, - variable, grid, attributte_name, attributte_value): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, + attributte_name, + attributte_value, + ): FixFile.__init__( - self, data_manager, startdate, member, chunk, - domain, variable, grid + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, ) self.attributte_name = attributte_name self.attributte_value = attributte_value def __str__(self): - return 'Write attributte output Startdate: {0.startdate} '\ - 'Member: {0.member} Chunk: {0.chunk} ' \ - 'Variable: {0.domain}:{0.variable} ' \ - 'Attributte: {0.attributte_name}:{0.attributte_value} ' \ - 'Grid: {0.grid}'.format(self) + return ( + "Write attributte output Startdate: {0.startdate} " + "Member: {0.member} Chunk: {0.chunk} " + "Variable: {0.domain}:{0.variable} " + "Attributte: {0.attributte_name}:{0.attributte_value} " + "Grid: {0.grid}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variable == other.variable \ - and self.grid == other.grid and \ - self.attributte_name == other.attributte_name and \ - self.attributte_value == other.attributte_value + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.grid == other.grid + and self.attributte_name == other.attributte_name + and self.attributte_value == other.attributte_value + ) @classmethod def generate_jobs(cls, diags, options): @@ -74,9 +100,9 @@ class Attribute(FixFile): options_available = ( DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticOption('name'), - DiagnosticComplexStrOption('value'), - DiagnosticOption('grid', '') + DiagnosticOption("name"), + DiagnosticComplexStrOption("value"), + DiagnosticOption("grid", ""), ) options = cls.process_options(options, options_available) job_list = list() @@ -84,9 +110,15 @@ class Attribute(FixFile): for startdate, member, chunk in chunk_list: job_list.append( Attribute( - diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], options['grid'], - options['name'], options['value'] + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["grid"], + options["name"], + options["value"], ) ) return job_list @@ -99,7 +131,7 @@ class Attribute(FixFile): handler.close() if not Utils.check_netcdf_file(variable_file): raise Exception( - 'Attribute {0} can not be set correctly to {1}'.format( + "Attribute {0} can not be set correctly to {1}".format( self.attributte_name, self.attributte_value ) ) diff --git a/earthdiagnostics/general/fix_file.py b/earthdiagnostics/general/fix_file.py index e2ef1703..e7fe2551 100644 --- a/earthdiagnostics/general/fix_file.py +++ b/earthdiagnostics/general/fix_file.py @@ -1,7 +1,11 @@ # coding=utf-8 """Base diagnostic for fixing files""" -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticDomainOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticVariableOption, +) class FixFile(Diagnostic): @@ -28,8 +32,9 @@ class FixFile(Diagnostic): :type domain: ModelingRealm """ - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, grid): + def __init__( + self, data_manager, startdate, member, chunk, domain, variable, grid + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -44,20 +49,24 @@ class FixFile(Diagnostic): _STR_PREFIX = None def __str__(self): - return '{0._STR_PREFIX} Startdate: {0.startdate} Member: {0.member} ' \ - 'Chunk: {0.chunk} Variable: {0.domain}:{0.variable} ' \ - 'Grid: {0.grid}'.format(self) + return ( + "{0._STR_PREFIX} Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Variable: {0.domain}:{0.variable} " + "Grid: {0.grid}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.grid == self.grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.grid == self.grid + ) @classmethod def generate_jobs(cls, diags, options): @@ -76,8 +85,15 @@ class FixFile(Diagnostic): chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: job_list.append( - cls(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], options['grid']) + cls( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["grid"], + ) ) return job_list @@ -86,21 +102,28 @@ class FixFile(Diagnostic): return [ DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticOption('grid', '') + DiagnosticOption("grid", ""), ] def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - grid=self.grid, to_modify=True + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + to_modify=True, ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.corrected = self.declare_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, ) diff --git a/earthdiagnostics/general/module.py b/earthdiagnostics/general/module.py index 01fde676..504d10ea 100644 --- a/earthdiagnostics/general/module.py +++ b/earthdiagnostics/general/module.py @@ -2,8 +2,12 @@ """Compute module of two variables""" import numpy as np -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticVariableOption, \ - DiagnosticDomainOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticVariableOption, + DiagnosticDomainOption, + DiagnosticOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -27,11 +31,21 @@ class Module(Diagnostic): :type domain: ModelingRealm """ - alias = 'module' + alias = "module" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, componentu, componentv, module_var, grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + componentu, + componentv, + module_var, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -48,23 +62,27 @@ class Module(Diagnostic): self.module_file = None def __str__(self): - return 'Calculate module Startdate: {0.startdate} ' \ - 'Member: {0.member} ' \ - 'Chunk: {0.chunk} ' \ - 'Variables: {0.domain}:{0.componentu},{0.componentv},{0.module} ' \ - 'Grid: {0.grid}'.format(self) + return ( + "Calculate module Startdate: {0.startdate} " + "Member: {0.member} " + "Chunk: {0.chunk} " + "Variables: {0.domain}:{0.componentu},{0.componentv},{0.module} " + "Grid: {0.grid}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.componentu == other.componentu and \ - self.componentv == other.componentv and \ - self.module == other.module and \ - self.grid == other.grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.componentu == other.componentu + and self.componentv == other.componentv + and self.module == other.module + and self.grid == other.grid + ) @classmethod def generate_jobs(cls, diags, options): @@ -80,10 +98,10 @@ class Module(Diagnostic): var_manager = diags.data_manager.config.var_manager options_available = ( DiagnosticDomainOption(), - DiagnosticVariableOption(var_manager, 'componentu'), - DiagnosticVariableOption(var_manager, 'componentv'), - DiagnosticVariableOption(var_manager, 'module'), - DiagnosticOption('grid', '') + DiagnosticVariableOption(var_manager, "componentu"), + DiagnosticVariableOption(var_manager, "componentv"), + DiagnosticVariableOption(var_manager, "module"), + DiagnosticOption("grid", ""), ) options = cls.process_options(options, options_available) job_list = list() @@ -91,10 +109,15 @@ class Module(Diagnostic): for startdate, member, chunk in chunk_list: job_list.append( Module( - diags.data_manager, startdate, member, chunk, - options['domain'], options['componentu'], - options['componentv'], options['module'], - options['grid'] + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["componentu"], + options["componentv"], + options["module"], + options["grid"], ) ) return job_list @@ -102,19 +125,31 @@ class Module(Diagnostic): def request_data(self): """Request data required by the diagnostic""" self.component_u_file = self.request_chunk( - self.domain, self.componentu, self.startdate, self.member, - self.chunk, grid=self.grid + self.domain, + self.componentu, + self.startdate, + self.member, + self.chunk, + grid=self.grid, ) self.component_v_file = self.request_chunk( - self.domain, self.componentv, self.startdate, self.member, - self.chunk, grid=self.grid + self.domain, + self.componentv, + self.startdate, + self.member, + self.chunk, + grid=self.grid, ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.module_file = self.declare_chunk( - self.domain, self.module, self.startdate, self.member, self.chunk, - grid=self.grid + self.domain, + self.module, + self.startdate, + self.member, + self.chunk, + grid=self.grid, ) def compute(self): @@ -128,9 +163,9 @@ class Module(Diagnostic): variable_u[:] = np.sqrt(variable_u[:] ** 2 + variable_v[:] ** 2) - if 'table' in variable_u.ncattrs(): + if "table" in variable_u.ncattrs(): del variable_u.table - if 'code' in variable_u.ncattrs(): + if "code" in variable_u.ncattrs(): del variable_u.code component_u.close() diff --git a/earthdiagnostics/general/relink.py b/earthdiagnostics/general/relink.py index cdad3d9d..bb277ce3 100644 --- a/earthdiagnostics/general/relink.py +++ b/earthdiagnostics/general/relink.py @@ -1,7 +1,12 @@ # coding=utf-8 """Create links for a variable""" -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticDomainOption, DiagnosticBoolOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticBoolOption, + DiagnosticVariableOption, +) class Relink(Diagnostic): @@ -29,11 +34,20 @@ class Relink(Diagnostic): :type move_old: bool """ - alias = 'relink' + alias = "relink" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, move_old, grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + move_old, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -45,9 +59,11 @@ class Relink(Diagnostic): self.var_manager = data_manager.config.var_manager def __str__(self): - return 'Relink output Startdate: {0.startdate} Member: {0.member} ' \ - 'Chunk: {0.chunk} Move old: {0.move_old} ' \ - 'Variable: {0.domain}:{0.variable} Grid: {0.grid}'.format(self) + return ( + "Relink output Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Move old: {0.move_old} " + "Variable: {0.domain}:{0.variable} Grid: {0.grid}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -55,13 +71,15 @@ class Relink(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.move_old == other.move_old and \ - self.grid == other.grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.move_old == other.move_old + and self.grid == other.grid + ) @classmethod def generate_jobs(cls, diags, options): @@ -77,8 +95,8 @@ class Relink(Diagnostic): options_available = ( DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticBoolOption('move_old', True), - DiagnosticOption('grid', '') + DiagnosticBoolOption("move_old", True), + DiagnosticOption("grid", ""), ) options = cls.process_options(options, options_available) job_list = list() @@ -86,9 +104,14 @@ class Relink(Diagnostic): for startdate, member, chunk in chunk_list: job_list.append( Relink( - diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], - options['move_old'], options['grid'] + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["move_old"], + options["grid"], ) ) return job_list @@ -102,8 +125,12 @@ class Relink(Diagnostic): def compute(self): """Run the diagnostic""" self.data_manager.link_file( - self.domain, self.variable, + self.domain, + self.variable, self.var_manager.get_variable(self.variable), - self.startdate, self.member, self.chunk, - move_old=self.move_old, grid=self.grid + self.startdate, + self.member, + self.chunk, + move_old=self.move_old, + grid=self.grid, ) diff --git a/earthdiagnostics/general/relinkall.py b/earthdiagnostics/general/relinkall.py index ef421842..107fa290 100644 --- a/earthdiagnostics/general/relinkall.py +++ b/earthdiagnostics/general/relinkall.py @@ -17,7 +17,7 @@ class RelinkAll(Diagnostic): :type startdate: str """ - alias = 'relinkall' + alias = "relinkall" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate): @@ -25,7 +25,7 @@ class RelinkAll(Diagnostic): self.startdate = startdate def __str__(self): - return 'Relink all output Startdate: {0}'.format(self.startdate) + return "Relink all output Startdate: {0}".format(self.startdate) def __hash__(self): return hash(str(self)) @@ -47,7 +47,7 @@ class RelinkAll(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The Relink All diagnostic has no options') + raise Exception("The Relink All diagnostic has no options") job_list = list() for startdate in diags.config.experiment.startdates: job_list.append(RelinkAll(diags.data_manager, startdate)) diff --git a/earthdiagnostics/general/rewrite.py b/earthdiagnostics/general/rewrite.py index 26cab98c..d681aa33 100644 --- a/earthdiagnostics/general/rewrite.py +++ b/earthdiagnostics/general/rewrite.py @@ -27,10 +27,10 @@ class Rewrite(FixFile): :type domain: ModelingRealm """ - alias = 'rewrite' + alias = "rewrite" "Diagnostic alias for the configuration file" - _STR_PREFIX = 'Rewrite output' + _STR_PREFIX = "Rewrite output" def compute(self): """Run the diagnostic""" diff --git a/earthdiagnostics/general/scale.py b/earthdiagnostics/general/scale.py index 6403ec41..b0470292 100644 --- a/earthdiagnostics/general/scale.py +++ b/earthdiagnostics/general/scale.py @@ -6,9 +6,14 @@ import numpy as np from earthdiagnostics.constants import Basins from earthdiagnostics.general.fix_file import FixFile -from earthdiagnostics.diagnostic import DiagnosticDomainOption, \ - DiagnosticVariableOption, DiagnosticFloatOption, DiagnosticBoolOption, \ - DiagnosticListFrequenciesOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + DiagnosticDomainOption, + DiagnosticVariableOption, + DiagnosticFloatOption, + DiagnosticBoolOption, + DiagnosticListFrequenciesOption, + DiagnosticOption, +) from earthdiagnostics.utils import Utils @@ -38,15 +43,34 @@ class Scale(FixFile): :type domain: ModelingRealm """ - alias = 'scale' + alias = "scale" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - value, offset, domain, variable, grid, - min_limit, max_limit, frequency, apply_mask): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + value, + offset, + domain, + variable, + grid, + min_limit, + max_limit, + frequency, + apply_mask, + ): FixFile.__init__( - self, data_manager, startdate, member, chunk, - domain, variable, grid + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, ) self.value = value self.offset = offset @@ -58,23 +82,27 @@ class Scale(FixFile): self.original_values = None def __str__(self): - return 'Scale output Startdate: {0.startdate} Member: {0.member} '\ - 'Chunk: {0.chunk} Scale value: {0.value} Offset: {0.offset} '\ - 'Variable: {0.domain}:{0.variable} ' \ - 'Frequency: {0.frequency} Apply mask: {0.apply_mask}'.format(self) + return ( + "Scale output Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Scale value: {0.value} Offset: {0.offset} " + "Variable: {0.domain}:{0.variable} " + "Frequency: {0.frequency} Apply mask: {0.apply_mask}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.frequency == other.frequency and \ - self.apply_mask == other.apply_mask and \ - self.value == other.value and \ - self.offset == other.offset + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.frequency == other.frequency + and self.apply_mask == other.apply_mask + and self.value == other.value + and self.offset == other.offset + ) @classmethod def generate_jobs(cls, diags, options): @@ -90,30 +118,36 @@ class Scale(FixFile): options_available = ( DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticFloatOption('value'), - DiagnosticFloatOption('offset'), - DiagnosticOption('grid', ''), - DiagnosticFloatOption('min_limit', float('nan')), - DiagnosticFloatOption('max_limit', float('nan')), + DiagnosticFloatOption("value"), + DiagnosticFloatOption("offset"), + DiagnosticOption("grid", ""), + DiagnosticFloatOption("min_limit", float("nan")), + DiagnosticFloatOption("max_limit", float("nan")), DiagnosticListFrequenciesOption( - 'frequencies', [diags.config.frequency] + "frequencies", [diags.config.frequency] ), - DiagnosticBoolOption('apply_mask', False) + DiagnosticBoolOption("apply_mask", False), ) options = cls.process_options(options, options_available) job_list = list() - for frequency in options['frequencies']: + for frequency in options["frequencies"]: chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: job_list.append( Scale( - diags.data_manager, startdate, member, chunk, - options['value'], options['offset'], options['domain'], - options['variable'], - options['grid'], - options['min_limit'], options['max_limit'], + diags.data_manager, + startdate, + member, + chunk, + options["value"], + options["offset"], + options["domain"], + options["variable"], + options["grid"], + options["min_limit"], + options["max_limit"], frequency, - options['apply_mask'] + options["apply_mask"], ) ) return job_list diff --git a/earthdiagnostics/general/select_levels.py b/earthdiagnostics/general/select_levels.py index 420299aa..a5c99a4f 100644 --- a/earthdiagnostics/general/select_levels.py +++ b/earthdiagnostics/general/select_levels.py @@ -1,8 +1,13 @@ # coding=utf-8 """Extract levels from variable""" from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticDomainOption, DiagnosticVariableListOption, DiagnosticIntOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticVariableListOption, + DiagnosticIntOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -23,11 +28,21 @@ class SelectLevels(Diagnostic): last_level: int """ - alias = 'selev' + alias = "selev" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, grid, first_level, last_level): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, + first_level, + last_level, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -43,21 +58,25 @@ class SelectLevels(Diagnostic): self.result = None def __str__(self): - return 'Select levels Startdate: {0.startdate} Member: {0.member} ' \ - 'Chunk: {0.chunk} Variable: {0.domain}:{0.variable} ' \ - 'Levels: {0.box.min_depth}-{0.box.max_depth} ' \ - 'Grid: {0.grid}'.format(self) + return ( + "Select levels Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Variable: {0.domain}:{0.variable} " + "Levels: {0.box.min_depth}-{0.box.max_depth} " + "Grid: {0.grid}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.box == other.box and \ - self.grid == self.grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.box == other.box + and self.grid == self.grid + ) @classmethod def generate_jobs(cls, diags, options): @@ -73,39 +92,54 @@ class SelectLevels(Diagnostic): options_available = ( DiagnosticDomainOption(), DiagnosticVariableListOption( - diags.data_manager.config.var_manager, 'variables' + diags.data_manager.config.var_manager, "variables" ), - DiagnosticIntOption('first_level'), - DiagnosticIntOption('last_level'), - DiagnosticOption('grid', '') + DiagnosticIntOption("first_level"), + DiagnosticIntOption("last_level"), + DiagnosticOption("grid", ""), ) options = cls.process_options(options, options_available) job_list = list() - variables = options['variables'] + variables = options["variables"] for var in variables: chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: - job_list.append(SelectLevels( - diags.data_manager, startdate, member, chunk, - options['domain'], var, options['grid'], - options['first_level'], options['last_level'] - )) + job_list.append( + SelectLevels( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + options["grid"], + options["first_level"], + options["last_level"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - grid=self.grid, to_modify=True + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + to_modify=True, ) def declare_data_generated(self): """Request data required by the diagnostic""" self.result = self.declare_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, ) def compute(self): @@ -113,16 +147,18 @@ class SelectLevels(Diagnostic): temp = TempFile.get() handler = Utils.open_cdf(self.variable_file.local_file) var_name = "" - for var in ('lev', 'plev'): + for var in ("lev", "plev"): if var in handler.variables: var_name = var continue handler.close() Utils.nco().ncks( - input=self.variable_file.local_file, output=temp, - options='-O -d {1},{0.min_depth},{0.max_depth}'.format( - self.box, var_name), + input=self.variable_file.local_file, + output=temp, + options="-O -d {1},{0.min_depth},{0.max_depth}".format( + self.box, var_name + ), ) self.result.set_local_file(temp) @@ -130,21 +166,23 @@ class SelectLevels(Diagnostic): def _create_var(var_name, var_values, source, destiny): old_var = source.variables[var_name] new_var = destiny.createVariable( - var_name, old_var.dtype, dimensions=(var_name, )) + var_name, old_var.dtype, dimensions=(var_name,) + ) new_var[:] = var_values Utils.copy_attributes(new_var, old_var) - vertices_name = '{0}_vertices'.format(var_name) + vertices_name = "{0}_vertices".format(var_name) if vertices_name in source.variables: var_vertices = source.variables[vertices_name] - if var_name == 'lon': + if var_name == "lon": vertices_values = var_vertices[0:1, ...] else: vertices_values = var_vertices[:, 0:1, :] new_lat_vertices = destiny.createVariable( - vertices_name, var_vertices.dtype, - dimensions=(var_name, 'vertices') + vertices_name, + var_vertices.dtype, + dimensions=(var_name, "vertices"), ) new_lat_vertices[:] = vertices_values Utils.copy_attributes(new_lat_vertices, var_vertices) diff --git a/earthdiagnostics/general/simplify_dimensions.py b/earthdiagnostics/general/simplify_dimensions.py index 9f9ab1b3..b86c52dd 100644 --- a/earthdiagnostics/general/simplify_dimensions.py +++ b/earthdiagnostics/general/simplify_dimensions.py @@ -2,8 +2,11 @@ """Convert i j files to lon lat when there is no interpolation required""" import numpy as np from earthdiagnostics.general.fix_file import FixFile -from earthdiagnostics.diagnostic import DiagnosticOption, \ - DiagnosticDomainOption, DiagnosticVariableListOption +from earthdiagnostics.diagnostic import ( + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticVariableListOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -25,36 +28,55 @@ class SimplifyDimensions(FixFile): data_convention: str """ - alias = 'simdim' + alias = "simdim" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, grid, data_convention): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, + data_convention, + ): FixFile.__init__( - self, data_manager, startdate, member, chunk, - domain, variable, grid + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid, ) - if data_convention in ('cmip6', 'primavera'): - self.lon_name = 'longitude' - self.lat_name = 'latitude' + if data_convention in ("cmip6", "primavera"): + self.lon_name = "longitude" + self.lat_name = "latitude" else: - self.lon_name = 'lon' - self.lat_name = 'lat' + self.lon_name = "lon" + self.lat_name = "lat" def __str__(self): - return 'Simplify dimension Startdate: {0.startdate} ' \ - 'Member: {0.member} Chunk: {0.chunk} ' \ - 'Variable: {0.domain}:{0.variable} Grid: {0.grid}'.format(self) + return ( + "Simplify dimension Startdate: {0.startdate} " + "Member: {0.member} Chunk: {0.chunk} " + "Variable: {0.domain}:{0.variable} Grid: {0.grid}".format(self) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.grid == self.grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.grid == self.grid + ) @classmethod def generate_jobs(cls, diags, options): @@ -70,37 +92,45 @@ class SimplifyDimensions(FixFile): options_available = ( DiagnosticDomainOption(), DiagnosticVariableListOption( - diags.data_manager.config.var_manager, 'variables' + diags.data_manager.config.var_manager, "variables" ), - DiagnosticOption('grid', '')) + DiagnosticOption("grid", ""), + ) options = cls.process_options(options, options_available) job_list = list() - variables = options['variables'] + variables = options["variables"] for var in variables: chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: - job_list.append(SimplifyDimensions( - diags.data_manager, startdate, member, chunk, - options['domain'], var, options['grid'], - diags.config.data_convention - )) + job_list.append( + SimplifyDimensions( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + options["grid"], + diags.config.data_convention, + ) + ) return job_list def compute(self): """Run the diagnostic""" handler = Utils.open_cdf(self.variable_file.local_file) - if 'i' not in handler.dimensions: + if "i" not in handler.dimensions: raise Exception( - 'Variable {0.domain}:{0.variable} does not have i,j ' - 'dimensions'.format(self) + "Variable {0.domain}:{0.variable} does not have i,j " + "dimensions".format(self) ) lat = handler.variables[self.lat_name] lat_values = lat[:, 0:1] # noinspection PyTypeChecker if np.any((lat[:] - lat_values) != 0): raise Exception( - 'Latitude is not constant over i dimension for variable ' - '{0.domain}:{0.variable}'.format(self) + "Latitude is not constant over i dimension for variable " + "{0.domain}:{0.variable}".format(self) ) lon = handler.variables[self.lon_name] @@ -108,31 +138,42 @@ class SimplifyDimensions(FixFile): # noinspection PyTypeChecker if np.any((lon[:] - lon) != 0): raise Exception( - 'Longitude is not constant over j dimension for variable ' - '{0.domain}:{0.variable}'.format(self) + "Longitude is not constant over j dimension for variable " + "{0.domain}:{0.variable}".format(self) ) temp = TempFile.get() - new_file = Utils.open_cdf(temp, 'w') + new_file = Utils.open_cdf(temp, "w") for dim in handler.dimensions.keys(): - if dim in (self.lon_name, self.lat_name, 'i', 'j', 'vertices'): + if dim in (self.lon_name, self.lat_name, "i", "j", "vertices"): continue Utils.copy_dimension( - handler, new_file, dim, - new_names={'i': self.lon_name, 'j': self.lat_name} + handler, + new_file, + dim, + new_names={"i": self.lon_name, "j": self.lat_name}, ) - new_file.createDimension(self.lon_name, handler.dimensions['i'].size) - new_file.createDimension(self.lat_name, handler.dimensions['j'].size) - new_file.createDimension('vertices', 2) + new_file.createDimension(self.lon_name, handler.dimensions["i"].size) + new_file.createDimension(self.lat_name, handler.dimensions["j"].size) + new_file.createDimension("vertices", 2) for var in handler.variables.keys(): - if var in (self.lon_name, self.lat_name, 'i', 'j', - '{0}_vertices'.format(self.lon_name), - '{0}_vertices'.format(self.lat_name)): + if var in ( + self.lon_name, + self.lat_name, + "i", + "j", + "{0}_vertices".format(self.lon_name), + "{0}_vertices".format(self.lat_name), + ): continue - Utils.copy_variable(handler, new_file, var, new_names={ - 'i': self.lon_name, 'j': self.lat_name}) + Utils.copy_variable( + handler, + new_file, + var, + new_names={"i": self.lon_name, "j": self.lat_name}, + ) self._create_var(self.lon_name, lon_values, handler, new_file) self._create_var(self.lat_name, lat_values, handler, new_file) @@ -145,21 +186,23 @@ class SimplifyDimensions(FixFile): def _create_var(var_name, var_values, source, destiny): old_var = source.variables[var_name] new_var = destiny.createVariable( - var_name, old_var.dtype, dimensions=(var_name, )) + var_name, old_var.dtype, dimensions=(var_name,) + ) new_var[:] = var_values Utils.copy_attributes(new_var, old_var) - vertices_name = '{0}_vertices'.format(var_name) + vertices_name = "{0}_vertices".format(var_name) if vertices_name in source.variables: var_vertices = source.variables[vertices_name] - if var_name == 'lon': + if var_name == "lon": vertices_values = var_vertices[0:1, :, 2:] else: vertices_values = var_vertices[:, 0:1, 1:3] new_lat_vertices = destiny.createVariable( - vertices_name, var_vertices.dtype, - dimensions=(var_name, 'vertices') + vertices_name, + var_vertices.dtype, + dimensions=(var_name, "vertices"), ) new_lat_vertices[:] = vertices_values Utils.copy_attributes(new_lat_vertices, var_vertices) diff --git a/earthdiagnostics/general/timemean.py b/earthdiagnostics/general/timemean.py index 2e90da89..a55310ac 100644 --- a/earthdiagnostics/general/timemean.py +++ b/earthdiagnostics/general/timemean.py @@ -7,8 +7,13 @@ import iris.exceptions import numpy as np -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticDomainOption, DiagnosticFrequencyOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticFrequencyOption, + DiagnosticVariableOption, +) from earthdiagnostics.frequency import Frequencies from earthdiagnostics.utils import TempFile, Utils @@ -35,8 +40,17 @@ class TimeMean(Diagnostic): :type grid: str """ - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, frequency, grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -51,10 +65,12 @@ class TimeMean(Diagnostic): self.mean_file = None def __str__(self): - return 'Calculate {0._target_frequency} mean '\ - 'Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Variable: {0.domain}:{0.variable} ' \ - 'Original frequency: {0.frequency} Grid: {0.grid}'.format(self) + return ( + "Calculate {0._target_frequency} mean " + "Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} " + "Variable: {0.domain}:{0.variable} " + "Original frequency: {0.frequency} Grid: {0.grid}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -62,14 +78,16 @@ class TimeMean(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.domain == other.domain and \ - self.variable == other.variable and \ - self.frequency == other.frequency and \ - self.grid == other.grid and \ - self._target_frequency == other._target_frequency + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.frequency == other.frequency + and self.grid == other.grid + and self._target_frequency == other._target_frequency + ) @classmethod def _process_options(cls, diags, options): @@ -77,7 +95,7 @@ class TimeMean(Diagnostic): DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), DiagnosticFrequencyOption(), - DiagnosticOption('grid', '') + DiagnosticOption("grid", ""), ) options = cls.process_options(options, options_available) return options @@ -97,19 +115,30 @@ class TimeMean(Diagnostic): job_list = list() chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: - job_list.append(cls( - diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], - options['frequency'], options['grid'] - )) + job_list.append( + cls( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["frequency"], + options["grid"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - frequency=self.frequency, grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + frequency=self.frequency, + grid=self.grid, ) def compute_mean(self, cube): @@ -130,20 +159,23 @@ class TimeMean(Diagnostic): """Run the diagnostic""" temp = TempFile.get() cube = iris.load_cube(self.variable_file.local_file) - time_centered = [coord for coord in cube.coords( - ) if coord.var_name == 'time_centered'] + time_centered = [ + coord + for coord in cube.coords() + if coord.var_name == "time_centered" + ] if time_centered: cube.remove_coord(time_centered[0]) - iris.coord_categorisation.add_day_of_month(cube, 'time') - iris.coord_categorisation.add_month_number(cube, 'time') - iris.coord_categorisation.add_year(cube, 'time') + iris.coord_categorisation.add_day_of_month(cube, "time") + iris.coord_categorisation.add_month_number(cube, "time") + iris.coord_categorisation.add_year(cube, "time") cube = self.compute_mean(cube) - cube.remove_coord('day_of_month') - cube.remove_coord('month_number') - cube.remove_coord('year') + cube.remove_coord("day_of_month") + cube.remove_coord("month_number") + cube.remove_coord("year") try: - region_coord = cube.coord('region') + region_coord = cube.coord("region") cube.remove_coord(region_coord) except iris.exceptions.CoordinateNotFoundError: region_coord = None @@ -151,14 +183,14 @@ class TimeMean(Diagnostic): iris.save(cube, temp) if region_coord: handler = Utils.open_cdf(temp) - region = handler.createVariable('region', str, ('dim0',)) + region = handler.createVariable("region", str, ("dim0",)) region.standard_name = region_coord.standard_name region[...] = region_coord.points.astype(np.dtype(str)) - handler.variables[self.variable].coordinates += ' region' + handler.variables[self.variable].coordinates += " region" handler.close() - Utils.rename_variable(temp, 'dim0', 'region', False) + Utils.rename_variable(temp, "dim0", "region", False) self.mean_file.set_local_file(temp) @@ -185,14 +217,32 @@ class DailyMean(TimeMean): :type grid: str """ - alias = 'daymean' + alias = "daymean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, frequency, grid): - TimeMean.__init__(self, data_manager, startdate, member, - chunk, domain, variable, frequency, grid) - self._target_frequency = 'daily' + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ): + TimeMean.__init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ) + self._target_frequency = "daily" def compute_mean(self, cube): """ @@ -207,16 +257,19 @@ class DailyMean(TimeMean): iris.cube.Cube """ return cube.aggregated_by( - ['day_of_month', 'month_number', 'year'], - iris.analysis.MEAN + ["day_of_month", "month_number", "year"], iris.analysis.MEAN ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.mean_file = self.declare_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - frequency=Frequencies.daily, grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.daily, + grid=self.grid, ) @@ -242,14 +295,32 @@ class MonthlyMean(TimeMean): :type grid: str """ - alias = 'monmean' + alias = "monmean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, frequency, grid): - TimeMean.__init__(self, data_manager, startdate, member, - chunk, domain, variable, frequency, grid) - self._target_frequency = 'monthly' + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ): + TimeMean.__init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ) + self._target_frequency = "monthly" def compute_mean(self, cube): """ @@ -263,14 +334,18 @@ class MonthlyMean(TimeMean): ------- iris.cube.Cube """ - return cube.aggregated_by(['month_number', 'year'], iris.analysis.MEAN) + return cube.aggregated_by(["month_number", "year"], iris.analysis.MEAN) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.mean_file = self.declare_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - frequency=Frequencies.monthly, grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.monthly, + grid=self.grid, ) @@ -296,14 +371,32 @@ class YearlyMean(TimeMean): :type grid: str """ - alias = 'yearmean' + alias = "yearmean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, frequency, grid): - TimeMean.__init__(self, data_manager, startdate, member, - chunk, domain, variable, frequency, grid) - self._target_frequency = 'yearly' + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ): + TimeMean.__init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + frequency, + grid, + ) + self._target_frequency = "yearly" def compute_mean(self, cube): """ @@ -317,12 +410,16 @@ class YearlyMean(TimeMean): ------- iris.cube.Cube """ - return cube.aggregated_by(['year'], iris.analysis.MEAN) + return cube.aggregated_by(["year"], iris.analysis.MEAN) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.mean_file = self.declare_chunk( - self.domain, self.variable, - self.startdate, self.member, self.chunk, - frequency=Frequencies.yearly, grid=self.grid + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.yearly, + grid=self.grid, ) diff --git a/earthdiagnostics/general/verticalmeanmetersiris.py b/earthdiagnostics/general/verticalmeanmetersiris.py index f88b6f83..093f6330 100644 --- a/earthdiagnostics/general/verticalmeanmetersiris.py +++ b/earthdiagnostics/general/verticalmeanmetersiris.py @@ -5,8 +5,12 @@ import iris.analysis import iris.exceptions from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticFloatOption, \ - DiagnosticDomainOption, DiagnosticVariableListOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticFloatOption, + DiagnosticDomainOption, + DiagnosticVariableListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import TempFile @@ -36,11 +40,12 @@ class VerticalMeanMetersIris(Diagnostic): """ - alias = 'vmean' + alias = "vmean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, - domain, variable, box): + def __init__( + self, data_manager, startdate, member, chunk, domain, variable, box + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -55,17 +60,21 @@ class VerticalMeanMetersIris(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and \ - self.chunk == other.chunk and \ - self.box == other.box and \ - self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + ) def __str__(self): - return 'Vertical mean meters Startdate: {0.startdate} ' \ - 'Member: {0.member} Chunk: {0.chunk} ' \ - 'Variable: {0.domain}:{0.variable} ' \ - 'Box: {0.box}'.format(self) + return ( + "Vertical mean meters Startdate: {0.startdate} " + "Member: {0.member} Chunk: {0.chunk} " + "Variable: {0.domain}:{0.variable} " + "Box: {0.box}".format(self) + ) @classmethod def generate_jobs(cls, diags, options): @@ -82,48 +91,62 @@ class VerticalMeanMetersIris(Diagnostic): options_available = ( DiagnosticDomainOption(), DiagnosticVariableListOption( - diags.data_manager.config.var_manager, 'variable' + diags.data_manager.config.var_manager, "variable" ), - DiagnosticFloatOption('min_depth', -1), - DiagnosticFloatOption('max_depth', -1) + DiagnosticFloatOption("min_depth", -1), + DiagnosticFloatOption("max_depth", -1), ) options = cls.process_options(options, options_available) box = Box(True) - if options['min_depth'] >= 0: - box.min_depth = options['min_depth'] - if options['max_depth'] >= 0: - box.max_depth = options['max_depth'] + if options["min_depth"] >= 0: + box.min_depth = options["min_depth"] + if options["max_depth"] >= 0: + box.max_depth = options["max_depth"] job_list = list() - for var in options['variable']: + for var in options["variable"]: chunk_list = diags.config.experiment.get_chunk_list() for startdate, member, chunk in chunk_list: - job_list.append(VerticalMeanMetersIris( - diags.data_manager, startdate, member, chunk, - options['domain'], var, box - )) + job_list.append( + VerticalMeanMetersIris( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + box, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk( - ModelingRealms.ocean, self.variable, self.startdate, self.member, - self.chunk + ModelingRealms.ocean, + self.variable, + self.startdate, + self.member, + self.chunk, ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.results = self.declare_chunk( - self.domain, self.variable + 'vmean', self.startdate, self.member, - self.chunk, box=self.box + self.domain, + self.variable + "vmean", + self.startdate, + self.member, + self.chunk, + box=self.box, ) def compute(self): """Run the diagnostic""" var_cube = iris.load_cube(self.variable_file.local_file) - lev_names = ('lev', 'depth', 'air_pressure') + lev_names = ("lev", "depth", "air_pressure") coord = None for coord_name in lev_names: try: diff --git a/earthdiagnostics/modelingrealm.py b/earthdiagnostics/modelingrealm.py index ce91670d..f05c703c 100644 --- a/earthdiagnostics/modelingrealm.py +++ b/earthdiagnostics/modelingrealm.py @@ -19,18 +19,20 @@ class ModelingRealm(object): def __init__(self, domain_name): lower_name = domain_name.lower() - if lower_name == 'seaice': - self.name = 'seaIce' - elif lower_name == 'landice': - self.name = 'landIce' - elif lower_name == 'atmoschem': - self.name = 'atmosChem' - elif lower_name == 'ocnbgchem': - self.name = 'ocnBgchem' - elif lower_name in ['ocean', 'atmos', 'land', 'aerosol']: + if lower_name == "seaice": + self.name = "seaIce" + elif lower_name == "landice": + self.name = "landIce" + elif lower_name == "atmoschem": + self.name = "atmosChem" + elif lower_name == "ocnbgchem": + self.name = "ocnBgchem" + elif lower_name in ["ocean", "atmos", "land", "aerosol"]: self.name = domain_name else: - raise ValueError('Modelling realm {0} not recognized!'.format(domain_name)) + raise ValueError( + "Modelling realm {0} not recognized!".format(domain_name) + ) def __eq__(self, other): return other.__class__ == ModelingRealm and self.name == other.name @@ -47,18 +49,26 @@ class ModelingRealm(object): def __repr__(self): return str(self) - def get_varfolder(self, var, ocean_timestep, atmos_timestep, grid=None, frequency=None): + def get_varfolder( + self, var, ocean_timestep, atmos_timestep, grid=None, frequency=None + ): """Get variable folder name for _ folder""" if grid: - var = '{0}-{1}'.format(var, grid) + var = "{0}-{1}".format(var, grid) - if self in [ModelingRealms.ocean, ModelingRealms.seaIce, ModelingRealms.ocnBgchem]: + if self in [ + ModelingRealms.ocean, + ModelingRealms.seaIce, + ModelingRealms.ocnBgchem, + ]: timestep = ocean_timestep else: timestep = atmos_timestep - is_base_frequency = frequency is not None and frequency.frequency.endswith('hr') + is_base_frequency = ( + frequency is not None and frequency.frequency.endswith("hr") + ) if not is_base_frequency and timestep > 0: - return '{0}_f{1}h'.format(var, timestep) + return "{0}_f{1}h".format(var, timestep) return var def get_table_name(self, frequency, data_convention): @@ -75,21 +85,22 @@ class ModelingRealm(object): str """ - if self.name == 'seaIce': - if data_convention.name in ('specs', 'preface'): - prefix = 'OI' + if self.name == "seaIce": + if data_convention.name in ("specs", "preface"): + prefix = "OI" else: - prefix = 'SI' - elif self.name == 'landIce': - prefix = 'LI' + prefix = "SI" + elif self.name == "landIce": + prefix = "LI" else: prefix = self.name[0].upper() if frequency == Frequencies.six_hourly: - table_name = '6hrPlev' + table_name = "6hrPlev" else: - if (frequency in (Frequencies.monthly, Frequencies.climatology)) or data_convention.name not in ('specs', - 'preface'): + if ( + frequency in (Frequencies.monthly, Frequencies.climatology) + ) or data_convention.name not in ("specs", "preface"): table_name = prefix + str(frequency) else: table_name = frequency.frequency @@ -120,28 +131,30 @@ class ModelingRealm(object): """ table_name = self.get_table_name(frequency, data_convention) from earthdiagnostics.variable import CMORTable - return CMORTable(table_name, frequency, 'December 2013', self) + + return CMORTable(table_name, frequency, "December 2013", self) class ModelingRealms(object): """Enumeration of supported modelling realms""" - seaIce = ModelingRealm('seaice') - ocean = ModelingRealm('ocean') - landIce = ModelingRealm('landIce') - atmos = ModelingRealm('atmos') - land = ModelingRealm('land') - aerosol = ModelingRealm('aerosol') - atmosChem = ModelingRealm('atmosChem') - ocnBgchem = ModelingRealm('ocnBgchem') + seaIce = ModelingRealm("seaice") + ocean = ModelingRealm("ocean") + landIce = ModelingRealm("landIce") + atmos = ModelingRealm("atmos") + land = ModelingRealm("land") + aerosol = ModelingRealm("aerosol") + atmosChem = ModelingRealm("atmosChem") + ocnBgchem = ModelingRealm("ocnBgchem") @classmethod def parse(cls, modelling_realm): """ Return the basin matching the given name. - If the parameter modelling_realm is a ModelingRealm instance, directly returns the same - instance. This behaviour is intended to facilitate the development of + If the parameter modelling_realm is a ModelingRealm instance, + directly returns the same instance. This behaviour is intended to + facilitate the development of methods that can either accept a name or a ModelingRealm instance to characterize the modelling realm. @@ -161,14 +174,16 @@ class ModelingRealms(object): """ if isinstance(modelling_realm, ModelingRealm): return modelling_realm - if modelling_realm == '': + if modelling_realm == "": return None for name in cls.__dict__.keys(): - if name.startswith('_'): + if name.startswith("_"): continue # noinspection PyCallByClass value = cls.__getattribute__(cls, name) if isinstance(value, ModelingRealm): if modelling_realm.lower() in [value.name.lower()]: return value - raise ValueError('Modelling realm {0} not recognized!'.format(modelling_realm)) + raise ValueError( + "Modelling realm {0} not recognized!".format(modelling_realm) + ) diff --git a/earthdiagnostics/obsreconmanager.py b/earthdiagnostics/obsreconmanager.py index 892676e2..ad74f5a9 100644 --- a/earthdiagnostics/obsreconmanager.py +++ b/earthdiagnostics/obsreconmanager.py @@ -2,8 +2,8 @@ """ Data management for BSC-Earth conventions -Focused on working with observations and reconstructions as well as with downloaded -but no cmorized models (like ECMWF System 4) +Focused on working with observations and reconstructions as well as +with downloaded but no cmorized models (like ECMWF System 4) """ import os @@ -25,22 +25,34 @@ class ObsReconManager(DataManager): def __init__(self, config): super(ObsReconManager, self).__init__(config) - data_folders = self.config.data_dir.split(':') + data_folders = self.config.data_dir.split(":") self.config.data_dir = None for data_folder in data_folders: - if os.path.isdir(os.path.join(data_folder, self.config.data_type, self.experiment.institute.lower(), - self.experiment.model.lower())): + if os.path.isdir( + os.path.join( + data_folder, + self.config.data_type, + self.experiment.institute.lower(), + self.experiment.model.lower(), + ) + ): self.config.data_dir = data_folder break if not self.config.data_dir: - raise Exception('Can not find model data') - - if self.config.data_type in ('obs', 'recon') and self.experiment.chunk_size != 1: - raise Exception('For obs and recon data chunk_size must be always 1') - - def get_file_path(self, startdate, domain, var, frequency, vartype, - box=None, grid=None): + raise Exception("Can not find model data") + + if ( + self.config.data_type in ("obs", "recon") + and self.experiment.chunk_size != 1 + ): + raise Exception( + "For obs and recon data chunk_size must be always 1" + ) + + def get_file_path( + self, startdate, domain, var, frequency, vartype, box=None, grid=None + ): """ Return the path to a concrete file @@ -64,34 +76,58 @@ class ObsReconManager(DataManager): if frequency is None: frequency = self.config.frequency - folder_path = self._get_folder_path(frequency, domain, var, grid, vartype) + folder_path = self._get_folder_path( + frequency, domain, var, grid, vartype + ) file_name = self._get_file_name(var, startdate, frequency) filepath = os.path.join(folder_path, file_name) return filepath def _get_folder_path(self, frequency, domain, variable, grid, vartype): - var_folder = domain.get_varfolder(variable, self.config.experiment.ocean_timestep, - self.config.experiment.atmos_timestep, grid=grid, frequency=frequency) - folder_path = os.path.join(self.config.data_dir, self.config.data_type, - self.experiment.institute.lower(), - self.experiment.model.lower(), - frequency.folder_name(vartype), - var_folder) + var_folder = domain.get_varfolder( + variable, + self.config.experiment.ocean_timestep, + self.config.experiment.atmos_timestep, + grid=grid, + frequency=frequency, + ) + folder_path = os.path.join( + self.config.data_dir, + self.config.data_type, + self.experiment.institute.lower(), + self.experiment.model.lower(), + frequency.folder_name(vartype), + var_folder, + ) return folder_path def _get_file_name(self, var, startdate, frequency): if startdate: - if self.config.data_type != 'exp' and frequency != Frequencies.weekly: + if ( + self.config.data_type != "exp" + and frequency != Frequencies.weekly + ): startdate = startdate[0:6] - return '{0}_{1}.nc'.format(var, startdate) + return "{0}_{1}.nc".format(var, startdate) else: - return '{0}.nc'.format(var) - - def request_chunk(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - vartype=VariableType.MEAN): + return "{0}.nc".format(var) + + def request_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + ): """ - Request a given file from the CMOR repository to the scratch folder and returns the path to the scratch's copy + Request a given file from the CMOR repository to the scratch folder + and returns the path to the scratch's copy Parameters ---------- @@ -111,13 +147,27 @@ class ObsReconManager(DataManager): """ var = self._get_final_var_name(box, var) - filepath = self.get_file_path(startdate, domain, var, frequency, vartype, box, grid) - Log.debug('{0} requested', filepath) + filepath = self.get_file_path( + startdate, domain, var, frequency, vartype, box, grid + ) + Log.debug("{0} requested", filepath) return self._get_file_from_storage(filepath) # noinspection PyUnusedLocal - def declare_chunk(self, domain, var, startdate, member, chunk, grid=None, region=None, box=None, frequency=None, - vartype=VariableType.MEAN, diagnostic=None): + def declare_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + region=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ Declare a variable chunk to be generated by a diagnostic @@ -147,9 +197,21 @@ class ObsReconManager(DataManager): var = cmor_var.short_name final_name = self._get_final_var_name(box, var) - filepath = self.get_file_path(startdate, domain, final_name, frequency, vartype, box, grid) - netcdf_file = self._declare_generated_file(filepath, domain, final_name, cmor_var, self.config.data_convention, - region, diagnostic, grid, vartype, original_name) + filepath = self.get_file_path( + startdate, domain, final_name, frequency, vartype, box, grid + ) + netcdf_file = self._declare_generated_file( + filepath, + domain, + final_name, + cmor_var, + self.config.data_convention, + region, + diagnostic, + grid, + vartype, + original_name, + ) netcdf_file.frequency = frequency - Log.debug('{0} will be generated', filepath) + Log.debug("{0} will be generated", filepath) return netcdf_file diff --git a/earthdiagnostics/ocean/areamoc.py b/earthdiagnostics/ocean/areamoc.py index cbc1ae83..d8c798f0 100644 --- a/earthdiagnostics/ocean/areamoc.py +++ b/earthdiagnostics/ocean/areamoc.py @@ -6,7 +6,11 @@ import numpy as np from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticIntOption, + DiagnosticBasinOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -40,10 +44,10 @@ class AreaMoc(Diagnostic): :type box: Box """ - alias = 'mocarea' + alias = "mocarea" "Diagnostic alias for the configuration file" - vsftmyz = 'vsftmyz' + vsftmyz = "vsftmyz" def __init__(self, data_manager, startdate, member, chunk, basin, box): Diagnostic.__init__(self, data_manager) @@ -51,19 +55,28 @@ class AreaMoc(Diagnostic): self.startdate = startdate self.member = member self.chunk = chunk - self.required_vars = ['vo'] - self.generated_vars = ['vsftmyz'] + self.required_vars = ["vo"] + self.generated_vars = ["vsftmyz"] self.box = box def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.basin == other.basin and self.box == other.box + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.basin == other.basin + and self.box == other.box + ) def __str__(self): - return 'Area MOC Startdate: {0} Member: {1} Chunk: {2} Box: {3} Basin: {4}'.format( - self.startdate, self.member, self.chunk, self.box, self.basin) + return ( + "Area MOC Startdate: {0} Member: {1} Chunk: {2} Box: {3} " + "Basin: {4}".format( + self.startdate, self.member, self.chunk, self.box, self.basin + ) + ) def __hash__(self): return hash(str(self)) @@ -75,36 +88,62 @@ class AreaMoc(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: minimum latitude, maximum latitude, minimum depth, maximum depth, basin=Global + :param options: minimum latitude, maximum latitude, minimum depth, + maximum depth, basin=Global :type options: list[str] :return: """ - options_available = (DiagnosticIntOption('min_lat'), - DiagnosticIntOption('max_lat'), - DiagnosticIntOption('min_depth'), - DiagnosticIntOption('max_depth'), - DiagnosticBasinOption('basin', Basins().Global)) + options_available = ( + DiagnosticIntOption("min_lat"), + DiagnosticIntOption("max_lat"), + DiagnosticIntOption("min_depth"), + DiagnosticIntOption("max_depth"), + DiagnosticBasinOption("basin", Basins().Global), + ) options = cls.process_options(options, options_available) box = Box() - box.min_lat = options['min_lat'] - box.max_lat = options['max_lat'] - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] + box.min_lat = options["min_lat"] + box.max_lat = options["max_lat"] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(AreaMoc(diags.data_manager, startdate, member, chunk, options['basin'], box)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + AreaMoc( + diags.data_manager, + startdate, + member, + chunk, + options["basin"], + box, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(ModelingRealms.ocean, AreaMoc.vsftmyz, - self.startdate, self.member, self.chunk) + self.variable_file = self.request_chunk( + ModelingRealms.ocean, + AreaMoc.vsftmyz, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.results = self.declare_chunk(ModelingRealms.ocean, AreaMoc.vsftmyz, - self.startdate, self.member, self.chunk, - box=self.box) + self.results = self.declare_chunk( + ModelingRealms.ocean, + AreaMoc.vsftmyz, + self.startdate, + self.member, + self.chunk, + box=self.box, + ) def compute(self): """Run the diagnostic""" @@ -117,16 +156,18 @@ class AreaMoc(Diagnostic): Utils.copy_file(self.variable_file.local_file, temp) handler = Utils.open_cdf(temp) - if 'i' in handler.dimensions: + if "i" in handler.dimensions: handler.close() - nco.ncwa(input=temp, output=temp, options=('-O -a i',)) + nco.ncwa(input=temp, output=temp, options=("-O -a i",)) handler = Utils.open_cdf(temp) - basin_index = np.where(handler.variables['basin'][:] == self.basin.name) - if 'lat' in handler.variables: - lat_name = 'lat' + basin_index = np.where( + handler.variables["basin"][:] == self.basin.name + ) + if "lat" in handler.variables: + lat_name = "lat" else: - lat_name = 'latitude' + lat_name = "latitude" var_lat = handler.variables[lat_name] lat_values = var_lat[:] lat_type = var_lat.dtype @@ -136,41 +177,65 @@ class AreaMoc(Diagnostic): handler.close() if len(basin_index) == 0: - raise Exception('Basin {0} not defined in file') + raise Exception("Basin {0} not defined in file") basin_index = basin_index[0][0] # To select basin and remove dimension - nco.ncwa(input=temp, output=temp, options=('-O -d basin,{0} -a basin'.format(basin_index),)) + nco.ncwa( + input=temp, + output=temp, + options=("-O -d basin,{0} -a basin".format(basin_index),), + ) source = Utils.open_cdf(temp) - destiny = Utils.open_cdf(temp2, 'w') + destiny = Utils.open_cdf(temp2, "w") - Utils.copy_dimension(source, destiny, 'time') - Utils.copy_dimension(source, destiny, 'lev') - Utils.copy_dimension(source, destiny, 'j', new_names={'j': lat_name}) + Utils.copy_dimension(source, destiny, "time") + Utils.copy_dimension(source, destiny, "lev") + Utils.copy_dimension(source, destiny, "j", new_names={"j": lat_name}) lat_variable = destiny.createVariable(lat_name, lat_type, lat_name) lat_variable[:] = lat_values[:] lat_variable.units = lat_units lat_variable.long_name = lat_long_name - Utils.copy_variable(source, destiny, 'lev') - Utils.copy_variable(source, destiny, 'time') - Utils.copy_variable(source, destiny, 'vsftmyz', new_names={'j': lat_name}) + Utils.copy_variable(source, destiny, "lev") + Utils.copy_variable(source, destiny, "time") + Utils.copy_variable( + source, destiny, "vsftmyz", new_names={"j": lat_name} + ) source.close() destiny.close() - nco.ncks(input=temp2, output=temp, - options=('-d lev,{0:.1f},{1:.1f} -d {4},{2:.1f},{3:.1f}'.format(self.box.min_depth, - self.box.max_depth, - self.box.min_lat, - self.box.max_lat, - lat_name),)) + nco.ncks( + input=temp2, + output=temp, + options=( + "-d lev,{0:.1f},{1:.1f} -d {4},{2:.1f},{3:.1f}".format( + self.box.min_depth, + self.box.max_depth, + self.box.min_lat, + self.box.max_lat, + lat_name, + ), + ), + ) cdo.vertmean(input=temp, output=temp2) os.remove(temp) - nco.ncap2(input=temp2, output=temp2, - options=('-s "coslat[{0}]=cos({0}[{0}]*3.141592657/180.0)"'.format(lat_name),)) - nco.ncwa(input=temp2, output=temp2, options=('-w coslat -a {0}'.format(lat_name),)) - nco.ncks(input=temp2, output=temp2, options=('-v vsftmyz,time',)) + nco.ncap2( + input=temp2, + output=temp2, + options=( + '-s "coslat[{0}]=cos({0}[{0}]*3.141592657/180.0)"'.format( + lat_name + ), + ), + ) + nco.ncwa( + input=temp2, + output=temp2, + options=("-w coslat -a {0}".format(lat_name),), + ) + nco.ncks(input=temp2, output=temp2, options=("-v vsftmyz,time",)) self.results.set_local_file(temp2) diff --git a/earthdiagnostics/ocean/averagesection.py b/earthdiagnostics/ocean/averagesection.py index 74f10b94..4b13792b 100644 --- a/earthdiagnostics/ocean/averagesection.py +++ b/earthdiagnostics/ocean/averagesection.py @@ -3,8 +3,13 @@ import os from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticDomainOption, DiagnosticVariableOption, \ - DiagnosticIntOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticDomainOption, + DiagnosticVariableOption, + DiagnosticIntOption, + DiagnosticOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -36,10 +41,20 @@ class AverageSection(Diagnostic): :type box: Box """ - alias = 'avgsection' + alias = "avgsection" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, box, grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + box, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -52,12 +67,21 @@ class AverageSection(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variable == other.variable and self.box == other.box + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.box == other.box + ) def __str__(self): - return 'Average section Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Box: {0.box} ' \ - 'Variable: {0.domain}:{0.variable} Grid: {0.grid}'.format(self) + return ( + "Average section Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Box: {0.box} " + "Variable: {0.domain}:{0.variable} Grid: {0.grid}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -69,49 +93,82 @@ class AverageSection(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: variable, minimum longitude, maximum longitude, minimum latitude, maximum latitude, domain=ocean + :param options: variable, minimum longitude, maximum longitude, + minimum latitude, maximum latitude, domain=ocean :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticIntOption('min_lon'), - DiagnosticIntOption('max_lon'), - DiagnosticIntOption('min_lat'), - DiagnosticIntOption('max_lat'), - DiagnosticOption('grid', '')) + options_available = ( + DiagnosticDomainOption(), + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticIntOption("min_lon"), + DiagnosticIntOption("max_lon"), + DiagnosticIntOption("min_lat"), + DiagnosticIntOption("max_lat"), + DiagnosticOption("grid", ""), + ) options = cls.process_options(options, options_available) box = Box() - box.min_lon = options['min_lon'] - box.max_lon = options['max_lon'] - box.min_lat = options['min_lat'] - box.max_lat = options['max_lat'] + box.min_lon = options["min_lon"] + box.max_lon = options["max_lon"] + box.min_lat = options["min_lat"] + box.max_lat = options["max_lat"] job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(AverageSection(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], box, options['grid'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + AverageSection( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + box, + options["grid"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.grid) + self.variable_file = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.mean = self.declare_chunk(self.domain, self.variable + 'mean', self.startdate, self.member, self.chunk, - box=self.box, grid=self.grid) + self.mean = self.declare_chunk( + self.domain, + self.variable + "mean", + self.startdate, + self.member, + self.chunk, + box=self.box, + grid=self.grid, + ) def compute(self): """Run the diagnostic""" temp = TempFile.get() variable_file = self.variable_file.local_file Utils.cdo().zonmean( - input='-sellonlatbox,{0},{1},{2},{3} {4}'.format( - self.box.min_lon, self.box.max_lon, self.box.min_lat, self.box.max_lat, - variable_file + input="-sellonlatbox,{0},{1},{2},{3} {4}".format( + self.box.min_lon, + self.box.max_lon, + self.box.min_lat, + self.box.max_lat, + variable_file, ), - output=temp + output=temp, ) os.remove(variable_file) - self.mean.set_local_file(temp, rename_var='tos') + self.mean.set_local_file(temp, rename_var="tos") diff --git a/earthdiagnostics/ocean/convectionsites.py b/earthdiagnostics/ocean/convectionsites.py index 8a82d654..321783e1 100644 --- a/earthdiagnostics/ocean/convectionsites.py +++ b/earthdiagnostics/ocean/convectionsites.py @@ -31,7 +31,7 @@ class ConvectionSites(Diagnostic): :type model_version: str """ - alias = 'convection' + alias = "convection" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, model_version): @@ -43,13 +43,19 @@ class ConvectionSites(Diagnostic): self.mlotst_handler = None def __str__(self): - return 'Convection sites Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk) + return "Convection sites Startdate: {0} Member: {1} Chunk: {2}".format( + self.startdate, self.member, self.chunk + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.model_version == other.model_version + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.model_version == other.model_version + ) @classmethod def generate_jobs(cls, diags, options): @@ -63,70 +69,119 @@ class ConvectionSites(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The convection sites diagnostic has no options') + raise Exception("The convection sites diagnostic has no options") job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(ConvectionSites(diags.data_manager, startdate, member, chunk, diags.model_version)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + ConvectionSites( + diags.data_manager, + startdate, + member, + chunk, + diags.model_version, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.mixed_layer = self.request_chunk(ModelingRealms.ocean, 'mlotst', self.startdate, self.member, self.chunk) + self.mixed_layer = self.request_chunk( + ModelingRealms.ocean, + "mlotst", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.sites = self.declare_chunk(ModelingRealms.ocean, 'site', self.startdate, self.member, self.chunk) + self.sites = self.declare_chunk( + ModelingRealms.ocean, + "site", + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - if self.model_version in [Models.ECEARTH_2_3_O1L42, Models.ECEARTH_3_0_O1L46, - Models.NEMO_3_2_O1L42, Models.NEMO_3_3_O1L46, - Models.NEMOVAR_O1L42]: + if self.model_version in [ + Models.ECEARTH_2_3_O1L42, + Models.ECEARTH_3_0_O1L46, + Models.NEMO_3_2_O1L42, + Models.NEMO_3_3_O1L46, + Models.NEMOVAR_O1L42, + ]: labrador = [225, 245, 215, 255] irminger = [245, 290, 215, 245] gin = [260, 310, 245, 291] wedell = [225, 280, 1, 50] - elif self.model_version in [Models.ECEARTH_3_0_O25L46, Models.ECEARTH_3_0_O25L75, - Models.ECEARTH_3_2_O25L75, Models.ECEARTH_3_2_O1L75, - Models.GLORYS2_V1_O25L75]: - raise Exception("Option convection not available yet for {0}".format(self.model_version)) + elif self.model_version in [ + Models.ECEARTH_3_0_O25L46, + Models.ECEARTH_3_0_O25L75, + Models.ECEARTH_3_2_O25L75, + Models.ECEARTH_3_2_O1L75, + Models.GLORYS2_V1_O25L75, + ]: + raise Exception( + "Option convection not available yet for {0}".format( + self.model_version + ) + ) else: - raise Exception("Input grid {0} not recognized".format(self.model_version)) + raise Exception( + "Input grid {0} not recognized".format(self.model_version) + ) mlotst_file = self.mixed_layer.local_file output = TempFile.get() self.mlotst_handler = Utils.open_cdf(mlotst_file) - handler = Utils.open_cdf(output, 'w') - handler.createDimension('time', self.mlotst_handler.variables['time'].shape[0]) - handler.createDimension('region', 4) - Utils.copy_variable(self.mlotst_handler, handler, 'time') - var_region = handler.createVariable('region', str, 'region') - var_gyre = handler.createVariable('site', 'f', ('time', 'region')) - var_gyre.short_name = 'site' - var_gyre.long_name = 'convection sites' - var_gyre.units = 'm^3/s' - - var_region[0] = 'labrador' + handler = Utils.open_cdf(output, "w") + handler.createDimension( + "time", self.mlotst_handler.variables["time"].shape[0] + ) + handler.createDimension("region", 4) + Utils.copy_variable(self.mlotst_handler, handler, "time") + var_region = handler.createVariable("region", str, "region") + var_gyre = handler.createVariable("site", "f", ("time", "region")) + var_gyre.short_name = "site" + var_gyre.long_name = "convection sites" + var_gyre.units = "m^3/s" + + var_region[0] = "labrador" var_gyre[:, 0] = self._convection_site(labrador) - var_region[1] = 'irminger' + var_region[1] = "irminger" var_gyre[:, 1] = self._convection_site(irminger) - var_region[2] = 'gin' + var_region[2] = "gin" var_gyre[:, 2] = self._convection_site(gin) - var_region[3] = 'wedell' + var_region[3] = "wedell" var_gyre[:, 3] = self._convection_site(wedell) self.mlotst_handler.close() handler.close() self.sites.set_local_file(output) - Log.info('Finished convection sites for startdate {0}, member {1}, chunk {2}', - self.startdate, self.member, self.chunk) + Log.info( + "Finished convection sites for startdate {0}, member {1}, " + "chunk {2}", + self.startdate, + self.member, + self.chunk, + ) def _convection_site(self, site): - return np.max(self.mlotst_handler.variables['mlotst'][:, site[2] - 1:site[3] - 1, site[0] - 1:site[1] - 1], - (1, 2)) + return np.max( + self.mlotst_handler.variables["mlotst"][ + :, site[2] - 1:site[3] - 1, site[0] - 1:site[1] - 1 + ], + (1, 2), + ) diff --git a/earthdiagnostics/ocean/cutsection.py b/earthdiagnostics/ocean/cutsection.py index b3a6e926..fb5a1cb3 100644 --- a/earthdiagnostics/ocean/cutsection.py +++ b/earthdiagnostics/ocean/cutsection.py @@ -4,8 +4,13 @@ import numpy as np from bscearth.utils.log import Log from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBoolOption, DiagnosticIntOption, \ - DiagnosticDomainOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBoolOption, + DiagnosticIntOption, + DiagnosticDomainOption, + DiagnosticVariableOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils @@ -38,10 +43,20 @@ class CutSection(Diagnostic): :type value: int """ - alias = 'cutsection' + alias = "cutsection" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, zonal, value): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + zonal, + value, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -62,14 +77,29 @@ class CutSection(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variable == other.variable and self.zonal == other.zonal and \ - self.value == other.value + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.zonal == other.zonal + and self.value == other.value + ) def __str__(self): - return 'Cut section Startdate: {0} Member: {1} Chunk: {2} Variable: {3}:{4} ' \ - 'Zonal: {5} Value: {6}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable, - self.zonal, self.value) + return ( + "Cut section Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3}:{4} Zonal: {5} Value: {6}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variable, + self.zonal, + self.value, + ) + ) @classmethod def generate_jobs(cls, diags, options): @@ -82,42 +112,66 @@ class CutSection(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticBoolOption('zonal'), - DiagnosticIntOption('value'), - DiagnosticDomainOption(default_value=ModelingRealms.ocean)) + options_available = ( + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticBoolOption("zonal"), + DiagnosticIntOption("value"), + DiagnosticDomainOption(default_value=ModelingRealms.ocean), + ) options = cls.process_options(options, options_available) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(CutSection(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], options['zonal'], options['value'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + CutSection( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["zonal"], + options["value"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk) + self.variable_file = self.request_chunk( + self.domain, self.variable, self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.section = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - box=self.box) + self.section = self.declare_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + box=self.box, + ) def compute(self): """Run the diagnostic""" nco = Utils.nco() - handler = Utils.open_cdf('mesh_hgr.nc') - dimi = handler.dimensions['i'].size - dimj = handler.dimensions['j'].size - dimlev = handler.dimensions['lev'].size + handler = Utils.open_cdf("mesh_hgr.nc") + dimi = handler.dimensions["i"].size + dimj = handler.dimensions["j"].size + dimlev = handler.dimensions["lev"].size - lon = handler.variables['lon'][:] - lat = handler.variables['lat'][:] + lon = handler.variables["lon"][:] + lat = handler.variables["lat"][:] handler.close() - handler = Utils.open_cdf('mask.nc') - mask_lev = handler.variables['tmask'][:] + handler = Utils.open_cdf("mask.nc") + mask_lev = handler.variables["tmask"][:] mask_lev = mask_lev.astype(float) # noinspection PyTypeChecker np.place(mask_lev, mask_lev == 0, [1e20]) @@ -153,14 +207,19 @@ class CutSection(Diagnostic): listi[jpt] = jpt listj[jpt] = pos[0][0] - temp = self.data_manager.get_file(self.domain, self.variable, self.startdate, self.member, self.chunk) + temp = self.data_manager.get_file( + self.domain, self.variable, self.startdate, self.member, self.chunk + ) handler = Utils.open_cdf(temp) - dimtime = handler.dimensions['time'].size + dimtime = handler.dimensions["time"].size var_array = handler.variables[self.variable][:] handler.close() - var = np.empty([dimtime, dimlev, size], dtype=handler.variables[self.variable].dtype) + var = np.empty( + [dimtime, dimlev, size], + dtype=handler.variables[self.variable].dtype, + ) new_coord = np.empty(size, dtype=float) if self.zonal: old_coord = lon @@ -168,26 +227,36 @@ class CutSection(Diagnostic): old_coord = lat for jpt in range(0, size): - var[:, :, jpt] = np.maximum(var_array[:, :, listj[jpt], listi[jpt]], - mask_lev[:, :, listj[jpt], listi[jpt]]) + var[:, :, jpt] = np.maximum( + var_array[:, :, listj[jpt], listi[jpt]], + mask_lev[:, :, listj[jpt], listi[jpt]], + ) new_coord[jpt] = old_coord[listj[jpt], listi[jpt]] - nco.ncks(input=temp, output=temp, options=('-O -v lev,time',)) + nco.ncks(input=temp, output=temp, options=("-O -v lev,time",)) handler = Utils.open_cdf(temp) if not self.zonal: - handler.createDimension('lat', size) - coord_var = handler.createVariable('lat', float, 'lat') - file_var = handler.createVariable(self.variable, float, ('time', 'lev', 'lat')) + handler.createDimension("lat", size) + coord_var = handler.createVariable("lat", float, "lat") + file_var = handler.createVariable( + self.variable, float, ("time", "lev", "lat") + ) else: - handler.createDimension('lon', size) - coord_var = handler.createVariable('lon', float, 'lon') - file_var = handler.createVariable(self.variable, float, ('time', 'lev', 'lon')) + handler.createDimension("lon", size) + coord_var = handler.createVariable("lon", float, "lon") + file_var = handler.createVariable( + self.variable, float, ("time", "lev", "lon") + ) coord_var[:] = new_coord[:] file_var[:] = var[:] file_var.missing_value = 1e20 handler.close() self.section.set_local_file(temp) - Log.info('Finished cut section for startdate {0}, member {1}, chunk {2}', - self.startdate, self.member, self.chunk) + Log.info( + "Finished cut section for startdate {0}, member {1}, chunk {2}", + self.startdate, + self.member, + self.chunk, + ) diff --git a/earthdiagnostics/ocean/gyres.py b/earthdiagnostics/ocean/gyres.py index 38d22f9f..f81bef33 100644 --- a/earthdiagnostics/ocean/gyres.py +++ b/earthdiagnostics/ocean/gyres.py @@ -31,7 +31,7 @@ class Gyres(Diagnostic): :type model_version: str """ - alias = 'gyres' + alias = "gyres" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, model_version): @@ -45,12 +45,18 @@ class Gyres(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.model_version == other.model_version + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.model_version == other.model_version + ) def __str__(self): - return 'Gyres Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Model version: {0.model_version}'.format(self) + return ( + "Gyres Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Model version: {0.model_version}".format(self) + ) @classmethod def generate_jobs(cls, diags, options): @@ -64,27 +70,51 @@ class Gyres(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The gyres diagnostic has no options') + raise Exception("The gyres diagnostic has no options") job_list = list() model_version = diags.config.experiment.model_version - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Gyres(diags.data_manager, startdate, member, chunk, model_version)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Gyres( + diags.data_manager, startdate, member, chunk, model_version + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.vsftbarot = self.request_chunk(ModelingRealms.ocean, 'vsftbarot', self.startdate, self.member, self.chunk) + self.vsftbarot = self.request_chunk( + ModelingRealms.ocean, + "vsftbarot", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.gyre = self.declare_chunk(ModelingRealms.ocean, 'gyre', self.startdate, self.member, self.chunk) + self.gyre = self.declare_chunk( + ModelingRealms.ocean, + "gyre", + self.startdate, + self.member, + self.chunk, + ) # noinspection PyPep8Naming def compute(self): """Run the diagnostic""" - if self.model_version in [Models.ECEARTH_2_3_O1L42, Models.ECEARTH_3_0_O1L46, - Models.NEMO_3_2_O1L42, Models.NEMO_3_3_O1L46, - Models.NEMOVAR_O1L42]: + if self.model_version in [ + Models.ECEARTH_2_3_O1L42, + Models.ECEARTH_3_0_O1L46, + Models.NEMO_3_2_O1L42, + Models.NEMO_3_3_O1L46, + Models.NEMOVAR_O1L42, + ]: subpol_n_atl = [230, 275, 215, 245] subpol_n_pac = [70, 145, 195, 235] @@ -95,66 +125,86 @@ class Gyres(Diagnostic): subtrop_ind = [320, 30, 110, 180] acc = [1, 361, 1, 65] - elif self.model_version in [Models.ECEARTH_3_0_O25L46, Models.ECEARTH_3_0_O25L75, Models.GLORYS2_V1_O25L75, - Models.ECEARTH_3_2_O1L75, Models.ECEARTH_3_2_O25L75]: - raise Exception("Option gyres not available yet for {0}".format(self.model_version)) + elif self.model_version in [ + Models.ECEARTH_3_0_O25L46, + Models.ECEARTH_3_0_O25L75, + Models.GLORYS2_V1_O25L75, + Models.ECEARTH_3_2_O1L75, + Models.ECEARTH_3_2_O25L75, + ]: + raise Exception( + "Option gyres not available yet for {0}".format( + self.model_version + ) + ) else: - raise Exception("Input grid {0} not recognized".format(self.model_version)) + raise Exception( + "Input grid {0} not recognized".format(self.model_version) + ) output = TempFile.get() vsftbarot_file = self.vsftbarot.local_file handler_original = Utils.open_cdf(vsftbarot_file) - self.var_vsftbarot = handler_original.variables['vsftbarot'] - handler = Utils.open_cdf(output, 'w') - handler.createDimension('time', handler_original.variables['time'].shape[0]) - handler.createDimension('region', 8) - Utils.copy_variable(handler_original, handler, 'time') - var_region = handler.createVariable('region', str, 'region') - - var_gyre = handler.createVariable('gyre', 'f', ('time', 'region'), fill_value=0.0) + self.var_vsftbarot = handler_original.variables["vsftbarot"] + handler = Utils.open_cdf(output, "w") + handler.createDimension( + "time", handler_original.variables["time"].shape[0] + ) + handler.createDimension("region", 8) + Utils.copy_variable(handler_original, handler, "time") + var_region = handler.createVariable("region", str, "region") + + var_gyre = handler.createVariable( + "gyre", "f", ("time", "region"), fill_value=0.0 + ) var_gyre.valid_max = 2e8 var_gyre.valid_min = 0.0 - var_gyre.short_name = 'gyre' - var_gyre.long_name = 'gyre' - var_gyre.units = 'm^3/s' + var_gyre.short_name = "gyre" + var_gyre.long_name = "gyre" + var_gyre.units = "m^3/s" - var_region[0] = 'subpolNAtl' + var_region[0] = "subpolNAtl" var_gyre[:, 0] = self._gyre(subpol_n_atl, True) - Log.debug('subpolNAtl: {0}', var_gyre[:, 0]) + Log.debug("subpolNAtl: {0}", var_gyre[:, 0]) - var_region[1] = 'subpolNPac' + var_region[1] = "subpolNPac" var_gyre[:, 1] = self._gyre(subpol_n_pac, True) - Log.debug('subpolNPac: {0}', var_gyre[:, 1]) + Log.debug("subpolNPac: {0}", var_gyre[:, 1]) - var_region[2] = 'subtropNPac' + var_region[2] = "subtropNPac" var_gyre[:, 2] = self._gyre(subtrop_n_pac) - Log.debug('subtropNPac: {0}', var_gyre[:, 2]) + Log.debug("subtropNPac: {0}", var_gyre[:, 2]) - var_region[3] = 'subtropSPac' + var_region[3] = "subtropSPac" var_gyre[:, 3] = self._gyre(subtrop_s_pac) - Log.debug('subtropSPac: {0}', var_gyre[:, 3]) + Log.debug("subtropSPac: {0}", var_gyre[:, 3]) - var_region[4] = 'subtropNAtl' + var_region[4] = "subtropNAtl" var_gyre[:, 4] = self._gyre(subtrop_n_atl) - Log.debug('subtropNAtl: {0}', var_gyre[:, 4]) + Log.debug("subtropNAtl: {0}", var_gyre[:, 4]) - var_region[5] = 'subtropSAtl' + var_region[5] = "subtropSAtl" var_gyre[:, 5] = self._gyre(subtrop_s_atl) - Log.debug('subtropSAtl: {0}', var_gyre[:, 5]) + Log.debug("subtropSAtl: {0}", var_gyre[:, 5]) - var_region[6] = 'subtropInd' + var_region[6] = "subtropInd" var_gyre[:, 6] = self._gyre(subtrop_ind) - Log.debug('subtropInd: {0}', var_gyre[:, 6]) + Log.debug("subtropInd: {0}", var_gyre[:, 6]) - var_region[7] = 'ACC' + var_region[7] = "ACC" var_gyre[:, 7] = self._gyre(acc) - Log.debug('ACC: {0}', var_gyre[:, 7]) + Log.debug("ACC: {0}", var_gyre[:, 7]) handler.close() handler_original.close() self.gyre.set_local_file(output) - Log.info('Finished gyres for startdate {0}, member {1}, chunk {2}', self.startdate, self.member, self.chunk) + Log.info( + "Finished gyres for startdate {0}, member {1}, chunk {2}", + self.startdate, + self.member, + self.chunk, + ) def _gyre(self, site, invert=False): if invert: @@ -165,17 +215,44 @@ class Gyres(Diagnostic): def _extract_section(self, site): if site[2] <= site[3]: if site[0] <= site[1]: - return self.var_vsftbarot[:, site[2] - 1:site[3] - 1, site[0] - 1:site[1] - 1] + return self.var_vsftbarot[ + :, site[2] - 1:site[3] - 1, site[0] - 1:site[1] - 1 + ] else: - return np.concatenate((self.var_vsftbarot[:, site[2] - 1:site[3] - 1, site[0] - 1:], - self.var_vsftbarot[:, site[2] - 1:site[3] - 1, :site[1] - 1]), axis=2) + return np.concatenate( + ( + self.var_vsftbarot[ + :, site[2] - 1:site[3] - 1, site[0] - 1: + ], + self.var_vsftbarot[ + :, site[2] - 1:site[3] - 1, :site[1] - 1 + ], + ), + axis=2, + ) else: if site[0] <= site[1]: - return np.concatenate((self.var_vsftbarot[:, site[2] - 1:, site[0] - 1: site[1] - 1], - self.var_vsftbarot[:, :site[3] - 1, site[0] - 1: site[1] - 1]), axis=1) + return np.concatenate( + ( + self.var_vsftbarot[ + :, site[2] - 1:, site[0] - 1:site[1] - 1 + ], + self.var_vsftbarot[ + :, :site[3] - 1, site[0] - 1:site[1] - 1 + ], + ), + axis=1, + ) else: - temp = np.concatenate((self.var_vsftbarot[:, site[2] - 1:, :], - self.var_vsftbarot[:, :site[3] - 1, :]), axis=1) - return np.concatenate((temp[:, :, site[0] - 1:], - temp[:, :, :site[1] - 1]), axis=2) + temp = np.concatenate( + ( + self.var_vsftbarot[:, site[2] - 1:, :], + self.var_vsftbarot[:, :site[3] - 1, :], + ), + axis=1, + ) + return np.concatenate( + (temp[:, :, site[0] - 1:], temp[:, :, :site[1] - 1]), + axis=2, + ) diff --git a/earthdiagnostics/ocean/heatcontent.py b/earthdiagnostics/ocean/heatcontent.py index de21b94c..6082290b 100644 --- a/earthdiagnostics/ocean/heatcontent.py +++ b/earthdiagnostics/ocean/heatcontent.py @@ -7,7 +7,11 @@ import numpy as np from earthdiagnostics import cdftools from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption, DiagnosticIntOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBasinOption, + DiagnosticIntOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -30,17 +34,29 @@ class HeatContent(Diagnostic): :type member: int :param chunk: chunk's number :type chunk: int - :param mixed_layer: If 1, restricts calculation to the mixed layer, if -1 exclude it. If 0, no effect + :param mixed_layer: If 1, restricts calculation to the mixed layer, + if -1 exclude it. If 0, no effect :type mixed_layer: int :param box: box to use for the average :type box: Box """ - alias = 'ohc' + alias = "ohc" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, basin, mixed_layer, box, min_level, max_level): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + basin, + mixed_layer, + box, + min_level, + max_level, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -54,12 +70,21 @@ class HeatContent(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.basin == other.basin and self.mxloption == other.mxloption + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.basin == other.basin + and self.mxloption == other.mxloption + ) def __str__(self): - return 'Heat content Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Mixed layer: {0.mxloption} ' \ - 'Box: {0.box} Basin: {0.basin}'.format(self) + return ( + "Heat content Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Mixed layer: {0.mxloption} " + "Box: {0.box} Basin: {0.basin}".format(self) + ) @classmethod def generate_jobs(cls, diags, options): @@ -68,19 +93,22 @@ class HeatContent(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: basin, mixed layer option (1 to only compute at the mixed layer, -1 to exclude it, 0 to ignore), - minimum depth, maximum depth + :param options: basin, mixed layer option (1 to only compute at + the mixed layer, -1 to exclude it, 0 to ignore), + minimum depth, maximum depth :type options: list[str] :return: """ - options_available = (DiagnosticBasinOption('basin'), - DiagnosticIntOption('mixed_layer', None, -1, 1), - DiagnosticIntOption('min_depth'), - DiagnosticIntOption('max_depth')) + options_available = ( + DiagnosticBasinOption("basin"), + DiagnosticIntOption("mixed_layer", None, -1, 1), + DiagnosticIntOption("min_depth"), + DiagnosticIntOption("max_depth"), + ) options = cls.process_options(options, options_available) box = Box() - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] min_level = 0 max_level = 0 @@ -89,9 +117,24 @@ class HeatContent(Diagnostic): max_level, min_level = cls._get_levels_from_meters(box) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(HeatContent(diags.data_manager, startdate, member, chunk, - options['basin'], options['mixed_layer'], box, min_level, max_level)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + HeatContent( + diags.data_manager, + startdate, + member, + chunk, + options["basin"], + options["mixed_layer"], + box, + min_level, + max_level, + ) + ) return job_list @classmethod @@ -134,27 +177,39 @@ class HeatContent(Diagnostic): @classmethod def _read_level_values(cls): - handler = Utils.open_cdf('mesh_zgr.nc') - if 'gdepw_1d' in handler.variables: - depth_w = handler.variables['gdepw_1d'][0, :] - elif 'gdepw_0' in handler.variables: - depth_w = handler.variables['gdepw_0'][0, :] + handler = Utils.open_cdf("mesh_zgr.nc") + if "gdepw_1d" in handler.variables: + depth_w = handler.variables["gdepw_1d"][0, :] + elif "gdepw_0" in handler.variables: + depth_w = handler.variables["gdepw_0"][0, :] else: - raise Exception('gdepw 1D variable can not be found') - if 'gdept_1d' in handler.variables: - depth_t = handler.variables['gdept_1d'][0, :] - elif 'gdept_0' in handler.variables: - depth_t = handler.variables['gdept_0'][0, :] + raise Exception("gdepw 1D variable can not be found") + if "gdept_1d" in handler.variables: + depth_t = handler.variables["gdept_1d"][0, :] + elif "gdept_0" in handler.variables: + depth_t = handler.variables["gdept_0"][0, :] else: - raise Exception('gdept 1D variable can not be found') + raise Exception("gdept 1D variable can not be found") handler.close() return depth_t, depth_w def request_data(self): """Request data required by the diagnostic""" - self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk) + self.thetao = self.request_chunk( + ModelingRealms.ocean, + "thetao", + self.startdate, + self.member, + self.chunk, + ) if self.mxloption != 0: - self.mlotst = self.request_chunk(ModelingRealms.ocean, 'mlotst', self.startdate, self.member, self.chunk) + self.mlotst = self.request_chunk( + ModelingRealms.ocean, + "mlotst", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" @@ -164,10 +219,24 @@ class HeatContent(Diagnostic): else: box_save = self.box - self.heatcsum = self.declare_chunk(ModelingRealms.ocean, 'heatcsum', self.startdate, self.member, self.chunk, - box=box_save, region=self.basin.name) - self.heatcmean = self.declare_chunk(ModelingRealms.ocean, 'heatcvmean', self.startdate, self.member, self.chunk, - box=box_save, region=self.basin.name) + self.heatcsum = self.declare_chunk( + ModelingRealms.ocean, + "heatcsum", + self.startdate, + self.member, + self.chunk, + box=box_save, + region=self.basin.name, + ) + self.heatcmean = self.declare_chunk( + ModelingRealms.ocean, + "heatcvmean", + self.startdate, + self.member, + self.chunk, + box=box_save, + region=self.basin.name, + ) def compute(self): """Run the diagnostic""" @@ -175,11 +244,15 @@ class HeatContent(Diagnostic): temperature_file = TempFile.get() Utils.copy_file(self.thetao.local_file, temperature_file) if self.mxloption != 0: - nco.ncks(input=self.mlotst.local_file, output=temperature_file, options=('-A -v mlotst',)) + nco.ncks( + input=self.mlotst.local_file, + output=temperature_file, + options=("-A -v mlotst",), + ) para = list() if self.min_level != 0: - para.append('-zoom') + para.append("-zoom") para.append(0) para.append(0) para.append(0) @@ -187,47 +260,59 @@ class HeatContent(Diagnostic): para.append(self.min_level) para.append(self.max_level) if self.mxloption != 0: - para.append('-mxloption') + para.append("-mxloption") para.append(str(self.mxloption)) if self.basin != Basins().Global: - handler = Utils.open_cdf('mask_regions.3d.nc') + handler = Utils.open_cdf("mask_regions.3d.nc") if self.basin.name not in handler.variables: - raise Exception('Basin {0} is not defined on mask_regions.nc'.format(self.basin.name)) + raise Exception( + "Basin {0} is not defined on mask_regions.nc".format( + self.basin.name + ) + ) handler.close() - para.append('-M') - para.append('mask_regions.3d.nc') + para.append("-M") + para.append("mask_regions.3d.nc") para.append(self.basin.name) temp2 = TempFile.get() - cdftools.run('cdfheatc', options=para, input_file=temperature_file, output_file=temp2, input_option='-f') + cdftools.run( + "cdfheatc", + options=para, + input_file=temperature_file, + output_file=temp2, + input_option="-f", + ) results = Utils.open_cdf(temp2) heatcsum_temp = TempFile.get() heatcvmean_temp = TempFile.get() - nco.ncks(input=temperature_file, output=heatcsum_temp, options=('-O -v time',)) + nco.ncks( + input=temperature_file, + output=heatcsum_temp, + options=("-O -v time",), + ) shutil.copy(heatcsum_temp, heatcvmean_temp) heatcsum_handler = Utils.open_cdf(heatcsum_temp) - thc = heatcsum_handler.createVariable('heatcsum', float, 'time') - thc.standard_name = "integral_of_sea_water_potential_temperature_expressed_as_heat_content" + thc = heatcsum_handler.createVariable("heatcsum", float, "time") thc.long_name = "Total heat content" thc.units = "J" - thc[:] = results.variables['heatc3d'][:, 0, 0] + thc[:] = results.variables["heatc3d"][:, 0, 0] heatcsum_handler.close() heatcvmean_handler = Utils.open_cdf(heatcvmean_temp) - uhc = heatcvmean_handler.createVariable('heatcvmean', float, 'time') - uhc.standard_name = "integral_of_sea_water_potential_temperature_expressed_as_heat_content" + uhc = heatcvmean_handler.createVariable("heatcvmean", float, "time") uhc.long_name = "Heat content per unit volume" uhc.units = "J*m^-3" - uhc[:] = results.variables['heatc3dpervol'][:, 0, 0] + uhc[:] = results.variables["heatc3dpervol"][:, 0, 0] heatcvmean_handler.close() results.close() - Utils.setminmax(heatcsum_temp, 'heatcsum') + Utils.setminmax(heatcsum_temp, "heatcsum") self.heatcsum.set_local_file(heatcsum_temp) - Utils.setminmax(heatcvmean_temp, 'heatcvmean') + Utils.setminmax(heatcvmean_temp, "heatcvmean") self.heatcmean.set_local_file(heatcvmean_temp) diff --git a/earthdiagnostics/ocean/heatcontentlayer.py b/earthdiagnostics/ocean/heatcontentlayer.py index f6b51ea2..178bac42 100644 --- a/earthdiagnostics/ocean/heatcontentlayer.py +++ b/earthdiagnostics/ocean/heatcontentlayer.py @@ -3,19 +3,18 @@ import numpy as np from earthdiagnostics.box import Box -from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinListOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticIntOption, + DiagnosticBasinListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile from bscearth.utils.log import Log -import diagonals from diagonals import ohc from diagonals.mesh_helpers.nemo import Nemo -import iris -import iris.coords - import netCDF4 @@ -43,11 +42,24 @@ class HeatContentLayer(Diagnostic): :type box: Box """ - alias = 'ohclayer' + alias = "ohclayer" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, box, areas, - weight, layers, basins, data_convention, min_level, max_level): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + box, + areas, + weight, + layers, + basins, + data_convention, + min_level, + max_level, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -62,8 +74,11 @@ class HeatContentLayer(Diagnostic): self.max_level = max_level def __str__(self): - return 'Heat content layer Startdate: {0} Member: {1} Chunk: {2} Box: {3}'.format(self.startdate, self.member, - self.chunk, self.box) + return ( + f"Heat content layer Startdate: {self.startdate} " + f"Member: {self.member} Chunk: {self.chunk} " + f"Box: {self.box}" + ) @classmethod def generate_jobs(cls, diags, options): @@ -76,21 +91,21 @@ class HeatContentLayer(Diagnostic): :type options: list[str] """ options_available = ( - DiagnosticIntOption('min_depth'), - DiagnosticIntOption('max_depth'), - DiagnosticBasinListOption('basins', 'global') + DiagnosticIntOption("min_depth"), + DiagnosticIntOption("max_depth"), + DiagnosticBasinListOption("basins", "global"), ) options = cls.process_options(options, options_available) box = Box(True) - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] - basins = options['basins'] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") layers = ((box.min_depth, box.max_depth),) @@ -106,38 +121,53 @@ class HeatContentLayer(Diagnostic): e3t = mesh.get_k_length() mask = mesh.get_landsea_mask() - depth = mesh.get_depth(cell_point='W') + depth = mesh.get_depth(cell_point="W") weight = ohc.get_weights(layers, mask, e3t, depth) max_level, min_level = cls._get_used_levels(weight) weight[0] = weight[0][:, min_level:max_level, :, :] del mask, depth, e3t - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(HeatContentLayer( - diags.data_manager, startdate, member, chunk, box, - areas, weight, layers, basins, - diags.config.data_convention, min_level, max_level - )) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + HeatContentLayer( + diags.data_manager, + startdate, + member, + chunk, + box, + areas, + weight, + layers, + basins, + diags.config.data_convention, + min_level, + max_level, + ) + ) return job_list @classmethod def _get_mesh_info(cls): - handler = Utils.open_cdf('mesh_zgr.nc') + handler = Utils.open_cdf("mesh_zgr.nc") # mask = Utils.get_mask(options['basin']) - mask = handler.variables['tmask'][:] - if 'e3t' in handler.variables: - e3t = handler.variables['e3t'][:] - elif 'e3t_0' in handler.variables: - e3t = handler.variables['e3t_0'][:] + mask = handler.variables["tmask"][:] + if "e3t" in handler.variables: + e3t = handler.variables["e3t"][:] + elif "e3t_0" in handler.variables: + e3t = handler.variables["e3t_0"][:] else: - raise Exception('e3t variable can not be found') - if 'gdepw' in handler.variables: - depth = handler.variables['gdepw'][:] - elif 'gdepw_0' in handler.variables: - depth = handler.variables['gdepw_0'][:] + raise Exception("e3t variable can not be found") + if "gdepw" in handler.variables: + depth = handler.variables["gdepw"][:] + elif "gdepw_0" in handler.variables: + depth = handler.variables["gdepw_0"][:] else: - raise Exception('gdepw variable can not be found') + raise Exception("gdepw variable can not be found") e3t_3d = e3t.shape != depth.shape if e3t_3d: mask = e3t_3d * mask @@ -150,28 +180,45 @@ class HeatContentLayer(Diagnostic): def request_data(self): """Request data required by the diagnostic""" - self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', - self.startdate, self.member, - self.chunk) + self.thetao = self.request_chunk( + ModelingRealms.ocean, + "thetao", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.heatc = self.declare_chunk(ModelingRealms.ocean, 'heatc', - self.startdate, self.member, - self.chunk, box=self.box) - self.heatcsum = self.declare_chunk(ModelingRealms.ocean, 'heatcsum', - self.startdate, self.member, - self.chunk, box=self.box) + self.heatc = self.declare_chunk( + ModelingRealms.ocean, + "heatc", + self.startdate, + self.member, + self.chunk, + box=self.box, + ) + self.heatcsum = self.declare_chunk( + ModelingRealms.ocean, + "heatcsum", + self.startdate, + self.member, + self.chunk, + box=self.box, + ) @classmethod def _get_used_levels(cls, weight): - # Now we will reduce to the levels with any weight != 0 to avoid loading too much data on memory + # Now we will reduce to the levels with any weight != 0 to avoid + # loading too much data on memory levels = weight[0].shape[1] min_level = 0 while min_level < levels and not weight[0][:, min_level, :].any(): min_level += 1 max_level = min_level - while max_level < (levels - 1) and weight[0][:, max_level + 1, :].any(): + while ( + max_level < (levels - 1) and weight[0][:, max_level + 1, :].any() + ): max_level += 1 return max_level, min_level @@ -179,54 +226,66 @@ class HeatContentLayer(Diagnostic): """Run the diagnostic""" thetao_file = TempFile.get() results = TempFile.get() - results1D = TempFile.get() + results_1d = TempFile.get() Utils.copy_file(self.thetao.local_file, thetao_file) handler = Utils.open_cdf(thetao_file) - Utils.convert_units(handler.variables['thetao'], 'K') - heatc_sl, heatc_sl1D = ohc.compute( - self.layers, self.weight, - handler.variables['thetao'][:, self.min_level:self.max_level, :, :], - self.areas) + Utils.convert_units(handler.variables["thetao"], "K") + heatc_sl, heatc_sl_1d = ohc.compute( + self.layers, + self.weight, + handler.variables["thetao"][ + :, self.min_level:self.max_level, :, : + ], + self.areas, + ) handler.sync() - handler.renameVariable('thetao', 'heatc_sl') + handler.renameVariable("thetao", "heatc_sl") results = TempFile.get() - handler_results = Utils.open_cdf(results, 'w') - lat_name = next(alias for alias in ('lat', 'latitude') - if alias in handler.variables.keys()) - lon_name = next(alias for alias in ('lon', 'longitude') - if alias in handler.variables.keys()) - Utils.copy_variable(handler, handler_results, 'time', True, True) - Utils.copy_variable(handler, handler_results, 'i', False, True) - Utils.copy_variable(handler, handler_results, 'j', False, True) + handler_results = Utils.open_cdf(results, "w") + lat_name = next( + alias + for alias in ("lat", "latitude") + if alias in handler.variables.keys() + ) + lon_name = next( + alias + for alias in ("lon", "longitude") + if alias in handler.variables.keys() + ) + Utils.copy_variable(handler, handler_results, "time", True, True) + Utils.copy_variable(handler, handler_results, "i", False, True) + Utils.copy_variable(handler, handler_results, "j", False, True) Utils.copy_variable(handler, handler_results, lat_name, True, True) Utils.copy_variable(handler, handler_results, lon_name, True, True) - var = handler_results.createVariable('heatc', float, - ('time', 'j', 'i'), - fill_value=1.e20) - var.units = 'J m-2''' - var.coordinates = ' '.join((lat_name, lon_name)) + var = handler_results.createVariable( + "heatc", float, ("time", "j", "i"), fill_value=1.0e20 + ) + var.units = "J m-2" "" + var.coordinates = " ".join((lat_name, lon_name)) handler_results.sync() # temporary fix, needs to loop over layers - handler_results.variables['heatc'][:] = heatc_sl[0] + handler_results.variables["heatc"][:] = heatc_sl[0] handler_results.close() - results1D = TempFile.get() - handler_results1D = Utils.open_cdf(results1D, 'w') - Utils.copy_variable(handler, handler_results1D, 'time', True, True) - handler_results1D.createDimension('region', len(self.basins)) - handler_results1D.createDimension('region_length', 50) - var_region = handler_results1D.createVariable( - 'region', 'S1', ('region', 'region_length')) - var_ohc1D = handler_results1D.createVariable( - 'heatcsum', float, ('time', 'region',),) - handler_results1D.sync() + results_1d = TempFile.get() + handler_results_1d = Utils.open_cdf(results_1d, "w") + Utils.copy_variable(handler, handler_results_1d, "time", True, True) + handler_results_1d.createDimension("region", len(self.basins)) + handler_results_1d.createDimension("region_length", 50) + var_region = handler_results_1d.createVariable( + "region", "S1", ("region", "region_length") + ) + var_ohc_1d = handler_results_1d.createVariable( + "heatcsum", float, ("time", "region",), + ) + handler_results_1d.sync() for i, basin in enumerate(self.basins): var_region[i, ...] = netCDF4.stringtoarr(basin.name, 50) - var_ohc1D[..., i] = heatc_sl1D[i] - handler_results1D.close() + var_ohc_1d[..., i] = heatc_sl_1d[i] + handler_results_1d.close() - Utils.setminmax(results, 'heatc') + Utils.setminmax(results, "heatc") self.heatc.set_local_file(results) - self.heatcsum.set_local_file(results1D) + self.heatcsum.set_local_file(results_1d) diff --git a/earthdiagnostics/ocean/interpolate.py b/earthdiagnostics/ocean/interpolate.py index 98ea91e4..cad9c105 100644 --- a/earthdiagnostics/ocean/interpolate.py +++ b/earthdiagnostics/ocean/interpolate.py @@ -6,8 +6,13 @@ import threading from bscearth.utils.log import Log -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticDomainOption, DiagnosticBoolOption, \ - DiagnosticVariableListOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticBoolOption, + DiagnosticVariableListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -40,13 +45,24 @@ class Interpolate(Diagnostic): :type model_version: str """ - alias = 'interp' + alias = "interp" "Diagnostic alias for the configuration file" lock = threading.Lock() - def __init__(self, data_manager, startdate, member, chunk, domain, variable, target_grid, model_version, - invert_lat, original_grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + target_grid, + model_version, + invert_lat, + original_grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -56,7 +72,7 @@ class Interpolate(Diagnostic): self.model_version = model_version self.required_vars = [variable] self.generated_vars = [variable] - self.tempTemplate = '' + self.tempTemplate = "" self.grid = target_grid self.invert_latitude = invert_lat self.original_grid = original_grid @@ -64,17 +80,34 @@ class Interpolate(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.model_version == other.model_version and self.domain == other.domain and \ - self.variable == other.variable and self.grid == other.grid and \ - self.invert_latitude == other.invert_latitude and self.original_grid == other.original_grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.model_version == other.model_version + and self.domain == other.domain + and self.variable == other.variable + and self.grid == other.grid + and self.invert_latitude == other.invert_latitude + and self.original_grid == other.original_grid + ) def __str__(self): - return 'Interpolate Startdate: {0} Member: {1} Chunk: {2} ' \ - 'Variable: {3}:{4} Target grid: {5} Invert lat: {6} ' \ - 'Model: {7} Original grid: {8}'.format(self.startdate, self.member, self.chunk, self.domain, - self.variable, self.grid, self.invert_latitude, - self.model_version, self.original_grid) + return ( + "Interpolate Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3}:{4} Target grid: {5} Invert lat: {6} " + "Model: {7} Original grid: {8}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variable, + self.grid, + self.invert_latitude, + self.model_version, + self.original_grid, + ) + ) def __hash__(self): return hash(str(self)) @@ -90,42 +123,74 @@ class Interpolate(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticOption('target_grid'), - DiagnosticVariableListOption(diags.data_manager.config.var_manager, 'variable'), - DiagnosticDomainOption(default_value=ModelingRealms.ocean), - DiagnosticBoolOption('invert_lat', False), - DiagnosticOption('original_grid', '')) + options_available = ( + DiagnosticOption("target_grid"), + DiagnosticVariableListOption( + diags.data_manager.config.var_manager, "variable" + ), + DiagnosticDomainOption(default_value=ModelingRealms.ocean), + DiagnosticBoolOption("invert_lat", False), + DiagnosticOption("original_grid", ""), + ) options = cls.process_options(options, options_available) job_list = list() - for var in options['variable']: - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): + for var in options["variable"]: + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): job_list.append( - Interpolate(diags.data_manager, startdate, member, chunk, - options['domain'], var, options['target_grid'], - diags.config.experiment.model_version, options['invert_lat'], options['original_grid'])) + Interpolate( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + options["target_grid"], + diags.config.experiment.model_version, + options["invert_lat"], + options["original_grid"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.original = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.original_grid) + self.original = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.original_grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.regridded = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.grid) + self.regridded = self.declare_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def compute(self): """Run the diagnostic""" variable_file = TempFile.get() Utils.copy_file(self.original.local_file, variable_file) - Utils.rename_variables(variable_file, {'i': 'x', 'j': 'y'}, must_exist=False) + Utils.rename_variables( + variable_file, {"i": "x", "j": "y"}, must_exist=False + ) cdo = Utils.cdo() nco = Utils.nco() handler = Utils.open_cdf(variable_file) - if 'lev' in handler.dimensions: - num_levels = handler.dimensions['lev'].size + if "lev" in handler.dimensions: + num_levels = handler.dimensions["lev"].size has_levels = True else: num_levels = 1 @@ -137,67 +202,95 @@ class Interpolate(Diagnostic): temp = TempFile.get() if has_levels: Interpolate.lock.acquire() - nco.ncrcat(input=self._get_level_file(0), output=temp, - options="-n {0},2,1 -v '{1}'".format(num_levels, self.variable)) + nco.ncrcat( + input=self._get_level_file(0), + output=temp, + options="-n {0},2,1 -v '{1}'".format( + num_levels, self.variable + ), + ) Interpolate.lock.release() else: Utils.move_file(self._get_level_file(0), temp) handler = Utils.open_cdf(temp) - if 'record' in handler.dimensions: - handler.renameDimension('record', 'lev') + if "record" in handler.dimensions: + handler.renameDimension("record", "lev") handler.close() - nco.ncpdq(input=temp, output=temp, options=('-h -a time,lev',)) + nco.ncpdq(input=temp, output=temp, options=("-h -a time,lev",)) if has_levels: - nco.ncks(input=variable_file, output=temp, options=('-A -v lev',)) + nco.ncks(input=variable_file, output=temp, options=("-A -v lev",)) for lev in range(0, num_levels): os.remove(self._get_level_file(lev)) temp2 = TempFile.get() - cdo.setgrid('t106grid', input=temp, output=temp2) + cdo.setgrid("t106grid", input=temp, output=temp2) os.remove(temp) if self.invert_latitude: cdo.invertlatdata(input=temp2, output=temp) shutil.move(temp, temp2) if not has_levels: - nco.ncks(input=temp2, output=temp2, options=('-v {0},lat,lon,time'.format(self.variable),)) + nco.ncks( + input=temp2, + output=temp2, + options=("-v {0},lat,lon,time".format(self.variable),), + ) self.regridded.set_local_file(temp2) def _get_level_file(self, lev): if not self.tempTemplate: - self.tempTemplate = TempFile.get(suffix='_01.nc') + self.tempTemplate = TempFile.get(suffix="_01.nc") # self.tempTemplate = 'temp_01.nc' - return self.tempTemplate.replace('_01.nc', '_{0:02d}.nc'.format(lev + 1)) + return self.tempTemplate.replace( + "_01.nc", "_{0:02d}.nc".format(lev + 1) + ) def _interpolate_level(self, lev, has_levels, input_file): nco = Utils.nco() temp = TempFile.get() if has_levels: - nco.ncks(input=input_file, output=temp, options='-O -d lev,{0} -v {1},lat,lon'.format(lev, self.variable)) - nco.ncwa(input=temp, output=temp, options=('-h -a lev',)) + nco.ncks( + input=input_file, + output=temp, + options="-O -d lev,{0} -v {1},lat,lon".format( + lev, self.variable + ), + ) + nco.ncwa(input=temp, output=temp, options=("-h -a lev",)) else: shutil.copy(input_file, temp) - weights_file = '/esnas/autosubmit/con_files/weigths/{0}/rmp_{0}_to_{1}_lev{2}.nc'.format(self.model_version, - self.grid, lev + 1) + weights_file = ( + f"/esnas/autosubmit/con_files/weigths/{self.model_version}/" + f"rmp_{self.model_version}_to_{self.grid}_lev{lev + 1}.nc" + ) if not os.path.isfile(weights_file): - raise Exception('Level {0} weights file does not exist for model {1} ' - 'and grid {2}'.format(lev + 1, self.model_version, self.grid)) - namelist_file = TempFile.get(suffix='') - scrip_use_in = open(namelist_file, 'w') + raise Exception( + "Level {0} weights file does not exist for model {1} " + "and grid {2}".format(lev + 1, self.model_version, self.grid) + ) + namelist_file = TempFile.get(suffix="") + scrip_use_in = open(namelist_file, "w") scrip_use_in.writelines("&remap_inputs\n") - scrip_use_in.writelines(" remap_wgt = '{0}'\n".format(weights_file)) + scrip_use_in.writelines( + " remap_wgt = '{0}'\n".format(weights_file) + ) scrip_use_in.writelines(" infile = '{0}'\n".format(temp)) scrip_use_in.writelines(" invertlat = FALSE\n") - scrip_use_in.writelines(" var = '{0}'\n".format(self.variable)) + scrip_use_in.writelines( + " var = '{0}'\n".format(self.variable) + ) scrip_use_in.writelines(" fromregular = FALSE\n") scrip_use_in.writelines(" outfile = '{0}'\n".format(temp)) scrip_use_in.writelines("/\n") scrip_use_in.close() - Utils.execute_shell_command('/home/Earth/jvegas/pyCharm/cfutools/interpolation/scrip_use ' - '{0}'.format(namelist_file), Log.INFO) + Utils.execute_shell_command( + "/home/Earth/jvegas/pyCharm/cfutools/interpolation/scrip_use " + "{0}".format(namelist_file), + Log.INFO, + ) os.remove(namelist_file) nco.ncecat(input=temp, output=temp, options=("-h",)) shutil.move(temp, self._get_level_file(lev)) diff --git a/earthdiagnostics/ocean/interpolatecdo.py b/earthdiagnostics/ocean/interpolatecdo.py index 76ba48b7..dd514e1d 100644 --- a/earthdiagnostics/ocean/interpolatecdo.py +++ b/earthdiagnostics/ocean/interpolatecdo.py @@ -4,11 +4,15 @@ import os import numpy as np -from bscearth.utils.log import Log - from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticDomainOption, DiagnosticVariableListOption, \ - DiagnosticChoiceOption, DiagnosticBoolOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticDomainOption, + DiagnosticVariableListOption, + DiagnosticChoiceOption, + DiagnosticBoolOption, + DiagnosticOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -39,18 +43,30 @@ class InterpolateCDO(Diagnostic): :type model_version: str """ - alias = 'interpcdo' + alias = "interpcdo" "Diagnostic alias for the configuration file" - BILINEAR = 'bilinear' - BICUBIC = 'bicubic' - CONSERVATIVE = 'conservative' - CONSERVATIVE2 = 'conservative2' + BILINEAR = "bilinear" + BICUBIC = "bicubic" + CONSERVATIVE = "conservative" + CONSERVATIVE2 = "conservative2" METHODS = [BILINEAR, BICUBIC, CONSERVATIVE, CONSERVATIVE2] - def __init__(self, data_manager, startdate, member, chunk, domain, variable, target_grid, model_version, - mask_oceans, original_grid, weights): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + target_grid, + model_version, + mask_oceans, + original_grid, + weights, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -60,7 +76,7 @@ class InterpolateCDO(Diagnostic): self.model_version = model_version self.required_vars = [variable] self.generated_vars = [variable] - self.tempTemplate = '' + self.tempTemplate = "" self.grid = target_grid self.mask_oceans = mask_oceans self.original_grid = original_grid @@ -70,18 +86,28 @@ class InterpolateCDO(Diagnostic): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.model_version == other.model_version and self.domain == other.domain and \ - self.variable == other.variable and self.mask_oceans == other.mask_oceans and self.grid == other.grid and \ - self.original_grid == other.original_grid + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.model_version == other.model_version + and self.domain == other.domain + and self.variable == other.variable + and self.mask_oceans == other.mask_oceans + and self.grid == other.grid + and self.original_grid == other.original_grid + ) def __hash__(self): return hash(str(self)) def __str__(self): - return 'Interpolate with CDO Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Variable: {0.domain}:{0.variable} Target grid: {0.grid} Original grid: {0.original_grid} ' \ - 'Mask ocean: {0.mask_oceans} Model: {0.model_version}'.format(self) + return ( + "Interpolate with CDO Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Variable: {0.domain}:{0.variable} " + "Target grid: {0.grid} Original grid: {0.original_grid} " + "Mask ocean: {0.mask_oceans} Model: {0.model_version}".format(self) + ) @classmethod def generate_jobs(cls, diags, options): @@ -94,37 +120,73 @@ class InterpolateCDO(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(default_value=ModelingRealms.ocean), - DiagnosticVariableListOption(diags.data_manager.config.var_manager, 'variables'), - DiagnosticOption('target_grid', diags.config.experiment.atmos_grid.lower()), - DiagnosticChoiceOption('method', InterpolateCDO.METHODS, InterpolateCDO.BILINEAR), - DiagnosticBoolOption('mask_oceans', True), - DiagnosticOption('original_grid', ''), - DiagnosticBoolOption('weights_from_mask', True) - ) + options_available = ( + DiagnosticDomainOption(default_value=ModelingRealms.ocean), + DiagnosticVariableListOption( + diags.data_manager.config.var_manager, "variables" + ), + DiagnosticOption( + "target_grid", diags.config.experiment.atmos_grid.lower() + ), + DiagnosticChoiceOption( + "method", InterpolateCDO.METHODS, InterpolateCDO.BILINEAR + ), + DiagnosticBoolOption("mask_oceans", True), + DiagnosticOption("original_grid", ""), + DiagnosticBoolOption("weights_from_mask", True), + ) options = cls.process_options(options, options_available) - target_grid = cls._translate_ifs_grids_to_cdo_names(options['target_grid']) + target_grid = cls._translate_ifs_grids_to_cdo_names( + options["target_grid"] + ) if not target_grid: - raise Exception('Target grid not provided') + raise Exception("Target grid not provided") job_list = list() weights = TempFile.get() - method = options['method'].lower() - if options['weights_from_mask']: + method = options["method"].lower() + if options["weights_from_mask"]: temp = cls.get_sample_grid_file() cls.compute_weights(method, target_grid, temp, weights) os.remove(temp) weights_job = None else: - startdate, member, chunk = diags.config.experiment.get_chunk_list()[0] - weights_job = ComputeWeights(diags.data_manager, startdate, member, chunk, options['domain'], - options['variables'][0], target_grid, options['original_grid'], weights, - options['method']) - - for var in options['variables']: - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job = InterpolateCDO(diags.data_manager, startdate, member, chunk, options['domain'], var, target_grid, - diags.config.experiment.model_version, options['mask_oceans'], - options['original_grid'], weights) + ( + startdate, + member, + chunk, + ) = diags.config.experiment.get_chunk_list()[0] + weights_job = ComputeWeights( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variables"][0], + target_grid, + options["original_grid"], + weights, + options["method"], + ) + + for var in options["variables"]: + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job = InterpolateCDO( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + target_grid, + diags.config.experiment.model_version, + options["mask_oceans"], + options["original_grid"], + weights, + ) if weights_job is not None: job.add_subjob(weights_job) job_list.append(job) @@ -169,16 +231,18 @@ class InterpolateCDO(Diagnostic): """ temp = TempFile.get() - handler = Utils.open_cdf('mask.nc') + handler = Utils.open_cdf("mask.nc") lat_name, lon_name = cls._get_lat_lon_alias(handler) - lon_bnds_name = '{0}_bnds'.format(lon_name) - lat_bnds_name = '{0}_bnds'.format(lat_name) + lon_bnds_name = "{0}_bnds".format(lon_name) + lat_bnds_name = "{0}_bnds".format(lat_name) Utils.nco().ncks( - input='mask.nc', + input="mask.nc", output=temp, - options=('-O -v tmask,{0},{1},gphif,glamf'.format(lat_name, lon_name),) + options=( + "-O -v tmask,{0},{1},gphif,glamf".format(lat_name, lon_name), + ), ) handler = Utils.open_cdf(temp) lon = handler.variables[lon_name] @@ -197,17 +261,21 @@ class InterpolateCDO(Diagnostic): lat.short_name = lat_name lat.bounds = lat_bnds_name - handler.createDimension('bounds', 4) + handler.createDimension("bounds", 4) - lon_bnds = handler.createVariable(lon_bnds_name, lon.datatype, ('j', 'i', 'bounds')) - corner_lat = handler.variables['glamf'][0, ...] + lon_bnds = handler.createVariable( + lon_bnds_name, lon.datatype, ("j", "i", "bounds") + ) + corner_lat = handler.variables["glamf"][0, ...] lon_bnds[:, :, 0] = corner_lat lon_bnds[:, :, 1] = np.roll(corner_lat, 1, 0) lon_bnds[:, :, 2] = np.roll(corner_lat, -1, 1) lon_bnds[:, :, 3] = np.roll(lon_bnds[:, :, 1], -1, 1) - lat_bnds = handler.createVariable(lat_bnds_name, lat.datatype, ('j', 'i', 'bounds')) - corner_lat = handler.variables['gphif'][0, ...] + lat_bnds = handler.createVariable( + lat_bnds_name, lat.datatype, ("j", "i", "bounds") + ) + corner_lat = handler.variables["gphif"][0, ...] lat_bnds[:, :, 0] = corner_lat lat_bnds[:, :, 1] = np.roll(corner_lat, 1, 0) lat_bnds[:, :, 2] = np.roll(corner_lat, 1, 1) @@ -215,23 +283,25 @@ class InterpolateCDO(Diagnostic): lat_bnds[0, :, 1] = lat_bnds[1, 0, 1] - 1 lat_bnds[0, :, 3] = lat_bnds[1, 0, 3] - 1 - tmask = handler.variables['tmask'] - tmask.coordinates = 'time lev {0} {1}'.format(lat_name, lon_name) + tmask = handler.variables["tmask"] + tmask.coordinates = "time lev {0} {1}".format(lat_name, lon_name) handler.close() - Utils.nco().ncks(input=temp, output=temp, options=('-O -x -v gphif,glamf',)) + Utils.nco().ncks( + input=temp, output=temp, options=("-O -x -v gphif,glamf",) + ) return temp @classmethod def _get_lat_lon_alias(cls, handler): lat_name = None - for lat_alias in ['lat', 'latitude']: + for lat_alias in ["lat", "latitude"]: if lat_alias in handler.variables: lat_name = lat_alias break lon_name = None - for lon_alias in ['lon', 'longitude']: + for lon_alias in ["lon", "longitude"]: if lon_alias in handler.variables: lon_name = lon_alias break @@ -239,23 +309,35 @@ class InterpolateCDO(Diagnostic): @classmethod def _translate_ifs_grids_to_cdo_names(cls, target_grid): - if target_grid.upper().startswith('T159L'): - target_grid = 't106grid' - if target_grid.upper().startswith('T255L'): - target_grid = 't170grid' - if target_grid.upper().startswith('T511L'): - target_grid = 't340grid' + if target_grid.upper().startswith("T159L"): + target_grid = "t106grid" + if target_grid.upper().startswith("T255L"): + target_grid = "t170grid" + if target_grid.upper().startswith("T511L"): + target_grid = "t340grid" return target_grid def request_data(self): """Request data required by the diagnostic""" - self.original = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.original_grid) + self.original = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.original_grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.regridded = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.grid) + self.regridded = self.declare_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def compute(self): """Run the diagnostic""" @@ -264,14 +346,23 @@ class InterpolateCDO(Diagnostic): Utils.rename_variables( variable_file, { - 'jpib': 'i', 'jpjb': 'j', 'x': 'i', 'y': 'j', - 'dim1': 'j', 'dim2': 'i', - 'time_counter': 'time', 't': 'time', - 'SSTK_ens0': 'tos', 'SSTK_ens1': 'tos', 'SSTK_ens2': 'tos', - 'nav_lat': 'lat', 'nav_lon': 'lon', - 'time_centered': None, 'time_centered_bnds': None + "jpib": "i", + "jpjb": "j", + "x": "i", + "y": "j", + "dim1": "j", + "dim2": "i", + "time_counter": "time", + "t": "time", + "SSTK_ens0": "tos", + "SSTK_ens1": "tos", + "SSTK_ens2": "tos", + "nav_lat": "lat", + "nav_lon": "lon", + "time_centered": None, + "time_centered_bnds": None, }, - must_exist=False + must_exist=False, ) handler = Utils.open_cdf(variable_file) lat_name, lon_name = self._get_lat_lon_alias(handler) @@ -282,13 +373,13 @@ class InterpolateCDO(Diagnostic): units = None coordinates = list() for dim in var.dimensions: - if dim == 'i': + if dim == "i": coordinates.append(lon_name) - elif dim == 'j': + elif dim == "j": coordinates.append(lat_name) else: coordinates.append(dim) - var.coordinates = ' '.join(coordinates) + var.coordinates = " ".join(coordinates) if self.mask_oceans: mask = Utils.get_mask(Basins().Global).astype(float) @@ -297,15 +388,21 @@ class InterpolateCDO(Diagnostic): handler.close() temp = TempFile.get() - Utils.cdo().remap(','.join((self.grid.split('_')[0], self.weights)), input=variable_file, output=temp) + Utils.cdo().remap( + ",".join((self.grid.split("_")[0], self.weights)), + input=variable_file, + output=temp, + ) handler = Utils.open_cdf(temp) if units: handler.variables[self.variable].units = units handler.close() - if lat_name != 'lat': - Utils.rename_variables(temp, {'lat': lat_name, 'lon': lon_name}, True) + if lat_name != "lat": + Utils.rename_variables( + temp, {"lat": lat_name, "lon": lon_name}, True + ) self.regridded.set_local_file(temp) @@ -329,20 +426,32 @@ class ComputeWeights(Diagnostic): """ - alias = 'computeinterpcdoweights' + alias = "computeinterpcdoweights" "Diagnostic alias for the configuration file" @classmethod def generate_jobs(cls, diags, options): """ - Generate the instances of the diagnostics that will be run by the manager + Generate the instances of the diagnostics that will be run - This method does not does anything as this diagnostic is not expected to be called by the users + This method does not does anything as this diagnostic is + not expected to be called by the users """ pass - def __init__(self, data_manager, startdate, member, chunk, domain, variable, target_grid, - original_grid, weights_file, method): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + target_grid, + original_grid, + weights_file, + method, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -355,16 +464,30 @@ class ComputeWeights(Diagnostic): self.method = method def __str__(self): - return 'Computing weights for CDO interpolation: Method {0.method} Target grid: {0.grid}'.format(self) + return ( + "Computing weights for CDO interpolation: Method {0.method} " + "Target grid: {0.grid}".format(self) + ) def compute(self): """Compute weights""" - InterpolateCDO.compute_weights(self.method, self.grid, self.sample_data.local_file, self.weights_file) + InterpolateCDO.compute_weights( + self.method, + self.grid, + self.sample_data.local_file, + self.weights_file, + ) def request_data(self): """Request data required by the diagnostic""" - self.sample_data = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.original_grid) + self.sample_data = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.original_grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" diff --git a/earthdiagnostics/ocean/mask_land.py b/earthdiagnostics/ocean/mask_land.py index 451d9d2a..3a4f18b4 100644 --- a/earthdiagnostics/ocean/mask_land.py +++ b/earthdiagnostics/ocean/mask_land.py @@ -2,8 +2,13 @@ """Changes values present in the mask for NaNs""" import numpy as np -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticVariableListOption, \ - DiagnosticDomainOption, DiagnosticChoiceOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticVariableListOption, + DiagnosticDomainOption, + DiagnosticChoiceOption, + DiagnosticOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -23,10 +28,20 @@ class MaskLand(Diagnostic): :type variable: str """ - alias = 'maskland' + alias = "maskland" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, mask, grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + mask, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -39,12 +54,26 @@ class MaskLand(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + ) def __str__(self): - return 'Land mask Startdate: {0} Member: {1} Chunk: {2} Variable: {3}:{4} ' \ - 'Grid: {5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable, self.grid) + return ( + "Land mask Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3}:{4} Grid: {5}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variable, + self.grid, + ) + ) @classmethod def generate_jobs(cls, diags, options): @@ -57,43 +86,75 @@ class MaskLand(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption('domain'), - DiagnosticVariableListOption(diags.data_manager.config.var_manager, 'variables'), - DiagnosticChoiceOption('cell', ('t', 'u', 'v', 'f', 'w'), 't'), - DiagnosticOption('grid', '')) + options_available = ( + DiagnosticDomainOption("domain"), + DiagnosticVariableListOption( + diags.data_manager.config.var_manager, "variables" + ), + DiagnosticChoiceOption("cell", ("t", "u", "v", "f", "w"), "t"), + DiagnosticOption("grid", ""), + ) options = cls.process_options(options, options_available) - cell_point = options['cell'] + cell_point = options["cell"] # W and T share the same mask - if cell_point == 'w': - cell_point = 't' + if cell_point == "w": + cell_point = "t" mask = cls._get_mask(cell_point) job_list = list() - for var in options['variables']: - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(MaskLand(diags.data_manager, startdate, member, chunk, - options['domain'], var, mask, options['grid'])) + for var in options["variables"]: + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + MaskLand( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + mask, + options["grid"], + ) + ) return job_list @classmethod def _get_mask(cls, cell_point): - mask_file = Utils.open_cdf('mask.nc') - mask = mask_file.variables['{0}mask'.format(cell_point)][:].astype(float) + mask_file = Utils.open_cdf("mask.nc") + mask = mask_file.variables["{0}mask".format(cell_point)][:].astype( + float + ) mask[mask == 0] = np.nan mask_file.close() return mask def request_data(self): """Request data required by the diagnostic""" - self.var_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.grid) + self.var_file = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.masked_file = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, - grid=self.grid) + self.masked_file = self.declare_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def compute(self): """Run the diagnostic""" @@ -101,7 +162,7 @@ class MaskLand(Diagnostic): Utils.copy_file(self.var_file.local_file, temp) handler = Utils.open_cdf(temp) - if 'lev' not in handler.dimensions: + if "lev" not in handler.dimensions: mask = self.mask[:, 0, ...] else: mask = self.mask diff --git a/earthdiagnostics/ocean/maxmoc.py b/earthdiagnostics/ocean/maxmoc.py index 7c0309c3..3947c672 100644 --- a/earthdiagnostics/ocean/maxmoc.py +++ b/earthdiagnostics/ocean/maxmoc.py @@ -6,7 +6,11 @@ from bscearth.utils.log import Log from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption, DiagnosticFloatOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBasinOption, + DiagnosticFloatOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile from earthdiagnostics.variable import VariableType @@ -38,10 +42,10 @@ class MaxMoc(Diagnostic): :type box: Box """ - alias = 'mocmax' + alias = "mocmax" "Diagnostic alias for the configuration file" - vsftmyz = 'vsftmyz' + vsftmyz = "vsftmyz" def __init__(self, data_manager, startdate, member, year, basin, box): Diagnostic.__init__(self, data_manager) @@ -52,14 +56,27 @@ class MaxMoc(Diagnostic): self.box = box def __str__(self): - return 'Max moc Startdate: {0} Member: {1} Year: {2} Box: {3} ' \ - 'Basin: {4}'.format(self.startdate, self.member, self.year, self.box, self.basin.name) + return ( + "Max moc Startdate: {0} Member: {1} Year: {2} Box: {3} " + "Basin: {4}".format( + self.startdate, + self.member, + self.year, + self.box, + self.basin.name, + ) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.year == other.year and \ - self.box == other.box and self.basin == other.basin + return ( + self.startdate == other.startdate + and self.member == other.member + and self.year == other.year + and self.box == other.box + and self.basin == other.basin + ) @classmethod def generate_jobs(cls, diags, options): @@ -68,56 +85,115 @@ class MaxMoc(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: minimum latitude, maximum latitude, minimum depth, maximum depth, basin=global + :param options: minimum latitude, maximum latitude, minimum depth, + maximum depth, basin=global :type options: list[str] :return: """ - options_available = (DiagnosticFloatOption('min_lat'), - DiagnosticFloatOption('max_lat'), - DiagnosticFloatOption('min_depth'), - DiagnosticFloatOption('max_depth'), - DiagnosticBasinOption('basin', Basins().Global)) + options_available = ( + DiagnosticFloatOption("min_lat"), + DiagnosticFloatOption("max_lat"), + DiagnosticFloatOption("min_depth"), + DiagnosticFloatOption("max_depth"), + DiagnosticBasinOption("basin", Basins().Global), + ) options = cls.process_options(options, options_available) box = Box(True) - box.min_lat = options['min_lat'] - box.max_lat = options['max_lat'] - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] + box.min_lat = options["min_lat"] + box.max_lat = options["max_lat"] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] job_list = list() for startdate in diags.config.experiment.startdates: for member in diags.config.experiment.members: years = diags.config.experiment.get_full_years(startdate) if len(years) == 0: - Log.user_warning('No complete years are available with the given configuration. ' - 'MaxMoc can not be computed') + Log.user_warning( + "No complete years are available with the given " + "configuration. MaxMoc can not be computed" + ) for year in years: - job_list.append(MaxMoc(diags.data_manager, startdate, member, year, options['basin'], box)) + job_list.append( + MaxMoc( + diags.data_manager, + startdate, + member, + year, + options["basin"], + box, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_year(ModelingRealms.ocean, MaxMoc.vsftmyz, - self.startdate, self.member, self.year) + self.variable_file = self.request_year( + ModelingRealms.ocean, + MaxMoc.vsftmyz, + self.startdate, + self.member, + self.year, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.results = {'vsftmyzmax': self.declare_year(ModelingRealms.ocean, 'vsftmyzmax', self.startdate, self.member, - self.year, box=self.box, vartype=VariableType.STATISTIC), - 'vsftmyzmaxlev': self.declare_year(ModelingRealms.ocean, 'vsftmyzmaxlev', self.startdate, - self.member, self.year, box=self.box, - vartype=VariableType.STATISTIC), - 'vsftmyzmaxlat': self.declare_year(ModelingRealms.ocean, 'vsftmyzmaxlat', self.startdate, - self.member, self.year, box=self.box, - vartype=VariableType.STATISTIC), - 'vsftmyzmin': self.declare_year(ModelingRealms.ocean, 'vsftmyzmin', self.startdate, self.member, - self.year, box=self.box, vartype=VariableType.STATISTIC), - 'vsftmyzminlev': self.declare_year(ModelingRealms.ocean, 'vsftmyzminlev', self.startdate, - self.member, self.year, box=self.box, - vartype=VariableType.STATISTIC), - 'vsftmyzminlat': self.declare_year(ModelingRealms.ocean, 'vsftmyzminlat', self.startdate, - self.member, self.year, box=self.box, - vartype=VariableType.STATISTIC)} + self.results = { + "vsftmyzmax": self.declare_year( + ModelingRealms.ocean, + "vsftmyzmax", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + "vsftmyzmaxlev": self.declare_year( + ModelingRealms.ocean, + "vsftmyzmaxlev", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + "vsftmyzmaxlat": self.declare_year( + ModelingRealms.ocean, + "vsftmyzmaxlat", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + "vsftmyzmin": self.declare_year( + ModelingRealms.ocean, + "vsftmyzmin", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + "vsftmyzminlev": self.declare_year( + ModelingRealms.ocean, + "vsftmyzminlev", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + "vsftmyzminlat": self.declare_year( + ModelingRealms.ocean, + "vsftmyzminlat", + self.startdate, + self.member, + self.year, + box=self.box, + vartype=VariableType.STATISTIC, + ), + } def compute(self): """Run the diagnostic""" @@ -127,35 +203,43 @@ class MaxMoc(Diagnostic): Utils.copy_file(self.variable_file.local_file, temp) handler = Utils.open_cdf(temp) - if 'i' in handler.dimensions: + if "i" in handler.dimensions: handler.close() - nco.ncwa(input=temp, output=temp, options=('-O -a i',)) + nco.ncwa(input=temp, output=temp, options=("-O -a i",)) else: handler.close() handler = Utils.open_cdf(temp) - basin_index = np.where(handler.variables['basin'][:] == self.basin.name) + basin_index = np.where( + handler.variables["basin"][:] == self.basin.name + ) if len(basin_index) == 0: - raise Exception("Basin {1} is not defined in {0}", temp, self.basin.name) + raise Exception( + "Basin {1} is not defined in {0}", temp, self.basin.name + ) basin_index = basin_index[0][0] - lev = handler.variables['lev'][:] + lev = handler.variables["lev"][:] if "lat" in handler.variables: - lat = handler.variables['lat'][:] + lat = handler.variables["lat"][:] else: - lat = handler.variables['latitude'][:] + lat = handler.variables["latitude"][:] if self.box.min_lat == self.box.max_lat: lat_inds = ((np.abs(lat - self.box.min_lat)).argmin(),) else: - lat_inds = np.where((lat > self.box.min_lat) & (lat < self.box.max_lat))[0] + lat_inds = np.where( + (lat > self.box.min_lat) & (lat < self.box.max_lat) + )[0] if self.box.min_depth == self.box.max_depth: lev_inds = ((np.abs(lev - self.box.min_depth)).argmin(),) else: - lev_inds = np.where((lev > self.box.min_depth) & (lev < self.box.max_depth))[0] + lev_inds = np.where( + (lev > self.box.min_depth) & (lev < self.box.max_depth) + )[0] - Log.info('Computing year {0}', str(self.year)) - moc = handler.variables['vsftmyz'][:, lev_inds, lat_inds, basin_index] + Log.info("Computing year {0}", str(self.year)) + moc = handler.variables["vsftmyz"][:, lev_inds, lat_inds, basin_index] handler.close() moc = np.mean(moc, 0) @@ -174,75 +258,85 @@ class MaxMoc(Diagnostic): # noinspection PyUnresolvedReferences min_lat = lat[lat_inds[minimum_index[1]]] - Log.info('Maximum {0} Sv, latitude {1} depth {2} m', maximum, max_lat, max_lev) - Log.info('Minimum {0} Sv, latitude {1} depth {2} m', minimum, min_lat, min_lev) + Log.info( + "Maximum {0} Sv, latitude {1} depth {2} m", + maximum, + max_lat, + max_lev, + ) + Log.info( + "Minimum {0} Sv, latitude {1} depth {2} m", + minimum, + min_lat, + min_lev, + ) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzmax', float, ('time',)) - var.long_name = 'Maximum_Overturning' - var.units = 'Sverdrup' - var.valid_min = -1000. - var.valid_max = 1000. + var = handler.createVariable("vsftmyzmax", float, ("time",)) + var.long_name = "Maximum_Overturning" + var.units = "Sverdrup" + var.valid_min = -1000.0 + var.valid_max = 1000.0 var[0] = maximum handler.close() - self.results['vsftmyzmax'].set_local_file(temp) + self.results["vsftmyzmax"].set_local_file(temp) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzmaxlat', float, ('time',)) - var.long_name = 'Latitude_of_Maximum_Overturning' - var.units = 'Degrees' - var.valid_min = -90. - var.valid_max = 90. + var = handler.createVariable("vsftmyzmaxlat", float, ("time",)) + var.long_name = "Latitude_of_Maximum_Overturning" + var.units = "Degrees" + var.valid_min = -90.0 + var.valid_max = 90.0 var[0] = max_lat handler.close() - self.results['vsftmyzmaxlat'].set_local_file(temp) + self.results["vsftmyzmaxlat"].set_local_file(temp) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzmaxlev', float, ('time',)) - var.long_name = 'Depth_of_Maximum_Overturning' - var.units = 'Meters' - var.valid_min = 0. - var.valid_max = 10000. + var = handler.createVariable("vsftmyzmaxlev", float, ("time",)) + var.long_name = "Depth_of_Maximum_Overturning" + var.units = "Meters" + var.valid_min = 0.0 + var.valid_max = 10000.0 var[0] = max_lev handler.close() - self.results['vsftmyzmaxlev'].set_local_file(temp) + self.results["vsftmyzmaxlev"].set_local_file(temp) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzmin', float, ('time',)) - var.long_name = 'Minimum_Overturning' - var.units = 'Sverdrup' - var.valid_min = -1000. - var.valid_max = 1000. + var = handler.createVariable("vsftmyzmin", float, ("time",)) + var.long_name = "Minimum_Overturning" + var.units = "Sverdrup" + var.valid_min = -1000.0 + var.valid_max = 1000.0 var[0] = minimum handler.close() - self.results['vsftmyzmin'].set_local_file(temp) + self.results["vsftmyzmin"].set_local_file(temp) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzminlat', float, ('time',)) - var.long_name = 'Latitude_of_Minimum_Overturning' - var.units = 'Degrees' - var.valid_min = -90. - var.valid_max = 90. + var = handler.createVariable("vsftmyzminlat", float, ("time",)) + var.long_name = "Latitude_of_Minimum_Overturning" + var.units = "Degrees" + var.valid_min = -90.0 + var.valid_max = 90.0 var[0] = min_lat handler.close() - self.results['vsftmyzminlat'].set_local_file(temp) + self.results["vsftmyzminlat"].set_local_file(temp) handler, temp = self._create_output_file() - var = handler.createVariable('vsftmyzminlev', float, ('time',)) - var.long_name = 'Depth_of_Minimum_Overturning' - var.units = 'Meters' - var.valid_min = 0. - var.valid_max = 10000. + var = handler.createVariable("vsftmyzminlev", float, ("time",)) + var.long_name = "Depth_of_Minimum_Overturning" + var.units = "Meters" + var.valid_min = 0.0 + var.valid_max = 10000.0 var[0] = min_lev handler.close() - self.results['vsftmyzminlev'].set_local_file(temp) + self.results["vsftmyzminlev"].set_local_file(temp) def _create_output_file(self): temp = TempFile.get() - handler = netCDF4.Dataset(temp, 'w') - handler.createDimension('time') + handler = netCDF4.Dataset(temp, "w") + handler.createDimension("time") - time = handler.createVariable('time', 'i2', ('time',)) - time.calendar = 'gregorian' - time.units = 'days since January 1, {0}'.format(self.year) + time = handler.createVariable("time", "i2", ("time",)) + time.calendar = "gregorian" + time.units = "days since January 1, {0}".format(self.year) return handler, temp diff --git a/earthdiagnostics/ocean/mixedlayerheatcontent.py b/earthdiagnostics/ocean/mixedlayerheatcontent.py index c316bdae..7bd6f95c 100644 --- a/earthdiagnostics/ocean/mixedlayerheatcontent.py +++ b/earthdiagnostics/ocean/mixedlayerheatcontent.py @@ -28,7 +28,7 @@ class MixedLayerHeatContent(Diagnostic): :type chunk: int """ - alias = 'mlotsthc' + alias = "mlotsthc" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk): @@ -36,17 +36,25 @@ class MixedLayerHeatContent(Diagnostic): self.startdate = startdate self.member = member self.chunk = chunk - self.required_vars = ['so', 'mlotst'] - self.generated_vars = ['scvertsum'] + self.required_vars = ["so", "mlotst"] + self.generated_vars = ["scvertsum"] def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + ) def __str__(self): - return 'Mixed layer heat content Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, - self.chunk) + return ( + "Mixed layer heat content Startdate: {0} Member: {1} " + "Chunk: {2}".format( + self.startdate, self.member, self.chunk + ) + ) @classmethod def generate_jobs(cls, diags, options): @@ -60,32 +68,68 @@ class MixedLayerHeatContent(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The mixed layer ocean heat content diagnostic has no options') + raise Exception( + "The mixed layer ocean heat content diagnostic has no options" + ) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(MixedLayerHeatContent(diags.data_manager, startdate, member, chunk)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + MixedLayerHeatContent( + diags.data_manager, startdate, member, chunk + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk) - self.mlotst = self.request_chunk(ModelingRealms.ocean, 'mlotst', self.startdate, self.member, self.chunk) + self.thetao = self.request_chunk( + ModelingRealms.ocean, + "thetao", + self.startdate, + self.member, + self.chunk, + ) + self.mlotst = self.request_chunk( + ModelingRealms.ocean, + "mlotst", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.ohcsum = self.declare_chunk(ModelingRealms.ocean, 'ohcvsumlotst', self.startdate, self.member, self.chunk) + self.ohcsum = self.declare_chunk( + ModelingRealms.ocean, + "ohcvsumlotst", + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" temperature_file = TempFile.get() Utils.copy_file(self.thetao.local_file, temperature_file) - Utils.nco().ncks(input=self.mlotst.local_file, output=temperature_file, options=('-A -v mlotst',)) + Utils.nco().ncks( + input=self.mlotst.local_file, + output=temperature_file, + options=("-A -v mlotst",), + ) temp = TempFile.get() - cdftools.run('cdfmxlheatc', input_file=temperature_file, output_file=temp) + cdftools.run( + "cdfmxlheatc", input_file=temperature_file, output_file=temp + ) os.remove(temperature_file) - Utils.rename_variables(temp, {'x': 'i', 'y': 'j', 'somxlheatc': 'ohcvsumlotst'}, False) - Utils.setminmax(temp, 'ohcvsumlotst') + Utils.rename_variables( + temp, {"x": "i", "y": "j", "somxlheatc": "ohcvsumlotst"}, False + ) + Utils.setminmax(temp, "ohcvsumlotst") self.ohcsum.set_local_file(temp) diff --git a/earthdiagnostics/ocean/mixedlayersaltcontent.py b/earthdiagnostics/ocean/mixedlayersaltcontent.py index 13654b9a..9634bd57 100644 --- a/earthdiagnostics/ocean/mixedlayersaltcontent.py +++ b/earthdiagnostics/ocean/mixedlayersaltcontent.py @@ -28,7 +28,7 @@ class MixedLayerSaltContent(Diagnostic): :type chunk: int """ - alias = 'mlotstsc' + alias = "mlotstsc" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk): @@ -36,17 +36,25 @@ class MixedLayerSaltContent(Diagnostic): self.startdate = startdate self.member = member self.chunk = chunk - self.required_vars = ['so', 'mlotst'] - self.generated_vars = ['scvertsum'] + self.required_vars = ["so", "mlotst"] + self.generated_vars = ["scvertsum"] def __str__(self): - return 'Mixed layer salt content Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, - self.chunk) + return ( + "Mixed layer salt content Startdate: {0} Member: {1} " + "Chunk: {2}".format( + self.startdate, self.member, self.chunk + ) + ) def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + ) @classmethod def generate_jobs(cls, diags, options): @@ -60,31 +68,61 @@ class MixedLayerSaltContent(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The mixed layer salt content diagnostic has no options') + raise Exception( + "The mixed layer salt content diagnostic has no options" + ) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(MixedLayerSaltContent(diags.data_manager, startdate, member, chunk)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + MixedLayerSaltContent( + diags.data_manager, startdate, member, chunk + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.so = self.request_chunk(ModelingRealms.ocean, 'so', self.startdate, self.member, self.chunk) - self.mlotst = self.request_chunk(ModelingRealms.ocean, 'mlotst', self.startdate, self.member, self.chunk) + self.so = self.request_chunk( + ModelingRealms.ocean, "so", self.startdate, self.member, self.chunk + ) + self.mlotst = self.request_chunk( + ModelingRealms.ocean, + "mlotst", + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.sosum = self.declare_chunk(ModelingRealms.ocean, 'scvsummlotst', self.startdate, self.member, self.chunk) + self.sosum = self.declare_chunk( + ModelingRealms.ocean, + "scvsummlotst", + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" salinity_file = TempFile.get() Utils.copy_file(self.so.local_file, salinity_file) - Utils.nco().ncks(input=self.mlotst.local_file, output=salinity_file, options=('-A -v mlotst',)) + Utils.nco().ncks( + input=self.mlotst.local_file, + output=salinity_file, + options=("-A -v mlotst",), + ) temp = TempFile.get() - cdftools.run('cdfmxlsaltc', input_file=salinity_file, output_file=temp) + cdftools.run("cdfmxlsaltc", input_file=salinity_file, output_file=temp) os.remove(salinity_file) - Utils.rename_variables(temp, {'x': 'i', 'y': 'j', 'somxlsaltc': 'scvsummlotst'}, False) - Utils.setminmax(temp, 'scvsummlotst') + Utils.rename_variables( + temp, {"x": "i", "y": "j", "somxlsaltc": "scvsummlotst"}, False + ) + Utils.setminmax(temp, "scvsummlotst") self.sosum.set_local_file(temp) diff --git a/earthdiagnostics/ocean/moc.py b/earthdiagnostics/ocean/moc.py index a294b824..b6f84506 100644 --- a/earthdiagnostics/ocean/moc.py +++ b/earthdiagnostics/ocean/moc.py @@ -1,12 +1,9 @@ # coding=utf-8 """Compute the MOC for oceanic basins""" import numpy as np -import six from bscearth.utils.log import Log import iris -from iris.coords import DimCoord, AuxCoord -from iris.cube import CubeList import iris.analysis @@ -41,10 +38,10 @@ class Moc(Diagnostic): :type chunk: int """ - alias = 'moc' + alias = "moc" "Diagnostic alias for the configuration file" - vsftmyz = 'vsftmyz' + vsftmyz = "vsftmyz" def __init__(self, data_manager, startdate, member, chunk, basins): Diagnostic.__init__(self, data_manager) @@ -58,8 +55,10 @@ class Moc(Diagnostic): def __str__(self): basins = [] basins.extend(self.basins) - return 'MOC Startdate: {0.startdate} Member: {0.member} ' \ - 'Chunk: {0.chunk} Basins: {1}'.format(self, basins) + return ( + "MOC Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Basins: {1}".format(self, basins) + ) def __hash__(self): return hash(str(self)) @@ -68,8 +67,9 @@ class Moc(Diagnostic): if self._different_type(other): return False return ( - self.startdate == other.startdate and - self.member == other.member and self.chunk == other.chunk + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk ) @classmethod @@ -84,50 +84,54 @@ class Moc(Diagnostic): :return: """ basins = Basins() - options_available = ( - DiagnosticBasinListOption( - 'basins', - 'glob' - ), - ) + options_available = (DiagnosticBasinListOption("basins", "glob"),) options = cls.process_options(options, options_available) - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Moc(diags.data_manager, startdate, member, chunk, - basins)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Moc(diags.data_manager, startdate, member, chunk, basins) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(ModelingRealms.ocean, 'vo', - self.startdate, self.member, - self.chunk) + self.variable_file = self.request_chunk( + ModelingRealms.ocean, "vo", self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.results = self.declare_chunk(ModelingRealms.ocean, Moc.vsftmyz, - self.startdate, self.member, - self.chunk) + self.results = self.declare_chunk( + ModelingRealms.ocean, + Moc.vsftmyz, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" vo_cube = iris.load_cube(self.variable_file.local_file) vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') - e1v = mesh.get_i_length(cell_point='V') - e3v = mesh.get_k_length(cell_point='V') + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") + e1v = mesh.get_i_length(cell_point="V") + e3v = mesh.get_k_length(cell_point="V") masks = {} self.basins.sort() for basin in self.basins: - if basin is 'Global': - global_mask = mesh.get_landsea_mask(cell_point='V') + if basin == "Global": + global_mask = mesh.get_landsea_mask(cell_point="V") global_mask[..., 0] = 0.0 global_mask[..., -1] = 0.0 masks[basin] = global_mask @@ -141,36 +145,37 @@ class Moc(Diagnostic): def _save_result(self, result, mesh): temp = TempFile.get() handler_source = Utils.open_cdf(self.variable_file.local_file) - handler_temp = Utils.open_cdf(temp, 'w') - gphiv = np.squeeze(mesh.get_grid_latitude(cell_point='V')) + handler_temp = Utils.open_cdf(temp, "w") + gphiv = np.squeeze(mesh.get_grid_latitude(cell_point="V")) max_gphiv = np.unravel_index(np.argmax(gphiv), gphiv.shape)[1] - Utils.copy_variable(handler_source, handler_temp, 'time', True, True) - Utils.copy_variable(handler_source, handler_temp, 'lev', True, True) - handler_temp.createDimension('i', 1) - handler_temp.createDimension('j', gphiv.shape[0]) - handler_temp.createDimension('region', len(result)) - handler_temp.createDimension('region_length', 50) + Utils.copy_variable(handler_source, handler_temp, "time", True, True) + Utils.copy_variable(handler_source, handler_temp, "lev", True, True) + handler_temp.createDimension("i", 1) + handler_temp.createDimension("j", gphiv.shape[0]) + handler_temp.createDimension("region", len(result)) + handler_temp.createDimension("region_length", 50) - var_region = handler_temp.createVariable('region', 'S1', - ('region', 'region_length')) + var_region = handler_temp.createVariable( + "region", "S1", ("region", "region_length") + ) - lat = handler_temp.createVariable('lat', float, ('j', 'i')) + lat = handler_temp.createVariable("lat", float, ("j", "i")) lat[...] = gphiv[:, max_gphiv] - lat.units = 'degrees_north' + lat.units = "degrees_north" lat.long_name = "Latitude" - lon = handler_temp.createVariable('lon', float, ('j', 'i')) + lon = handler_temp.createVariable("lon", float, ("j", "i")) lon[...] = 0 - lon.units = 'degrees_east' + lon.units = "degrees_east" lon.long_name = "Longitude" - var = handler_temp.createVariable('vsftmyz', float, ('time', 'lev', - 'i', 'j', - 'region')) - var.units = 'Sverdrup' - var.coordinates = 'lev time' - var.long_name = 'Ocean meridional overturning volume streamfunction' + var = handler_temp.createVariable( + "vsftmyz", float, ("time", "lev", "i", "j", "region") + ) + var.units = "Sverdrup" + var.coordinates = "lev time" + var.long_name = "Ocean meridional overturning volume streamfunction" var.missing_value = 1e20 var.fill_value = 1e20 diff --git a/earthdiagnostics/ocean/mxl.py b/earthdiagnostics/ocean/mxl.py index 2022daaa..a453cd0f 100644 --- a/earthdiagnostics/ocean/mxl.py +++ b/earthdiagnostics/ocean/mxl.py @@ -22,7 +22,7 @@ class Mxl(Diagnostic): :type chunk: int """ - alias = 'mxl' + alias = "mxl" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk): @@ -34,10 +34,16 @@ class Mxl(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + ) def __str__(self): - return 'Mixed layer Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk) + return "Mixed layer Startdate: {0} Member: {1} Chunk: {2}".format( + self.startdate, self.member, self.chunk + ) @classmethod def generate_jobs(cls, diags, options): @@ -51,34 +57,59 @@ class Mxl(Diagnostic): :return: """ if len(options) > 1: - raise Exception('The mxl diagnostic has no options') + raise Exception("The mxl diagnostic has no options") job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): job_list.append(Mxl(diags.data_manager, startdate, member, chunk)) return job_list def request_data(self): """Request data required by the diagnostic""" - self.thetao_file = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk) - self.so_file = self.request_chunk(ModelingRealms.ocean, 'so', self.startdate, self.member, self.chunk) + self.thetao_file = self.request_chunk( + ModelingRealms.ocean, + "thetao", + self.startdate, + self.member, + self.chunk, + ) + self.so_file = self.request_chunk( + ModelingRealms.ocean, "so", self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.mlotst_file = self.declare_chunk(ModelingRealms.ocean, 'mlotst', self.startdate, self.member, self.chunk) + self.mlotst_file = self.declare_chunk( + ModelingRealms.ocean, + "mlotst", + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" temp = TempFile.get() - cdftools.run('cdfmxl', input_file=[self.thetao_file, self.so_file], output_file=temp, options='-nc4') + cdftools.run( + "cdfmxl", + input_file=[self.thetao_file, self.so_file], + output_file=temp, + options="-nc4", + ) temp2 = TempFile.get() source = Utils.open_cdf(temp) - destiny = Utils.open_cdf(temp2, 'w') - Utils.copy_variable(source, destiny, 'somxl010', must_exist=True, add_dimensions=True) - Utils.copy_variable(source, destiny, 'lat', must_exist=False) - Utils.copy_variable(source, destiny, 'latitude', must_exist=False) - Utils.copy_variable(source, destiny, 'lon', must_exist=False) - Utils.copy_variable(source, destiny, 'longitude', must_exist=False) + destiny = Utils.open_cdf(temp2, "w") + Utils.copy_variable( + source, destiny, "somxl010", must_exist=True, add_dimensions=True + ) + Utils.copy_variable(source, destiny, "lat", must_exist=False) + Utils.copy_variable(source, destiny, "latitude", must_exist=False) + Utils.copy_variable(source, destiny, "lon", must_exist=False) + Utils.copy_variable(source, destiny, "longitude", must_exist=False) source.close() destiny.close() - self.mlotst_file.set_local_file(temp2, rename_var='somxl010') + self.mlotst_file.set_local_file(temp2, rename_var="somxl010") os.remove(temp) diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index e13a20be..24d8b322 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -1,8 +1,4 @@ # coding=utf-8 -"""Compute the barotropic stream function""" -import os -import datetime - import numpy as np import netCDF4 @@ -13,9 +9,7 @@ import iris.coords import iris.util from bscearth.utils.log import Log -from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, \ - DiagnosticBasinListOption +from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinListOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -43,13 +37,14 @@ class Psi(Diagnostic): :type chunk: int """ - alias = 'psi' + alias = "psi" "Diagnostic alias for the configuration file" - vsftbarot = 'vsftbarot' + vsftbarot = "vsftbarot" def __init__(self, data_manager, startdate, member, chunk, masks): Diagnostic.__init__(self, data_manager) + self.data_convention = self.data_manager.config.data_convention self.startdate = startdate self.member = member self.chunk = chunk @@ -58,11 +53,19 @@ class Psi(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + ) def __str__(self): - return 'PSI Startdate: {0} Member: {1} Chunk: {2} Basins: {3}'.format( - self.startdate, self.member, self.chunk, ','.join(str(basin) for basin in self.masks.keys())) + return "PSI Startdate: {0} Member: {1} Chunk: {2} Basins: {3}".format( + self.startdate, + self.member, + self.chunk, + ",".join(str(basin) for basin in self.masks.keys()), + ) def __hash__(self): return hash(str(self)) @@ -78,12 +81,12 @@ class Psi(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption('basins', 'global'),) + options_available = (DiagnosticBasinListOption("basins", "global"),) options = cls.process_options(options, options_available) - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () masks = {} @@ -92,38 +95,50 @@ class Psi(Diagnostic): masks[basin] = Utils.get_mask(basin) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Psi(diags.data_manager, startdate, member, chunk, - masks)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Psi(diags.data_manager, startdate, member, chunk, masks) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.uo = self.request_chunk(ModelingRealms.ocean, 'uo', - self.startdate, self.member, self.chunk) - self.vo = self.request_chunk(ModelingRealms.ocean, 'vo', - self.startdate, self.member, self.chunk) + self.uo = self.request_chunk( + ModelingRealms.ocean, "uo", self.startdate, self.member, self.chunk + ) + self.vo = self.request_chunk( + ModelingRealms.ocean, "vo", self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.psi = self.declare_chunk(ModelingRealms.ocean, Psi.vsftbarot, - self.startdate, self.member, self.chunk) + self.psi = self.declare_chunk( + ModelingRealms.ocean, + Psi.vsftbarot, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - self._fix_coordinates_attribute(self.uo.local_file, 'uo') - self._fix_coordinates_attribute(self.vo.local_file, 'vo') + self._fix_coordinates_attribute(self.uo.local_file, "uo") + self._fix_coordinates_attribute(self.vo.local_file, "vo") uo_cube = iris.load_cube(self.uo.local_file) vo_cube = iris.load_cube(self.vo.local_file) uo = np.ma.filled(uo_cube.data, 0.0).astype(np.float32) vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') - e2u = mesh.get_j_length(cell_point='U') - e3u = mesh.get_k_length(cell_point='U') - e1v = mesh.get_i_length(cell_point='V') - e3v = mesh.get_k_length(cell_point='V') + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") + e2u = mesh.get_j_length(cell_point="U") + e3u = mesh.get_k_length(cell_point="U") + e1v = mesh.get_i_length(cell_point="V") + e3v = mesh.get_k_length(cell_point="V") results = psi.compute(self.masks, e2u, e1v, e3u, e3v, uo, vo) @@ -132,30 +147,30 @@ class Psi(Diagnostic): def save(self, result): temp = TempFile.get() handler_source = Utils.open_cdf(self.uo.local_file) - handler_temp = Utils.open_cdf(temp, 'w') - lat_name = next(alias for alias in ('lat', 'latitude') - if alias in handler_source.variables.keys()) - lon_name = next(alias for alias in ('lon', 'longitude') - if alias in handler_source.variables.keys()) - Utils.copy_variable(handler_source, handler_temp, 'time', True, True) - Utils.copy_variable(handler_source, handler_temp, 'i', False, True) - Utils.copy_variable(handler_source, handler_temp, 'j', False, True) + handler_temp = Utils.open_cdf(temp, "w") + lat_name = self.data_convention.lat_name + lon_name = self.data_convention.lon_name + Utils.copy_variable(handler_source, handler_temp, "time", True, True) + Utils.copy_variable(handler_source, handler_temp, "i", False, True) + Utils.copy_variable(handler_source, handler_temp, "j", False, True) Utils.copy_variable(handler_source, handler_temp, lat_name, True, True) Utils.copy_variable(handler_source, handler_temp, lon_name, True, True) - handler_temp.createDimension('region', len(result)) - handler_temp.createDimension('region_length', 50) - var_region = handler_temp.createVariable('region', 'S1', - ('region', 'region_length')) + handler_temp.createDimension("region", len(result)) + handler_temp.createDimension("region_length", 50) + var_region = handler_temp.createVariable( + "region", "S1", ("region", "region_length") + ) var = handler_temp.createVariable( - 'vsftbarot', float, ('time', 'j', 'i', 'region')) - var.units = 'm3/s' - var.coordinates = ' '.join((lat_name, lon_name)) + "vsftbarot", float, ("time", "j", "i", "region") + ) + var.units = "m3/s" + var.coordinates = " ".join((lat_name, lon_name)) var.missing_value = 1e20 var.fill_value = 1e20 var.valid_min = -300e6 var.valid_max = 300e6 - var.long_name = 'Barotropic_Stream_Function' + var.long_name = "Barotropic_Stream_Function" for i, basin in enumerate(result): var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) @@ -165,11 +180,10 @@ class Psi(Diagnostic): self.psi.set_local_file(temp, diagnostic=self) def _fix_coordinates_attribute(self, filepath, var_name): - add_coordinates = { - 'time_centered', 'leadtime' - } + add_coordinates = {"time_centered", "leadtime"} handler = Utils.open_cdf(filepath) coordinates = handler.variables[var_name].coordinates.split() - handler.variables[var_name].coordinates = \ - ' '.join(set(coordinates) | add_coordinates) + handler.variables[var_name].coordinates = " ".join( + set(coordinates) | add_coordinates + ) handler.close() diff --git a/earthdiagnostics/ocean/regionmean.py b/earthdiagnostics/ocean/regionmean.py index b05c1e18..2eb41416 100644 --- a/earthdiagnostics/ocean/regionmean.py +++ b/earthdiagnostics/ocean/regionmean.py @@ -5,17 +5,23 @@ import iris.util import iris.coords import iris.analysis import iris.exceptions +from bscearth.utils.log import Log import numpy as np import netCDF4 from earthdiagnostics.box import Box -from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticIntOption, DiagnosticDomainOption, \ - DiagnosticBoolOption, DiagnosticBasinListOption, \ - DiagnosticVariableListOption, DiagnosticFrequencyOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticIntOption, + DiagnosticDomainOption, + DiagnosticBoolOption, + DiagnosticBasinListOption, + DiagnosticVariableListOption, + DiagnosticFrequencyOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -49,12 +55,24 @@ class RegionMean(Diagnostic): :type box: Box """ - alias = 'regmean' + alias = "regmean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, - variable, box, save3d, variance, basins, grid_point, - frequency): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + box, + save3d, + variance, + basins, + grid_point, + frequency, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -70,21 +88,27 @@ class RegionMean(Diagnostic): self.declared = {} - self.lat_name = 'lat' - self.lon_name = 'lon' + self.lat_name = "lat" + self.lon_name = "lon" def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + ) def __str__(self): - return ('Region mean Startdate: {0.startdate} Member: {0.member} ' - 'Chunk: {0.chunk} Variable: {0.variable} ' - 'Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance} ' - 'Grid point: {0.grid_point}'.format(self)) + return ( + "Region mean Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Variable: {0.variable} " + "Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance} " + "Grid point: {0.grid_point}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -103,37 +127,48 @@ class RegionMean(Diagnostic): options_available = ( DiagnosticDomainOption(), DiagnosticVariableListOption( - diags.data_manager.config.var_manager, 'variable'), - DiagnosticBasinListOption('basins', 'global'), - DiagnosticOption('grid_point', 'T'), - DiagnosticIntOption('min_depth', -1), - DiagnosticIntOption('max_depth', -1), - DiagnosticBoolOption('save3D', True), - DiagnosticBoolOption('variance', False), - DiagnosticOption('grid', ''), - DiagnosticFrequencyOption('frequency', diags.config.frequency), + diags.data_manager.config.var_manager, "variable" + ), + DiagnosticBasinListOption("basins", "global"), + DiagnosticOption("grid_point", "T"), + DiagnosticIntOption("min_depth", -1), + DiagnosticIntOption("max_depth", -1), + DiagnosticBoolOption("save3D", True), + DiagnosticBoolOption("variance", False), + DiagnosticOption("grid", ""), + DiagnosticFrequencyOption("frequency", diags.config.frequency), ) options = cls.process_options(options, options_available) box = Box() - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') - return() + Log.error("Basins not recognized") + return () job_list = list() - for var in options['variable']: - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): + for var in options["variable"]: + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): job = RegionMean( - diags.data_manager, startdate, member, chunk, - options['domain'], var, box, - options['save3D'], options['variance'], - options['basins'], - options['grid_point'].lower(), - options['frequency'], + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + box, + options["save3D"], + options["variance"], + options["basins"], + options["grid_point"].lower(), + options["frequency"], ) job_list.append(job) @@ -141,9 +176,9 @@ class RegionMean(Diagnostic): def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(self.domain, self.variable, - self.startdate, self.member, - self.chunk) + self.variable_file = self.request_chunk( + self.domain, self.variable, self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" @@ -152,8 +187,8 @@ class RegionMean(Diagnostic): else: box_save = self.box - self._declare_var('mean', False, box_save) - self._declare_var('mean', True, box_save) + self._declare_var("mean", False, box_save) + self._declare_var("mean", True, box_save) def compute(self): """Run the diagnostic""" @@ -163,7 +198,7 @@ class RegionMean(Diagnostic): self.basins.sort() for basin in self.basins: masks[basin] = Utils.get_mask(basin) - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") if has_levels: self._meand_3d_variable(data, mesh, masks) else: @@ -172,13 +207,13 @@ class RegionMean(Diagnostic): def _mean_2d_var(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) mean = regmean.compute_regmean_2D(data.data, masks, areacello) - self._save_result_2D('mean', mean, data) + self._save_result_2d("mean", mean, data) def _meand_3d_variable(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) e3 = self._try_load_cube(3) e3 = self._rename_depth(e3) - e3.coord('depth').bounds = data.coord('depth').bounds + e3.coord("depth").bounds = data.coord("depth").bounds if self.box.min_depth is not -1 and self.box.max_depth is not -1: depth_constraint = iris.Constraint( depth=lambda c: self.box.min_depth <= c <= self.box.max_depth @@ -187,66 +222,91 @@ class RegionMean(Diagnostic): data = data.extract(depth_constraint) volcello = areacello * e3.data.astype(np.float32) mean = regmean.compute_regmean_3D(data.data, masks, volcello) - self._save_result_2D('mean', mean, data) + self._save_result_2d("mean", mean, data) if self.save3d: mean3d = regmean.compute_regmean_levels(data.data, masks, volcello) - self._save_result_3D('mean', mean3d, data) + self._save_result_3d("mean", mean3d, data) def _try_load_cube(self, number): try: - cube = iris.load_cube('mesh_hgr.nc', f'e{number}{self.grid_point}') + cube = iris.load_cube("mesh_hgr.nc", f"e{number}{self.grid_point}") except iris.exceptions.ConstraintMismatchError: cube = iris.load_cube( - 'mesh_hgr.nc', f'e{number}{self.grid_point}_0') + "mesh_hgr.nc", f"e{number}{self.grid_point}_0" + ) cube = iris.util.squeeze(cube) dims = len(cube.shape) try: - cube.coord('i') + cube.coord("i") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), - var_name='i'), dims - 1) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 1]), var_name="i" + ), + dims - 1, + ) try: - cube.coord('j') + cube.coord("j") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), - var_name='j'), dims - 2) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 2]), var_name="j" + ), + dims - 2, + ) return cube def _load_data(self): coords = [] handler = Utils.open_cdf(self.variable_file.local_file) for variable in handler.variables: - if variable in ('time', 'lev', 'lat', 'lon', 'latitude', - 'longitude', 'leadtime', 'time_centered'): + if variable in ( + "time", + "lev", + "lat", + "lon", + "latitude", + "longitude", + "leadtime", + "time_centered", + ): coords.append(variable) - if variable == 'time_centered': - handler.variables[variable].standard_name = '' + if variable == "time_centered": + handler.variables[variable].standard_name = "" - handler.variables[self.variable].coordinates = ' '.join(coords) + handler.variables[self.variable].coordinates = " ".join(coords) handler.close() data = iris.load_cube(self.variable_file.local_file) return self._rename_depth(data) def _rename_depth(self, data): - for coord_name in ('model_level_number', 'Vertical T levels', 'lev'): + for coord_name in ("model_level_number", "Vertical T levels", "lev"): if data.coords(coord_name): coord = data.coord(coord_name) - coord.standard_name = 'depth' - coord.long_name = 'depth' + coord.standard_name = "depth" + coord.long_name = "depth" break return data def _fix_file_metadata(self): handler = Utils.open_cdf(self.variable_file.local_file) var = handler.variables[self.variable] - coordinates = '' + coordinates = "" has_levels = False for dimension in handler.variables.keys(): - if dimension in ['time', 'lev', 'lat', 'latitude', - 'lon', 'longitude', 'i', 'j']: - coordinates += ' {0}'.format(dimension) - if dimension == 'lev': + if dimension in [ + "time", + "lev", + "lat", + "latitude", + "lon", + "longitude", + "i", + "j", + ]: + coordinates += " {0}".format(dimension) + if dimension == "lev": has_levels = True var.coordinates = coordinates handler.close() @@ -256,60 +316,68 @@ class RegionMean(Diagnostic): if threed: if not self.save3d: return False - final_name = '{1}3d{0}'.format(var, self.variable) + final_name = "{1}3d{0}".format(var, self.variable) else: - final_name = '{1}{0}'.format(var, self.variable) - - self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, - final_name, - self.startdate, - self.member, - self.chunk, - box=box_save, - region=self.basins) - - def _save_result_2D(self, var, result, data): - final_name = '{1}{0}'.format(var, self.variable) + final_name = "{1}{0}".format(var, self.variable) + + self.declared[final_name] = self.declare_chunk( + ModelingRealms.ocean, + final_name, + self.startdate, + self.member, + self.chunk, + box=box_save, + region=self.basins, + ) + + def _save_result_2d(self, var, result, data): + final_name = "{1}{0}".format(var, 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) - var_region = handler_temp.createVariable('region', 'S1', - ('region', 'region_length')) - var = handler_temp.createVariable('{1}{0}'.format(var, self.variable), - float, ('time', 'region',),) - var.units = '{0}'.format(data.units) + 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) + var_region = handler_temp.createVariable( + "region", "S1", ("region", "region_length") + ) + var = handler_temp.createVariable( + "{1}{0}".format(var, self.variable), float, ("time", "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) - def _save_result_3D(self, var, result, data): - final_name = '{1}3d{0}'.format(var, self.variable) + def _save_result_3d(self, var, result, data): + final_name = "{1}3d{0}".format(var, 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]) - 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('{1}3d{0}'.format(var, self.variable), - float, ('time', 'lev', 'region',),) - - var.units = '{0}'.format(data.units) + 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]) + 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( + "{1}3d{0}".format(var, self.variable), + float, + ("time", "lev", "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] diff --git a/earthdiagnostics/ocean/regionsum.py b/earthdiagnostics/ocean/regionsum.py index 99b46ce1..adadf303 100644 --- a/earthdiagnostics/ocean/regionsum.py +++ b/earthdiagnostics/ocean/regionsum.py @@ -1,23 +1,26 @@ # coding=utf-8 """Diagnostic to calculate a region total""" -import os - import iris import iris.util import iris.coords import iris.analysis import iris.exceptions +from bscearth.utils.log import Log import numpy as np import netCDF4 -from earthdiagnostics import cdftools from earthdiagnostics.box import Box -from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \ - DiagnosticIntOption, DiagnosticDomainOption, \ - DiagnosticBoolOption, DiagnosticBasinListOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticIntOption, + DiagnosticDomainOption, + DiagnosticBoolOption, + DiagnosticBasinListOption, + DiagnosticVariableOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -51,12 +54,23 @@ class RegionSum(Diagnostic): :type box: Box """ - alias = 'regsum' + alias = "regsum" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, - variable, grid_point, box, save3d, basins, - grid): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + grid_point, + box, + save3d, + basins, + grid, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -70,24 +84,31 @@ class RegionSum(Diagnostic): self.grid = grid self.declared = {} - self.lat_name = 'lat' - self.lon_name = 'lon' + self.lat_name = "lat" + self.lon_name = "lon" def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and \ - self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.variable == other.variable and \ - self.grid_point == other.grid_point and self.grid == other.grid \ + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + and self.grid_point == other.grid_point + and self.grid == other.grid and self.basin == other.basin + ) def __str__(self): - return 'Region sum Startdate: {0.startdate} Member: {0.member}' \ - 'Chunk: {0.chunk} Variable: {0.variable} ' \ - 'Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d}' \ - 'Original grid: {0.grid} Basin: {0.basins}'.format(self) + return ( + "Region sum Startdate: {0.startdate} Member: {0.member}" + "Chunk: {0.chunk} Variable: {0.variable} " + "Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d}" + "Original grid: {0.grid} Basin: {0.basins}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -103,48 +124,68 @@ class RegionSum(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticOption('grid_point', 'T'), - DiagnosticBasinListOption('basins', Basins().Global), - DiagnosticIntOption('min_depth', -1), - DiagnosticIntOption('max_depth', -1), - DiagnosticIntOption('min_lat', -1), - DiagnosticIntOption('max_lat', -1), - DiagnosticIntOption('min_lon', -1), - DiagnosticIntOption('max_lon', -1), - DiagnosticBoolOption('save3D', True), - DiagnosticOption('grid', '')) + options_available = ( + DiagnosticDomainOption(), + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticOption("grid_point", "T"), + DiagnosticBasinListOption("basins", "global"), + DiagnosticIntOption("min_depth", -1), + DiagnosticIntOption("max_depth", -1), + DiagnosticIntOption("min_lat", -1), + DiagnosticIntOption("max_lat", -1), + DiagnosticIntOption("min_lon", -1), + DiagnosticIntOption("max_lon", -1), + DiagnosticBoolOption("save3D", True), + DiagnosticOption("grid", ""), + ) options = cls.process_options(options, options_available) box = Box() - box.min_depth = options['min_depth'] - box.max_depth = options['max_depth'] - box.min_lat = options['min_lat'] - box.max_lat = options['max_lat'] - box.min_lon = options['min_lon'] - box.max_lon = options['max_lon'] - - basins = options['basins'] + box.min_depth = options["min_depth"] + box.max_depth = options["max_depth"] + box.min_lat = options["min_lat"] + box.max_lat = options["max_lat"] + box.min_lon = options["min_lon"] + box.max_lon = options["max_lon"] + + basins = options["basins"] if not basins: - Log.error('Basins not recognized') - return() + Log.error("Basins not recognized") + return () job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(RegionSum(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], - options['grid_point'].lower(), box, - options['save3D'], options['basins'], - options['grid'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + RegionSum( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["grid_point"].lower(), + box, + options["save3D"], + options["basins"], + options["grid"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(self.domain, self.variable, - self.startdate, - self.member, self.chunk, - grid=self.grid) + self.variable_file = self.request_chunk( + self.domain, + self.variable, + self.startdate, + self.member, + self.chunk, + grid=self.grid, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" @@ -153,8 +194,8 @@ class RegionSum(Diagnostic): else: box_save = self.box - self._declare_var('sum', False, box_save) - self._declare_var('sum', True, box_save) + self._declare_var("sum", False, box_save) + self._declare_var("sum", True, box_save) def compute(self): """Run the diagnostic""" @@ -165,15 +206,23 @@ class RegionSum(Diagnostic): for basin in self.basins: masks[basin] = Utils.get_mask(basin) - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') - if self.box.min_lat is not -1 and self.box.max_lat is not -1 and \ - self.box.min_lon is not -1 and self.box.max_lat is not -1: - name = '{0}_{1}'.format(Box.get_lat_str(self.box), Box.get_lon_str(self.box)) - - masks[name] = mesh.get_region_mask(self.box.min_lat, - self.box.max_lat, - self.box.min_lon, - self.box.max_lon) + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") + if ( + self.box.min_lat is not -1 + and self.box.max_lat is not -1 + and self.box.min_lon is not -1 + and self.box.max_lat is not -1 + ): + name = "{0}_{1}".format( + Box.get_lat_str(self.box), Box.get_lon_str(self.box) + ) + + masks[name] = mesh.get_region_mask( + self.box.min_lat, + self.box.max_lat, + self.box.min_lon, + self.box.max_lon, + ) if has_levels: self._sum_3d_variable(data, mesh, masks) else: @@ -182,173 +231,223 @@ class RegionSum(Diagnostic): def _sum_2d_var(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) varsum = regsum.compute_regsum_2D(data.data, masks, areacello) - self._save_result_2D('sum', varsum, data) + self._save_result_2d("sum", varsum, data) def _sum_3d_variable(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) - tmask = iris.load_cube('mesh_hgr.nc', '{0}mask'.format(self.grid_point)) + tmask = iris.load_cube( + "mesh_hgr.nc", "{0}mask".format(self.grid_point) + ) e3 = self._try_load_cube(3) e3 = self._rename_depth(e3) - e3.coord('depth').bounds = data.coord('depth').bounds + e3.coord("depth").bounds = data.coord("depth").bounds if self.box.min_depth is not -1 and self.box.max_depth is not -1: - depth_constraint = iris.Constraint(depth=lambda c: self.box.min_depth <= c <= self.box.max_depth) + depth_constraint = iris.Constraint( + depth=lambda c: self.box.min_depth <= c <= self.box.max_depth + ) e3 = e3.extract(depth_constraint) data = data.extract(depth_constraint) tmask = tmask.extract(depth_constraint) - volcello = areacello*e3.data.astype(np.float32) - varsum = regsum.compute_regsum_3D(data.data, masks, volcello, - tmask.data) - self._save_result_2D('sum', varsum, data) + volcello = areacello * e3.data.astype(np.float32) + varsum = regsum.compute_regsum_3D( + data.data, masks, volcello, tmask.data + ) + self._save_result_2d("sum", varsum, data) if self.save3d: - varsum3d = regsum.compute_regsum_levels(data.data, masks, - volcello, tmask) - self._save_result_3D('sum', varsum3d, data) + varsum3d = regsum.compute_regsum_levels( + data.data, masks, volcello, tmask + ) + self._save_result_3d("sum", varsum3d, data) def _try_load_cube(self, number): try: - cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}'.format(number, self.grid_point)) + cube = iris.load_cube( + "mesh_hgr.nc", "e{0}{1}".format(number, self.grid_point) + ) except iris.exceptions.ConstraintMismatchError: - cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}_0'.format(number, self.grid_point)) + cube = iris.load_cube( + "mesh_hgr.nc", "e{0}{1}_0".format(number, self.grid_point) + ) cube = iris.util.squeeze(cube) dims = len(cube.shape) try: - cube.coord('i') + cube.coord("i") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), var_name='i'), dims - 1) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 1]), var_name="i" + ), + dims - 1, + ) try: - cube.coord('j') + cube.coord("j") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), var_name='j'), dims - 2) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 2]), var_name="j" + ), + dims - 2, + ) return cube def _load_data(self): coords = [] handler = Utils.open_cdf(self.variable_file.local_file) for variable in handler.variables: - if variable in ('time', 'lev', 'lat', 'lon', 'latitude', 'longitude', 'leadtime', 'time_centered'): + if variable in ( + "time", + "lev", + "lat", + "lon", + "latitude", + "longitude", + "leadtime", + "time_centered", + ): coords.append(variable) - if variable == 'time_centered': - handler.variables[variable].standard_name = '' + if variable == "time_centered": + handler.variables[variable].standard_name = "" - handler.variables[self.variable].coordinates = ' '.join(coords) + handler.variables[self.variable].coordinates = " ".join(coords) handler.close() data = iris.load_cube(self.variable_file.local_file) return self._rename_depth(data) def _rename_depth(self, data): - for coord_name in ('model_level_number', 'Vertical T levels', 'lev'): + for coord_name in ("model_level_number", "Vertical T levels", "lev"): if data.coords(coord_name): coord = data.coord(coord_name) - coord.standard_name = 'depth' - coord.long_name = 'depth' + coord.standard_name = "depth" + coord.long_name = "depth" break return data def _fix_file_metadata(self): handler = Utils.open_cdf(self.variable_file.local_file) var = handler.variables[self.variable] - coordinates = '' + coordinates = "" has_levels = False for dimension in handler.variables.keys(): - if dimension in ['time', 'lev', 'lat', 'latitude', 'lon', 'longitude', 'i', 'j']: - coordinates += ' {0}'.format(dimension) - if dimension == 'lev': + if dimension in [ + "time", + "lev", + "lat", + "latitude", + "lon", + "longitude", + "i", + "j", + ]: + coordinates += " {0}".format(dimension) + if dimension == "lev": has_levels = True var.coordinates = coordinates handler.close() return has_levels - def _declare_var(self, var, threed, box_save): - if threed: - if not self.save3d: - return False - final_name = '{1}3d{0}'.format(var, self.variable) - else: - final_name = '{1}{0}'.format(var, self.variable) - - self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member, - self.chunk, box=box_save, region=self.basins) - def _send_var(self, threed, mean_file): - var = 'sum' + var = "sum" if threed: if not self.save3d: return False - original_name = '{0}_{1}'.format(var, self.variable) - final_name = '{1}3d{0}'.format(var, self.variable) - levels = ',lev' + original_name = "{0}_{1}".format(var, self.variable) + final_name = "{1}3d{0}".format(var, self.variable) + levels = ",lev" else: - original_name = '{0}_3D{1}'.format(var, self.variable) - final_name = '{1}{0}'.format(var, self.variable) - levels = '' + original_name = "{0}_3D{1}".format(var, self.variable) + final_name = "{1}{0}".format(var, self.variable) + levels = "" temp2 = TempFile.get() Utils.nco().ncks( input=mean_file, output=temp2, - options=('-v {0},{2.lat_name},{2.lon_name}{1}'.format(original_name, levels, self),) + options=( + "-v {0},{2.lat_name},{2.lon_name}{1}".format( + original_name, levels, self + ), + ), ) handler = Utils.open_cdf(temp2) var_handler = handler.variables[original_name] - if hasattr(var_handler, 'valid_min'): + if hasattr(var_handler, "valid_min"): del var_handler.valid_min - if hasattr(var_handler, 'valid_max'): + if hasattr(var_handler, "valid_max"): del var_handler.valid_max handler.close() - self.declared[final_name].set_local_file(temp2, diagnostic=self, rename_var=original_name, region=self.basins) + self.declared[final_name].set_local_file( + temp2, + diagnostic=self, + rename_var=original_name, + region=self.basins, + ) def _declare_var(self, var, threed, box_save): if threed: if not self.save3d: return False - final_name = '{1}3d{0}'.format(var, self.variable) + final_name = "{1}3d{0}".format(var, self.variable) else: - final_name = '{1}{0}'.format(var, self.variable) - - self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member, - self.chunk, box=box_save, region=self.basins, grid=self.grid) + final_name = "{1}{0}".format(var, self.variable) + + self.declared[final_name] = self.declare_chunk( + ModelingRealms.ocean, + final_name, + self.startdate, + self.member, + self.chunk, + box=box_save, + region=self.basins, + grid=self.grid, + ) - def _save_result_2D(self, var, result, data): - final_name = '{1}{0}'.format(var, self.variable) + def _save_result_2d(self, var, result, data): + final_name = "{1}{0}".format(var, 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) - var_region = handler_temp.createVariable('region', 'S1', - ('region', 'region_length')) - var = handler_temp.createVariable('{1}{0}'.format(var, self.variable), - float, ('time', 'region',),) - var.units = '{0}'.format(data.units) + 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) + var_region = handler_temp.createVariable( + "region", "S1", ("region", "region_length") + ) + var = handler_temp.createVariable( + "{1}{0}".format(var, self.variable), float, ("time", "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) - def _save_result_3D(self, var, result, data): - final_name = '{1}3d{0}'.format(var, self.variable) + def _save_result_3d(self, var, result, data): + final_name = "{1}3d{0}".format(var, 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]) - 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('{1}3d{0}'.format(var, self.variable), - float, ('time', 'lev', 'region',),) - var.units = '{0}'.format(data.units) + 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]) + 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( + "{1}3d{0}".format(var, self.variable), + float, + ("time", "lev", "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] diff --git a/earthdiagnostics/ocean/rotation.py b/earthdiagnostics/ocean/rotation.py index 22fdd405..5aa8a0be 100644 --- a/earthdiagnostics/ocean/rotation.py +++ b/earthdiagnostics/ocean/rotation.py @@ -4,7 +4,12 @@ import shutil from bscearth.utils.log import Log -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticDomainOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticVariableOption, +) from earthdiagnostics.utils import Utils, TempFile @@ -30,10 +35,20 @@ class Rotation(Diagnostic): :type domain: Domain """ - alias = 'rotate' + alias = "rotate" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variableu, variablev, executable): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variableu, + variablev, + executable, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -47,14 +62,29 @@ class Rotation(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variableu == other.variableu and \ - self.variablev == other.variablev and self.executable == other.executable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variableu == other.variableu + and self.variablev == other.variablev + and self.executable == other.executable + ) def __str__(self): - return 'Rotate variables Startdate: {0} Member: {1} Chunk: {2} Variables: {3}:{4} , ' \ - '{3}:{5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variableu, - self.variablev) + return ( + "Rotate variables Startdate: {0} Member: {1} Chunk: {2} " + "Variables: {3}:{4} , " + "{3}:{5}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variableu, + self.variablev, + ) + ) def __hash__(self): return hash(str(self)) @@ -70,37 +100,83 @@ class Rotation(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager, 'variableu'), - DiagnosticVariableOption(diags.data_manager.config.var_manager, 'variablev'), - DiagnosticOption('executable', - '/home/Earth/jvegas/pyCharm/cfutools/interpolation/rotateUVorca')) + options_available = ( + DiagnosticDomainOption(), + DiagnosticVariableOption( + diags.data_manager.config.var_manager, "variableu" + ), + DiagnosticVariableOption( + diags.data_manager.config.var_manager, "variablev" + ), + DiagnosticOption( + "executable", + "/home/Earth/jvegas/pyCharm/cfutools/interpolation/" + "rotateUVorca", + ), + ) options = cls.process_options(options, options_available) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Rotation(diags.data_manager, startdate, member, chunk, - options['domain'], options['variableu'], options['variablev'], - options['executable'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Rotation( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variableu"], + options["variablev"], + options["executable"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.ufile = self.request_chunk(self.domain, self.variableu, self.startdate, self.member, self.chunk) - self.vfile = self.request_chunk(self.domain, self.variablev, self.startdate, self.member, self.chunk) + self.ufile = self.request_chunk( + self.domain, + self.variableu, + self.startdate, + self.member, + self.chunk, + ) + self.vfile = self.request_chunk( + self.domain, + self.variablev, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.urotated_file = self.declare_chunk(self.domain, self.variableu, self.startdate, self.member, self.chunk, - grid='rotated') - self.vrotated_file = self.declare_chunk(self.domain, self.variablev, self.startdate, self.member, self.chunk, - grid='rotated') + self.urotated_file = self.declare_chunk( + self.domain, + self.variableu, + self.startdate, + self.member, + self.chunk, + grid="rotated", + ) + self.vrotated_file = self.declare_chunk( + self.domain, + self.variablev, + self.startdate, + self.member, + self.chunk, + grid="rotated", + ) def compute(self): """Run the diagnostic""" handler = Utils.open_cdf(self.ufile.local_file) - if 'lev' in handler.dimensions: - self.num_levels = handler.dimensions['lev'].size + if "lev" in handler.dimensions: + self.num_levels = handler.dimensions["lev"].size self.has_levels = True else: self.num_levels = 1 @@ -110,8 +186,8 @@ class Rotation(Diagnostic): for lev in range(0, self.num_levels): self._rotate_level(lev) - urotated = self._merge_levels(self.variableu, 'u') - vrotated = self._merge_levels(self.variablev, 'v') + urotated = self._merge_levels(self.variableu, "u") + vrotated = self._merge_levels(self.variablev, "v") ufile_handler = Utils.open_cdf(self.ufile.local_file) self._add_metadata_and_vars(ufile_handler, urotated, self.variableu) @@ -129,14 +205,18 @@ class Rotation(Diagnostic): Utils.nco().ncecat( input=self._get_level_file(0, direction), output=temp, - options=("-n {0},2,1 -v '{1}'".format(self.num_levels, var),) + options=("-n {0},2,1 -v '{1}'".format(self.num_levels, var),), ) handler = Utils.open_cdf(temp) - if 'record' in handler.dimensions: - handler.renameDimension('record', 'lev') + if "record" in handler.dimensions: + handler.renameDimension("record", "lev") handler.close() - Utils.nco().ncpdq(input=temp, output=temp, options=('-O -h -a time,lev',)) - Utils.rename_variables(temp, {'x': 'i', 'y': 'j'}, must_exist=False) + Utils.nco().ncpdq( + input=temp, output=temp, options=("-O -h -a time,lev",) + ) + Utils.rename_variables( + temp, {"x": "i", "y": "j"}, must_exist=False + ) else: Utils.move_file(self._get_level_file(0, direction), temp) return temp @@ -144,9 +224,15 @@ class Rotation(Diagnostic): def _rotate_level(self, lev): ufile = self._extract_level(self.ufile.local_file, self.variableu, lev) vfile = self._extract_level(self.vfile.local_file, self.variablev, lev) - namelist_file = self._create_namelist(ufile, self._get_level_file(lev, 'u'), - vfile, self._get_level_file(lev, 'v')) - Utils.execute_shell_command('{0} {1}'.format(self.executable, namelist_file), Log.INFO) + namelist_file = self._create_namelist( + ufile, + self._get_level_file(lev, "u"), + vfile, + self._get_level_file(lev, "v"), + ) + Utils.execute_shell_command( + "{0} {1}".format(self.executable, namelist_file), Log.INFO + ) def _extract_level(self, input_file, var, level): temp = TempFile.get() @@ -154,17 +240,19 @@ class Rotation(Diagnostic): Utils.nco().ncks( input=input_file, output=temp, - options=('-O -d lev,{0} -v {1},lat,lon'.format(level, var),) + options=("-O -d lev,{0} -v {1},lat,lon".format(level, var),), + ) + Utils.nco().ncwa( + input=temp, output=temp, options=("-O -h -a lev",) ) - Utils.nco().ncwa(input=temp, output=temp, options=('-O -h -a lev',)) else: shutil.copy(input_file, temp) return temp def _create_namelist(self, ufile, urotated, vfile, vrotated): - namelist_file = TempFile.get(suffix='') - rotate_namelist = open(namelist_file, 'w') - rotate_namelist.write('&nam_rotUV\n') + namelist_file = TempFile.get(suffix="") + rotate_namelist = open(namelist_file, "w") + rotate_namelist.write("&nam_rotUV\n") rotate_namelist.write(' Ufilein = "{0}"\n'.format(ufile)) rotate_namelist.write(' Uvarin = "{0}"\n'.format(self.variableu)) rotate_namelist.write(' Vfilein = "{0}"\n'.format(vfile)) @@ -176,20 +264,32 @@ class Rotation(Diagnostic): rotate_namelist.close() return namelist_file - def _add_metadata_and_vars(self, reference_file_handler, rotaded_file, var_name): + def _add_metadata_and_vars( + self, reference_file_handler, rotaded_file, var_name + ): rotated_handler = Utils.open_cdf(rotaded_file) self._copy_extra_variables(reference_file_handler, rotated_handler) - Utils.copy_attributes(rotated_handler.variables[var_name], reference_file_handler.variables[var_name], - ('_FillValue',)) + Utils.copy_attributes( + rotated_handler.variables[var_name], + reference_file_handler.variables[var_name], + ("_FillValue",), + ) rotated_handler.close() def _copy_extra_variables(self, reference_file_handler, rotated_handler): for var in reference_file_handler.variables.keys(): - if var not in rotated_handler.variables.keys() and var not in [self.variableu, self.variablev]: - Utils.copy_variable(reference_file_handler, rotated_handler, var, True, True) + if var not in rotated_handler.variables.keys() and var not in [ + self.variableu, + self.variablev, + ]: + Utils.copy_variable( + reference_file_handler, rotated_handler, var, True, True + ) def _get_level_file(self, lev, direction): if not self.tempTemplate: - self.tempTemplate = TempFile.get(suffix='_01.nc') + self.tempTemplate = TempFile.get(suffix="_01.nc") # self.tempTemplate = 'temp_01.nc' - return self.tempTemplate.replace('_01.nc', '_{1}_{0:02d}.nc'.format(lev + 1, direction)) + return self.tempTemplate.replace( + "_01.nc", "_{1}_{0:02d}.nc".format(lev + 1, direction) + ) diff --git a/earthdiagnostics/ocean/siarea.py b/earthdiagnostics/ocean/siarea.py index 9c867f54..bf00d185 100644 --- a/earthdiagnostics/ocean/siarea.py +++ b/earthdiagnostics/ocean/siarea.py @@ -1,5 +1,5 @@ # coding=utf-8 -"""Compute the sea ice extent , area and volume in both hemispheres or a specified region""" +"""Compute the sea ice area and extent""" import os import six @@ -12,7 +12,10 @@ import iris.util from bscearth.utils.log import Log from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinListOption, DiagnosticBoolOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBasinListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -22,7 +25,7 @@ from earthdiagnostics.utils import Utils, TempFile class Siarea(Diagnostic): """ - Compute the sea ice extent and area in both hemispheres or a specified region. + Compute the sea ice extent and area in both hemispheres or a region. Parameters ---------- @@ -36,14 +39,23 @@ class Siarea(Diagnostic): mask: numpy.array """ - alias = 'siarea' + alias = "siarea" "Diagnostic alias for the configuration file" e1t = None e2t = None gphit = None - def __init__(self, data_manager, startdate, member, chunk, masks, var_manager, data_convention): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + masks, + var_manager, + data_convention, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -51,16 +63,20 @@ class Siarea(Diagnostic): self.masks = masks self.generated = {} self.var_manager = var_manager - self.sic_varname = self.var_manager.get_variable('sic').short_name + self.sic_varname = self.var_manager.get_variable("sic").short_name self.data_convention = data_convention self.results = {} - for var in ('siarean', 'siareas', 'siextentn', 'siextents'): + for var in ("siarean", "siareas", "siextentn", "siextents"): self.results[var] = {} def __str__(self): - return 'Siarea Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Basins: {1}'.format(self, ','.join(str(basin) for basin in self.masks.keys())) + return ( + "Siarea Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Basins: {1}".format( + self, ",".join(str(basin) for basin in self.masks.keys()) + ) + ) def __hash__(self): return hash(str(self)) @@ -76,16 +92,18 @@ class Siarea(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption('basins', [Basins().Global])) + options_available = DiagnosticBasinListOption( + "basins", [Basins().Global] + ) options = cls.process_options(options, options_available) - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () - e1t = iris.load_cube('mesh_hgr.nc', 'e1t') - e2t = iris.load_cube('mesh_hgr.nc', 'e2t') + e1t = iris.load_cube("mesh_hgr.nc", "e1t") + e2t = iris.load_cube("mesh_hgr.nc", "e2t") area = e1t * e2t masks = {} @@ -94,102 +112,150 @@ class Siarea(Diagnostic): masks[basin] = Utils.get_mask(basin) * area.data job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Siarea( - diags.data_manager, startdate, member, chunk, masks, - diags.config.var_manager, diags.config.data_convention - )) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Siarea( + diags.data_manager, + startdate, + member, + chunk, + masks, + diags.config.var_manager, + diags.config.data_convention, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.sic = self.request_chunk(ModelingRealms.seaIce, self.sic_varname, - self.startdate, self.member, self.chunk) + self.sic = self.request_chunk( + ModelingRealms.seaIce, + self.sic_varname, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self._declare_var('siareas') - self._declare_var('siextents') + self._declare_var("siareas") + self._declare_var("siextents") - self._declare_var('siarean') - self._declare_var('siextentn') + self._declare_var("siarean") + self._declare_var("siextentn") def _declare_var(self, var_name): - self.generated[var_name] = self.declare_chunk(ModelingRealms.seaIce, var_name, - self.startdate, self.member, self.chunk) + self.generated[var_name] = self.declare_chunk( + ModelingRealms.seaIce, + var_name, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - coordinates = ' '.join(('time', 'leadtime', 'time_centered', - self.data_convention.lon_name, self.data_convention.lat_name)) + coordinates = " ".join( + ( + "time", + "leadtime", + "time_centered", + self.data_convention.lon_name, + self.data_convention.lat_name, + ) + ) handler = Utils.open_cdf(self.sic.local_file) handler.variables[self.sic_varname].coordinates = coordinates handler.close() sic = iris.load_cube(self.sic.local_file) - if sic.units.origin == '%' and sic.data.max() < 2: - sic.units = '1.0' + if sic.units.origin == "%" and sic.data.max() < 2: + sic.units = "1.0" extent = sic.copy((sic.data >= 0.15).astype(np.int8)) for basin, mask in six.iteritems(self.masks): - self.results['siarean'][basin] = self.sum(sic, mask, north=True) - self.results['siareas'][basin] = self.sum(sic, mask, north=False) - self.results['siextentn'][basin] = self.sum(extent, mask, north=True) - self.results['siextents'][basin] = self.sum(extent, mask, north=False) + self.results["siarean"][basin] = self.sum(sic, mask, north=True) + self.results["siareas"][basin] = self.sum(sic, mask, north=False) + self.results["siextentn"][basin] = self.sum( + extent, mask, north=True + ) + self.results["siextents"][basin] = self.sum( + extent, mask, north=False + ) self.save() def sum(self, data, mask, north=True): if north: - condition = data.coord('latitude').points > 0 + condition = data.coord("latitude").points > 0 else: - condition = data.coord('latitude').points < 0 - weights = iris.util.broadcast_to_shape(condition, data.shape, data.coord_dims('latitude')) * mask - return data.collapsed(('latitude', 'longitude'), iris.analysis.SUM, weights=weights) + condition = data.coord("latitude").points < 0 + weights = ( + iris.util.broadcast_to_shape( + condition, data.shape, data.coord_dims("latitude") + ) + * mask + ) + return data.collapsed( + ("latitude", "longitude"), iris.analysis.SUM, weights=weights + ) def save(self): for var, basins in six.iteritems(self.results): results = iris.cube.CubeList() for basin, result in six.iteritems(basins): result.var_name = var - result.units = 'm^2' - result.add_aux_coord(iris.coords.AuxCoord(basin.name, var_name='region')) + result.units = "m^2" + result.add_aux_coord( + iris.coords.AuxCoord(basin.name, var_name="region") + ) results.append(result) self._save_file(results.merge_cube(), var) def _save_file(self, data, var): generated_file = self.generated[var] temp = TempFile.get() - region = data.coord('region').points - data.remove_coord('region') + region = data.coord("region").points + data.remove_coord("region") iris.save(data, temp, zlib=True) if len(region) > 1: - Utils.rename_variable(temp, 'dim0', 'region', False) + Utils.rename_variable(temp, "dim0", "region", False) handler = Utils.open_cdf(temp) - var = handler.createVariable('region2', str, ('region',)) + var = handler.createVariable("region2", str, ("region",)) var[...] = region handler.close() - Utils.rename_variable(temp, 'region2', 'region', True) + Utils.rename_variable(temp, "region2", "region", True) else: handler = Utils.open_cdf(temp) - if 'region' not in handler.dimensions: + if "region" not in handler.dimensions: new_file = TempFile.get() - new_handler = Utils.open_cdf(new_file, 'w') + new_handler = Utils.open_cdf(new_file, "w") - new_handler.createDimension('region', 1) + new_handler.createDimension("region", 1) for dimension in handler.dimensions: Utils.copy_dimension(handler, new_handler, dimension) for variable in handler.variables.keys(): - if variable in (var, 'region'): + if variable in (var, "region"): continue Utils.copy_variable(handler, new_handler, variable) old_var = handler.variables[var] - new_var = new_handler.createVariable(var, old_var.dtype, ('region',) + old_var.dimensions, - zlib=True, fill_value=1.0e20) + new_var = new_handler.createVariable( + var, + old_var.dtype, + ("region",) + old_var.dimensions, + zlib=True, + fill_value=1.0e20, + ) Utils.copy_attributes(new_var, old_var) new_var[0, :] = old_var[:] - new_var = new_handler.createVariable('region', str, ('region',)) + new_var = new_handler.createVariable( + "region", str, ("region",) + ) new_var[0] = region[0] new_handler.close() diff --git a/earthdiagnostics/ocean/siasiesiv.py b/earthdiagnostics/ocean/siasiesiv.py index 0b81ff6c..94c1e78b 100644 --- a/earthdiagnostics/ocean/siasiesiv.py +++ b/earthdiagnostics/ocean/siasiesiv.py @@ -1,8 +1,5 @@ # coding=utf-8 -"""Compute the sea ice extent , area and volume in both hemispheres or a specified region""" -import os -import six - +"""Compute the sea ice extent, area and volume""" import numpy as np import netCDF4 @@ -14,8 +11,11 @@ import iris.util from bscearth.utils.log import Log from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, \ - DiagnosticBasinListOption, DiagnosticBoolOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBasinListOption, + DiagnosticBoolOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -44,11 +44,20 @@ class Siasiesiv(Diagnostic): omit_vol: bool """ - alias = 'siasiesiv' + alias = "siasiesiv" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, masks, - var_manager, data_convention, omit_vol): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + masks, + var_manager, + data_convention, + omit_vol, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -57,18 +66,21 @@ class Siasiesiv(Diagnostic): self.generated = {} self.var_manager = var_manager self.omit_volume = omit_vol - self.sic_varname = self.var_manager.get_variable('sic').short_name - self.sit_varname = self.var_manager.get_variable('sit').short_name + self.sic_varname = self.var_manager.get_variable("sic").short_name + self.sit_varname = self.var_manager.get_variable("sit").short_name self.data_convention = data_convention self.results = {} - for var in ('siarean', 'siareas', 'siextentn', 'siextents'): + for var in ("siarean", "siareas", "siextentn", "siextents"): self.results[var] = {} def __str__(self): - return 'Siasiesiv Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Basins: {1} Omit volume: {0.omit_volume}'.format(self, - ','.join(str(basin) for basin in self.masks.keys())) + return ( + "Siasiesiv Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Basins: {1} Omit volume: {0.omit_volume}".format( + self, ",".join(str(basin) for basin in self.masks.keys()) + ) + ) def __hash__(self): return hash(str(self)) @@ -84,14 +96,15 @@ class Siasiesiv(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption('basins', - [Basins().Global]), - DiagnosticBoolOption('omit_volume', True)) + options_available = ( + DiagnosticBasinListOption("basins", [Basins().Global]), + DiagnosticBoolOption("omit_volume", True), + ) options = cls.process_options(options, options_available) - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () masks = {} @@ -100,94 +113,123 @@ class Siasiesiv(Diagnostic): masks[basin] = Utils.get_mask(basin) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Siasiesiv(diags.data_manager, startdate, member, - chunk, masks, diags.config.var_manager, - diags.config.data_convention, - options['omit_volume'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + Siasiesiv( + diags.data_manager, + startdate, + member, + chunk, + masks, + diags.config.var_manager, + diags.config.data_convention, + options["omit_volume"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" if not self.omit_volume: - self.sit = self.request_chunk(ModelingRealms.seaIce, - self.sit_varname, - self.startdate, - self.member, - self.chunk) - self.sic = self.request_chunk(ModelingRealms.seaIce, self.sic_varname, - self.startdate, self.member, self.chunk) + self.sit = self.request_chunk( + ModelingRealms.seaIce, + self.sit_varname, + self.startdate, + self.member, + self.chunk, + ) + self.sic = self.request_chunk( + ModelingRealms.seaIce, + self.sic_varname, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" if not self.omit_volume: - self._declare_var('sivols') - self._declare_var('sivoln') + self._declare_var("sivols") + self._declare_var("sivoln") - self._declare_var('siareas') - self._declare_var('siextents') + self._declare_var("siareas") + self._declare_var("siextents") - self._declare_var('siarean') - self._declare_var('siextentn') + self._declare_var("siarean") + self._declare_var("siextentn") def _declare_var(self, var_name): - self.generated[var_name] = self.declare_chunk(ModelingRealms.seaIce, - var_name, self.startdate, - self.member, self.chunk) + self.generated[var_name] = self.declare_chunk( + ModelingRealms.seaIce, + var_name, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - self._fix_coordinates_attribute( - self.sic.local_file, self.sic_varname - ) + self._fix_coordinates_attribute(self.sic.local_file, self.sic_varname) sic = iris.load_cube(self.sic.local_file) - # if sic.units.origin == '%' and sic.data.max() < 2: - # sic.units = '1.0' - if sic.units.origin is not '%': - sic.convert_units('%') + if sic.units.origin != "%": + sic.convert_units("%") sic_slices = [] - for sic_data in sic.slices_over('time'): + for sic_data in sic.slices_over("time"): sic_data.data = np.ma.filled(sic_data.data, 0.0).astype(np.float32) sic_data.data[np.isinf(sic_data.data)] = 0.0 sic_slices.append(sic_data) - mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') - areacello = mesh.get_areacello(cell_point='T') - gphit = mesh.get_grid_latitude(cell_point='T') + mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") + areacello = mesh.get_areacello(cell_point="T") + gphit = mesh.get_grid_latitude(cell_point="T") if not self.omit_volume: sit = iris.load_cube(self.sit.local_file) sit_slices = [] - for sit_data in sit.slices_over('time'): - sit_data.data = np.ma.filled(sit_data.data, 0.0).astype(np.float32) + for sit_data in sit.slices_over("time"): + sit_data.data = np.ma.filled(sit_data.data, 0.0).astype( + np.float32 + ) sit_data.data[np.isinf(sit_data.data)] = 0.0 sit_slices.append(sit_data) - results = siasie.compute(gphit, areacello, sic_slices, self.masks, sit_slices) - self.results['siextentn'] = results[0] - self.results['siextents'] = results[1] - self.results['siarean'] = results[2] - self.results['siareas'] = results[3] - self.results['sivoln'] = results[4] - self.results['sivols'] = results[5] + results = siasie.compute( + gphit, areacello, sic_slices, self.masks, sit_slices + ) + self.results["siextentn"] = results[0] + self.results["siextents"] = results[1] + self.results["siarean"] = results[2] + self.results["siareas"] = results[3] + self.results["sivoln"] = results[4] + self.results["sivols"] = results[5] else: - results = siasie.compute(gphit, areacello, sic_slices, self.masks, None) - self.results['siextentn'] = results[0] - self.results['siextents'] = results[1] - self.results['siarean'] = results[2] - self.results['siareas'] = results[3] + results = siasie.compute( + gphit, areacello, sic_slices, self.masks, None + ) + self.results["siextentn"] = results[0] + self.results["siextents"] = results[1] + self.results["siarean"] = results[2] + self.results["siareas"] = results[3] self.save() def _fix_coordinates_attribute(self, filepath, var_name): add_coordinates = { - 'time', 'leadtime', 'time_centered', - self.data_convention.lon_name, self.data_convention.lat_name + "time", + "leadtime", + "time_centered", + self.data_convention.lon_name, + self.data_convention.lat_name, } handler = Utils.open_cdf(filepath) coordinates = handler.variables[var_name].coordinates.split() - handler.variables[var_name].coordinates = \ - ' '.join(set(coordinates) | add_coordinates) + handler.variables[var_name].coordinates = " ".join( + set(coordinates) | add_coordinates + ) handler.close() def save(self): @@ -195,20 +237,22 @@ class Siasiesiv(Diagnostic): res = self.results[var] temp = TempFile.get() handler_source = Utils.open_cdf(self.sic.local_file) - handler_temp = Utils.open_cdf(temp, 'w') - Utils.copy_variable(handler_source, handler_temp, 'time', True, True) - handler_temp.createDimension('region', len(self.masks)) - handler_temp.createDimension('region_length', 50) + handler_temp = Utils.open_cdf(temp, "w") + Utils.copy_variable( + handler_source, handler_temp, "time", True, True + ) + handler_temp.createDimension("region", len(self.masks)) + handler_temp.createDimension("region_length", 50) var_region = handler_temp.createVariable( - 'region', 'S1', ('region', 'region_length') + "region", "S1", ("region", "region_length") ) var_res = handler_temp.createVariable( - '{0}'.format(var), float, ('time', 'region',) + "{0}".format(var), float, ("time", "region",) ) - if var in ('sivoln', 'sivols'): - var_res.units = 'm^3' + if var in ("sivoln", "sivols"): + var_res.units = "m^3" else: - var_res.units = 'm^2' + var_res.units = "m^2" for i, basin in enumerate(self.masks): if not np.all(res[i, ...] == 0): var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) diff --git a/earthdiagnostics/ocean/sivol2d.py b/earthdiagnostics/ocean/sivol2d.py index fe0e9dc7..e74f4a80 100644 --- a/earthdiagnostics/ocean/sivol2d.py +++ b/earthdiagnostics/ocean/sivol2d.py @@ -1,25 +1,14 @@ # coding=utf-8 -"""Compute the sea ice extent , area and volume in both hemispheres or a specified region""" -import os -import six - -import numpy as np - +"""Compute the sea ice volume""" import iris import iris.analysis import iris.coords import iris.util -from bscearth.utils.log import Log - -from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinListOption, DiagnosticBoolOption +from earthdiagnostics.diagnostic import Diagnostic from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile -# noinspection PyUnresolvedReferences - - class Sivol2d(Diagnostic): """ Compute the sea ice volumen from sic and sit @@ -36,27 +25,38 @@ class Sivol2d(Diagnostic): mask: numpy.array """ - alias = 'sivol2d' + alias = "sivol2d" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, var_manager, data_convention): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + var_manager, + data_convention, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.generated = {} self.var_manager = var_manager - self.sic_varname = self.var_manager.get_variable('sic').short_name - self.sit_varname = self.var_manager.get_variable('sit').short_name - self.sivol_varname = self.var_manager.get_variable('sivol').short_name + self.sic_varname = self.var_manager.get_variable("sic").short_name + self.sit_varname = self.var_manager.get_variable("sit").short_name + self.sivol_varname = self.var_manager.get_variable("sivol").short_name self.data_convention = data_convention self.results = {} - for var in ('siarean', 'siareas', 'siextentn', 'siextents'): + for var in ("siarean", "siareas", "siextentn", "siextents"): self.results[var] = {} def __str__(self): - return 'Sivol2d Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk}'.format(self) + return ( + f"Sivol2d Startdate: {self.startdate} Member: {self.member} " + f"Chunk: {self.chunk}" + ) def __hash__(self): return hash(str(self)) @@ -76,38 +76,68 @@ class Sivol2d(Diagnostic): options = cls.process_options(options, options_available) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): job_list.append( Sivol2d( - diags.data_manager, startdate, member, chunk, - diags.config.var_manager, diags.config.data_convention + diags.data_manager, + startdate, + member, + chunk, + diags.config.var_manager, + diags.config.data_convention, ) ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.sic = self.request_chunk(ModelingRealms.seaIce, self.sic_varname, - self.startdate, self.member, self.chunk) - self.sit = self.request_chunk(ModelingRealms.seaIce, self.sit_varname, - self.startdate, self.member, self.chunk) + self.sic = self.request_chunk( + ModelingRealms.seaIce, + self.sic_varname, + self.startdate, + self.member, + self.chunk, + ) + self.sit = self.request_chunk( + ModelingRealms.seaIce, + self.sit_varname, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.sivol = self.declare_chunk(ModelingRealms.seaIce, self.sivol_varname, - self.startdate, self.member, self.chunk) + self.sivol = self.declare_chunk( + ModelingRealms.seaIce, + self.sivol_varname, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - coordinates = ' '.join(('time', 'leadtime', 'time_centered', - self.data_convention.lon_name, self.data_convention.lat_name)) + coordinates = " ".join( + ( + "time", + "leadtime", + "time_centered", + self.data_convention.lon_name, + self.data_convention.lat_name, + ) + ) handler = Utils.open_cdf(self.sic.local_file) handler.variables[self.sic_varname].coordinates = coordinates handler.close() sic = iris.load_cube(self.sic.local_file) - if sic.units.origin == '%' and sic.data.max() < 2: - sic.units = '1.0' - sic.convert_units('1.0') + if sic.units.origin == "%" and sic.data.max() < 2: + sic.units = "1.0" + sic.convert_units("1.0") handler = Utils.open_cdf(self.sit.local_file) handler.variables[self.sit_varname].coordinates = coordinates @@ -119,8 +149,10 @@ class Sivol2d(Diagnostic): del sic sivol.var_name = self.sivol_varname sivol.standard_name = "sea_ice_thickness" - sivol.long_name = "Total volume of sea ice divided by grid-cell area" \ - " (this used to be called ice thickness in CMIP5)" + sivol.long_name = ( + "Total volume of sea ice divided by grid-cell area" + " (this used to be called ice thickness in CMIP5)" + ) temp = TempFile.get() iris.save(sivol, temp, zlib=True) del sivol diff --git a/earthdiagnostics/ocean/sivolume.py b/earthdiagnostics/ocean/sivolume.py index 75c71e23..75ba01cb 100644 --- a/earthdiagnostics/ocean/sivolume.py +++ b/earthdiagnostics/ocean/sivolume.py @@ -1,10 +1,8 @@ # coding=utf-8 -"""Compute the sea ice extent , area and volume in both hemispheres or a specified region""" +"""Computes the seaice extent, area and volume""" import os import six -import numpy as np - import iris import iris.analysis import iris.coords @@ -12,17 +10,17 @@ import iris.util from bscearth.utils.log import Log from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinListOption, DiagnosticBoolOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticBasinListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile -# noinspection PyUnresolvedReferences - - class Sivolume(Diagnostic): """ - Compute the sea ice volume from sivol in both hemispheres or a specified region. + Compute the sea ice volume from sivol in both hemispheres or a region. Parameters ---------- @@ -35,10 +33,19 @@ class Sivolume(Diagnostic): mask: numpy.array """ - alias = 'sivolume' + alias = "sivolume" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, masks, var_manager, data_convention): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + masks, + var_manager, + data_convention, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -47,15 +54,19 @@ class Sivolume(Diagnostic): self.generated = {} self.var_manager = var_manager self.data_convention = data_convention - self.sivol_varname = self.var_manager.get_variable('sivol').short_name + self.sivol_varname = self.var_manager.get_variable("sivol").short_name self.results = {} self.sivoln = {} self.sivols = {} def __str__(self): - return 'Sivolume Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \ - 'Basins: {1} '.format(self, ','.join(str(basin) for basin in self.masks.keys())) + return ( + "Sivolume Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Basins: {1} ".format( + self, ",".join(str(basin) for basin in self.masks.keys()) + ) + ) def __hash__(self): return hash(str(self)) @@ -71,53 +82,79 @@ class Sivolume(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption('basins', [Basins().Global]),) + options_available = ( + DiagnosticBasinListOption("basins", [Basins().Global]), + ) options = cls.process_options(options, options_available) - basins = options['basins'] + basins = options["basins"] if not basins: - Log.error('Basins not recognized') + Log.error("Basins not recognized") return () masks = {} basins.sort() - e1t = iris.load_cube('mesh_hgr.nc', 'e1t') - e2t = iris.load_cube('mesh_hgr.nc', 'e2t') + e1t = iris.load_cube("mesh_hgr.nc", "e1t") + e2t = iris.load_cube("mesh_hgr.nc", "e2t") area = e1t * e2t for basin in basins: masks[basin] = Utils.get_mask(basin) * area.data job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): job_list.append( Sivolume( - diags.data_manager, startdate, member, chunk, masks, - diags.config.var_manager, diags.config.data_convention, - ) + diags.data_manager, + startdate, + member, + chunk, + masks, + diags.config.var_manager, + diags.config.data_convention, + ) ) return job_list def request_data(self): """Request data required by the diagnostic""" self.sivol = self.request_chunk( - ModelingRealms.seaIce, self.sivol_varname, - self.startdate, self.member, self.chunk + ModelingRealms.seaIce, + self.sivol_varname, + self.startdate, + self.member, + self.chunk, ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self._declare_var('sivols') - self._declare_var('sivoln') + self._declare_var("sivols") + self._declare_var("sivoln") def _declare_var(self, var_name): - self.generated[var_name] = self.declare_chunk(ModelingRealms.seaIce, var_name, - self.startdate, self.member, self.chunk) + self.generated[var_name] = self.declare_chunk( + ModelingRealms.seaIce, + var_name, + self.startdate, + self.member, + self.chunk, + ) def compute(self): """Run the diagnostic""" - coordinates = ' '.join(('time', 'leadtime', 'time_centered', - self.data_convention.lon_name, self.data_convention.lat_name)) + coordinates = " ".join( + ( + "time", + "leadtime", + "time_centered", + self.data_convention.lon_name, + self.data_convention.lat_name, + ) + ) handler = Utils.open_cdf(self.sivol.local_file) handler.variables[self.sivol_varname].coordinates = coordinates @@ -129,62 +166,78 @@ class Sivolume(Diagnostic): self.sivols[basin] = self.sum(sivol, mask, north=False) del sivol - self.save('sivoln', self.sivoln) - self.save('sivols', self.sivols) + self.save("sivoln", self.sivoln) + self.save("sivols", self.sivols) del self.sivols del self.sivoln def sum(self, data, mask, north=True): if north: - condition = data.coord('latitude').points > 0 + condition = data.coord("latitude").points > 0 else: - condition = data.coord('latitude').points < 0 - weights = iris.util.broadcast_to_shape(condition, data.shape, data.coord_dims('latitude')) * mask - return data.collapsed(('latitude', 'longitude'), iris.analysis.SUM, weights=weights) + condition = data.coord("latitude").points < 0 + weights = ( + iris.util.broadcast_to_shape( + condition, data.shape, data.coord_dims("latitude") + ) + * mask + ) + return data.collapsed( + ("latitude", "longitude"), iris.analysis.SUM, weights=weights + ) def save(self, var, results): cubes = iris.cube.CubeList() for basin, result in six.iteritems(results): result.var_name = var - result.units = 'm^3' - result.add_aux_coord(iris.coords.AuxCoord(basin.name, var_name='region')) + result.units = "m^3" + result.add_aux_coord( + iris.coords.AuxCoord(basin.name, var_name="region") + ) cubes.append(result) self._save_file(cubes.merge_cube(), var) def _save_file(self, data, var): generated_file = self.generated[var] temp = TempFile.get() - region = data.coord('region').points - data.remove_coord('region') + region = data.coord("region").points + data.remove_coord("region") iris.save(data, temp, zlib=True) if len(region) > 1: - Utils.rename_variable(temp, 'dim0', 'region', False) + Utils.rename_variable(temp, "dim0", "region", False) handler = Utils.open_cdf(temp) - var = handler.createVariable('region2', str, ('region',)) + var = handler.createVariable("region2", str, ("region",)) var[...] = region handler.close() - Utils.rename_variable(temp, 'region2', 'region', True) + Utils.rename_variable(temp, "region2", "region", True) else: handler = Utils.open_cdf(temp) - if 'region' not in handler.dimensions: + if "region" not in handler.dimensions: new_file = TempFile.get() - new_handler = Utils.open_cdf(new_file, 'w') + new_handler = Utils.open_cdf(new_file, "w") - new_handler.createDimension('region', 1) + new_handler.createDimension("region", 1) for dimension in handler.dimensions: Utils.copy_dimension(handler, new_handler, dimension) for variable in handler.variables.keys(): - if variable in (var, 'region'): + if variable in (var, "region"): continue Utils.copy_variable(handler, new_handler, variable) old_var = handler.variables[var] - new_var = new_handler.createVariable(var, old_var.dtype, ('region',) + old_var.dimensions, - zlib=True, fill_value=1.0e20) + new_var = new_handler.createVariable( + var, + old_var.dtype, + ("region",) + old_var.dimensions, + zlib=True, + fill_value=1.0e20, + ) Utils.copy_attributes(new_var, old_var) new_var[0, :] = old_var[:] - new_var = new_handler.createVariable('region', str, ('region',)) + new_var = new_handler.createVariable( + "region", str, ("region",) + ) new_var[0] = region[0] new_handler.close() diff --git a/earthdiagnostics/ocean/verticalgradient.py b/earthdiagnostics/ocean/verticalgradient.py index f1b1af24..133cc29b 100644 --- a/earthdiagnostics/ocean/verticalgradient.py +++ b/earthdiagnostics/ocean/verticalgradient.py @@ -1,7 +1,11 @@ # coding=utf-8 """Calculate the gradient between 2 ocean levels""" from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticIntOption, + DiagnosticVariableOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -31,7 +35,7 @@ class VerticalGradient(Diagnostic): :type box: Box """ - alias = 'vgrad' + alias = "vgrad" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, variable, box): @@ -45,12 +49,25 @@ class VerticalGradient(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + ) def __str__(self): - return 'Vertical gradient Startdate: {0} Member: {1} Chunk: {2} Variable: {3} ' \ - 'Box: {4}'.format(self.startdate, self.member, self.chunk, self.variable, self.box) + return ( + "Vertical gradient Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3} Box: {4}".format( + self.startdate, + self.member, + self.chunk, + self.variable, + self.box, + ) + ) def __hash__(self): return hash(str(self)) @@ -66,54 +83,83 @@ class VerticalGradient(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticIntOption('upper_level', 1), - DiagnosticIntOption('low_level', 2)) + options_available = ( + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticIntOption("upper_level", 1), + DiagnosticIntOption("low_level", 2), + ) options = cls.process_options(options, options_available) box = Box(False) - if options['upper_level'] >= 0: - box.min_depth = options['upper_level'] - if options['low_level'] >= 0: - box.max_depth = options['low_level'] + if options["upper_level"] >= 0: + box.min_depth = options["upper_level"] + if options["low_level"] >= 0: + box.max_depth = options["low_level"] job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(VerticalGradient(diags.data_manager, startdate, member, chunk, - options['variable'], box)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + VerticalGradient( + diags.data_manager, + startdate, + member, + chunk, + options["variable"], + box, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(ModelingRealms.ocean, self.variable, self.startdate, - self.member, self.chunk) + self.variable_file = self.request_chunk( + ModelingRealms.ocean, + self.variable, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.gradient_file = self.declare_chunk(ModelingRealms.ocean, self.variable + 'vgrad', - self.startdate, self.member, self.chunk, box=self.box) + self.gradient_file = self.declare_chunk( + ModelingRealms.ocean, + self.variable + "vgrad", + self.startdate, + self.member, + self.chunk, + box=self.box, + ) def compute(self): """Run the diagnostic""" handler = Utils.open_cdf(self.variable_file) - if 'lev' not in handler.dimensions: - raise Exception('Variable {0} does not have a level dimension') + if "lev" not in handler.dimensions: + raise Exception("Variable {0} does not have a level dimension") var_handler = handler.variables[self.variable] upper_level = var_handler[:, self.box.min_depth - 1, ...] lower_level = var_handler[:, self.box.max_depth - 1, ...] gradient = upper_level - lower_level temp = TempFile.get() - new_file = Utils.open_cdf(temp, 'w') + new_file = Utils.open_cdf(temp, "w") for var in handler.variables.keys(): - if var in (self.variable, 'lev', 'lev_bnds'): + if var in (self.variable, "lev", "lev_bnds"): continue Utils.copy_variable(handler, new_file, var, add_dimensions=True) - new_var = new_file.createVariable(self.variable + 'vgrad', var_handler.dtype, - dimensions=('time', 'j', 'i'), zlib=True) + new_var = new_file.createVariable( + self.variable + "vgrad", + var_handler.dtype, + dimensions=("time", "j", "i"), + zlib=True, + ) Utils.copy_attributes(new_var, var_handler) new_var[...] = gradient[...] - new_var.long_name += ' Vertical gradient' - new_var.standard_name += '_vertical_gradient' + new_var.long_name += " Vertical gradient" + new_var.standard_name += "_vertical_gradient" self.gradient_file.set_local_file(temp) diff --git a/earthdiagnostics/ocean/verticalmean.py b/earthdiagnostics/ocean/verticalmean.py index 9e4e130c..b9fa4bf4 100644 --- a/earthdiagnostics/ocean/verticalmean.py +++ b/earthdiagnostics/ocean/verticalmean.py @@ -1,15 +1,19 @@ # coding=utf-8 -"""Chooses vertical level in ocean, or vertically averages between 2 or more ocean levels""" +"""Chooses vertical level in ocean, or vertically averages between 2+ levels""" from earthdiagnostics import cdftools from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticIntOption, + DiagnosticVariableOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile class VerticalMean(Diagnostic): """ - Chooses vertical level in ocean, or vertically averages between 2 or more ocean levels + Chooses vertical level in ocean, or vertically averages between 2+ levels :original author: Virginie Guemas :contributor: Eleftheria Exarchou @@ -32,7 +36,7 @@ class VerticalMean(Diagnostic): :type box: Box """ - alias = 'vertmean' + alias = "vertmean" "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, variable, box): @@ -46,12 +50,25 @@ class VerticalMean(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + ) def __str__(self): - return 'Vertical mean Startdate: {0} Member: {1} Chunk: {2} Variable: {3} ' \ - 'Box: {4}'.format(self.startdate, self.member, self.chunk, self.variable, self.box) + return ( + "Vertical mean Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3} Box: {4}".format( + self.startdate, + self.member, + self.chunk, + self.variable, + self.box, + ) + ) def __hash__(self): return hash(str(self)) @@ -67,49 +84,80 @@ class VerticalMean(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticIntOption('min_depth', -1), - DiagnosticIntOption('max_depth', -1)) + options_available = ( + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticIntOption("min_depth", -1), + DiagnosticIntOption("max_depth", -1), + ) options = cls.process_options(options, options_available) box = Box() - if options['min_depth'] >= 0: - box.min_depth = options['min_depth'] - if options['max_depth'] >= 0: - box.max_depth = options['max_depth'] + if options["min_depth"] >= 0: + box.min_depth = options["min_depth"] + if options["max_depth"] >= 0: + box.max_depth = options["max_depth"] job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(VerticalMean(diags.data_manager, startdate, member, chunk, - options['variable'], box)) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + VerticalMean( + diags.data_manager, + startdate, + member, + chunk, + options["variable"], + box, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(ModelingRealms.ocean, self.variable, self.startdate, self.member, - self.chunk) + self.variable_file = self.request_chunk( + ModelingRealms.ocean, + self.variable, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.results = self.declare_chunk(ModelingRealms.ocean, self.variable + 'vmean', self.startdate, self.member, - self.chunk, box=self.box) + self.results = self.declare_chunk( + ModelingRealms.ocean, + self.variable + "vmean", + self.startdate, + self.member, + self.chunk, + box=self.box, + ) def compute(self): """Run the diagnostic""" temp = TempFile.get() handler = Utils.open_cdf(self.variable_file.local_file) if self.box.min_depth is None: - lev_min = handler.variables['lev'][0] + lev_min = handler.variables["lev"][0] else: lev_min = self.box.min_depth if self.box.max_depth is None: - lev_max = handler.variables['lev'][-1] + lev_max = handler.variables["lev"][-1] else: lev_max = self.box.max_depth handler.close() - cdftools.run('cdfvertmean', input_file=self.variable_file.local_file, output_file=temp, - options=[self.variable, 'T', lev_min, lev_max, '-debug']) - Utils.setminmax(temp, '{0}_vert_mean'.format(self.variable)) - self.results.set_local_file(temp, rename_var='{0}_vert_mean'.format(self.variable)) + cdftools.run( + "cdfvertmean", + input_file=self.variable_file.local_file, + output_file=temp, + options=[self.variable, "T", lev_min, lev_max, "-debug"], + ) + Utils.setminmax(temp, "{0}_vert_mean".format(self.variable)) + self.results.set_local_file( + temp, rename_var="{0}_vert_mean".format(self.variable) + ) diff --git a/earthdiagnostics/ocean/verticalmeanmeters.py b/earthdiagnostics/ocean/verticalmeanmeters.py index d1f1cbfd..8c476e4a 100644 --- a/earthdiagnostics/ocean/verticalmeanmeters.py +++ b/earthdiagnostics/ocean/verticalmeanmeters.py @@ -2,8 +2,13 @@ """Averages vertically any given variable""" from earthdiagnostics import cdftools from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticFloatOption, DiagnosticDomainOption, \ - DiagnosticVariableOption, DiagnosticChoiceOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticFloatOption, + DiagnosticDomainOption, + DiagnosticVariableOption, + DiagnosticChoiceOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -33,10 +38,20 @@ class VerticalMeanMeters(Diagnostic): """ - alias = 'vertmeanmeters' + alias = "vertmeanmeters" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, box, grid_point): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + box, + grid_point, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -49,12 +64,26 @@ class VerticalMeanMeters(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.box == other.box and self.variable == other.variable + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.box == other.box + and self.variable == other.variable + ) def __str__(self): - return 'Vertical mean meters Startdate: {0} Member: {1} Chunk: {2} Variable: {3}:{4} ' \ - 'Box: {5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable, self.box) + return ( + "Vertical mean meters Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3}:{4} Box: {5}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variable, + self.box, + ) + ) def __hash__(self): return hash(str(self)) @@ -66,55 +95,94 @@ class VerticalMeanMeters(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: variable, minimum depth (meters), maximum depth (meters) + :param options: variable, minimum depth (m), maximum depth (m) :type options: list[str] :return: """ - options_available = (DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticFloatOption('min_depth', -1), - DiagnosticFloatOption('max_depth', -1), - DiagnosticDomainOption(default_value=ModelingRealms.ocean), - DiagnosticChoiceOption('grid_point', ('T', 'U', 'V'), 'T')) + options_available = ( + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticFloatOption("min_depth", -1), + DiagnosticFloatOption("max_depth", -1), + DiagnosticDomainOption(default_value=ModelingRealms.ocean), + DiagnosticChoiceOption("grid_point", ("T", "U", "V"), "T"), + ) options = cls.process_options(options, options_available) box = Box(True) - if options['min_depth'] >= 0: - box.min_depth = options['min_depth'] - if options['max_depth'] >= 0: - box.max_depth = options['max_depth'] + if options["min_depth"] >= 0: + box.min_depth = options["min_depth"] + if options["max_depth"] >= 0: + box.max_depth = options["max_depth"] job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(VerticalMeanMeters(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], box, options['grid_point'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + VerticalMeanMeters( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + box, + options["grid_point"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(ModelingRealms.ocean, self.variable, self.startdate, self.member, - self.chunk) + self.variable_file = self.request_chunk( + ModelingRealms.ocean, + self.variable, + self.startdate, + self.member, + self.chunk, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.results = self.declare_chunk(self.domain, self.variable + 'vmean', self.startdate, self.member, - self.chunk, box=self.box) + self.results = self.declare_chunk( + self.domain, + self.variable + "vmean", + self.startdate, + self.member, + self.chunk, + box=self.box, + ) def compute(self): """Run the diagnostic""" temp = TempFile.get() handler = Utils.open_cdf(self.variable_file.local_file) if self.box.min_depth is None: - lev_min = handler.variables['lev'][0] + lev_min = handler.variables["lev"][0] else: lev_min = self.box.min_depth if self.box.max_depth is None: - lev_max = handler.variables['lev'][-1] + lev_max = handler.variables["lev"][-1] else: lev_max = self.box.max_depth handler.close() - cdftools.run('cdfvertmean', input_file=self.variable_file.local_file, output_file=temp, - options=[self.variable, self.grid_point, lev_min, lev_max, '-debug']) - Utils.setminmax(temp, '{0}_vert_mean'.format(self.variable)) - self.results.set_local_file(temp, rename_var='{0}_vert_mean'.format(self.variable)) + cdftools.run( + "cdfvertmean", + input_file=self.variable_file.local_file, + output_file=temp, + options=[ + self.variable, + self.grid_point, + lev_min, + lev_max, + "-debug", + ], + ) + Utils.setminmax(temp, "{0}_vert_mean".format(self.variable)) + self.results.set_local_file( + temp, rename_var="{0}_vert_mean".format(self.variable) + ) diff --git a/earthdiagnostics/ocean/zonalmean.py b/earthdiagnostics/ocean/zonalmean.py index ed2c1907..8e00db5f 100644 --- a/earthdiagnostics/ocean/zonalmean.py +++ b/earthdiagnostics/ocean/zonalmean.py @@ -5,18 +5,18 @@ import iris.util import iris.coords import iris.analysis import iris.exceptions -from iris.coord_categorisation import add_categorised_coord -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, DiagnosticBasinListOption, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticBasinListOption, + DiagnosticVariableOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -42,10 +42,20 @@ class ZonalMean(Diagnostic): :type box: Box """ - alias = 'zonmean' + alias = "zonmean" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, basins, 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 @@ -57,18 +67,27 @@ class ZonalMean(Diagnostic): self.declared = {} - self.lat_name = 'lat' - self.lon_name = 'lon' + self.lat_name = "lat" + self.lon_name = "lon" def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.variable == other.variable and self.grid_point == other.grid_point and self.basin == other.basin + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.variable == other.variable + and self.grid_point == other.grid_point + and self.basins == other.basins + ) def __str__(self): - return 'Zonal mean Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Variable: {0.variable} ' \ - 'Grid point: {0.grid_point}'.format(self) + return ( + "Zonal mean Startdate: {0.startdate} Member: {0.member} " + "Chunk: {0.chunk} Variable: {0.variable} " + "Grid point: {0.grid_point}".format(self) + ) def __hash__(self): return hash(str(self)) @@ -87,16 +106,27 @@ class ZonalMean(Diagnostic): options_available = ( DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticOption('grid_point', 'T'), - DiagnosticBasinListOption('basins', Basins().Global), + DiagnosticBasinListOption("basins", "global"), + DiagnosticOption("grid_point", "T"), ) 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['basins'], - options['grid_point'].lower()) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job = ZonalMean( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["basins"], + options["grid_point"].lower(), + ) job_list.append(job) return job_list @@ -110,16 +140,20 @@ class ZonalMean(Diagnostic): def declare_data_generated(self): """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.basins + ModelingRealms.ocean, + self.variable + "zonal", + 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') + 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) + lats = mesh.get_mesh_var("latitude", dtype=np.float32) data = self._load_data() masks = {} self.basins.sort() @@ -131,93 +165,132 @@ class ZonalMean(Diagnostic): def _try_load_cube(self, number): try: - cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}'.format(number, self.grid_point)) + cube = iris.load_cube( + "mesh_hgr.nc", "e{0}{1}".format(number, self.grid_point) + ) except iris.exceptions.ConstraintMismatchError: - cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}_0'.format(number, self.grid_point)) + cube = iris.load_cube( + "mesh_hgr.nc", "e{0}{1}_0".format(number, self.grid_point) + ) cube = iris.util.squeeze(cube) dims = len(cube.shape) try: - cube.coord('i') + cube.coord("i") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), var_name='i'), dims - 1) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 1]), var_name="i" + ), + dims - 1, + ) try: - cube.coord('j') + cube.coord("j") except iris.exceptions.CoordinateNotFoundError: - cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), var_name='j'), dims - 2) + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(cube.shape[dims - 2]), var_name="j" + ), + dims - 2, + ) return cube def _load_data(self): coords = [] handler = Utils.open_cdf(self.variable_file.local_file) for variable in handler.variables: - if variable in ('time', 'lev', 'lat', 'lon', 'latitude', 'longitude', 'leadtime', 'time_centered'): + if variable in ( + "time", + "lev", + "lat", + "lon", + "latitude", + "longitude", + "leadtime", + "time_centered", + ): coords.append(variable) - if variable == 'time_centered': - handler.variables[variable].standard_name = '' + if variable == "time_centered": + handler.variables[variable].standard_name = "" - handler.variables[self.variable].coordinates = ' '.join(coords) + handler.variables[self.variable].coordinates = " ".join(coords) handler.close() data = iris.load_cube(self.variable_file.local_file) return self._rename_depth(data) def _rename_depth(self, data): - for coord_name in ('model_level_number', 'Vertical T levels', 'lev'): + for coord_name in ("model_level_number", "Vertical T levels", "lev"): if data.coords(coord_name): coord = data.coord(coord_name) - coord.standard_name = 'depth' - coord.long_name = 'depth' + coord.standard_name = "depth" + coord.long_name = "depth" break return data def _fix_file_metadata(self): handler = Utils.open_cdf(self.variable_file.local_file) var = handler.variables[self.variable] - coordinates = '' + coordinates = "" has_levels = False for dimension in handler.variables.keys(): - if dimension in ['time', 'lev', 'lat', 'latitude', 'lon', 'longitude', 'i', 'j']: - coordinates += ' {0}'.format(dimension) - if dimension == 'lev': + if dimension in [ + "time", + "lev", + "lat", + "latitude", + "lon", + "longitude", + "i", + "j", + ]: + coordinates += " {0}".format(dimension) + if dimension == "lev": has_levels = True 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') + 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_lat.units = "degrees_north" + var_lat.long_name = "latitude" + var_lat.standard_name = "latitude" + + var_region = handler_temp.createVariable( + "region", "S1", ("region", "region_length") + ) + + try: + data.coord('depth') + except iris.exceptions.CoordinateNotFoundError: + dims = ("region", "time", "lat",) + else: + 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" + dims = ("region", "time", "lev", "lat",) var = handler_temp.createVariable( - '{0}zonal'.format(self.variable), + "{0}zonal".format(self.variable), float, - ('time', 'lev', 'lat', 'region',), + dims, ) - var.units = '{0}'.format(data.units) + var.units = "{0}".format(data.units) for i, basin in enumerate(result): var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) - var[..., i] = result[basin] + var[i, ...] = result[basin] handler_temp.close() - self.declared[final_name].set_local_file(temp, diagnostic=self) + self.zonal_mean.set_local_file(temp, diagnostic=self) diff --git a/earthdiagnostics/publisher.py b/earthdiagnostics/publisher.py index 42a8d3ce..4f808731 100644 --- a/earthdiagnostics/publisher.py +++ b/earthdiagnostics/publisher.py @@ -3,7 +3,7 @@ class Publisher(object): - """Base class to provide functionality to notify updates to other objects""" + """Base class to provide functionality to notify updates to objects""" def __init__(self): self._subscribers = dict() @@ -20,7 +20,7 @@ class Publisher(object): Callback to call """ if callback is None: - callback = getattr(who, 'update') + callback = getattr(who, "update") self._subscribers[who] = callback def unsubscribe(self, who): diff --git a/earthdiagnostics/statistics/climatologicalpercentile.py b/earthdiagnostics/statistics/climatologicalpercentile.py index 865a2b9e..20e97303 100644 --- a/earthdiagnostics/statistics/climatologicalpercentile.py +++ b/earthdiagnostics/statistics/climatologicalpercentile.py @@ -8,8 +8,13 @@ import numpy as np import six from bscearth.utils.log import Log -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticVariableOption, DiagnosticDomainOption, \ - DiagnosticIntOption, DiagnosticListIntOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticVariableOption, + DiagnosticDomainOption, + DiagnosticIntOption, + DiagnosticListIntOption, +) from earthdiagnostics.frequency import Frequencies from earthdiagnostics.utils import TempFile from earthdiagnostics.variable import VariableType @@ -27,13 +32,21 @@ class ClimatologicalPercentile(Diagnostic): :type experiment_config: ExperimentConfig """ - alias = 'climpercent' + alias = "climpercent" "Diagnostic alias for the configuration file" Percentiles = np.array([0.1, 0.25, 0.33, 0.5, 0.66, 0.75, 0.9]) - def __init__(self, data_manager, domain, variable, start_year, end_year, - forecast_month, experiment_config): + def __init__( + self, + data_manager, + domain, + variable, + start_year, + end_year, + forecast_month, + experiment_config, + ): Diagnostic.__init__(self, data_manager) self.variable = variable self.domain = domain @@ -49,13 +62,20 @@ class ClimatologicalPercentile(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.domain == other.domain and self.variable == other.variable and \ - self.start_year == other.start_year and self.end_year == other.end_year and \ - self.forecast_month == other.forecast_month + return ( + self.domain == other.domain + and self.variable == other.variable + and self.start_year == other.start_year + and self.end_year == other.end_year + and self.forecast_month == other.forecast_month + ) def __str__(self): - return 'Climatological percentile Variable: {0.domain}:{0.variable} Period: {0.start_year}-{0.end_year} ' \ - 'Forecast month: {0.forecast_month}'.format(self) + return ( + "Climatological percentile Variable: {0.domain}:{0.variable} " + "Period: {0.start_year}-{0.end_year} " + "Forecast month: {0.forecast_month}".format(self) + ) @classmethod def generate_jobs(cls, diags, options): @@ -64,23 +84,32 @@ class ClimatologicalPercentile(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: domain, variable, percentil number, maximum depth (level) + :param options: domain, variable, percentil number, max depth (level) :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticIntOption('start_year'), - DiagnosticIntOption('end_year'), - DiagnosticListIntOption('forecast_month'), - ) + options_available = ( + DiagnosticDomainOption(), + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticIntOption("start_year"), + DiagnosticIntOption("end_year"), + DiagnosticListIntOption("forecast_month"), + ) options = cls.process_options(options, options_available) job_list = list() - for forecast_month in options['forecast_month']: - job_list.append(ClimatologicalPercentile(diags.data_manager, options['domain'], options['variable'], - options['start_year'], options['end_year'], - forecast_month, diags.config.experiment)) + for forecast_month in options["forecast_month"]: + job_list.append( + ClimatologicalPercentile( + diags.data_manager, + options["domain"], + options["variable"], + options["start_year"], + options["end_year"], + forecast_month, + diags.config.experiment, + ) + ) return job_list def requested_startdates(self): @@ -92,22 +121,41 @@ class ClimatologicalPercentile(Diagnostic): list of str """ - return ['{0}{1:02}01'.format(year, self.forecast_month) for year in range(self.start_year, self.end_year + 1)] + return [ + "{0}{1:02}01".format(year, self.forecast_month) + for year in range(self.start_year, self.end_year + 1) + ] def request_data(self): """Request data required by the diagnostic""" for startdate in self.requested_startdates(): if startdate not in self.leadtime_files: self.leadtime_files[startdate] = {} - Log.debug('Retrieving startdate {0}', startdate) - self.leadtime_files[startdate] = self.request_chunk(self.domain, '{0}_dis'.format(self.variable), startdate, - None, None, vartype=VariableType.STATISTIC) + Log.debug("Retrieving startdate {0}", startdate) + self.leadtime_files[startdate] = self.request_chunk( + self.domain, + "{0}_dis".format(self.variable), + startdate, + None, + None, + vartype=VariableType.STATISTIC, + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - var_name = '{0.variable}prct{0.start_year}{0.forecast_month}-{0.end_year}{0.forecast_month:02d}'.format(self) - self.percentiles_file = self.declare_chunk(self.domain, var_name, None, None, None, - frequency=Frequencies.climatology, vartype=VariableType.STATISTIC) + var_name = ( + f"{self.variable}prct{self.start_year}{self.forecast_month}-" + f"{self.end_year}{self.forecast_month:02d}" + ) + self.percentiles_file = self.declare_chunk( + self.domain, + var_name, + None, + None, + None, + frequency=Frequencies.climatology, + vartype=VariableType.STATISTIC, + ) def compute(self): """Run the diagnostic""" @@ -120,12 +168,12 @@ class ClimatologicalPercentile(Diagnostic): temp = TempFile.get() iris.save(percentile_values.merge_cube(), temp, zlib=True) - self.percentiles_file.set_local_file(temp, rename_var='percent') + self.percentiles_file.set_local_file(temp, rename_var="percent") def _calculate_percentiles(self): - Log.debug('Calculating percentiles') + Log.debug("Calculating percentiles") - bins = self.distribution.coord('bin').points + bins = self.distribution.coord("bin").points def calculate(point_distribution): cs = np.cumsum(point_distribution) @@ -135,25 +183,32 @@ class ClimatologicalPercentile(Diagnostic): return [bins[i] for i in index] results = iris.cube.CubeList() - percentile_coord = iris.coords.DimCoord(ClimatologicalPercentile.Percentiles, long_name='percentile') + percentile_coord = iris.coords.DimCoord( + ClimatologicalPercentile.Percentiles, long_name="percentile" + ) print(self.distribution) - for leadtime_slice in self.distribution.slices_over('leadtime'): - result = iris.cube.Cube(np.apply_along_axis(calculate, 0, leadtime_slice.data), var_name='percent', - units=self.distribution.coord('bin').units) + for leadtime_slice in self.distribution.slices_over("leadtime"): + result = iris.cube.Cube( + np.apply_along_axis(calculate, 0, leadtime_slice.data), + var_name="percent", + units=self.distribution.coord("bin").units, + ) result.add_dim_coord(percentile_coord, 0) - result.add_dim_coord(leadtime_slice.coord('latitude'), 1) - result.add_dim_coord(leadtime_slice.coord('longitude'), 2) - result.add_aux_coord(leadtime_slice.coord('leadtime')) + result.add_dim_coord(leadtime_slice.coord("latitude"), 1) + result.add_dim_coord(leadtime_slice.coord("longitude"), 2) + result.add_aux_coord(leadtime_slice.coord("leadtime")) results.append(result) return results def _get_distribution(self): - for startdate, startdate_file in six.iteritems(self.leadtime_files): - Log.info('Getting data for startdate {0}', startdate) + for startdate, startdate_file in six.iteritems(self.leadtime_files): + Log.info("Getting data for startdate {0}", startdate) data_cube = iris.load_cube(startdate_file.local_file) if self.distribution is None: self.distribution = data_cube else: self.distribution += data_cube - if len(self.distribution.coords('leadtime')) == 0: - self.distribution.add_aux_coord(iris.coords.AuxCoord(1, var_name='leadtime', units='months')) + if len(self.distribution.coords("leadtime")) == 0: + self.distribution.add_aux_coord( + iris.coords.AuxCoord(1, var_name="leadtime", units="months") + ) diff --git a/earthdiagnostics/statistics/daysoverpercentile.py b/earthdiagnostics/statistics/daysoverpercentile.py index 88ed3189..dcacabc1 100644 --- a/earthdiagnostics/statistics/daysoverpercentile.py +++ b/earthdiagnostics/statistics/daysoverpercentile.py @@ -1,5 +1,5 @@ # coding=utf-8 -"""Diagnostic to coompute statistics about days above and below a set ofprecomputed percentiles""" +"""Diagnostic to compute days above and below a set of percentiles""" import os import iris @@ -12,10 +12,17 @@ from bscearth.utils.date import parse_date, add_months from bscearth.utils.log import Log from iris.time import PartialDateTime -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticDomainOption, DiagnosticIntOption, \ - DiagnosticListIntOption, DiagnosticOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticDomainOption, + DiagnosticIntOption, + DiagnosticListIntOption, + DiagnosticOption, +) from earthdiagnostics.frequency import Frequencies -from earthdiagnostics.statistics.climatologicalpercentile import ClimatologicalPercentile +from earthdiagnostics.statistics.climatologicalpercentile import ( + ClimatologicalPercentile, +) from earthdiagnostics.utils import Utils, TempFile from earthdiagnostics.variable import VariableType @@ -32,10 +39,19 @@ class DaysOverPercentile(Diagnostic): :type variable: str """ - alias = 'daysover' + alias = "daysover" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, domain, variable, start_year, end_year, startdate, forecast_month): + def __init__( + self, + data_manager, + domain, + variable, + start_year, + end_year, + startdate, + forecast_month, + ): Diagnostic.__init__(self, data_manager) self.variable = variable self.domain = domain @@ -50,13 +66,20 @@ class DaysOverPercentile(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.domain == other.domain and \ - self.variable == other.variable and self.start_year == other.start_year and \ - self.end_year == other.end_year + return ( + self.startdate == other.startdate + and self.domain == other.domain + and self.variable == other.variable + and self.start_year == other.start_year + and self.end_year == other.end_year + ) def __str__(self): - return 'Days over percentile Startdate: {0.startdate} Variable: {0.domain}:{0.variable} ' \ - 'Climatology: {0.start_year}-{0.end_year}'.format(self) + return ( + f"Days over percentile Startdate: {self.startdate} " + f"Variable: {self.domain}:{self.variable} " + f"Climatology: {self.start_year}-{self.end_year}" + ) @classmethod def generate_jobs(cls, diags, options): @@ -65,104 +88,180 @@ class DaysOverPercentile(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: domain, variable, percentil number, maximum depth (level) + :param options: domain, variable, percentil number, maximum depth (lev) :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticOption('variable'), - DiagnosticIntOption('start_year'), - DiagnosticIntOption('end_year'), - DiagnosticListIntOption('forecast_month'),) + options_available = ( + DiagnosticDomainOption(), + DiagnosticOption("variable"), + DiagnosticIntOption("start_year"), + DiagnosticIntOption("end_year"), + DiagnosticListIntOption("forecast_month"), + ) options = cls.process_options(options, options_available) job_list = list() for startdate in diags.config.experiment.startdates: - for forecast_month in options['forecast_month']: - job_list.append(DaysOverPercentile(diags.data_manager, options['domain'], options['variable'], - options['start_year'], options['end_year'], - startdate, forecast_month)) + for forecast_month in options["forecast_month"]: + job_list.append( + DaysOverPercentile( + diags.data_manager, + options["domain"], + options["variable"], + options["start_year"], + options["end_year"], + startdate, + forecast_month, + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - var_name = '{0.variable}prct{0.start_year}{0.forecast_month}-{0.end_year}{0.forecast_month:02d}'.format(self) - self.percentiles_file = self.request_chunk(self.domain, var_name, None, None, None, - frequency=Frequencies.climatology) - - self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, None, None) + var_name = ( + f"{self.variable}prct{self.start_year}{self.forecast_month}-" + f"{self.end_year}{self.forecast_month:02d}" + ) + self.percentiles_file = self.request_chunk( + self.domain, + var_name, + None, + None, + None, + frequency=Frequencies.climatology, + ) + + self.variable_file = self.request_chunk( + self.domain, self.variable, self.startdate, None, None + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - var_over = self.variable + '_daysover_q{0}_{1.start_year}-{1.end_year}' - var_below = self.variable + '_daysbelow_q{0}_{1.start_year}-{1.end_year}' + var_over = self.variable + "_daysover_q{0}_{1.start_year}-{1.end_year}" + var_below = ( + self.variable + "_daysbelow_q{0}_{1.start_year}-{1.end_year}" + ) self.days_over_file = {} self.days_below_file = {} for perc in ClimatologicalPercentile.Percentiles: - self.days_over_file[perc] = self.declare_chunk(self.domain, var_over.format(int(perc * 100), self), - self.startdate, None, - None, frequency=Frequencies.monthly, - vartype=VariableType.STATISTIC) - - self.days_below_file[perc] = self.declare_chunk(self.domain, var_below.format(int(perc * 100), self), - self.startdate, None, - None, frequency=Frequencies.monthly, - vartype=VariableType.STATISTIC) + self.days_over_file[perc] = self.declare_chunk( + self.domain, + var_over.format(int(perc * 100), self), + self.startdate, + None, + None, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) + + self.days_below_file[perc] = self.declare_chunk( + self.domain, + var_below.format(int(perc * 100), self), + self.startdate, + None, + None, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) def compute(self): """Run the diagnostic""" leadtimes, percentiles, var = self._load_data() try: - realization_coord = var.coord('realization') + realization_coord = var.coord("realization") except iris.exceptions.CoordinateNotFoundError: realization_coord = None - self.lat_coord = var.coord('latitude') - self.lon_coord = var.coord('longitude') - results_over = {perc: iris.cube.CubeList() for perc in ClimatologicalPercentile.Percentiles} - results_below = {perc: iris.cube.CubeList() for perc in ClimatologicalPercentile.Percentiles} - - var_daysover = 'days_over' - var_days_below = 'days_below' - long_name_days_over = 'Proportion of days over a given percentile for {0.start_year}-{0.end_year} ' \ - 'climatology'.format(self) - long_name_days_below = 'Proportion of days below a given percentile for {0.start_year}-{0.end_year} ' \ - 'climatology'.format(self) + self.lat_coord = var.coord("latitude") + self.lon_coord = var.coord("longitude") + results_over = { + perc: iris.cube.CubeList() + for perc in ClimatologicalPercentile.Percentiles + } + results_below = { + perc: iris.cube.CubeList() + for perc in ClimatologicalPercentile.Percentiles + } + + var_daysover = "days_over" + var_days_below = "days_below" + long_name_days_over = ( + "Proportion of days over a given percentile for " + f"{self.start_year}-{self.end_year} climatology" + ) + long_name_days_below = ( + "Proportion of days below a given percentile for " + f"{self.start_year}-{self.end_year} climatology" + ) for leadtime in leadtimes.keys(): - Log.debug('Computing startdate {0} leadtime {1}', self.startdate, leadtime) + Log.debug( + "Computing startdate {0} leadtime {1}", + self.startdate, + leadtime, + ) leadtime_slice = var.extract(iris.Constraint(leadtime=leadtime)) - if len(percentiles.coords('leadtime')) > 0: - percentiles_leadtime = percentiles.extract(iris.Constraint(leadtime=leadtime)) + if len(percentiles.coords("leadtime")) > 0: + percentiles_leadtime = percentiles.extract( + iris.Constraint(leadtime=leadtime) + ) else: percentiles_leadtime = percentiles - time_coord = iris.coords.AuxCoord.from_coord(leadtime_slice.coord('time')) + time_coord = iris.coords.AuxCoord.from_coord( + leadtime_slice.coord("time") + ) first_time = time_coord.points[0] last_time = time_coord.points[-1] - timesteps = leadtime_slice.coord('time').shape[0] - time_coord = time_coord.copy(first_time + (last_time - first_time) / 2, (first_time, last_time)) - for percentile_slice in percentiles_leadtime.slices_over('percentile'): - percentile = percentile_slice.coord('percentile').points[0] + timesteps = leadtime_slice.coord("time").shape[0] + time_coord = time_coord.copy( + first_time + (last_time - first_time) / 2, + (first_time, last_time), + ) + for percentile_slice in percentiles_leadtime.slices_over( + "percentile" + ): + percentile = percentile_slice.coord("percentile").points[0] # noinspection PyTypeChecker - days_over = np.sum(leadtime_slice.data > percentile_slice.data, 0) / float(timesteps) - result = self._create_results_cube(days_over, percentile, realization_coord, - time_coord, var_daysover, long_name_days_over) + days_over = np.sum( + leadtime_slice.data > percentile_slice.data, 0 + ) / float(timesteps) + result = self._create_results_cube( + days_over, + percentile, + realization_coord, + time_coord, + var_daysover, + long_name_days_over, + ) results_over[percentile].append(result) # noinspection PyTypeChecker - days_below = np.sum(leadtime_slice.data < percentile_slice.data, 0) / float(timesteps) - result = self._create_results_cube(days_below, percentile, realization_coord, - time_coord, var_days_below, long_name_days_below) + days_below = np.sum( + leadtime_slice.data < percentile_slice.data, 0 + ) / float(timesteps) + result = self._create_results_cube( + days_below, + percentile, + realization_coord, + time_coord, + var_days_below, + long_name_days_below, + ) results_below[percentile].append(result) - Log.debug('Saving percentiles startdate {0}', self.startdate) + Log.debug("Saving percentiles startdate {0}", self.startdate) for perc in ClimatologicalPercentile.Percentiles: - self.days_over_file[perc].set_local_file(self._save_to_file(perc, results_over, var_daysover), - rename_var=var_daysover) - self.days_below_file[perc].set_local_file(self._save_to_file(perc, results_below, var_days_below), - rename_var=var_days_below) + self.days_over_file[perc].set_local_file( + self._save_to_file(perc, results_over, var_daysover), + rename_var=var_daysover, + ) + self.days_below_file[perc].set_local_file( + self._save_to_file(perc, results_below, var_days_below), + rename_var=var_days_below, + ) del self.days_over_file del self.days_below_file @@ -173,13 +272,17 @@ class DaysOverPercentile(Diagnostic): percentiles = iris.load_cube(self.percentiles_file.local_file) handler = Utils.open_cdf(self.variable_file.local_file) - if 'realization' in handler.variables: - handler.variables[self.variable].coordinates = 'realization' + if "realization" in handler.variables: + handler.variables[self.variable].coordinates = "realization" handler.close() var = iris.load_cube(self.variable_file.local_file) date = parse_date(self.startdate) - lead_date = add_months(date, 1, self.data_manager.config.experiment.calendar) - leadtimes = {1: PartialDateTime(lead_date.year, lead_date.month, lead_date.day)} + lead_date = add_months( + date, 1, self.data_manager.config.experiment.calendar + ) + leadtimes = { + 1: PartialDateTime(lead_date.year, lead_date.month, lead_date.day) + } def assign_leadtime(coord, x): leadtime_month = 1 @@ -189,47 +292,71 @@ class DaysOverPercentile(Diagnostic): try: partial_date = leadtimes[leadtime_month] except KeyError: - new_date = add_months(date, leadtime_month, self.data_manager.config.experiment.calendar) - partial_date = PartialDateTime(new_date.year, new_date.month, new_date.day) + new_date = add_months( + date, + leadtime_month, + self.data_manager.config.experiment.calendar, + ) + partial_date = PartialDateTime( + new_date.year, new_date.month, new_date.day + ) leadtimes[leadtime_month] = partial_date return leadtime_month - iris.coord_categorisation.add_categorised_coord(var, 'leadtime', 'time', assign_leadtime) - iris.coord_categorisation.add_year(var, 'time') - iris.coord_categorisation.add_day_of_year(var, 'time') + iris.coord_categorisation.add_categorised_coord( + var, "leadtime", "time", assign_leadtime + ) + iris.coord_categorisation.add_year(var, "time") + iris.coord_categorisation.add_day_of_year(var, "time") return leadtimes, percentiles, var @staticmethod def _save_to_file(perc, results_over, var_daysover): temp = TempFile.get() - iris.save(results_over[perc].merge_cube(), temp, zlib=True, unlimited_dimensions=['time']) - Utils.rename_variables(temp, {'dim2': 'ensemble', 'dim1': 'ensemble'}, must_exist=False) + iris.save( + results_over[perc].merge_cube(), + temp, + zlib=True, + unlimited_dimensions=["time"], + ) + Utils.rename_variables( + temp, {"dim2": "ensemble", "dim1": "ensemble"}, must_exist=False + ) handler = Utils.open_cdf(temp) - if 'time' not in handler.dimensions: + if "time" not in handler.dimensions: new_file = TempFile.get() - new_handler = Utils.open_cdf(new_file, 'w') + new_handler = Utils.open_cdf(new_file, "w") - new_handler.createDimension('time', 1) + new_handler.createDimension("time", 1) for dimension in handler.dimensions: Utils.copy_dimension(handler, new_handler, dimension) for variable in handler.variables.keys(): - if variable in (var_daysover, 'time', 'time_bnds'): + if variable in (var_daysover, "time", "time_bnds"): continue Utils.copy_variable(handler, new_handler, variable) old_var = handler.variables[var_daysover] - new_var = new_handler.createVariable(var_daysover, old_var.dtype, ('time',) + old_var.dimensions, - zlib=True, fill_value=1.0e20) + new_var = new_handler.createVariable( + var_daysover, + old_var.dtype, + ("time",) + old_var.dimensions, + zlib=True, + fill_value=1.0e20, + ) Utils.copy_attributes(new_var, old_var) new_var[0, :] = old_var[:] - old_var = handler.variables['time'] - new_var = new_handler.createVariable('time', old_var.dtype, ('time',)) + old_var = handler.variables["time"] + new_var = new_handler.createVariable( + "time", old_var.dtype, ("time",) + ) Utils.copy_attributes(new_var, old_var) new_var[0] = old_var[0] - old_var = handler.variables['time_bnds'] - new_var = new_handler.createVariable('time_bnds', old_var.dtype, ('time',) + old_var.dimensions) + old_var = handler.variables["time_bnds"] + new_var = new_handler.createVariable( + "time_bnds", old_var.dtype, ("time",) + old_var.dimensions + ) Utils.copy_attributes(new_var, old_var) new_var[0, :] = old_var[:] @@ -239,9 +366,21 @@ class DaysOverPercentile(Diagnostic): handler.close() return temp - def _create_results_cube(self, days_over, percentile, realization_coord, time_coord, - var_name, long_name): - result = iris.cube.Cube(days_over.astype(np.float32), var_name=var_name, long_name=long_name, units=1.0) + def _create_results_cube( + self, + days_over, + percentile, + realization_coord, + time_coord, + var_name, + long_name, + ): + result = iris.cube.Cube( + days_over.astype(np.float32), + var_name=var_name, + long_name=long_name, + units=1.0, + ) if realization_coord is not None: result.add_aux_coord(realization_coord, 0) result.add_dim_coord(self.lat_coord, 1) @@ -249,6 +388,8 @@ class DaysOverPercentile(Diagnostic): else: result.add_dim_coord(self.lat_coord, 0) result.add_dim_coord(self.lon_coord, 1) - result.add_aux_coord(iris.coords.AuxCoord(percentile, long_name='percentile')) + result.add_aux_coord( + iris.coords.AuxCoord(percentile, long_name="percentile") + ) result.add_aux_coord(time_coord) return result diff --git a/earthdiagnostics/statistics/discretize.py b/earthdiagnostics/statistics/discretize.py index 29e119d4..f379e979 100644 --- a/earthdiagnostics/statistics/discretize.py +++ b/earthdiagnostics/statistics/discretize.py @@ -15,8 +15,13 @@ from bscearth.utils.log import Log from iris.cube import Cube from iris.time import PartialDateTime -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticVariableOption, DiagnosticDomainOption, \ - DiagnosticIntOption, DiagnosticFloatOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticVariableOption, + DiagnosticDomainOption, + DiagnosticIntOption, + DiagnosticFloatOption, +) from earthdiagnostics.utils import Utils, TempFile from earthdiagnostics.variable import VariableType @@ -31,12 +36,21 @@ class Discretize(Diagnostic): :type variable: str """ - alias = 'discretize' + alias = "discretize" "Diagnostic alias for the configuration file" Percentiles = np.array([0.1, 0.25, 0.33, 0.5, 0.66, 0.75, 0.9]) - def __init__(self, data_manager, startdate, domain, variable, num_bins, min_value, max_value): + def __init__( + self, + data_manager, + startdate, + domain, + variable, + num_bins, + min_value, + max_value, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate @@ -47,7 +61,9 @@ class Discretize(Diagnostic): self.num_bins = num_bins self._bins = None self.discretized_data = None - self.cmor_var = data_manager.variable_list.get_variable(variable, silent=True) + self.cmor_var = data_manager.variable_list.get_variable( + variable, silent=True + ) if not math.isnan(min_value): self.min_value = min_value @@ -72,7 +88,11 @@ class Discretize(Diagnostic): self.process = psutil.Process() def _print_memory_used(self): - Log.debug('Memory: {0:.2f} GB'.format(self.process.memory_info().rss / 1024.0**3)) + Log.debug( + "Memory: {0:.2f} GB".format( + self.process.memory_info().rss / 1024.0 ** 3 + ) + ) @property def bins(self): @@ -94,16 +114,26 @@ class Discretize(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.domain == other.domain and self.variable == other.variable and self.num_bins == other.num_bins and \ - self.min_value == other.min_value and self.max_value == other.max_value and \ - self.startdate == other.startdate + return ( + self.domain == other.domain + and self.variable == other.variable + and self.num_bins == other.num_bins + and self.min_value == other.min_value + and self.max_value == other.max_value + and self.startdate == other.startdate + ) def __hash__(self): return hash(str(self)) def __str__(self): - return 'Discretizing variable: {0.domain}:{0.variable} Startdate: {0.startdate} ' \ - 'Bins: {0.num_bins} Range: [{0.min_value}, {0.max_value}]'.format(self) + return ( + "Discretizing variable: {0.domain}:{0.variable} " + "Startdate: {0.startdate} " + "Bins: {0.num_bins} Range: [{0.min_value}, {0.max_value}]".format( + self + ) + ) @classmethod def generate_jobs(cls, diags, options): @@ -112,33 +142,51 @@ class Discretize(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: domain, variable, percentil number, maximum depth (level) + :param options: domain, variable, percentil number, maximum depth (lev) :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticIntOption('bins', 2000), - DiagnosticFloatOption('min_value', float('nan')), - DiagnosticFloatOption('max_value', float('nan')), - ) + options_available = ( + DiagnosticDomainOption(), + DiagnosticVariableOption(diags.data_manager.config.var_manager), + DiagnosticIntOption("bins", 2000), + DiagnosticFloatOption("min_value", float("nan")), + DiagnosticFloatOption("max_value", float("nan")), + ) options = cls.process_options(options, options_available) job_list = list() for startdate in diags.config.experiment.startdates: - job_list.append(Discretize(diags.data_manager, startdate, options['domain'], options['variable'], - options['bins'], options['min_value'], options['max_value'])) + job_list.append( + Discretize( + diags.data_manager, + startdate, + options["domain"], + options["variable"], + options["bins"], + options["min_value"], + options["max_value"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.original_data = self.request_chunk(self.domain, self.variable, self.startdate, None, None) + self.original_data = self.request_chunk( + self.domain, self.variable, self.startdate, None, None + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - var_name = '{0.variable}_dis'.format(self) - self.discretized_data = self.declare_chunk(self.domain, var_name, self.startdate, None, None, - vartype=VariableType.STATISTIC) + var_name = "{0.variable}_dis".format(self) + self.discretized_data = self.declare_chunk( + self.domain, + var_name, + self.startdate, + None, + None, + vartype=VariableType.STATISTIC, + ) def compute(self): """Run the diagnostic""" @@ -148,7 +196,7 @@ class Discretize(Diagnostic): self._print_memory_used() self._get_value_interval() self._print_memory_used() - Log.info('Range: [{0}, {1}]', self.min_value, self.max_value) + Log.info("Range: [{0}, {1}]", self.min_value, self.max_value) self._get_distribution() self._print_memory_used() self._save_results() @@ -160,14 +208,18 @@ class Discretize(Diagnostic): def _load_cube(self): handler = Utils.open_cdf(self.original_data.local_file) - if 'realization' in handler.variables: - handler.variables[self.variable].coordinates = 'realization' + if "realization" in handler.variables: + handler.variables[self.variable].coordinates = "realization" handler.close() data_cube = iris.load_cube(self.original_data.local_file) date = parse_date(self.startdate) - lead_date = add_months(date, 1, self.data_manager.config.experiment.calendar) - leadtimes = {1: PartialDateTime(lead_date.year, lead_date.month, lead_date.day)} + lead_date = add_months( + date, 1, self.data_manager.config.experiment.calendar + ) + leadtimes = { + 1: PartialDateTime(lead_date.year, lead_date.month, lead_date.day) + } def assign_leadtime(coord, x): leadtime_month = 1 @@ -177,16 +229,24 @@ class Discretize(Diagnostic): try: partial_date = leadtimes[leadtime_month] except KeyError: - new_date = add_months(date, leadtime_month, self.data_manager.config.experiment.calendar) - partial_date = PartialDateTime(new_date.year, new_date.month, new_date.day) + new_date = add_months( + date, + leadtime_month, + self.data_manager.config.experiment.calendar, + ) + partial_date = PartialDateTime( + new_date.year, new_date.month, new_date.day + ) leadtimes[leadtime_month] = partial_date return leadtime_month - iris.coord_categorisation.add_categorised_coord(data_cube, 'leadtime', 'time', assign_leadtime) + iris.coord_categorisation.add_categorised_coord( + data_cube, "leadtime", "time", assign_leadtime + ) self.data_cube = data_cube def _save_results(self): - Log.debug('Saving results...') + Log.debug("Saving results...") bins = np.zeros(self.num_bins) bins_bounds = np.zeros((self.num_bins, 2)) @@ -196,7 +256,12 @@ class Discretize(Diagnostic): bins_bounds[x, 0] = self.bins[x] bins_bounds[x, 1] = self.bins[x + 1] - bins_coord = iris.coords.DimCoord(bins, var_name='bin', units=self.data_cube.units, bounds=bins_bounds) + bins_coord = iris.coords.DimCoord( + bins, + var_name="bin", + units=self.data_cube.units, + bounds=bins_bounds, + ) cubes = iris.cube.CubeList() @@ -204,50 +269,81 @@ class Discretize(Diagnostic): date = add_days(date, 14, self.data_manager.config.experiment.calendar) for leadtime, distribution in six.iteritems(self.distribution): - leadtime_cube = Cube(distribution.astype(np.uint32), var_name=self.data_cube.var_name, - standard_name=self.data_cube.standard_name, units='1') + leadtime_cube = Cube( + distribution.astype(np.uint32), + var_name=self.data_cube.var_name, + standard_name=self.data_cube.standard_name, + units="1", + ) leadtime_cube.add_dim_coord(bins_coord, 0) - leadtime_cube.add_dim_coord(self.data_cube.coord('latitude'), 1) - leadtime_cube.add_dim_coord(self.data_cube.coord('longitude'), 2) - leadtime_cube.add_aux_coord(iris.coords.AuxCoord(leadtime, - var_name='leadtime', - units='months')) - lead_date = add_months(date, leadtime - 1, self.data_manager.config.experiment.calendar) - leadtime_cube.add_aux_coord(iris.coords.AuxCoord(cf_units.date2num(lead_date, - unit='days since 1950-01-01', - calendar="standard"), - var_name='time', - units='days since 1950-01-01')) + leadtime_cube.add_dim_coord(self.data_cube.coord("latitude"), 1) + leadtime_cube.add_dim_coord(self.data_cube.coord("longitude"), 2) + leadtime_cube.add_aux_coord( + iris.coords.AuxCoord( + leadtime, var_name="leadtime", units="months" + ) + ) + lead_date = add_months( + date, + leadtime - 1, + self.data_manager.config.experiment.calendar, + ) + leadtime_cube.add_aux_coord( + iris.coords.AuxCoord( + cf_units.date2num( + lead_date, + unit="days since 1950-01-01", + calendar="standard", + ), + var_name="time", + units="days since 1950-01-01", + ) + ) cubes.append(leadtime_cube) temp = TempFile.get() iris.save(cubes.merge_cube(), temp, zlib=True) - self.discretized_data.set_local_file(temp, rename_var=self.data_cube.var_name) + self.discretized_data.set_local_file( + temp, rename_var=self.data_cube.var_name + ) def _get_distribution(self): self.distribution = {} - Log.debug('Discretizing...') - for leadtime in set(self.data_cube.coord('leadtime').points): - Log.debug('Discretizing leadtime {0}', leadtime) - leadtime_cube = self.data_cube.extract(iris.Constraint(leadtime=leadtime)) - if 'realization' in leadtime_cube.coords(): - for realization_cube in self.data_cube.slices_over('realization'): - Log.debug('Discretizing realization {0}', realization_cube.coord('realization').points[0]) + Log.debug("Discretizing...") + for leadtime in set(self.data_cube.coord("leadtime").points): + Log.debug("Discretizing leadtime {0}", leadtime) + leadtime_cube = self.data_cube.extract( + iris.Constraint(leadtime=leadtime) + ) + if "realization" in leadtime_cube.coords(): + for realization_cube in self.data_cube.slices_over( + "realization" + ): + Log.debug( + "Discretizing realization {0}", + realization_cube.coord("realization").points[0], + ) self._print_memory_used() if leadtime not in self.distribution: - self.distribution[leadtime] = self._calculate_distribution(realization_cube) + self.distribution[ + leadtime + ] = self._calculate_distribution(realization_cube) else: - self.distribution[leadtime] += self._calculate_distribution(realization_cube) + self.distribution[ + leadtime + ] += self._calculate_distribution(realization_cube) else: self._print_memory_used() - self.distribution[leadtime] = self._calculate_distribution(leadtime_cube) + self.distribution[leadtime] = self._calculate_distribution( + leadtime_cube + ) # noinspection PyTypeChecker def _get_value_interval(self): if self.check_min_value or self.check_max_value: - Log.debug('Calculating max and min values...') - for time_slice in self.data_cube.slices_over('time'): + Log.debug("Calculating max and min values...") + for time_slice in self.data_cube.slices_over("time"): if self.check_min_value: file_min = np.amin(time_slice.data) if self.min_value is None: @@ -260,8 +356,11 @@ class Discretize(Diagnostic): def _calculate_distribution(self, data_cube): def calculate_histogram(time_series): - histogram, self.bins = np.histogram(time_series, bins=self.bins, - range=(self.min_value, self.max_value)) + histogram, self.bins = np.histogram( + time_series, + bins=self.bins, + range=(self.min_value, self.max_value), + ) return histogram return np.apply_along_axis(calculate_histogram, 0, data_cube.data) diff --git a/earthdiagnostics/statistics/monthlypercentile.py b/earthdiagnostics/statistics/monthlypercentile.py index 26a10174..b46b09a0 100644 --- a/earthdiagnostics/statistics/monthlypercentile.py +++ b/earthdiagnostics/statistics/monthlypercentile.py @@ -4,7 +4,12 @@ from calendar import monthrange from bscearth.utils.log import Log -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticDomainOption, DiagnosticListIntOption +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticDomainOption, + DiagnosticListIntOption, +) from earthdiagnostics.frequency import Frequencies from earthdiagnostics.utils import Utils, TempFile from earthdiagnostics.variable import VariableType @@ -26,10 +31,19 @@ class MonthlyPercentile(Diagnostic): :type variable: str """ - alias = 'monpercent' + alias = "monpercent" "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, domain, variable, percentiles): + def __init__( + self, + data_manager, + startdate, + member, + chunk, + domain, + variable, + percentiles, + ): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member @@ -41,14 +55,27 @@ class MonthlyPercentile(Diagnostic): def __eq__(self, other): if self._different_type(other): return False - return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ - self.domain == other.domain and self.variable == other.variable and self.percentiles == other.percentiles + return ( + self.startdate == other.startdate + and self.member == other.member + and self.chunk == other.chunk + and self.domain == other.domain + and self.variable == other.variable + and self.percentiles == other.percentiles + ) def __str__(self): - return 'Monthly percentile Startdate: {0} Member: {1} Chunk: {2} ' \ - 'Variable: {3}:{4} Percentiles: {5}'.format(self.startdate, self.member, self.chunk, - self.domain, self.variable, - ', '.join(str(i) for i in self.percentiles)) + return ( + "Monthly percentile Startdate: {0} Member: {1} Chunk: {2} " + "Variable: {3}:{4} Percentiles: {5}".format( + self.startdate, + self.member, + self.chunk, + self.domain, + self.variable, + ", ".join(str(i) for i in self.percentiles), + ) + ) @classmethod def generate_jobs(cls, diags, options): @@ -57,40 +84,83 @@ class MonthlyPercentile(Diagnostic): :param diags: Diagnostics manager class :type diags: Diags - :param options: domain, variable, percentil number, maximum depth (level) + :param options: domain, variable, percentil number, maximum depth (lev) :type options: list[str] :return: """ - options_available = (DiagnosticDomainOption(), - DiagnosticOption('variable'), - DiagnosticListIntOption('percentiles', [], 0, 100)) + options_available = ( + DiagnosticDomainOption(), + DiagnosticOption("variable"), + DiagnosticListIntOption("percentiles", [], 0, 100), + ) options = cls.process_options(options, options_available) job_list = list() - for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(MonthlyPercentile(diags.data_manager, startdate, member, chunk, - options['domain'], options['variable'], options['percentiles'])) + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + MonthlyPercentile( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + options["variable"], + options["percentiles"], + ) + ) return job_list def request_data(self): """Request data required by the diagnostic""" - self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk) + self.variable_file = self.request_chunk( + self.domain, self.variable, self.startdate, self.member, self.chunk + ) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.max_file = self.declare_chunk(self.domain, self.variable_max, self.startdate, self.member, self.chunk, - frequency=Frequencies.monthly, vartype=VariableType.STATISTIC) - self.min_file = self.declare_chunk(self.domain, self.variable_min, self.startdate, self.member, self.chunk, - frequency=Frequencies.monthly, vartype=VariableType.STATISTIC) + self.max_file = self.declare_chunk( + self.domain, + self.variable_max, + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) + self.min_file = self.declare_chunk( + self.domain, + self.variable_min, + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) self.percentile_file = {} for percentile in self.percentiles: - self.percentile_file[percentile] = self.declare_chunk(self.domain, self.percentile(percentile), - self.startdate, self.member, self.chunk, - frequency=Frequencies.monthly, - vartype=VariableType.STATISTIC) + self.percentile_file[percentile] = self.declare_chunk( + self.domain, + self.percentile(percentile), + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) - self.declare_chunk(self.domain, '{0}_q{1}'.format(self.variable, percentile), self.startdate, - self.member, self.chunk, frequency=Frequencies.monthly, vartype=VariableType.STATISTIC) + self.declare_chunk( + self.domain, + "{0}_q{1}".format(self.variable, percentile), + self.startdate, + self.member, + self.chunk, + frequency=Frequencies.monthly, + vartype=VariableType.STATISTIC, + ) @property def variable_max(self): @@ -102,7 +172,7 @@ class MonthlyPercentile(Diagnostic): str """ - return '{0}max'.format(self.variable) + return "{0}max".format(self.variable) @property def variable_min(self): @@ -114,7 +184,7 @@ class MonthlyPercentile(Diagnostic): str """ - return '{0}min'.format(self.variable) + return "{0}min".format(self.variable) def percentile(self, percentile): """ @@ -129,7 +199,7 @@ class MonthlyPercentile(Diagnostic): str """ - return '{0}_q{1}'.format(self.variable, percentile) + return "{0}_q{1}".format(self.variable, percentile) def compute(self): """Run the diagnostic""" @@ -142,53 +212,74 @@ class MonthlyPercentile(Diagnostic): while datetimes[start_index].day != 1 and start_index < handler.size: start_index += 1 if start_index == datetimes.size: - raise Exception('No complete month for diagnostic {0}'.format(self)) + raise Exception( + "No complete month for diagnostic {0}".format(self) + ) end_index = datetimes.size - 1 - while datetimes[end_index].day != monthrange(datetimes[end_index].year, datetimes[end_index].month)[1] \ - and end_index >= 0: + while ( + datetimes[end_index].day + != monthrange( + datetimes[end_index].year, datetimes[end_index].month + )[1] + and end_index >= 0 + ): end_index -= 1 if end_index < 0: - raise Exception('No complete month for diagnostic {0}'.format(self)) + raise Exception( + "No complete month for diagnostic {0}".format(self) + ) if start_index != 0 or end_index != datetimes.size - 1: - start_date = '{0.year}-{0.month}-{0.day}'.format(datetimes[start_index]) - end_date = '{0.year}-{0.month}-{0.day}'.format(datetimes[end_index]) + start_date = "{0.year}-{0.month}-{0.day}".format( + datetimes[start_index] + ) + end_date = "{0.year}-{0.month}-{0.day}".format( + datetimes[end_index] + ) Utils.cdo().seldate( - '{0},{1}'.format(start_date, end_date), + "{0},{1}".format(start_date, end_date), input=self.variable_file.local_file, - output=temp + output=temp, ) - Utils.rename_variable(temp, 'lev', 'ensemble', False) + Utils.rename_variable(temp, "lev", "ensemble", False) else: Utils.copy_file(self.variable_file.local_file, temp) - Log.debug('Computing minimum') + Log.debug("Computing minimum") monmin_file = TempFile.get() Utils.cdo().monmin(input=temp, output=monmin_file) - Log.debug('Computing maximum') + Log.debug("Computing maximum") monmax_file = TempFile.get() Utils.cdo().monmax(input=temp, output=monmax_file) for percentile in self.percentiles: - Log.debug('Computing percentile {0}', percentile) - Utils.cdo().monpctl(str(percentile), input=[temp, monmin_file, monmax_file], output=temp) - Utils.rename_variable(temp, 'lev', 'ensemble', False) + Log.debug("Computing percentile {0}", percentile) + Utils.cdo().monpctl( + str(percentile), + input=[temp, monmin_file, monmax_file], + output=temp, + ) + Utils.rename_variable(temp, "lev", "ensemble", False) handler = Utils.open_cdf(monmax_file) - handler.variables[self.variable].long_name += ' {0} Percentile'.format(percentile) + handler.variables[ + self.variable + ].long_name += " {0} Percentile".format(percentile) handler.close() - self.percentiles[percentile].set_local_file(temp, rename_var=self.variable) + self.percentiles[percentile].set_local_file( + temp, rename_var=self.variable + ) - Utils.rename_variable(monmax_file, 'lev', 'ensemble', False) + Utils.rename_variable(monmax_file, "lev", "ensemble", False) handler = Utils.open_cdf(monmax_file) - handler.variables[self.variable].long_name += ' Monthly Maximum' + handler.variables[self.variable].long_name += " Monthly Maximum" handler.close() self.max_file.set_local_file(monmax_file, rename_var=self.variable) - Utils.rename_variable(monmin_file, 'lev', 'ensemble', False) + Utils.rename_variable(monmin_file, "lev", "ensemble", False) handler = Utils.open_cdf(monmin_file) - handler.variables[self.variable].long_name += ' Monthly Minimum' + handler.variables[self.variable].long_name += " Monthly Minimum" handler.close() self.min_file.set_local_file(monmin_file, rename_var=self.variable) diff --git a/earthdiagnostics/threddsmanager.py b/earthdiagnostics/threddsmanager.py index 5869f9b1..aeec37cb 100644 --- a/earthdiagnostics/threddsmanager.py +++ b/earthdiagnostics/threddsmanager.py @@ -32,23 +32,45 @@ class THREDDSManager(DataManager): def __init__(self, config): super(THREDDSManager, self).__init__(config) self.server_url = config.thredds.server_url - data_folders = self.config.data_dir.split(':') + data_folders = self.config.data_dir.split(":") self.config.data_dir = None for data_folder in data_folders: - if os.path.isdir(os.path.join(data_folder, self.config.data_type, self.experiment.institute.lower(), - self.experiment.model.lower())): + if os.path.isdir( + os.path.join( + data_folder, + self.config.data_type, + self.experiment.institute.lower(), + self.experiment.model.lower(), + ) + ): self.config.data_dir = data_folder break if not self.config.data_dir: - raise Exception('Can not find model data') + raise Exception("Can not find model data") - if self.config.data_type in ('obs', 'recon') and self.experiment.chunk_size != 1: - raise Exception('For obs and recon data chunk_size must be always 1') + if ( + self.config.data_type in ("obs", "recon") + and self.experiment.chunk_size != 1 + ): + raise Exception( + "For obs and recon data chunk_size must be always 1" + ) # noinspection PyUnusedLocal - def file_exists(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - vartype=VariableType.MEAN, possible_versions=None): + def file_exists( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + possible_versions=None, + ): """ Check if a file exists in the storage @@ -71,17 +93,32 @@ class THREDDSManager(DataManager): THREDDSSubset """ - aggregation_path = self.get_var_url(var, startdate, frequency, box, vartype) - - start_chunk = chunk_start_date(parse_date(startdate), chunk, self.experiment.chunk_size, 'month', - self.experiment.calendar) - end_chunk = chunk_end_date(start_chunk, self.experiment.chunk_size, 'month', self.experiment.calendar) - - thredds_subset = THREDDSSubset(aggregation_path, "", var, start_chunk, end_chunk) + aggregation_path = self.get_var_url( + var, startdate, frequency, box, vartype + ) + + start_chunk = chunk_start_date( + parse_date(startdate), + chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) + end_chunk = chunk_end_date( + start_chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) + + thredds_subset = THREDDSSubset( + aggregation_path, "", var, start_chunk, end_chunk + ) return thredds_subset - def get_file_path(self, startdate, domain, var, frequency, vartype, - box=None, grid=None): + def get_file_path( + self, startdate, domain, var, frequency, vartype, box=None, grid=None + ): """ Return the path to a concrete file @@ -104,7 +141,9 @@ class THREDDSManager(DataManager): frequency = self.config.frequency var = self._get_final_var_name(box, var) - folder_path = self._get_folder_path(frequency, domain, var, grid, vartype) + folder_path = self._get_folder_path( + frequency, domain, var, grid, vartype + ) file_name = self._get_file_name(var, startdate) filepath = os.path.join(folder_path, file_name) @@ -112,21 +151,38 @@ class THREDDSManager(DataManager): def _get_folder_path(self, frequency, domain, variable, grid, vartype): - if self.config.data_type == 'exp': - var_folder = domain.get_varfolder(variable, self.config.experiment.ocean_timestep, - self.config.experiment.atmos_timestep, grid=grid) + if self.config.data_type == "exp": + var_folder = domain.get_varfolder( + variable, + self.config.experiment.ocean_timestep, + self.config.experiment.atmos_timestep, + grid=grid, + ) else: var_folder = variable - folder_path = os.path.join(self.config.data_dir, self.config.data_type, - self.experiment.institute.lower(), - self.experiment.model.lower(), - frequency.folder_name(vartype), - var_folder) + folder_path = os.path.join( + self.config.data_dir, + self.config.data_type, + self.experiment.institute.lower(), + self.experiment.model.lower(), + frequency.folder_name(vartype), + var_folder, + ) return folder_path # noinspection PyUnusedLocal - def get_year(self, domain, var, startdate, member, year, grid=None, box=None, vartype=VariableType.MEAN): + def get_year( + self, + domain, + var, + startdate, + member, + year, + grid=None, + box=None, + vartype=VariableType.MEAN, + ): """ Ge a file containing all the data for one year for one variable @@ -149,7 +205,13 @@ class THREDDSManager(DataManager): :return: """ aggregation_path = self.get_var_url(var, startdate, None, box, vartype) - thredds_subset = THREDDSSubset(aggregation_path, "", var, datetime(year, 1, 1), datetime(year + 1, 1, 1)) + thredds_subset = THREDDSSubset( + aggregation_path, + "", + var, + datetime(year, 1, 1), + datetime(year + 1, 1, 1), + ) return thredds_subset.download() def get_var_url(self, var, startdate, frequency, box, vartype): @@ -171,26 +233,45 @@ class THREDDSManager(DataManager): if frequency is None: frequency = self.config.frequency var = self._get_final_var_name(box, var) - full_path = os.path.join(self.server_url, 'dodsC', self.config.data_type, self.experiment.institute, - self.experiment.model, frequency.folder_name(vartype)) - if self.config.data_type == 'exp': - full_path = os.path.join(full_path, var, self._get_file_name(var, startdate)) + full_path = os.path.join( + self.server_url, + "dodsC", + self.config.data_type, + self.experiment.institute, + self.experiment.model, + frequency.folder_name(vartype), + ) + if self.config.data_type == "exp": + full_path = os.path.join( + full_path, var, self._get_file_name(var, startdate) + ) else: full_path = os.path.join(full_path, self._get_file_name(var, None)) return full_path def _get_file_name(self, var, startdate): if startdate: - if self.config.data_type != 'exp': + if self.config.data_type != "exp": startdate = startdate[0:6] - return '{0}_{1}.nc'.format(var, startdate) + return "{0}_{1}.nc".format(var, startdate) else: - return '{0}.nc'.format(var) - - def request_chunk(self, domain, var, startdate, member, chunk, grid=None, box=None, frequency=None, - vartype=VariableType.MEAN): + return "{0}.nc".format(var) + + def request_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + ): """ - Request a given file from the CMOR repository to the scratch folder and returns the path to the scratch's copy + Request a given file from the CMOR repository to the scratch folder + and returns the path to the scratch's copy Parameters ---------- @@ -209,23 +290,52 @@ class THREDDSManager(DataManager): DataFile """ - aggregation_path = self.get_var_url(var, startdate, frequency, box, vartype) - file_path = self.get_file_path(startdate, domain, var, frequency, vartype, box=box) - - start_chunk = chunk_start_date(parse_date(startdate), chunk, self.experiment.chunk_size, 'month', - self.experiment.calendar) - end_chunk = chunk_end_date(start_chunk, self.experiment.chunk_size, 'month', self.experiment.calendar) - - thredds_subset = THREDDSSubset(aggregation_path, file_path, var, start_chunk, end_chunk) + aggregation_path = self.get_var_url( + var, startdate, frequency, box, vartype + ) + file_path = self.get_file_path( + startdate, domain, var, frequency, vartype, box=box + ) + + start_chunk = chunk_start_date( + parse_date(startdate), + chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) + end_chunk = chunk_end_date( + start_chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) + + thredds_subset = THREDDSSubset( + aggregation_path, file_path, var, start_chunk, end_chunk + ) thredds_subset.local_status = LocalStatus.PENDING self.requested_files[file_path] = thredds_subset return thredds_subset # noinspection PyUnusedLocal - def declare_chunk(self, domain, var, startdate, member, chunk, grid=None, region=None, box=None, frequency=None, - vartype=VariableType.MEAN, diagnostic=None): + def declare_chunk( + self, + domain, + var, + startdate, + member, + chunk, + grid=None, + region=None, + box=None, + frequency=None, + vartype=VariableType.MEAN, + diagnostic=None, + ): """ - Copy a given file from the CMOR repository to the scratch folder and returns the path to the scratch's copy + Copy a given file from the CMOR repository to the scratch folder + and returns the path to the scratch's copy :param diagnostic: :param region: @@ -243,24 +353,41 @@ class THREDDSManager(DataManager): :type grid: str|None :param box: file's box (only needed to retrieve sections or averages) :type box: Box - :param frequency: file's frequency (only needed if it is different from the default) + :param frequency: file's frequency (only needed if it is different + from the default) :type frequency: Frequency|None :param vartype: Variable type (mean, statistic) :type vartype: VariableType :return: path to the copy created on the scratch folder :rtype: str """ - aggregation_path = self.get_var_url(var, startdate, frequency, box, vartype) - file_path = self.get_file_path(startdate, domain, var, frequency, vartype, box=box) - - start_chunk = chunk_start_date(parse_date(startdate), chunk, self.experiment.chunk_size, 'month', - self.experiment.calendar) - end_chunk = chunk_end_date(start_chunk, self.experiment.chunk_size, 'month', self.experiment.calendar) + aggregation_path = self.get_var_url( + var, startdate, frequency, box, vartype + ) + file_path = self.get_file_path( + startdate, domain, var, frequency, vartype, box=box + ) + + start_chunk = chunk_start_date( + parse_date(startdate), + chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) + end_chunk = chunk_end_date( + start_chunk, + self.experiment.chunk_size, + "month", + self.experiment.calendar, + ) final_name = self._get_final_var_name(box, var) if file_path in self.requested_files: thredds_subset = self.requested_files[file_path] else: - thredds_subset = THREDDSSubset(aggregation_path, file_path, var, start_chunk, end_chunk) + thredds_subset = THREDDSSubset( + aggregation_path, file_path, var, start_chunk, end_chunk + ) self.requested_files[file_path] = thredds_subset thredds_subset.final_name = final_name @@ -294,19 +421,21 @@ class THREDDSSubset(DataFile): self.thredds_path = thredds_path self.remote_file = file_path self.local_file = None - if '_f' in var: - self.var = var[:var.index('_f')] - self.hourly = var[var.index('_f'):] + if "_f" in var: + self.var = var[: var.index("_f")] + self.hourly = var[var.index("_f"):] else: self.var = var - self.hourly = '' + self.hourly = "" self.dimension_indexes = {} self.handler = None self.start_time = start_time self.end_time = end_time def __str__(self): - return 'THREDDS {0.thredds_path} ({0.start_time}-{0.end_time})'.format(self) + return "THREDDS {0.thredds_path} ({0.start_time}-{0.end_time})".format( + self + ) def download(self): """ @@ -319,39 +448,63 @@ class THREDDSSubset(DataFile): """ try: - Log.debug('Downloading thredds subset {0}...', self) + Log.debug("Downloading thredds subset {0}...", self) with iris.FUTURE.context(cell_datetime_objects=True): - time_constraint = iris.Constraint(time=lambda cell: self.start_time <= cell.point <= self.end_time) - var_cube = iris.load_cube(self.thredds_path, constraint=time_constraint, callback=self._correct_cube) + time_constraint = iris.Constraint( + time=lambda cell: self.start_time + <= cell.point + <= self.end_time + ) + var_cube = iris.load_cube( + self.thredds_path, + constraint=time_constraint, + callback=self._correct_cube, + ) if not self.local_file: self.local_file = TempFile.get() iris.save(var_cube, self.local_file, zlib=True) if not Utils.check_netcdf_file(self.local_file): - raise THREDDSError('netcdf check for downloaded file failed') - Log.info('Request {0} ready!', self) + raise THREDDSError("netcdf check for downloaded file failed") + Log.info("Request {0} ready!", self) self.local_status = LocalStatus.READY except THREDDSError as ex: - Log.error('Can not retrieve {0} from server: {1}'.format(self, ex)) + Log.error("Can not retrieve {0} from server: {1}".format(self, ex)) self.local_status = LocalStatus.FAILED # noinspection PyUnusedLocal @staticmethod def _correct_cube(cube, field, filename): - if not cube.coords('time'): + if not cube.coords("time"): return - time = cube.coord('time') - if time.units.origin.startswith('month'): - ref = strptime(time.units.origin[time.units.origin.index(' since ') + 7:], '%Y-%m-%d %H:%M:%S') - helper = np.vectorize(lambda x: datetime(year=ref.tm_year + int(x) / 12, - month=int(x - 1) % 12 + 1, - day=ref.tm_mday)) + time = cube.coord("time") + if time.units.origin.startswith("month"): + ref = strptime( + time.units.origin[time.units.origin.index(" since ") + 7:], + "%Y-%m-%d %H:%M:%S", + ) + helper = np.vectorize( + lambda x: datetime( + year=ref.tm_year + int(x) / 12, + month=int(x - 1) % 12 + 1, + day=ref.tm_mday, + ) + ) times = np.round(time.points + ref.tm_mon) dates = helper(times) - dates = netCDF4.date2num(dates, units='days since 1850-01-01', calendar=time.units.calendar) - new_time = DimCoord(dates, standard_name=time.standard_name, long_name=time.long_name, - var_name=time.var_name, attributes=time.attributes, - units=Unit('days since 1850-01-01', time.units.calendar)) + dates = netCDF4.date2num( + dates, + units="days since 1850-01-01", + calendar=time.units.calendar, + ) + new_time = DimCoord( + dates, + standard_name=time.standard_name, + long_name=time.long_name, + var_name=time.var_name, + attributes=time.attributes, + units=Unit("days since 1850-01-01", time.units.calendar), + ) [dimension] = cube.coord_dims(time) cube.remove_coord(time) cube.add_dim_coord(new_time, dimension) diff --git a/earthdiagnostics/utils.py b/earthdiagnostics/utils.py index d9a69c48..ce898db1 100644 --- a/earthdiagnostics/utils.py +++ b/earthdiagnostics/utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -"""Common utilities for multiple topics that are not big enough to have their own module""" +"""Common utilities for multiple topics that do not have their own module""" import datetime import os import time @@ -57,7 +57,7 @@ class Utils(object): """An instance of Cdo class ready to be used""" if not Utils._cdo: Utils._cdo = Cdo(env=os.environ) - Utils._cdo.CDO = find_executable('cdo') + Utils._cdo.CDO = find_executable("cdo") Utils._cdo.debug = Log.console_handler.level <= Log.DEBUG return Utils._cdo @@ -82,21 +82,22 @@ class Utils(object): basin = Basins().parse(basin) if basin != Basins().Global: try: - basins = iris.load_cube('basins.nc') + basins = iris.load_cube("basins.nc") return basins.extract(iris.Constraint(region=basin.name)).data except IOError: raise Exception( - 'File basins.nc is required for basin {0}'.format(basin)) + "File basins.nc is required for basin {0}".format(basin) + ) else: - mask_handler = Utils.open_cdf('mask.nc') - mask = mask_handler.variables['tmask'][0, 0, :] + mask_handler = Utils.open_cdf("mask.nc") + mask = mask_handler.variables["tmask"][0, 0, :] mask_handler.close() return np.squeeze(mask) @staticmethod def setminmax(filename, variable_list): """ - Set the valid_max and valid_min values to the current max and min values on the file + Set the valid_max and valid_min values to current max and min values Parameters ---------- @@ -108,7 +109,7 @@ class Utils(object): if isinstance(variable_list, six.string_types): variable_list = variable_list.split() - Log.info('Getting max and min values for {0}', ' '.join(variable_list)) + Log.info("Getting max and min values for {0}", " ".join(variable_list)) handler = Utils.open_cdf(filename) for variable in variable_list: @@ -118,18 +119,22 @@ class Utils(object): input=filename, output=filename, options=( - '-h -a valid_max,{0},m,f,{1}'.format(variable, values[0]),) + "-h -a valid_max,{0},m,f,{1}".format(variable, values[0]), + ), ) Utils.nco().ncatted( input=filename, output=filename, options=( - '-h -a valid_min,{0},m,f,{1}'.format(variable, values[1]),) + "-h -a valid_min,{0},m,f,{1}".format(variable, values[1]), + ), ) handler.close() @staticmethod - def rename_variable(filepath, old_name, new_name, must_exist=True, rename_dimension=True): + def rename_variable( + filepath, old_name, new_name, must_exist=True, rename_dimension=True + ): """ Rename variable from a NetCDF file @@ -148,10 +153,13 @@ class Utils(object): """ Utils.rename_variables( - filepath, {old_name: new_name}, must_exist, rename_dimension) + filepath, {old_name: new_name}, must_exist, rename_dimension + ) @staticmethod - def rename_variables(filepath, dic_names, must_exist=True, rename_dimension=True): + def rename_variables( + filepath, dic_names, must_exist=True, rename_dimension=True + ): """ Rename multiple variables from a NetCDF file @@ -171,30 +179,47 @@ class Utils(object): """ for old, new in six.iteritems(dic_names): if old == new: - raise ValueError('{0} original name is the same as the new') + raise ValueError("{0} original name is the same as the new") original_handler = Utils.open_cdf(filepath) original_names = set(original_handler.variables.keys()).union( - original_handler.dimensions.keys()) + original_handler.dimensions.keys() + ) if not any((True for x in dic_names.keys() if x in original_names)): original_handler.close() if must_exist: - raise Exception("Variables {0} does not exist in file {1}".format( - ','.join(dic_names.keys()), filepath)) + raise Exception( + "Variables {0} does not exist in file {1}".format( + ",".join(dic_names.keys()), filepath + ) + ) return temp = TempFile.get() - new_handler = Utils.open_cdf(temp, 'w') + new_handler = Utils.open_cdf(temp, "w") for attribute in original_handler.ncattrs(): original = getattr(original_handler, attribute) - setattr(new_handler, attribute, - Utils.convert_to_ascii_if_possible(original)) + setattr( + new_handler, + attribute, + Utils.convert_to_ascii_if_possible(original), + ) for dimension in original_handler.dimensions.keys(): - Utils.copy_dimension(original_handler, new_handler, dimension, - new_names=dic_names, rename_dimension=rename_dimension) + Utils.copy_dimension( + original_handler, + new_handler, + dimension, + new_names=dic_names, + rename_dimension=rename_dimension, + ) for variable in original_handler.variables.keys(): - Utils.copy_variable(original_handler, new_handler, variable, - new_names=dic_names, rename_dimension=rename_dimension) + Utils.copy_variable( + original_handler, + new_handler, + variable, + new_names=dic_names, + rename_dimension=rename_dimension, + ) original_handler.close() new_handler.close() @@ -205,8 +230,8 @@ class Utils(object): """ Check if a NetCDF file is well stored - This functions is used to check if a NetCDF file is corrupted. It prefers to raise a false postive than - to have false negatives. + This functions is used to check if a NetCDF file is corrupted. + It prefers to raise a false postive than to have false negatives. Parameters ---------- @@ -219,8 +244,8 @@ class Utils(object): with suppress_stdout(): try: handler = Utils.open_cdf(filepath) - if 'time' in handler.variables: - if handler.variables['time'].dimensions != ('time', ): + if "time" in handler.variables: + if handler.variables["time"].dimensions != ("time",): handler.close() return False handler.close() @@ -228,7 +253,7 @@ class Utils(object): if len(cubes) == 0: return False except (iris.exceptions.IrisError, RuntimeError, OSError) as ex: - Log.debug('netCDF checks failed: {0}', ex) + Log.debug("netCDF checks failed: {0}", ex) return False return True @@ -253,12 +278,12 @@ class Utils(object): # noinspection PyPep8Naming @staticmethod - def convert_to_ascii_if_possible(string, encoding='ascii'): + def convert_to_ascii_if_possible(string, encoding="ascii"): u""" Convert an Unicode string to ASCII if all characters can be translated. - If a string can not be translated it is unchanged. It also automatically - replaces Bretonnière with Bretonniere + If a string can not be translated it is unchanged. It also + automatically replaces Bretonnière with Bretonniere Parameters ---------- @@ -275,39 +300,59 @@ class Utils(object): try: return string.encode(encoding) except UnicodeEncodeError: - if u'Bretonnière' in string: - string = string.replace(u'Bretonnière', 'Bretonniere') + if u"Bretonnière" in string: + string = string.replace(u"Bretonnière", "Bretonniere") return Utils.convert_to_ascii_if_possible(string, encoding) return string @staticmethod - def _rename_vars_directly(dic_names, filepath, handler, must_exist, rename_dimension): + def _rename_vars_directly( + dic_names, filepath, handler, must_exist, rename_dimension + ): for old_name, new_name in dic_names.items(): if rename_dimension: if old_name in handler.dimensions: handler.renameDimension(old_name, new_name) elif must_exist: raise Exception( - "Dimension {0} does not exist in file {1}".format(old_name, filepath)) + "Dimension {0} does not exist in file {1}".format( + old_name, filepath + ) + ) if old_name in handler.variables: if new_name not in handler.variables: handler.renameVariable(old_name, new_name) for var in handler.variables: - if hasattr(var, 'coordinates') and " {0} ".format(old_name) in var.coordinates: - new_coordinates = var.coordinates.replace(" {0} ".format(old_name), - " {0} ".format(new_name)) - var.coordinates = Utils.convert_to_ascii_if_possible( - new_coordinates) + Utils._rename_coords_attr(var, old_name, new_name) elif must_exist: raise Exception( - "Variable {0} does not exist in file {1}".format(old_name, filepath)) + "Variable {0} does not exist in file {1}".format( + old_name, filepath + ) + ) handler.sync() @staticmethod - def copy_file(source, destiny, save_hash=False, use_stored_hash=True, retrials=3): + def _rename_coords_attr(var, old_name, new_name): + if ( + hasattr(var, "coordinates") + and " {0} ".format(old_name) in var.coordinates + ): + new_coordinates = var.coordinates.replace( + " {0} ".format(old_name), + " {0} ".format(new_name), + ) + var.coordinates = Utils.convert_to_ascii_if_possible( + new_coordinates + ) + + @staticmethod + def copy_file( + source, destiny, save_hash=False, use_stored_hash=True, retrials=3 + ): """ - Copy a file and compute a hash to check if the copy is equal to the source + Copy a file and compute a hash to check if the copy is equal to source Parameters ---------- @@ -316,7 +361,8 @@ class Utils(object): save_hash: bool, optional If True, stores a copy of the hash use_stored_hash: bool, optional - If True, try to use the stored value of the source hash instead of computing it + If True, try to use the stored value of the source hash instead + of computing it retrials: int, optional Minimum value is 1 @@ -330,32 +376,20 @@ class Utils(object): os.makedirs(dirname_path) Utils.give_group_write_permissions(dirname_path) except OSError as ex: - # This can be due to a race condition. If directory already exists, we don have to do nothing + # This can be due to a race condition. If directory already + # exists, we do nothing if not os.path.exists(dirname_path): raise ex - # hash_destiny = None - # Log.debug('Hashing original file... {0}', datetime.datetime.now()) - # hash_original = Utils.get_file_hash(source, use_stored=use_stored_hash) - - # if retrials < 1: - # retrials = 1 - - # while hash_original != hash_destiny: - # if retrials == 0: - # raise Utils.CopyException('Can not copy {0} to {1}'.format(source, destiny)) - Log.debug('Copying... {0}', datetime.datetime.now()) + Log.debug("Copying... {0}", datetime.datetime.now()) shutil.copyfile(source, destiny) - # Log.debug('Hashing copy ... {0}', datetime.datetime.now()) - # hash_destiny = Utils.get_file_hash(destiny, save=save_hash) - # retrials -= 1 - Log.debug('Finished {0}', datetime.datetime.now()) + Log.debug("Finished {0}", datetime.datetime.now()) @staticmethod def move_file(source, destiny, save_hash=False, retrials=3): """ - Move a file and compute a hash to check if the copy is equal to the source + Move a file and compute hash to check - It is just a call to Utils.copy_file followed bu + It is just a call to Utils.copy_file followed by a remove Parameters ---------- @@ -447,19 +481,19 @@ class Utils(object): if use_stored: hash_file = Utils._get_hash_filename(filepath) if os.path.isfile(hash_file): - hash_value = open(hash_file, 'r').readline() + hash_value = open(hash_file, "r").readline() return hash_value blocksize = 104857600 hasher = xxhash.xxh64() - with open(filepath, 'rb') as afile: + with open(filepath, "rb") as afile: buf = afile.read(blocksize) while len(buf) > 0: hasher.update(buf) buf = afile.read(blocksize) hash_value = hasher.hexdigest() if save: - hash_file = open(Utils._get_hash_filename(filepath), 'w') + hash_file = open(Utils._get_hash_filename(filepath), "w") hash_file.write(hash_value) hash_file.close() @@ -469,7 +503,7 @@ class Utils(object): def _get_hash_filename(filepath): folder = os.path.dirname(filepath) filename = os.path.basename(filepath) - hash_file = os.path.join(folder, '.{0}.xxhash64.hash'.format(filename)) + hash_file = os.path.join(folder, ".{0}.xxhash64.hash".format(filename)) return hash_file @staticmethod @@ -505,13 +539,16 @@ class Utils(object): if not line: continue if six.PY3: - line = str(line, encoding='UTF-8') + line = str(line, encoding="UTF-8") if log_level != Log.NO_LOG: Log.log.log(log_level, line) output.append(line) if process.returncode != 0: - raise Utils.ExecutionError('Error executing {0}\n Return code: {1}'.format(' '.join(command), - str(process.returncode))) + raise Utils.ExecutionError( + "Error executing {0}\n Return code: {1}".format( + " ".join(command), str(process.returncode) + ) + ) return output _cpu_count = None @@ -521,16 +558,20 @@ class Utils(object): """Number of available virtual or physical CPUs on this system""" if Utils._cpu_count is None: try: - match = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', - open('/proc/self/status').read()) + match = re.search( + r"(?m)^Cpus_allowed:\s*(.*)$", + open("/proc/self/status").read(), + ) if match: - res = bin( - int(match.group(1).replace(',', ''), 16)).count('1') + res = bin(int(match.group(1).replace(",", ""), 16)).count( + "1" + ) if res > 0: Utils._cpu_count = res except IOError: try: import multiprocessing + Utils._cpu_count = multiprocessing.cpu_count() return Utils._cpu_count except (ImportError, NotImplementedError): @@ -542,7 +583,8 @@ class Utils(object): """ Convert a file to NetCDF4 - Conversion only performed if required. Deflation level set to 4 and shuffle activated. + Conversion only performed if required. Deflation level set to 4 + and shuffle activated. Parameters ---------- @@ -551,25 +593,28 @@ class Utils(object): if Utils._is_compressed_netcdf4(filetoconvert): return - Log.debug('Reformatting to netCDF-4') + Log.debug("Reformatting to netCDF-4") temp = TempFile.get() Utils.execute_shell_command( - ["nccopy", "-4", "-d4", "-s", filetoconvert, temp]) + ["nccopy", "-4", "-d4", "-s", filetoconvert, temp] + ) shutil.move(temp, filetoconvert) @classmethod def _is_compressed_netcdf4(cls, filetoconvert): ncdump_result = Utils.execute_shell_command( - 'ncdump -hs {0}'.format(filetoconvert), Log.NO_LOG) - ncdump_result = ncdump_result[0].replace('\t', '').split('\n') - if any(':_Shuffle = "true"' in line for line in ncdump_result) and \ - any(':_DeflateLevel = 4' in line for line in ncdump_result): + "ncdump -hs {0}".format(filetoconvert), Log.NO_LOG + ) + ncdump_result = ncdump_result[0].replace("\t", "").split("\n") + if any(':_Shuffle = "true"' in line for line in ncdump_result) and any( + ":_DeflateLevel = 4" in line for line in ncdump_result + ): return True return False -# noinspection PyPep8Naming + # noinspection PyPep8Naming @staticmethod - def open_cdf(filepath, mode='a'): + def open_cdf(filepath, mode="a"): """ Open a NetCDF file @@ -585,7 +630,7 @@ class Utils(object): return netCDF4.Dataset(filepath, mode) @staticmethod - def get_datetime_from_netcdf(handler, time_variable='time'): + def get_datetime_from_netcdf(handler, time_variable="time"): """ Get time from NetCDF files @@ -609,8 +654,15 @@ class Utils(object): return netCDF4.num2date(nctime, units=units, calendar=cal_temps) @staticmethod - def copy_variable(source, destiny, variable, must_exist=True, add_dimensions=False, new_names=None, - rename_dimension=True): + def copy_variable( + source, + destiny, + variable, + must_exist=True, + add_dimensions=False, + new_names=None, + rename_dimension=True, + ): """ Copy the given variable from source to destiny @@ -626,7 +678,8 @@ class Utils(object): Raises ------ Exception - If dimensions are not correct in the destiny file and add_dimensions is False + If dimensions are not correct in the destiny file and + add_dimensions is False """ if not must_exist and variable not in source.variables.keys(): @@ -641,38 +694,68 @@ class Utils(object): if not new_name or new_name in destiny.variables.keys(): return - translated_dimensions = Utils._copy_dimensions(add_dimensions, destiny, must_exist, new_names, rename_dimension, - source, variable) + translated_dimensions = Utils._copy_dimensions( + add_dimensions, + destiny, + must_exist, + new_names, + rename_dimension, + source, + variable, + ) if new_name in destiny.variables.keys(): # Just in case the variable we are copying match a dimension name return original_var = source.variables[variable] new_var = destiny.createVariable( - new_name, original_var.datatype, translated_dimensions) + new_name, original_var.datatype, translated_dimensions + ) Utils.copy_attributes(new_var, original_var) - if hasattr(new_var, 'coordinates'): + if hasattr(new_var, "coordinates"): coords = [ - new_names[coord] if coord in new_names else coord for coord in new_var.coordinates.split(' ')] + new_names[coord] if coord in new_names else coord + for coord in new_var.coordinates.split(" ") + ] coords = [coord for coord in coords if coord] if coords: new_var.coordinates = Utils.convert_to_ascii_if_possible( - ' '.join(coords)) + " ".join(coords) + ) new_var[:] = original_var[:] @staticmethod - def _copy_dimensions(add_dimensions, destiny, must_exist, new_names, rename_dimension, source, variable): + def _copy_dimensions( + add_dimensions, + destiny, + must_exist, + new_names, + rename_dimension, + source, + variable, + ): if rename_dimension: translated_dimensions = Utils._translate( - source.variables[variable].dimensions, new_names) + source.variables[variable].dimensions, new_names + ) else: translated_dimensions = list(source.variables[variable].dimensions) if not set(translated_dimensions).issubset(destiny.dimensions): if not add_dimensions: - raise Exception('Variable {0} can not be added because dimensions does not match: ' - '{1} {2}'.format(variable, translated_dimensions, destiny.dimensions)) + raise Exception( + "Variable {0} can not be added because dimensions do " + "not match: {1} {2}".format( + variable, translated_dimensions, destiny.dimensions + ) + ) for dimension in source.variables[variable].dimensions: Utils.copy_dimension( - source, destiny, dimension, must_exist, new_names, rename_dimension) + source, + destiny, + dimension, + must_exist, + new_names, + rename_dimension, + ) return translated_dimensions @staticmethod @@ -689,13 +772,27 @@ class Utils(object): """ if omitted_attributtes is None: omitted_attributtes = [] - new_var.setncatts({k: Utils.convert_to_ascii_if_possible(original_var.getncattr(k)) - for k in original_var.ncattrs() if k not in omitted_attributtes}) + new_var.setncatts( + { + k: Utils.convert_to_ascii_if_possible( + original_var.getncattr(k) + ) + for k in original_var.ncattrs() + if k not in omitted_attributtes + } + ) @staticmethod - def copy_dimension(source, destiny, dimension, must_exist=True, new_names=None, rename_dimension=False): + def copy_dimension( + source, + destiny, + dimension, + must_exist=True, + new_names=None, + rename_dimension=False, + ): """ - Copy the given dimension from source to destiny, including dimension variables if present + Copy dimension from source to destiny, including dimension variables Parameters ---------- @@ -719,8 +816,14 @@ class Utils(object): new_name = dimension destiny.createDimension(new_name, source.dimensions[dimension].size) if dimension in source.variables: - Utils.copy_variable(source, destiny, dimension, - new_names=new_names, rename_dimension=rename_dimension, add_dimensions=True) + Utils.copy_variable( + source, + destiny, + dimension, + new_names=new_names, + rename_dimension=rename_dimension, + add_dimensions=True, + ) @staticmethod def concat_variables(source, destiny, remove_source=False): @@ -742,13 +845,19 @@ class Utils(object): for var in handler_variable.variables: if var not in handler_total.variables: Utils.copy_variable( - handler_variable, handler_total, var, add_dimensions=True) + handler_variable, + handler_total, + var, + add_dimensions=True, + ) else: variable = handler_variable.variables[var] - if 'time' not in variable.dimensions: + if "time" not in variable.dimensions: continue - concatenated[var] = np.concatenate((handler_total.variables[var][:], variable[:]), - axis=variable.dimensions.index('time')) + concatenated[var] = np.concatenate( + (handler_total.variables[var][:], variable[:]), + axis=variable.dimensions.index("time"), + ) for var, array in six.iteritems(concatenated): handler_total.variables[var][:] = array @@ -793,8 +902,9 @@ class Utils(object): try: os.makedirs(path) except OSError: - # This could happen if two threads are tying to create the folder. - # Let's check again for existence and rethrow if still not exists + # This could happen if two threads are tying to create + # the folder. Let's check again for existence and rethrow + # if still not exists if not os.path.isdir(path): raise @@ -807,7 +917,9 @@ class Utils(object): os.chmod(path, stats.st_mode | stat.S_IWGRP) @staticmethod - def convert_units(var_handler, new_units, calendar=None, old_calendar=None): + def convert_units( + var_handler, new_units, calendar=None, old_calendar=None + ): """ Convert units @@ -821,19 +933,23 @@ class Utils(object): if new_units == var_handler.units: return - if hasattr(var_handler, 'calendar'): + if hasattr(var_handler, "calendar"): old_calendar = var_handler.calendar import cf_units + new_unit = cf_units.Unit(new_units, calendar=calendar) old_unit = cf_units.Unit(var_handler.units, calendar=old_calendar) var_handler[:] = old_unit.convert( - np.array(var_handler[:]), new_unit, inplace=True) - if 'valid_min' in var_handler.ncattrs(): - var_handler.valid_min = old_unit.convert(float(var_handler.valid_min), new_unit, - inplace=True) - if 'valid_max' in var_handler.ncattrs(): - var_handler.valid_max = old_unit.convert(float(var_handler.valid_max), new_unit, - inplace=True) + np.array(var_handler[:]), new_unit, inplace=True + ) + if "valid_min" in var_handler.ncattrs(): + var_handler.valid_min = old_unit.convert( + float(var_handler.valid_min), new_unit, inplace=True + ) + if "valid_max" in var_handler.ncattrs(): + var_handler.valid_max = old_unit.convert( + float(var_handler.valid_max), new_unit, inplace=True + ) var_handler.units = new_units @staticmethod @@ -848,28 +964,37 @@ class Utils(object): """ for filepath in files: - Log.debug('Unpacking {0}', filepath) + Log.debug("Unpacking {0}", filepath) extracted = False tar = tarfile.open(filepath) for file_compressed in tar.getmembers(): if file_compressed.isdir(): - if os.path.isdir(os.path.join(destiny_path, file_compressed.name)): + if os.path.isdir( + os.path.join(destiny_path, file_compressed.name) + ): continue else: if filter_files: - if not any(f in file_compressed.name for f in filter_files): + if not any( + f in file_compressed.name for f in filter_files + ): continue - if os.path.exists(os.path.join(destiny_path, file_compressed.name)): - os.remove(os.path.join( - destiny_path, file_compressed.name)) + if os.path.exists( + os.path.join(destiny_path, file_compressed.name) + ): + os.remove( + os.path.join(destiny_path, file_compressed.name) + ) extracted = True tar.extract(file_compressed, destiny_path) - Log.debug('File {0} extracted', - os.path.basename(file_compressed.name)) + Log.debug( + "File {0} extracted", + os.path.basename(file_compressed.name), + ) tar.close() if filter_files and not extracted: - ValueError('No files extracted. Check filters') + ValueError("No files extracted. Check filters") @staticmethod def unzip(files, force=False): @@ -885,17 +1010,19 @@ class Utils(object): if isinstance(files, six.string_types): files = [files] for filepath in files: - Log.debug('Unzipping {0}', filepath) + Log.debug("Unzipping {0}", filepath) if force: - option = ' -f' + option = " -f" else: - option = '' + option = "" try: Utils.execute_shell_command( - 'gunzip{1} {0}'.format(filepath, option)) + "gunzip{1} {0}".format(filepath, option) + ) except Exception as ex: raise Utils.UnzipException( - 'Can not unzip {0}: {1}'.format(filepath, ex)) + "Can not unzip {0}: {1}".format(filepath, ex) + ) class UnzipException(Exception): """Exception raised when unzip fails""" @@ -919,22 +1046,23 @@ class TempFile(object): """ List of files to clean automatically """ - scratch_folder = '' + scratch_folder = "" """ Scratch folder to create temporary files on it """ - prefix = 'temp' + prefix = "temp" """ Prefix for temporary filenames """ @staticmethod - def get(filename=None, clean=None, suffix='.nc'): + def get(filename=None, clean=None, suffix=".nc"): """ Get a new temporal filename, storing it for automated cleaning :param suffix: - :param filename: if it is not none, the function will use this filename instead of a random one + :param filename: if it is not none, the function will use this filename + instead of a random one :type filename: str :param clean: if true, stores filename for cleaning :type clean: bool @@ -948,7 +1076,10 @@ class TempFile(object): path = os.path.join(TempFile.scratch_folder, filename) else: file_descriptor, path = tempfile.mkstemp( - dir=TempFile.scratch_folder, prefix=TempFile.prefix, suffix=suffix) + dir=TempFile.scratch_folder, + prefix=TempFile.prefix, + suffix=suffix, + ) path = str(path) os.close(file_descriptor) diff --git a/earthdiagnostics/variable.py b/earthdiagnostics/variable.py index bb1ccb42..668b6b16 100644 --- a/earthdiagnostics/variable.py +++ b/earthdiagnostics/variable.py @@ -12,23 +12,24 @@ from bscearth.utils.log import Log from earthdiagnostics.constants import Basins from earthdiagnostics.frequency import Frequency from earthdiagnostics.modelingrealm import ModelingRealms -from concurrent.futures import ThreadPoolExecutor class VariableJsonException(Exception): - """Exception to be raised when an error related to the json reading is encountered""" + """Json reading ERROR""" pass class VariableManager(object): - """Class for translating variable alias into standard names and provide the correct description for them""" + """Class to manage CMOR variables info""" def __init__(self): self._cmor_tables_folder = os.path.join( - os.path.dirname(os.path.realpath(__file__)), 'cmor_tables') - self._aliases_folder = os.path.join(os.path.dirname( - os.path.realpath(__file__)), 'variable_alias') + os.path.dirname(os.path.realpath(__file__)), "cmor_tables" + ) + self._aliases_folder = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "variable_alias" + ) self.clean() def clean(self): @@ -53,8 +54,10 @@ class VariableManager(object): return self._dict_aliases[original_name.lower()][1] except KeyError: if not silent: - Log.warning('Variable {0} is not defined in the CMOR table. Please add it'.format( - original_name)) + Log.warning( + f"Variable {original_name} is not defined in the CMOR" + " table. Please add it" + ) return None def get_all_variables(self): @@ -82,8 +85,10 @@ class VariableManager(object): return self._dict_aliases[original_name.lower()] except KeyError: if not silent: - Log.warning('Variable {0} is not defined in the CMOR table. Please add it'.format( - original_name)) + Log.warning( + f"Variable {original_name} is not defined in the CMOR" + " table. Please add it" + ) return None, None def load_variables(self, table_name): @@ -119,28 +124,33 @@ class VariableManager(object): self._load_file(self.table_name) return - raise Exception('Data convention {0} unknown'.format(self.table_name)) + raise Exception("Data convention {0} unknown".format(self.table_name)) def _get_csv_path(self, table_name): csv_table_path = os.path.join( - self._cmor_tables_folder, '{0}.csv'.format(table_name)) + self._cmor_tables_folder, "{0}.csv".format(table_name) + ) return csv_table_path def _get_json_folder(self): json_folder = os.path.join( - self._cmor_tables_folder, '{0}/Tables'.format(self.table_name)) + self._cmor_tables_folder, "{0}/Tables".format(self.table_name) + ) return json_folder def _load_file(self, csv_table_path, default=False): - with open(self._get_csv_path(csv_table_path), 'r') as csvfile: - reader = csv.reader(csvfile, dialect='excel') + with open(self._get_csv_path(csv_table_path), "r") as csvfile: + reader = csv.reader(csvfile, dialect="excel") for line in reader: - if line[0] == 'Variable': + if line[0] == "Variable": continue var = Variable() var.parse_csv(line) - if not var.short_name or var.short_name.lower() in self._dict_variables: + if ( + not var.short_name + or var.short_name.lower() in self._dict_variables + ): continue var.default = default @@ -158,13 +168,10 @@ class VariableManager(object): self._dict_variables[var.short_name.lower()] = var def _load_json(self, json_folder): - # executor = ThreadPoolExecutor(max_workers=1) for file_name in os.listdir(json_folder): - if file_name in ('CMIP6_grids.json', 'CMIP6_formula_terms.json'): + if file_name in ("CMIP6_grids.json", "CMIP6_formula_terms.json"): continue self._load_json_file(os.path.join(json_folder, file_name)) - # executor.submit(self._load_json_file, os.path.join(json_folder, file_name)) - # executor.shutdown(True) def _load_json_file(self, json_path): with open(json_path) as json_file: @@ -173,27 +180,30 @@ class VariableManager(object): data = json.loads(json_data) except ValueError: return - if 'variable_entry' in data: - Log.debug('Parsing file {0}'.format(json_path)) - table_id = data['Header']['table_id'][6:] + if "variable_entry" in data: + Log.debug("Parsing file {0}".format(json_path)) + table_id = data["Header"]["table_id"][6:] - var_freqs = (var['frequency'] - for var in data['variable_entry'].values()) + var_freqs = ( + var["frequency"] for var in data["variable_entry"].values() + ) table_freq, _ = Counter(var_freqs).most_common(1)[0] table = CMORTable( table_id, Frequency(table_freq), - data['Header']['table_date'], - ModelingRealms.parse(data['Header']['realm'].split(' ')[0]) + data["Header"]["table_date"], + ModelingRealms.parse( + data["Header"]["realm"].split(" ")[0] + ), ) self.tables[table_id] = table - self._load_json_variables(data['variable_entry'], table) + self._load_json_variables(data["variable_entry"], table) else: - Log.debug('Skipping file {0}'.format(json_path)) + Log.debug("Skipping file {0}".format(json_path)) def _load_json_variables(self, json_data, table): for short_name in json_data.keys(): - if short_name == 'ta19': + if short_name == "ta19": pass short_name = str.strip(str(short_name)) if short_name.lower() in self._dict_variables: @@ -205,19 +215,19 @@ class VariableManager(object): variable.add_table(table) self.register_variable(variable) except VariableJsonException: - Log.error('Could not read variable {0}'.format(short_name)) + Log.error("Could not read variable {0}".format(short_name)) def _load_known_aliases(self): - self._load_alias_csv('default') + self._load_alias_csv("default") self._load_alias_csv(self.table_name) def _load_alias_csv(self, filename): file_path = self._get_aliases_csv_path(filename) - with open(file_path, 'r') as csvfile: - reader = csv.reader(csvfile, dialect='excel') + with open(file_path, "r") as csvfile: + reader = csv.reader(csvfile, dialect="excel") for line in reader: - if line[0] == 'Aliases': + if line[0] == "Aliases": continue aliases = self._get_aliases(line) @@ -234,19 +244,28 @@ class VariableManager(object): if len(cmor_vars) == 0: Log.warning( - 'Aliases {0} could not be mapped to any variable'.format(aliases)) + "Aliases {0} could not be mapped to any variable", + aliases + ) continue elif len(cmor_vars) > 1: non_default = [var for var in cmor_vars if not var.default] if len(non_default) == 1: - for default in [var for var in cmor_vars if var not in non_default]: - del self._dict_variables[default.short_name.lower( - )] + for default in [ + var for var in cmor_vars if var not in non_default + ]: + del self._dict_variables[ + default.short_name.lower() + ] cmor_vars = non_default else: - Log.error('Aliases {0} can be be mapped to multiple variables ' - '[{1}]'.format(aliases, ', '.join(map(str, cmor_vars)))) + Log.error( + "Aliases {0} can be be mapped to multiple " + "variables [{1}]".format( + aliases, ", ".join(map(str, cmor_vars)) + ) + ) continue cmor_var = cmor_vars[0] @@ -254,14 +273,16 @@ class VariableManager(object): @staticmethod def _get_aliases(line): - aliases = line[0].split(':') + aliases = line[0].split(":") return aliases def _register_aliases(self, aliases, cmor_var, line): for alias in aliases: if alias != cmor_var.short_name and alias in self._dict_variables: - Log.error('Alias {0} for variable {1} is already a different ' - 'variable!'.format(alias, cmor_var.short_name)) + Log.error( + "Alias {0} for variable {1} is already a different " + "variable!".format(alias, cmor_var.short_name) + ) continue alias_object = VariableAlias(alias) if line[2]: @@ -272,14 +293,15 @@ class VariableManager(object): def _get_aliases_csv_path(self, filename): csv_table_path = os.path.join( - self._aliases_folder, '{0}.csv'.format(filename)) + self._aliases_folder, "{0}.csv".format(filename) + ) return csv_table_path def create_aliases_dict(self): """Create aliases dictionary for the registered variables""" self._dict_aliases = {} for cmor_var_name in self._dict_variables: - if cmor_var_name == 'tos': + if cmor_var_name == "tos": pass cmor_var = self._dict_variables[cmor_var_name] base_alias = VariableAlias(cmor_var_name) @@ -290,12 +312,14 @@ class VariableManager(object): def _get_xlsx_path(self): xlsx_table_path = os.path.join( - self._cmor_tables_folder, '{0}.xlsx'.format(self.table_name)) + self._cmor_tables_folder, "{0}.xlsx".format(self.table_name) + ) if os.path.isfile(xlsx_table_path): return xlsx_table_path xlsx_table_path = os.path.join( - self._cmor_tables_folder, self.table_name, 'etc', '*.xlsx') + self._cmor_tables_folder, self.table_name, "etc", "*.xlsx" + ) xlsx_table_path = glob.glob(xlsx_table_path) if len(xlsx_table_path) == 1: return xlsx_table_path[0] @@ -309,12 +333,14 @@ class VariableManager(object): for row in data_sheet.rows: if row[1].value in excel.sheetnames: table_data[row[1].value] = ( - Frequency(row[2].value), 'Date missing') + Frequency(row[2].value), + "Date missing", + ) for sheet_name in excel.sheetnames: sheet = excel[sheet_name] - if sheet.title == 'Primday': + if sheet.title == "Primday": pass - if sheet['A1'].value not in ['Priority', 'rm']: + if sheet["A1"].value not in ["Priority", "rm"]: continue self._load_xlsx_table(sheet, table_data) @@ -325,25 +351,29 @@ class VariableManager(object): table = CMORTable(sheet.title, table_frequency, table_date, realm) self.tables[sheet.title] = table for row in sheet.rows: - if row[0].value in ('Priority', 'rm') or not row[5].value: + if row[0].value in ("Priority", "rm") or not row[5].value: continue self._parse_xlsx_var_row(row, table) except Exception as ex: - Log.error('Table {0} can not be loaded: {1}', sheet.title, ex) + Log.error("Table {0} can not be loaded: {1}", sheet.title, ex) import traceback + traceback.print_exc() def _read_realm_from_json(self, table_name): - for prefix in ('CMIP6', 'PRIMAVERA'): - json_path = os.path.join(self._get_json_folder( - ), '{0}_{1}.json'.format(prefix, table_name)) + for prefix in ("CMIP6", "PRIMAVERA"): + json_path = os.path.join( + self._get_json_folder(), + "{0}_{1}.json".format(prefix, table_name), + ) if os.path.isfile(json_path): with open(json_path) as json_file: json_data = json_file.read() data = json.loads(json_data) - # Cogemos el primer realm para las tablas que tienen varios - # Solo se usa al generar los links para una startdate concreta - return ModelingRealms.parse(data['Header']['realm'].split(' ')[0]) + # Take the first realm for those with several ones + return ModelingRealms.parse( + data["Header"]["realm"].split(" ")[0] + ) return None def _parse_xlsx_var_row(self, row, table): @@ -352,7 +382,10 @@ class VariableManager(object): cmor_name = row[5].value priority = int(row[0].value) bsc_commitment = row[30].value - if bsc_commitment is not None and bsc_commitment.strip().lower() == 'false': + if ( + bsc_commitment is not None + and bsc_commitment.strip().lower() == "false" + ): priority = priority + 3 if cmor_name.lower() in self._dict_variables: var = self._dict_variables[cmor_name.lower()] @@ -371,12 +404,12 @@ class VariableManager(object): @staticmethod def _process_modelling_realm(var, value): if value is None: - value = '' - modelling_realm = value.split(' ') + value = "" + modelling_realm = value.split(" ") return var.get_modelling_realm(modelling_realm) def _load_missing_defaults(self): - self._load_file('default', True) + self._load_file("default", True) class Variable(object): @@ -384,14 +417,15 @@ class Variable(object): Class to characterize a CMOR variable. It also contains the static method to make the match between the original - name and the standard name. Requires data _convetion to be available in cmor_tables to work. + name and the standard name. Requires data _convetion to be available + in cmor_tables to work. """ def __str__(self): - return '{0} ({1})'.format(self.standard_name, self.short_name) + return "{0} ({1})".format(self.standard_name, self.short_name) def __repr__(self): - return '{0} ({1})'.format(self.standard_name, self.short_name) + return "{0} ({1})".format(self.standard_name, self.short_name) def __init__(self): self.short_name = None @@ -431,24 +465,25 @@ class Variable(object): ------- """ - if 'out_name' in json_var: - self.short_name = json_var['out_name'].strip() + if "out_name" in json_var: + self.short_name = json_var["out_name"].strip() else: raise VariableJsonException( - 'Variable {0} has no out name defined'.format(variable)) - self.standard_name = json_var['standard_name'].strip() - self.long_name = json_var['long_name'].strip() + "Variable {0} has no out name defined".format(variable) + ) + self.standard_name = json_var["standard_name"].strip() + self.long_name = json_var["long_name"].strip() - domain = json_var['modeling_realm'].split(' ') + domain = json_var["modeling_realm"].split(" ") self.domain = self.get_modelling_realm(domain) - self.valid_min = json_var['valid_min'].strip() - self.valid_max = json_var['valid_max'].strip() - self.units = json_var['units'].strip() - if 'priority' in json_var: - self.priority = int(json_var['priority'].strip()) - elif 'primavera_priority' in json_var: - self.priority = int(json_var['primavera_priority'].strip()) + self.valid_min = json_var["valid_min"].strip() + self.valid_max = json_var["valid_max"].strip() + self.units = json_var["units"].strip() + if "priority" in json_var: + self.priority = int(json_var["priority"].strip()) + elif "primavera_priority" in json_var: + self.priority = int(json_var["primavera_priority"].strip()) else: self.priority = 1 @@ -467,7 +502,10 @@ class Variable(object): """ if len(domains) > 1: Log.warning( - 'Multiple modeling realms assigned to variable {0}: {1}. ', self, domains) + "Multiple modeling realms assigned to variable {0}: {1}. ", + self, + domains, + ) parsed = [] for domain in domains: parsed.append(ModelingRealms.parse(domain)) @@ -475,16 +513,23 @@ class Variable(object): selected = self._select_most_specific(parsed) if selected: Log.warning( - 'We will use {0} as it is the most specific', selected) + "We will use {0} as it is the most specific", selected + ) return selected - Log.warning('We will use {0} as it is the first on the list and there is no one that is more specific', - parsed[0]) + Log.warning( + "We will use {0} as it is the first on the list and there is " + "no one that is more specific", + parsed[0], + ) return parsed[0] elif len(domains) == 0: Log.warning( - 'Variable {0} has no modeling realm defined'.format(self.short_name)) + "Variable {0} has no modeling realm defined".format( + self.short_name + ) + ) return None else: return ModelingRealms.parse(domains[0]) @@ -512,7 +557,8 @@ class Variable(object): """ Get a table object given the frequency and data_covention - If the variable does not contain the table information, it uses the domain to make a guess + If the variable does not contain the table information, it uses + the domain to make a guess Parameters ---------- @@ -533,17 +579,24 @@ class Variable(object): if len(tables) == 1: return tables[0] elif len(tables) > 1: - for priority_table in ['day', 'Amon', 'Oday', 'Omon']: + for priority_table in ["day", "Amon", "Oday", "Omon"]: try: - return next((t for t in tables if priority_table == t.name)) + return next( + (t for t in tables if priority_table == t.name) + ) except StopIteration: pass return tables[0] if self.domain: table_name = self.domain.get_table_name(frequency, data_convention) - return CMORTable(table_name, frequency, 'December 2013', self.domain) + return CMORTable( + table_name, frequency, "December 2013", self.domain + ) raise ValueError( - 'Can not get table for {0} and frequency {1}'.format(self, frequency)) + "Can not get table for {0} and frequency {1}".format( + self, frequency + ) + ) @staticmethod def _select_most_specific(parsed): @@ -568,7 +621,8 @@ class VariableAlias(object): Class to characterize a CMOR variable. It also contains the static method to make the match between the original - name and the standard name. Requires data _convetion to be available in cmor_tables to work. + name and the standard name. Requires data _convetion to be available + in cmor_tables to work. Parameters ---------- @@ -584,15 +638,19 @@ class VariableAlias(object): def __str__(self): string = self.alias if self.basin: - string += ' Basin: {0}'.format(self.basin) + string += " Basin: {0}".format(self.basin) if self.grid: - string += ' Grid: {0}'.format(self.grid) + string += " Grid: {0}".format(self.grid) return string def __eq__(self, other): if other is None: return False - return self.alias == other.alias and self.grid == other.grid and self.basin == other.basin + return ( + self.alias == other.alias + and self.grid == other.grid + and self.basin == other.basin + ) def __ne__(self, other): return not self == other @@ -619,7 +677,7 @@ class CMORTable(object): return self.name def __repr__(self): - return '{0.name} ({0.domain} {0.frequency}, {0.date})'.format(self) + return "{0.name} ({0.domain} {0.frequency}, {0.date})".format(self) def __lt__(self, other): return self.name < other.name @@ -635,8 +693,8 @@ class VariableType(object): def to_str(vartype): """Get str representation of vartype for the folder convention""" if vartype == VariableType.MEAN: - return 'mean' + return "mean" elif vartype == VariableType.STATISTIC: - return 'statistics' + return "statistics" else: - raise ValueError('Variable type {0} not supported'.format(vartype)) + raise ValueError("Variable type {0} not supported".format(vartype)) diff --git a/earthdiagnostics/work_manager.py b/earthdiagnostics/work_manager.py index c3d2ccd9..048388f7 100644 --- a/earthdiagnostics/work_manager.py +++ b/earthdiagnostics/work_manager.py @@ -6,7 +6,7 @@ import operator import sys import threading import traceback -import resource + # noinspection PyCompatibility from concurrent.futures import ThreadPoolExecutor from functools import cmp_to_key @@ -16,7 +16,11 @@ import six from bscearth.utils.log import Log from earthdiagnostics.datafile import StorageStatus, LocalStatus -from earthdiagnostics.diagnostic import DiagnosticStatus, Diagnostic, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticStatus, + Diagnostic, + DiagnosticOptionError, +) from earthdiagnostics.utils import Utils, TempFile @@ -49,7 +53,7 @@ class WorkManager(object): self._register_diagnostics() for fulldiag in self.config.get_commands(): Log.info("Adding {0} to diagnostic list", fulldiag) - diag_options = fulldiag.split(',') + diag_options = fulldiag.split(",") diag_class = Diagnostic.get_diagnostic(diag_options[0]) if diag_class: try: @@ -58,10 +62,16 @@ class WorkManager(object): for subjob in job.subjobs: self.add_job(subjob) except DiagnosticOptionError as ex: - Log.error('Can not configure diagnostic {0}: {1}', diag_options[0], ex) + Log.error( + "Can not configure diagnostic {0}: {1}", + diag_options[0], + ex, + ) self.had_errors = True else: - Log.error('{0} is not an available diagnostic', diag_options[0]) + Log.error( + "{0} is not an available diagnostic", diag_options[0] + ) self.had_errors = True def add_job(self, job, old_status=None): @@ -79,8 +89,11 @@ class WorkManager(object): bool Only True if all diagnostic were correctly executed """ - if not self.jobs[DiagnosticStatus.WAITING] and not self.jobs[DiagnosticStatus.READY]: - Log.result('No diagnostics to run') + if ( + not self.jobs[DiagnosticStatus.WAITING] + and not self.jobs[DiagnosticStatus.READY] + ): + Log.result("No diagnostics to run") return True start_time = datetime.datetime.now() @@ -88,7 +101,7 @@ class WorkManager(object): self.threads = Utils.available_cpu_count() if 0 < self.config.max_cores < self.threads: self.threads = self.config.max_cores - Log.info('Using {0} threads', self.threads) + Log.info("Using {0} threads", self.threads) self.downloader = Downloader() self.uploader = ThreadPoolExecutor(self.config.parallel_uploads) @@ -103,7 +116,7 @@ class WorkManager(object): job.subscribe(self, self._job_status_changed) if self.config.skip_diags_done: if job.can_skip_run(): - Log.info('Diagnostic {0} already done. Skipping !', job) + Log.info("Diagnostic {0} already done. Skipping !", job) job.status = DiagnosticStatus.COMPLETED continue job.check_is_ready() @@ -116,7 +129,9 @@ class WorkManager(object): self.downloader.start() printer_lock = threading.Lock() printer_lock.acquire() - printer = Thread(target=self._print_status, name="Printer", args=(printer_lock,)) + printer = Thread( + target=self._print_status, name="Printer", args=(printer_lock,) + ) printer.start() self.lock.acquire() @@ -147,14 +162,25 @@ class WorkManager(object): return time.sleep(step) - Log.info('Current status:') - Log.info('===============') - Log.info('Waiting: {0:4}', len(self.jobs[DiagnosticStatus.WAITING])) - Log.info('Ready to run: {0:4}', len(self.jobs[DiagnosticStatus.READY])) - Log.info('Running: {0:4}', len(self.jobs[DiagnosticStatus.RUNNING])) - Log.info('Completed: {0:4}', len(self.jobs[DiagnosticStatus.COMPLETED])) - Log.info('Failed: {0:4}', len(self.jobs[DiagnosticStatus.FAILED])) - Log.info('===============') + Log.info("Current status:") + Log.info("===============") + Log.info( + "Waiting: {0:4}", len(self.jobs[DiagnosticStatus.WAITING]) + ) + Log.info( + "Ready to run: {0:4}", len(self.jobs[DiagnosticStatus.READY]) + ) + Log.info( + "Running: {0:4}", len(self.jobs[DiagnosticStatus.RUNNING]) + ) + Log.info( + "Completed: {0:4}", + len(self.jobs[DiagnosticStatus.COMPLETED]), + ) + Log.info( + "Failed: {0:4}", len(self.jobs[DiagnosticStatus.FAILED]) + ) + Log.info("===============") def _job_status_changed(self, job, old_status): self.add_job(job, old_status) @@ -168,7 +194,11 @@ class WorkManager(object): self._check_completion() def _file_object_status_changed(self, file_object): - Log.debug('Checking file {0}. Local status {0.local_status} Storage status{0.storage_status}', file_object) + Log.debug( + "Checking file {0}. Local status {0.local_status} " + "Storage status{0.storage_status}", + file_object, + ) if file_object.download_required(): self.downloader.submit(file_object) return @@ -178,9 +208,11 @@ class WorkManager(object): return if file_object.storage_status == StorageStatus.FAILED: self.uploads_failed.append(file_object) - if file_object.local_status != LocalStatus.COMPUTING and \ - file_object.storage_status != StorageStatus.UPLOADING and \ - file_object.only_suscriber(self): + if ( + file_object.local_status != LocalStatus.COMPUTING + and file_object.storage_status != StorageStatus.UPLOADING + and file_object.only_suscriber(self) + ): del self.data_manager.requested_files[file_object.remote_file] file_object.unsubscribe(self) file_object.clean_local() @@ -201,7 +233,10 @@ class WorkManager(object): return True def _jobs_running_or_ready(self): - if self.jobs[DiagnosticStatus.READY] or self.jobs[DiagnosticStatus.RUNNING]: + if ( + self.jobs[DiagnosticStatus.READY] + or self.jobs[DiagnosticStatus.RUNNING] + ): return True for job in self.jobs[DiagnosticStatus.WAITING]: @@ -225,11 +260,14 @@ class WorkManager(object): self.downloader.on_hold = len(self.jobs[DiagnosticStatus.READY]) > 20 def _print_stats(self): - Log.info('Time consumed by each diagnostic class') - Log.info('--------------------------------------') + Log.info("Time consumed by each diagnostic class") + Log.info("--------------------------------------") times = {} - for job in self.jobs[DiagnosticStatus.COMPLETED] + self.jobs[DiagnosticStatus.FAILED]: + for job in ( + self.jobs[DiagnosticStatus.COMPLETED] + + self.jobs[DiagnosticStatus.FAILED] + ): job_type = job.alias if job_type in times.keys(): times[job_type] += job.consumed_time @@ -237,46 +275,57 @@ class WorkManager(object): times[job_type] = job.consumed_time for diag in sorted(times, key=operator.itemgetter(1)): - Log.info('{0:23} {1:}', diag, times[diag]) + Log.info("{0:23} {1:}", diag, times[diag]) def _print_errors(self): if self.jobs[DiagnosticStatus.FAILED]: self.had_errors = True - Log.error('Failed jobs') - Log.error('-----------') + Log.error("Failed jobs") + Log.error("-----------") for job in self.jobs[DiagnosticStatus.FAILED]: - Log.error('{0}: {0.message}', job) - Log.error('Total wasted time: {0}', sum([job.consumed_time for job in self.jobs[DiagnosticStatus.FAILED]], - datetime.timedelta())) - Log.info('') + Log.error("{0}: {0.message}", job) + Log.error( + "Total wasted time: {0}", + sum( + [ + job.consumed_time + for job in self.jobs[DiagnosticStatus.FAILED] + ], + datetime.timedelta(), + ), + ) + Log.info("") if self.uploads_failed: self.had_errors = True - Log.error('Failed uploads') - Log.error('--------------') + Log.error("Failed uploads") + Log.error("--------------") for file_object in self.uploads_failed: - Log.error('{0}', file_object.remote_file) - Log.info('') + Log.error("{0}", file_object.remote_file) + Log.info("") @staticmethod def _run_job(job): start_time = datetime.datetime.now() try: - Log.info('Starting {0}', job) + Log.info("Starting {0}", job) job.status = DiagnosticStatus.RUNNING job.compute() except Exception as ex: job.consumed_time = datetime.datetime.now() - start_time exc_type, _, exc_traceback = sys.exc_info() - job.message = '{0}\n{1}'.format(ex, ''.join(traceback.format_tb(exc_traceback))) - Log.error('Job {0} failed ({2}): {1}', job, job.message, exc_type) + job.message = "{0}\n{1}".format( + ex, "".join(traceback.format_tb(exc_traceback)) + ) + Log.error("Job {0} failed ({2}): {1}", job, job.message, exc_type) job.status = DiagnosticStatus.FAILED return False job.consumed_time = datetime.datetime.now() - start_time - Log.result('Finished {0}', job) + Log.result("Finished {0}", job) job.status = DiagnosticStatus.COMPLETED return True + count = 0 @staticmethod @@ -287,9 +336,15 @@ class WorkManager(object): @staticmethod def _register_stats_diagnostics(): - from earthdiagnostics.statistics.monthlypercentile import MonthlyPercentile - from earthdiagnostics.statistics.climatologicalpercentile import ClimatologicalPercentile - from earthdiagnostics.statistics.daysoverpercentile import DaysOverPercentile + from earthdiagnostics.statistics.monthlypercentile import ( + MonthlyPercentile, + ) + from earthdiagnostics.statistics.climatologicalpercentile import ( + ClimatologicalPercentile, + ) + from earthdiagnostics.statistics.daysoverpercentile import ( + DaysOverPercentile, + ) from earthdiagnostics.statistics.discretize import Discretize Diagnostic.register(MonthlyPercentile) @@ -300,14 +355,22 @@ class WorkManager(object): @staticmethod def _register_general_diagnostics(): from earthdiagnostics.general.attribute import Attribute - from earthdiagnostics.general.timemean import DailyMean, MonthlyMean, YearlyMean + from earthdiagnostics.general.timemean import ( + DailyMean, + MonthlyMean, + YearlyMean, + ) from earthdiagnostics.general.module import Module from earthdiagnostics.general.rewrite import Rewrite from earthdiagnostics.general.relink import Relink from earthdiagnostics.general.relinkall import RelinkAll from earthdiagnostics.general.scale import Scale - from earthdiagnostics.general.verticalmeanmetersiris import VerticalMeanMetersIris - from earthdiagnostics.general.simplify_dimensions import SimplifyDimensions + from earthdiagnostics.general.verticalmeanmetersiris import ( + VerticalMeanMetersIris, + ) + from earthdiagnostics.general.simplify_dimensions import ( + SimplifyDimensions, + ) Diagnostic.register(DailyMean) Diagnostic.register(MonthlyMean) @@ -378,7 +441,8 @@ class Downloader(object): """ Download manager for EarthDiagnostics - We are not using a ThreadPoolExecutor because we want to be able to control priorities in the download + We are not using a ThreadPoolExecutor because we want to be able to + control priorities in the download """ def __init__(self): @@ -407,12 +471,18 @@ class Downloader(object): return time.sleep(0.1) continue - self._downloads.sort(key=cmp_to_key(Downloader._prioritize)) + self._downloads.sort( + key=cmp_to_key(Downloader._prioritize) + ) datafile = self._downloads[0] self._downloads.remove(datafile) datafile.download() except Exception as ex: - Log.critical('Unhandled error at downloader: {0}\n{1}', ex, traceback.print_exc()) + Log.critical( + "Unhandled error at downloader: {0}\n{1}", + ex, + traceback.print_exc(), + ) @staticmethod def _suscribers_waiting(datafile): @@ -426,7 +496,9 @@ class Downloader(object): @staticmethod def _prioritize(datafile1, datafile2): - waiting = Downloader._suscribers_waiting(datafile1) - Downloader._suscribers_waiting(datafile2) + waiting = Downloader._suscribers_waiting( + datafile1 + ) - Downloader._suscribers_waiting(datafile2) if waiting: return -waiting diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000..5566188f --- /dev/null +++ b/setup.cfg @@ -0,0 +1,31 @@ +[aliases] +test=pytest + +[build_sphinx] +source-dir = doc/ +build-dir = doc/build +all_files = 1 +builder = html + +[tool:pytest] +addopts = + --flake8 + --doctest-modules + --ignore=earthdiagnostics/cmor_tables/ + --cov=earthdiagnostics + --cov-report=term + --cov-report=xml:test/report/python3/coverage.xml + --cov-report=html:test/report/python3/coverage_html + --junit-xml=test/report/python3/report.xml + --html=test/report/python3/report.html +env = + MPLBACKEND = Agg +flake8-ignore = + doc/source/conf.py ALL +log_level = WARNING + +[coverage:run] +parallel = true + +[pydocstyle] +convention = numpy diff --git a/setup.py b/setup.py index 44c44346..b5a06b99 100644 --- a/setup.py +++ b/setup.py @@ -3,101 +3,65 @@ """Installation script for EarthDiagnostics package""" from os import path -import sys -from setuptools import setup, Command, find_packages +from setuptools import setup, find_packages here = path.abspath(path.dirname(__file__)) # Get the version number from the relevant file -with open(path.join(here, 'VERSION')) as f: +with open(path.join(here, "VERSION")) as f: version = f.read().strip() -class RunTests(Command): - """Class to run tests and generate reports.""" - - user_options = [] - - def initialize_options(self): - """Do nothing.""" - - def finalize_options(self): - """Do nothing.""" - - def run(self): - """Run tests and generate a coverage report.""" - import pytest - - version = sys.version_info[0] - report_dir = 'test/report/python{}'.format(version) - args = [ - 'test', - 'earthdiagnostics', # for doctests - '--ignore=test/report', - '--doctest-modules', - '--cov=earthdiagnostics', - '--cov-report=term', - '--cov-report=html:{}/coverage_html'.format(report_dir), - '--cov-report=xml:{}/coverage.xml'.format(report_dir), - '--junit-xml={}/report.xml'.format(report_dir), - '--html={}/report.html'.format(report_dir), - ] - errno = pytest.main(args) - sys.exit(errno) - - REQUIREMENTS = { - 'setup': [ - 'pyproj', + "setup": [ + "pyproj", + 'pytest-runner', + 'setuptools_scm', ], - 'install': [ - 'bscearth.utils', - 'cdo>=1.3.4', - 'cfgrib', - 'dask[array]', - 'diagonals', - 'exrex', - 'netCDF4', - 'nco>=0.0.3', - 'numba', - 'numpy', - 'psutil', - 'openpyxl', - 'scitools-iris>=2.2', - 'six', - 'xxhash', + "install": [ + "bscearth.utils", + "cdo>=1.3.4", + "cfgrib", + "dask[array]", + "diagonals", + "exrex", + "netCDF4", + "nco>=0.0.3", + "numba", + "numpy", + "psutil", + "openpyxl", + "scitools-iris>=2.2", + "six", + "xxhash", ], - 'test': [ - 'coverage', - 'mock', - 'pycodestyle', - 'pytest', + "test": [ + "mock", + 'pytest>=3.9', 'pytest-cov', + 'pytest-env', + 'pytest-flake8', 'pytest-html', - ] + 'pytest-metadata>=1.5.1', + ], } setup( - name='earthdiagnostics', - license='GNU GPL v3', - platforms=['GNU/Linux Debian'], + name="earthdiagnostics", + license="GNU GPL v3", + platforms=["GNU/Linux Debian"], version=version, - description='EarthDiagnostics', - author='BSC-CNS Earth Sciences Department', - author_email='javier.vegas@bsc.es', - url='http://www.bsc.es/projects/earthsciences/autosubmit/', - keywords=['climate', 'weather', 'diagnostic'], - setup_requires=REQUIREMENTS['setup'], - install_requires=REQUIREMENTS['install'], - test_requires=REQUIREMENTS['test'], - extras_require={ - 'develop': REQUIREMENTS['test'], - }, + description="EarthDiagnostics", + author="BSC-CNS Earth Sciences Department", + author_email="javier.vegas@bsc.es", + url="http://www.bsc.es/projects/earthsciences/autosubmit/", + keywords=["climate", "weather", "diagnostic"], + setup_requires=REQUIREMENTS["setup"], + install_requires=REQUIREMENTS["install"], + tests_requires=REQUIREMENTS["test"], + extras_require={"develop": REQUIREMENTS["test"], }, packages=find_packages(), include_package_data=True, - scripts=['bin/earthdiags'], - cmdclass={ - 'test': RunTests, - }, + scripts=["bin/earthdiags"], ) diff --git a/test/integration/test_cmorizer.py b/test/integration/test_cmorizer.py index b6b35e4a..74a04024 100644 --- a/test/integration/test_cmorizer.py +++ b/test/integration/test_cmorizer.py @@ -33,15 +33,17 @@ class TestCmorizer(TestCase): def _get_variable(self, variable, silent=False): mock_variable = Mock() mock_variable.short_name = variable - mock_variable.domain = 'domain' + mock_variable.domain = "domain" return mock_variable def _get_file_path(self, *args, **kwargs): - return os.path.join(self.tmp_dir, args[3], '{0[3]}.nc'.format(args)) + return os.path.join(self.tmp_dir, args[3], "{0[3]}.nc".format(args)) def _get_file_path_grib(self, *args, **kwargs): - return os.path.join(self.tmp_dir, args[3], str(args[6]), '{0[3]}.nc'.format(args)) + return os.path.join( + self.tmp_dir, args[3], str(args[6]), "{0[3]}.nc".format(args) + ) def setUp(self): """Prepare tests""" @@ -50,26 +52,36 @@ class TestCmorizer(TestCase): self.data_manager = Mock() self.data_manager.is_cmorized.return_value = False - self.data_manager.config.data_dir = os.path.join(self.tmp_dir, 'data') - self.data_manager.config.scratch_dir = os.path.join(self.tmp_dir, 'scratch') + self.data_manager.config.data_dir = os.path.join(self.tmp_dir, "data") + self.data_manager.config.scratch_dir = os.path.join( + self.tmp_dir, "scratch" + ) TempFile.scratch_folder = self.data_manager.config.scratch_dir - self.data_manager.config.data_convention = create_autospec(DataConvention) - self.data_manager.config.data_convention.name = 'data_convention' - self.data_manager.config.data_convention.lat_name = 'lat' - self.data_manager.config.data_convention.lon_name = 'lon' - self.data_manager.config.data_convention.get_file_path = self._get_file_path + self.data_manager.config.data_convention = create_autospec( + DataConvention + ) + self.data_manager.config.data_convention.name = "data_convention" + self.data_manager.config.data_convention.lat_name = "lat" + self.data_manager.config.data_convention.lon_name = "lon" + self.data_manager.config.data_convention.get_file_path = ( + self._get_file_path + ) - self.data_manager.config.var_manager.get_variable_and_alias = self._get_variable_and_alias + self.data_manager.config.var_manager.get_variable_and_alias = ( + self._get_variable_and_alias + ) self.data_manager.config.var_manager.get_variable = self._get_variable self.data_manager.variable_list = self.data_manager.config.var_manager - self.data_manager.config.experiment.expid = 'expid' - self.data_manager.config.experiment.model = 'model' - self.data_manager.config.experiment.experiment_name = 'experiment_name' + self.data_manager.config.experiment.expid = "expid" + self.data_manager.config.experiment.model = "model" + self.data_manager.config.experiment.experiment_name = "experiment_name" self.data_manager.config.experiment.num_chunks = 1 self.data_manager.config.experiment.chunk_size = 1 - self.data_manager.config.experiment.institute = 'institute' - self.data_manager.config.experiment.get_member_str.return_value = 'member' + self.data_manager.config.experiment.institute = "institute" + self.data_manager.config.experiment.get_member_str.return_value = ( + "member" + ) self.data_manager.config.experiment.atmos_timestep = 6 self.data_manager.config.cmor.force = False @@ -77,41 +89,89 @@ class TestCmorizer(TestCase): self.data_manager.config.cmor.atmosphere = True self.data_manager.config.cmor.use_grib = True self.data_manager.config.cmor.filter_files = [] - self.data_manager.config.cmor.associated_experiment = 'associated_experiment' - self.data_manager.config.cmor.initialization_method = 'initialization_method' - self.data_manager.config.cmor.initialization_description = 'initialization_description' - self.data_manager.config.cmor.physics_version = 'physics_version' - self.data_manager.config.cmor.physics_description = 'physics_description' - self.data_manager.config.cmor.initialization_description = 'initialization_description' - self.data_manager.config.cmor.associated_model = 'initialization_description' - self.data_manager.config.cmor.source = 'source' - self.data_manager.config.cmor.get_requested_codes.return_value = {228, 142, 143, 201, 202, 129, 169, 180} - self.data_manager.config.cmor.get_variables.return_value = {228, 142, 143, 201, 202, 129, 169, 180} + self.data_manager.config.cmor.associated_experiment = ( + "associated_experiment" + ) + self.data_manager.config.cmor.initialization_method = ( + "initialization_method" + ) + self.data_manager.config.cmor.initialization_description = ( + "initialization_description" + ) + self.data_manager.config.cmor.physics_version = "physics_version" + self.data_manager.config.cmor.physics_description = ( + "physics_description" + ) + self.data_manager.config.cmor.initialization_description = ( + "initialization_description" + ) + self.data_manager.config.cmor.associated_model = ( + "initialization_description" + ) + self.data_manager.config.cmor.source = "source" + self.data_manager.config.cmor.get_requested_codes.return_value = { + 228, + 142, + 143, + 201, + 202, + 129, + 169, + 180, + } + self.data_manager.config.cmor.get_variables.return_value = { + 228, + 142, + 143, + 201, + 202, + 129, + 169, + 180, + } self.data_manager.config.cmor.get_levels.return_value = None os.makedirs(self.data_manager.config.data_dir) os.makedirs(self.data_manager.config.scratch_dir) - def _create_ocean_files(self, filename, tar_name, gzip=False, backup=False): - folder_path = os.path.join(self.data_manager.config.data_dir, 'expid', 'original_files', '19900101', 'member', - 'outputs') + def _create_ocean_files( + self, filename, tar_name, gzip=False, backup=False + ): + folder_path = os.path.join( + self.data_manager.config.data_dir, + "expid", + "original_files", + "19900101", + "member", + "outputs", + ) file_path, filename = self._create_file(folder_path, filename, gzip) if backup: - filename = os.path.join('backup', filename) + filename = os.path.join("backup", filename) - tar = tarfile.TarFile(os.path.join(folder_path, tar_name), mode='w') + tar = tarfile.TarFile(os.path.join(folder_path, tar_name), mode="w") tar.add(file_path, arcname=filename, recursive=False) tar.close() os.remove(file_path) def _create_mma_files(self, filename, tar_name, gzip=False): - folder_path = os.path.join(self.data_manager.config.data_dir, 'expid', 'original_files', '19900101', 'member', - 'outputs') - filepath_gg, filename_gg = self._create_file(folder_path, filename.replace('??', 'GG'), gzip) - filepath_sh, filename_sh = self._create_file(folder_path, filename.replace('??', 'SH'), gzip) + folder_path = os.path.join( + self.data_manager.config.data_dir, + "expid", + "original_files", + "19900101", + "member", + "outputs", + ) + filepath_gg, filename_gg = self._create_file( + folder_path, filename.replace("??", "GG"), gzip + ) + filepath_sh, filename_sh = self._create_file( + folder_path, filename.replace("??", "SH"), gzip + ) - tar = tarfile.TarFile(os.path.join(folder_path, tar_name), mode='w') + tar = tarfile.TarFile(os.path.join(folder_path, tar_name), mode="w") tar.add(filepath_gg, arcname=filename_gg, recursive=False) tar.add(filepath_sh, arcname=filename_sh, recursive=False) tar.close() @@ -119,36 +179,64 @@ class TestCmorizer(TestCase): os.remove(filepath_sh) def _create_file(self, folder_path, filename, gzip): - var1 = self._create_sample_cube('Variable 1', 'var1', threed=False, time_bounds=True) - var2 = self._create_sample_cube('Variable 2', 'var2', threed=True, time_bounds=True) + var1 = self._create_sample_cube( + "Variable 1", "var1", threed=False, time_bounds=True + ) + var2 = self._create_sample_cube( + "Variable 2", "var2", threed=True, time_bounds=True + ) if not os.path.isdir(folder_path): os.makedirs(folder_path) file_path = os.path.join(folder_path, filename) iris.save((var1, var2), file_path, zlib=True) if gzip: import subprocess - process = subprocess.Popen(('gzip', file_path), stdout=subprocess.PIPE) + + process = subprocess.Popen( + ("gzip", file_path), stdout=subprocess.PIPE + ) comunicate = process.communicate() file_path = "{0}.gz".format(file_path) filename = "{0}.gz".format(filename) if process.returncode != 0: - raise Exception('Can not compress: {0}'.format(comunicate)) + raise Exception("Can not compress: {0}".format(comunicate)) return file_path, filename def _create_sample_cube(self, long_name, var_name, threed, time_bounds): coord_data = np.array([1, 2], np.float) - lat = DimCoord(coord_data, standard_name='latitude', long_name='latitude', var_name='lat', - units='degrees_north') - lon = DimCoord(coord_data, standard_name='longitude', long_name='longitude', var_name='lon', - units='degrees_east') - time = DimCoord(coord_data, standard_name='time', long_name='time', var_name='time', - units='days since 1950-01-01') + lat = DimCoord( + coord_data, + standard_name="latitude", + long_name="latitude", + var_name="lat", + units="degrees_north", + ) + lon = DimCoord( + coord_data, + standard_name="longitude", + long_name="longitude", + var_name="lon", + units="degrees_east", + ) + time = DimCoord( + coord_data, + standard_name="time", + long_name="time", + var_name="time", + units="days since 1950-01-01", + ) if time_bounds: time.bounds = np.array([[0.5, 1.5], [1.5, 2.5]], np.float) if threed: data = np.random.rand(2, 2, 2, 2).astype(np.float) - depth = DimCoord(coord_data, standard_name='depth', long_name='Depth', var_name='lev', units='m') + depth = DimCoord( + coord_data, + standard_name="depth", + long_name="Depth", + var_name="lev", + units="m", + ) else: data = np.random.rand(2, 2, 2).astype(np.float) @@ -164,43 +252,121 @@ class TestCmorizer(TestCase): """Clean up after tests""" shutil.rmtree(self.tmp_dir) - def _test_ocean_cmor(self, success=True, error=False, critical=False, warnings=False, message='', check_vars=None): - self._test_cmorization(success=success, error=error, critical=critical, warnings=warnings, message=message, - ocean=True, check_vars=check_vars) + def _test_ocean_cmor( + self, + success=True, + error=False, + critical=False, + warnings=False, + message="", + check_vars=None, + ): + self._test_cmorization( + success=success, + error=error, + critical=critical, + warnings=warnings, + message=message, + ocean=True, + check_vars=check_vars, + ) - def _test_atmos_cmor(self, success=True, error=False, critical=False, warnings=False, message='', check_vars=None): - self._test_cmorization(success=success, error=error, critical=critical, warnings=warnings, message=message, - ocean=False, check_vars=check_vars) + def _test_atmos_cmor( + self, + success=True, + error=False, + critical=False, + warnings=False, + message="", + check_vars=None, + ): + self._test_cmorization( + success=success, + error=error, + critical=critical, + warnings=warnings, + message=message, + ocean=False, + check_vars=check_vars, + ) - def _test_cmorization(self, success=True, error=False, critical=False, warnings=False, message='', ocean=True, - check_vars=None): + def _test_cmorization( + self, + success=True, + error=False, + critical=False, + warnings=False, + message="", + ocean=True, + check_vars=None, + ): self._check_logs(critical, error, message, ocean, success, warnings) if check_vars: for variable, status in six.iteritems(check_vars): if status: - self.assertTrue(os.path.isfile(os.path.join(self.tmp_dir, variable, '{}.nc'.format(variable)))) + self.assertTrue( + os.path.isfile( + os.path.join( + self.tmp_dir, + variable, + "{}.nc".format(variable), + ) + ) + ) else: - self.assertFalse(os.path.isfile(os.path.join(self.tmp_dir, variable, '{}.nc'.format(variable)))) + self.assertFalse( + os.path.isfile( + os.path.join( + self.tmp_dir, + variable, + "{}.nc".format(variable), + ) + ) + ) def _check_logs(self, critical, error, message, ocean, success, warnings): if six.PY3: with self.assertLogs(log.Log.log) as cmd: - cmorizer = Cmorizer(self.data_manager, '19900101', 0) + cmorizer = Cmorizer(self.data_manager, "19900101", 0) if ocean: cmorizer.cmorize_ocean() else: cmorizer.cmorize_atmos() if message: - self.assertTrue([record for record in cmd.records if record.message == message]) + self.assertTrue( + [ + record + for record in cmd.records + if record.message == message + ] + ) else: - for level, value in six.iteritems({log.Log.RESULT: success, log.Log.ERROR: error, - log.Log.CRITICAL: critical, log.Log.WARNING: warnings}): + for level, value in six.iteritems( + { + log.Log.RESULT: success, + log.Log.ERROR: error, + log.Log.CRITICAL: critical, + log.Log.WARNING: warnings, + } + ): if value: - self.assertTrue([record for record in cmd.records if record.levelno == level]) + self.assertTrue( + [ + record + for record in cmd.records + if record.levelno == level + ] + ) else: - self.assertFalse([record for record in cmd.records if record.levelno == level]) + self.assertFalse( + [ + record + for record in cmd.records + if record.levelno == level + ] + ) else: - cmorizer = Cmorizer(self.data_manager, '19900101', 0) + cmorizer = Cmorizer(self.data_manager, "19900101", 0) if ocean: cmorizer.cmorize_ocean() else: @@ -209,36 +375,51 @@ class TestCmorizer(TestCase): def test_skip_ocean_cmorization(self): """Test ocean cmorization flag disabled option""" self.data_manager.config.cmor.ocean = False - self._test_ocean_cmor(message='Skipping ocean cmorization due to configuration') + self._test_ocean_cmor( + message="Skipping ocean cmorization due to configuration" + ) def test_skip_atmos_cmorization(self): """Test atmos cmorization flag disabled option""" self.data_manager.config.cmor.atmosphere = False if six.PY3: with self.assertLogs(log.Log.log) as cmd: - cmorizer = Cmorizer(self.data_manager, '19900101', 0) + cmorizer = Cmorizer(self.data_manager, "19900101", 0) cmorizer.cmorize_atmos() - self.assertTrue([record for record in cmd.records if - record.message == 'Skipping atmosphere cmorization due to configuration']) + self.assertTrue( + [ + record + for record in cmd.records + if record.message + == "Skipping atmosphere cmorization due to configuration" + ] + ) else: - cmorizer = Cmorizer(self.data_manager, '19900101', 0) + cmorizer = Cmorizer(self.data_manager, "19900101", 0) cmorizer.cmorize_ocean() def test_skip_when_cmorized(self): """Test cmorization skipped if already done""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) self.data_manager.is_cmorized.return_value = True - self._test_ocean_cmor(message='No need to unpack file 1/1') + self._test_ocean_cmor(message="No need to unpack file 1/1") def test_skip_when_not_requested(self): """Test cmorization skipped if chunk is not requested""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self.data_manager.config.cmor.chunk_cmorization_requested.return_value = False - self._test_ocean_cmor(message='No need to unpack file 1/1') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) + chunk = self.data_manager.config.cmor.chunk_cmorization_requested + chunk.return_value = False + self._test_ocean_cmor(message="No need to unpack file 1/1") def test_force(self): """Test cmorization force works""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) self.data_manager.is_cmorized.return_value = True self.data_manager.config.cmor.force = True self._test_ocean_cmor() @@ -249,80 +430,123 @@ class TestCmorizer(TestCase): def test_ocean_cmorization_not_vars_requested(self): """Test ocean cmorization report success if no vars qhere requested""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) self.data_manager.config.cmor.any_required.return_value = False - self._test_ocean_cmor(check_vars={'var1': False, 'var2': False}) + self._test_ocean_cmor(check_vars={"var1": False, "var2": False}) def test_ocean_cmorization_no_vars_recognized(self): """Test ocean cmorization report success if no vars where recognized""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) def not_recognized(*args): return None, None - self.data_manager.config.var_manager.get_variable_and_alias = not_recognized - self._test_ocean_cmor(check_vars={'var1': False, 'var2': False}) + + self.data_manager.config.var_manager.get_variable_and_alias = ( + not_recognized + ) + self._test_ocean_cmor(check_vars={"var1": False, "var2": False}) def test_ocean_cmorization_var2_not_requested(self): """Test ocean cmorization with var2 not recognized""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) def _reject_var2(cmor_var): - return cmor_var.short_name != 'var2' + return cmor_var.short_name != "var2" self.data_manager.config.cmor.cmorize = _reject_var2 - self._test_ocean_cmor(check_vars={'var1': True, 'var2': False}) + self._test_ocean_cmor(check_vars={"var1": True, "var2": False}) def test_ocean_cmorization(self): """Test basic ocean cmorization""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) def test_ocean_cmorization_with_filter(self): """Test ocean cmorization filtering files""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self.data_manager.config.cmor.filter_files = ['expid'] - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) + self.data_manager.config.cmor.filter_files = ["expid"] + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) def test_ocean_cmorization_with_bad_filter(self): """Test ocean cmorization fails if a bad filter is added""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar') - self.data_manager.config.cmor.filter_files = ['badfilter'] - self._test_ocean_cmor(success=False, error=True, check_vars={'var1': False, 'var2': False}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", "MMO_19900101-19900131.tar" + ) + self.data_manager.config.cmor.filter_files = ["badfilter"] + self._test_ocean_cmor( + success=False, + error=True, + check_vars={"var1": False, "var2": False}, + ) def test_ocean_cmorization_gzip(self): """Test ocean cmorization if tars are also zipped""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar', gzip=True) - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", + "MMO_19900101-19900131.tar", + gzip=True, + ) + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) def test_ocean_cmorization_backup(self): """Test ocean cmorization when files are in backup path""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'MMO_19900101-19900131.tar', backup=True) - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", + "MMO_19900101-19900131.tar", + backup=True, + ) + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) - def test_ocean_cmorization_PPO(self): + def test_ocean_cmorization_ppo(self): """Test ocean cmorization when files are PPO""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'PPO_expid_1D_xx_19900101_19900131.tar') - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", + "PPO_expid_1D_xx_19900101_19900131.tar", + ) + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) def test_ocean_cmorization_diags(self): """Test ocean cmorization when files are diags""" - self._create_ocean_files('expid_1d_19900101_19900131.nc', 'diags_expid_1D_xx_19900101_19900131.tar') - self._test_ocean_cmor(check_vars={'var1': True, 'var2': True}) + self._create_ocean_files( + "expid_1d_19900101_19900131.nc", + "diags_expid_1D_xx_19900101_19900131.tar", + ) + self._test_ocean_cmor(check_vars={"var1": True, "var2": True}) def test_atmos_cmorization(self): """Test basic atmos cmorization from nc""" - self._create_mma_files('MMA_1d_??_19900101_19900131.nc', 'MMA_expid_19901101_fc0_19900101-19900131.tar') - self._test_atmos_cmor(check_vars={'var1': True, 'var2': True}) + self._create_mma_files( + "MMA_1d_??_19900101_19900131.nc", + "MMA_expid_19901101_fc0_19900101-19900131.tar", + ) + self._test_atmos_cmor(check_vars={"var1": True, "var2": True}) def test_skip_when_not_requested_mma(self): """Test atmos cmorization is skipped if chunk is not requested""" - self._create_mma_files('MMA_1d_??_19900101_19900131.nc', 'MMA_expid_19901101_fc0_19900101-19900131.tar') - self.data_manager.config.cmor.chunk_cmorization_requested.return_value = False - self._test_atmos_cmor(message='No need to unpack file 1/1') + self._create_mma_files( + "MMA_1d_??_19900101_19900131.nc", + "MMA_expid_19901101_fc0_19900101-19900131.tar", + ) + chunk = self.data_manager.config.cmor.chunk_cmorization_requested + chunk.return_value = (False) + self._test_atmos_cmor(message="No need to unpack file 1/1") def test_force_mma(self): """Test force atmos cmorization""" - self._create_mma_files('MMA_1d_??_19900101_19900131.nc', 'MMA_expid_19901101_fc0_19900101-19900131.tar') + self._create_mma_files( + "MMA_1d_??_19900101_19900131.nc", + "MMA_expid_19901101_fc0_19900101-19900131.tar", + ) self.data_manager.is_cmorized.return_value = True self.data_manager.config.cmor.force = True self._test_atmos_cmor() @@ -334,76 +558,117 @@ class TestCmorizer(TestCase): def _create_grib_files(self, filename, month): filename = filename.format(month) coord_data = np.array([0, 1], np.float) - folder_path = os.path.join(self.data_manager.config.data_dir, 'expid', 'original_files', '19900101', 'member', - 'outputs') - self._create_file_for_grib(coord_data, folder_path, filename.replace('??', 'GG'), [142, 143, 129, 169, 180], - month) - self._create_file_for_grib(coord_data, folder_path, filename.replace('??', 'SH'), [201, 202], - month) - - def _create_file_for_grib(self, coord_data, folder_path, filename, codes, month): - lat = DimCoord(coord_data, standard_name='latitude', long_name='latitude', var_name='lat', - units='degrees_north') - lon = DimCoord(coord_data, standard_name='longitude', long_name='longitude', var_name='lon', - units='degrees_east') + folder_path = os.path.join( + self.data_manager.config.data_dir, + "expid", + "original_files", + "19900101", + "member", + "outputs", + ) + self._create_file_for_grib( + coord_data, + folder_path, + filename.replace("??", "GG"), + [142, 143, 129, 169, 180], + month, + ) + self._create_file_for_grib( + coord_data, + folder_path, + filename.replace("??", "SH"), + [201, 202], + month, + ) + + def _create_file_for_grib( + self, coord_data, folder_path, filename, codes, month + ): + lat = DimCoord( + coord_data, + standard_name="latitude", + long_name="latitude", + var_name="lat", + units="degrees_north", + ) + lon = DimCoord( + coord_data, + standard_name="longitude", + long_name="longitude", + var_name="lon", + units="degrees_east", + ) month_days = calendar.monthrange(1990, month)[1] month -= 1 - time_data = np.arange(0.25, month_days + 0.25, 0.25, np.float) + month * 31 - time = DimCoord(time_data, standard_name='time', long_name='time', var_name='time', - units='days since 1990-01-01 00:00:00') + time_data = ( + np.arange(0.25, month_days + 0.25, 0.25, np.float) + month * 31 + ) + time = DimCoord( + time_data, + standard_name="time", + long_name="time", + var_name="time", + units="days since 1990-01-01 00:00:00", + ) variables = [] for code in codes: - var = iris.cube.Cube(np.ones((month_days * 4, 2, 2), np.float) * code, - long_name='Variable {}'.format(code), - var_name='var{}'.format(code)) + var = iris.cube.Cube( + np.ones((month_days * 4, 2, 2), np.float) * code, + long_name="Variable {}".format(code), + var_name="var{}".format(code), + ) for x, data in enumerate(time_data): var.data[x, ...] += data var.add_dim_coord(time, 0) var.add_dim_coord(lat, 1) var.add_dim_coord(lon, 2) - var.attributes['table'] = np.int32(128) - var.attributes['code'] = np.int32(code) + var.attributes["table"] = np.int32(128) + var.attributes["code"] = np.int32(code) variables.append(var) if not os.path.isdir(folder_path): os.makedirs(folder_path) file_path = os.path.join(folder_path, filename) - iris.save(variables, file_path, zlib=True, local_keys=('table', 'code')) + iris.save( + variables, file_path, zlib=True, local_keys=("table", "code") + ) Utils.cdo().settaxis( - '1990-0{}-01,06:00,6hour'.format(month + 1), + "1990-0{}-01,06:00,6hour".format(month + 1), input=file_path, - output=file_path.replace('.nc', '.grb'), - options='-f grb2' + output=file_path.replace(".nc", ".grb"), + options="-f grb2", ) os.remove(file_path) def test_grib_cmorization(self): """Test atmos cmorization from grib""" - self.data_manager.config.data_convention.get_file_path = self._get_file_path_grib + self.data_manager.config.data_convention.get_file_path = ( + self._get_file_path_grib + ) self.data_manager.config.experiment.chunk_size = 2 self.data_manager.get_file_path = self._get_file_path_grib - self._create_grib_files('ICM??expid+19900{}.nc', 1) - self._create_grib_files('ICM??expid+19900{}.nc', 2) + self._create_grib_files("ICM??expid+19900{}.nc", 1) + self._create_grib_files("ICM??expid+19900{}.nc", 2) self._test_atmos_cmor() variables = { - 'CP': 143, - 'EWSS': 180, - 'LSP': 142, - 'MN2T': 202, - 'MX2T': 201, - 'SSRD': 169, - 'TP': 228, - 'Z': 129 + "CP": 143, + "EWSS": 180, + "LSP": 142, + "MN2T": 202, + "MX2T": 201, + "SSRD": 169, + "TP": 228, + "Z": 129, } for var, code in six.iteritems(variables): self.assertTrue(os.path.isdir(os.path.join(self.tmp_dir, var))) base_data = np.ones((2, 2), np.float) * code - if var in ('EWSS', 'TP', 'MN2T', 'MX2T', 'SSRD', 'TP'): - if var == 'MX2T': + if var in ("EWSS", "TP", "MN2T", "MX2T", "SSRD", "TP"): + if var == "MX2T": month_offsets = np.array((16, 45.5)) daily_offsets = np.arange(1.0, 60.0) - elif var == 'MN2T': + elif var == "MN2T": month_offsets = np.array((15.25, 44.75)) daily_offsets = np.arange(0.25, 59.25) else: @@ -434,17 +699,27 @@ class TestCmorizer(TestCase): daily_offsets /= factor hourly_offsets /= factor - monthly = iris.load_cube(os.path.join(self.tmp_dir, var, 'mon', '{}.nc'.format(var))) - self._test_data(monthly, base_data, month_offsets, var, 'Month') + monthly = iris.load_cube( + os.path.join(self.tmp_dir, var, "mon", "{}.nc".format(var)) + ) + self._test_data(monthly, base_data, month_offsets, var, "Month") - daily = iris.load_cube(os.path.join(self.tmp_dir, var, 'day', '{}.nc'.format(var))) - self._test_data(daily, base_data, daily_offsets, var, 'Day') + daily = iris.load_cube( + os.path.join(self.tmp_dir, var, "day", "{}.nc".format(var)) + ) + self._test_data(daily, base_data, daily_offsets, var, "Day") - hourly = iris.load_cube(os.path.join(self.tmp_dir, var, '6hr', '{}.nc'.format(var))) - self._test_data(hourly, base_data, hourly_offsets, var, 'Hour') + hourly = iris.load_cube( + os.path.join(self.tmp_dir, var, "6hr", "{}.nc".format(var)) + ) + self._test_data(hourly, base_data, hourly_offsets, var, "Hour") def _test_data(self, data, base_data, offsets, var, freq): - self.assertEqual(data.coord('time').shape, (len(offsets),)) + self.assertEqual(data.coord("time").shape, (len(offsets),)) for x, offset in enumerate(offsets): - self.assertTrue(np.allclose(data.data[x, ...], base_data + offset), - '{} {} data wrong for {}: {}'.format(freq, x, var, data.data[x, ...] - base_data)) + self.assertTrue( + np.allclose(data.data[x, ...], base_data + offset), + "{} {} data wrong for {}: {}".format( + freq, x, var, data.data[x, ...] - base_data + ), + ) diff --git a/test/unit/data_convention/test_data_convention.py b/test/unit/data_convention/test_data_convention.py index 152198ec..26abd6c4 100644 --- a/test/unit/data_convention/test_data_convention.py +++ b/test/unit/data_convention/test_data_convention.py @@ -11,48 +11,58 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestDataConvention(TestCase): - def setUp(self): self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.config = Mock() - self.config.frequency = 'mon' - self.config.experiment.experiment_name = 'expname' + self.config.frequency = "mon" + self.config.experiment.experiment_name = "expname" self.config.data_dir = self.tmp_dir - self.config.experiment.expid = 'expid' - self.config.cmor.activity = 'activity' + self.config.experiment.expid = "expid" + self.config.cmor.activity = "activity" - self.config.experiment.institute = 'institute' - self.config.experiment.model = 'model' + self.config.experiment.institute = "institute" + self.config.experiment.model = "model" self.config.experiment.atmos_timestep = 3 self.config.experiment.ocean_timestep = 6 self.config.cmor.append_startdate = False - self.convention = DataConvention('name', self.config) + self.convention = DataConvention("name", self.config) def tearDown(self): shutil.rmtree(self.tmp_dir) def test_get_filename(self): with self.assertRaises(NotImplementedError): - self.convention.get_file_name(None, None, None, None, None, None, None, None, None, None) + self.convention.get_file_name( + None, None, None, None, None, None, None, None, None, None + ) def test_get_startdate_path(self): - self.assertEqual(self.convention.get_startdate_path('19900101'), - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/expname')) + self.assertEqual( + self.convention.get_startdate_path("19900101"), + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/expname", + ), + ) def test_experiment_name(self): - self.assertEqual(self.convention.experiment_name('19900101'), - 'expname') + self.assertEqual( + self.convention.experiment_name("19900101"), "expname" + ) def test_experiment_name_append(self): self.config.cmor.append_startdate = True - self.assertEqual(self.convention.experiment_name('19900101'), - 'expnameS19900101') + self.assertEqual( + self.convention.experiment_name("19900101"), "expnameS19900101" + ) def test_get_cmor_folder_path(self): with self.assertRaises(NotImplementedError): - self.convention.get_cmor_folder_path(None, None, None, None, None, None, None) + self.convention.get_cmor_folder_path( + None, None, None, None, None, None, None + ) def test_create_links(self): with self.assertRaises(NotImplementedError): @@ -66,131 +76,266 @@ class TestDataConvention(TestCase): with self.assertRaises(NotImplementedError): self.convention.is_cmorized(None, None, None, None) - @mock.patch('earthdiagnostics.data_convention.DataConvention.get_cmor_folder_path', autospec=True) - @mock.patch('earthdiagnostics.data_convention.DataConvention.get_file_name', autospec=True) + @mock.patch( + "earthdiagnostics.data_convention.DataConvention.get_cmor_folder_path", + autospec=True, + ) + @mock.patch( + "earthdiagnostics.data_convention.DataConvention.get_file_name", + autospec=True, + ) def test_get_file_path(self, mock_filename, mock_cmor_path): - mock_cmor_path.return_value = 'path' - mock_filename.return_value = 'filename' + mock_cmor_path.return_value = "path" + mock_filename.return_value = "filename" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - filepath = self.convention.get_file_path('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, 1, 'mon') - self.assertEqual('path/filename', filepath) - - @mock.patch('earthdiagnostics.data_convention.DataConvention.get_cmor_folder_path', autospec=True) - @mock.patch('earthdiagnostics.data_convention.DataConvention.get_file_name', autospec=True) + filepath = self.convention.get_file_path( + "19900101", 1, ModelingRealms.ocean, "var", cmor_var, 1, "mon" + ) + self.assertEqual("path/filename", filepath) + + @mock.patch( + "earthdiagnostics.data_convention.DataConvention.get_cmor_folder_path", + autospec=True, + ) + @mock.patch( + "earthdiagnostics.data_convention.DataConvention.get_file_name", + autospec=True, + ) def test_get_file_path_no_freq(self, mock_filename, mock_cmor_path): - mock_cmor_path.return_value = 'path' - mock_filename.return_value = 'filename' + mock_cmor_path.return_value = "path" + mock_filename.return_value = "filename" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - filepath = self.convention.get_file_path('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, 1, None) - self.assertEqual('path/filename', filepath) + filepath = self.convention.get_file_path( + "19900101", 1, ModelingRealms.ocean, "var", cmor_var, 1, None + ) + self.assertEqual("path/filename", filepath) def test_create_link(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', None, False, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", None, False, None + ) filename = os.path.basename(path) - self.assertTrue(os.path.islink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', filename))) + self.assertTrue( + os.path.islink( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "var_f6h", + filename, + ) + ) + ) def test_create_link_no_source(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) os.remove(path) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' + frequency.folder_name.return_value = "frequency_folder" with self.assertRaises(ValueError): - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', None, False, None) + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", None, False, None + ) def test_create_link_exists(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' + frequency.folder_name.return_value = "frequency_folder" filename = os.path.basename(path) - os.makedirs(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h')) - os.symlink(path, os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', filename)) - - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', None, False, None) - self.assertTrue(os.path.islink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', filename))) + os.makedirs( + os.path.join(self.tmp_dir, "expid", "frequency_folder", "var_f6h") + ) + os.symlink( + path, + os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var_f6h", filename + ), + ) + + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", None, False, None + ) + self.assertTrue( + os.path.islink( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "var_f6h", + filename, + ) + ) + ) def test_create_link_default_is_link(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) - os.makedirs(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-original_f6h')) - os.symlink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-original_f6h'), - os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h')) + os.makedirs( + os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-original_f6h" + ) + ) + os.symlink( + os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-original_f6h" + ), + os.path.join(self.tmp_dir, "expid", "frequency_folder", "var_f6h"), + ) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', None, False, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", None, False, None + ) filename = os.path.basename(path) - self.assertTrue(os.path.islink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', filename))) + self.assertTrue( + os.path.islink( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "var_f6h", + filename, + ) + ) + ) def test_create_link_move_old(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) - os.makedirs(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h')) - fd = open(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', 'var_20001101.nc'), 'w') + os.makedirs( + os.path.join(self.tmp_dir, "expid", "frequency_folder", "var_f6h") + ) + fd = open( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "var_f6h", + "var_20001101.nc", + ), + "w", + ) fd.close() frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', None, True, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", None, True, None + ) filename = os.path.basename(path) - self.assertTrue(os.path.islink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h', filename))) - self.assertTrue(os.path.isfile(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'old_var_f6h', - 'var_20001101.nc'))) + self.assertTrue( + os.path.islink( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "var_f6h", + filename, + ) + ) + ) + self.assertTrue( + os.path.isfile( + os.path.join( + self.tmp_dir, + "expid", + "frequency_folder", + "old_var_f6h", + "var_20001101.nc", + ) + ) + ) def test_create_link_with_grid(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', 'grid', False, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", "grid", False, None + ) filename = os.path.basename(path) - var_mainfolder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h') - var_grid_folder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-grid_f6h') + var_mainfolder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var_f6h" + ) + var_grid_folder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-grid_f6h" + ) self.assertTrue(os.path.islink(var_mainfolder)) - self.assertTrue(os.path.islink(os.path.join(var_grid_folder, filename))) + self.assertTrue( + os.path.islink(os.path.join(var_grid_folder, filename)) + ) self.assertTrue(os.readlink(var_mainfolder), var_grid_folder) def test_create_link_with_grid_default_is_link(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) - os.makedirs(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-original_f6h')) - os.symlink(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-original_f6h'), - os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h')) + os.makedirs( + os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-original_f6h" + ) + ) + os.symlink( + os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-original_f6h" + ), + os.path.join(self.tmp_dir, "expid", "frequency_folder", "var_f6h"), + ) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', 'grid', False, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", "grid", False, None + ) filename = os.path.basename(path) - var_mainfolder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h') - var_grid_folder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-grid_f6h') + var_mainfolder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var_f6h" + ) + var_grid_folder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-grid_f6h" + ) self.assertTrue(os.path.islink(var_mainfolder)) - self.assertTrue(os.path.islink(os.path.join(var_grid_folder, filename))) + self.assertTrue( + os.path.islink(os.path.join(var_grid_folder, filename)) + ) self.assertTrue(os.readlink(var_mainfolder), var_grid_folder) def test_create_link_with_grid_default_is_folder(self): file_descriptor, path = tempfile.mkstemp(dir=self.tmp_dir) os.close(file_descriptor) - os.makedirs(os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h')) + os.makedirs( + os.path.join(self.tmp_dir, "expid", "frequency_folder", "var_f6h") + ) frequency = Mock() - frequency.folder_name.return_value = 'frequency_folder' - self.convention.create_link(ModelingRealms.ocean, path, frequency, 'var', 'grid', False, None) + frequency.folder_name.return_value = "frequency_folder" + self.convention.create_link( + ModelingRealms.ocean, path, frequency, "var", "grid", False, None + ) filename = os.path.basename(path) - var_mainfolder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var_f6h') - var_grid_folder = os.path.join(self.tmp_dir, 'expid', 'frequency_folder', 'var-grid_f6h') + var_mainfolder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var_f6h" + ) + var_grid_folder = os.path.join( + self.tmp_dir, "expid", "frequency_folder", "var-grid_f6h" + ) self.assertTrue(os.path.islink(var_mainfolder)) - self.assertTrue(os.path.islink(os.path.join(var_grid_folder, filename))) + self.assertTrue( + os.path.islink(os.path.join(var_grid_folder, filename)) + ) self.assertTrue(os.readlink(var_mainfolder), var_grid_folder) diff --git a/test/unit/data_convention/test_meteofrance.py b/test/unit/data_convention/test_meteofrance.py index 825977f2..c0833f72 100644 --- a/test/unit/data_convention/test_meteofrance.py +++ b/test/unit/data_convention/test_meteofrance.py @@ -4,7 +4,6 @@ import tempfile from unittest import TestCase from mock import Mock -import mock from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.frequency import Frequencies @@ -12,153 +11,249 @@ from earthdiagnostics.data_convention import MeteoFranceConvention class TestMeteoFranceConvention(TestCase): - def setUp(self): self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.config = Mock() self.config.data_dir = self.tmp_dir - self.config.experiment.experiment_name = 'experiment_name' - self.config.experiment.expid = 'expid' - self.config.experiment.institute = 'institute' - self.config.experiment.model = 'model' + self.config.experiment.experiment_name = "experiment_name" + self.config.experiment.expid = "expid" + self.config.experiment.institute = "institute" + self.config.experiment.model = "model" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" - self.config.cmor.activity = 'activity' + self.config.cmor.activity = "activity" self.config.cmor.append_startdate = False self.config.cmor.initialization_number = 1 - self.config.cmor.version = 'version' + self.config.cmor.version = "version" - self.convention = MeteoFranceConvention('name', self.config) + self.convention = MeteoFranceConvention("name", self.config) def tearDown(self): shutil.rmtree(self.tmp_dir) def test_get_startdate_path(self): - self.assertEqual(self.convention.get_startdate_path('19900101'), - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name')) + self.assertEqual( + self.convention.get_startdate_path("19900101"), + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name", + ), + ) def test_experiment_name(self): - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_name') + self.assertEqual( + self.convention.experiment_name("19900101"), "experiment_name" + ) def test_experiment_name_append(self): self.config.cmor.append_startdate = True - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_nameS19900101') + self.assertEqual( + self.convention.experiment_name("19900101"), + "experiment_nameS19900101", + ) def test_get_cmor_folder_path(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'experiment_name/HA/1990')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, os.path.join(self.tmp_dir, "experiment_name/HA/1990") + ) def test_get_cmor_folder_path_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.atmos, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'experiment_name/HA/1990')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.atmos, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, os.path.join(self.tmp_dir, "experiment_name/HA/1990") + ) def test_get_cmor_folder_path_custom_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - 'grid', cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'experiment_name/HA/1990')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", "grid", cmor_var + ) + self.assertEqual( + file_path, os.path.join(self.tmp_dir, "experiment_name/HA/1990") + ) def test_get_cmor_folder_path_no_cmor(self): frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'mon' - frequency.frequency = 'mon' - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'experiment_name/HA/1990')) + frequency.__str__.return_value = "mon" + frequency.frequency = "mon" + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, os.path.join(self.tmp_dir, "experiment_name/HA/1990") + ) def test_get_file_path_no_version(self): - self.config.cmor.version = '' + self.config.cmor.version = "" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'experiment_name/HA/1990')) + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, os.path.join(self.tmp_dir, "experiment_name/HA/1990") + ) def test_get_filename(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_mon_199001_01.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual(file_path, "var_mon_199001_01.nc") def test_get_filename_daily(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.daily, 1, - None, None, None) - self.assertEqual(file_path, - 'var_day_199001_01.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.daily, + 1, + None, + None, + None, + ) + self.assertEqual(file_path, "var_day_199001_01.nc") def test_get_filename_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Amon' + omon.name = "Amon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.atmos, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_mon_199001_01.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.atmos, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual(file_path, "var_mon_199001_01.nc") def test_get_filename_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, 'grid') - self.assertEqual(file_path, - 'var_mon_199001_01.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + "grid", + ) + self.assertEqual(file_path, "var_mon_199001_01.nc") def test_get_filename_year(self): with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', None, - Frequencies.yearly, None, 1990, None, None) + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + None, + Frequencies.yearly, + None, + 1990, + None, + None, + ) - def test_get_filename_date_Str(self): + def test_get_filename_date_str(self): with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', None, - Frequencies.yearly, None, None, 'date_str', None) + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + None, + Frequencies.yearly, + None, + None, + "date_str", + None, + ) def test_get_filename_no_date_info(self): with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', None, - Frequencies.monthly, None, None, None, None) + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + None, + Frequencies.monthly, + None, + None, + None, + None, + ) def test_is_cmorized(self): - self.assertTrue(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + self.assertTrue( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_create_links_meteofrance(self): self.convention.create_links(None, None) diff --git a/test/unit/data_convention/test_preface.py b/test/unit/data_convention/test_preface.py index 878428fe..0b81eb00 100644 --- a/test/unit/data_convention/test_preface.py +++ b/test/unit/data_convention/test_preface.py @@ -12,269 +12,452 @@ from earthdiagnostics.data_convention import PrefaceConvention class TestPrefaceConvention(TestCase): - def setUp(self): self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.config = Mock() self.config.data_dir = self.tmp_dir - self.config.experiment.experiment_name = 'experiment_name' - self.config.experiment.expid = 'expid' - self.config.experiment.institute = 'institute' - self.config.experiment.model = 'model' + self.config.experiment.experiment_name = "experiment_name" + self.config.experiment.expid = "expid" + self.config.experiment.institute = "institute" + self.config.experiment.model = "model" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" - self.config.cmor.activity = 'activity' + self.config.cmor.activity = "activity" self.config.cmor.append_startdate = False self.config.cmor.initialization_number = 1 - self.config.cmor.version = 'version' + self.config.cmor.version = "version" - self.convention = PrefaceConvention('name', self.config) + self.convention = PrefaceConvention("name", self.config) def tearDown(self): shutil.rmtree(self.tmp_dir) def test_get_startdate_path(self): - self.assertEqual(self.convention.get_startdate_path('19900101'), - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101')) + self.assertEqual( + self.convention.get_startdate_path("19900101"), + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101", + ), + ) def test_experiment_name(self): - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_name') + self.assertEqual( + self.convention.experiment_name("19900101"), "experiment_name" + ) def test_experiment_name_append(self): self.config.cmor.append_startdate = True - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_nameS19900101') + self.assertEqual( + self.convention.experiment_name("19900101"), + "experiment_nameS19900101", + ) def test_get_cmor_folder_path(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101/mon/" + "ocean/var/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.atmos, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101/mon/' - 'atmos/var/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.atmos, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101/mon/" + "atmos/var/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_custom_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - 'grid', cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101/mon/' - 'ocean/var/grid/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", "grid", cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101/mon/" + "ocean/var/grid/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_no_cmor(self): frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'mon' - frequency.frequency = 'mon' - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1/version')) + frequency.__str__.return_value = "mon" + frequency.frequency = "mon" + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101/mon/" + "ocean/var/r2i1p1/version", + ), + ) def test_get_file_path_no_version(self): - self.config.cmor.version = '' + self.config.cmor.version = "" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1')) + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S19900101/mon/" + "ocean/var/r2i1p1", + ), + ) def test_get_filename(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc", + ) def test_get_filename_daily(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.daily, 1, - None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.daily, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc", + ) def test_get_filename_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Amon' + omon.name = "Amon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.atmos, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Amon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.atmos, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Amon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc", + ) def test_get_filename_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, 'grid') - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + "grid", + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc", + ) def test_get_filename_year(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.yearly, None, 1990, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_1990.nc') - - def test_get_filename_date_Str(self): + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.yearly, + None, + 1990, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_1990.nc", + ) + + def test_get_filename_date_str(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, 'date_str', None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_date_str.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + "date_str", + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_date_str.nc", + ) def test_get_filename_no_date_info(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, None, None) - - @mock.patch('os.path.isfile') + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + None, + None, + ) + + @mock.patch("os.path.isfile") def test_is_cmorized(self, mock_is_file): mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20000101/mon/' - 'ocean/var/r2i1p1/version')) - self.assertTrue(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101/mon/" + "ocean/var/r2i1p1/version", + ) + ) + self.assertTrue( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_not_cmorized(self, mock_is_file): mock_is_file.return_value = False cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20000101/mon/ocean/var/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101/mon/" + "ocean/var/r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_false_not_member_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20000101/mon/ocean/var/')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101/mon" + "/ocean/var/", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_cmorized_not_enough_vars(self, mock_is_file): mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 os.makedirs( - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S20000101/mon/ocean/var/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101/mon/" + "ocean/var/r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_not_domain_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 - os.makedirs(os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S20000101/mon/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101/mon/" + "r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_not_freq_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 - os.makedirs(os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/experiment_name/S20000101')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('earthdiagnostics.data_convention.PrefaceConvention.create_link') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20000101", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch( + "earthdiagnostics.data_convention.PrefaceConvention.create_link" + ) def test_create_links(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20010101/mon/ocean/var/r2i1p1') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20010101/mon/ocean/" + "var/r2i1p1", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() - @mock.patch('earthdiagnostics.data_convention.PrefaceConvention.create_link') + @mock.patch( + "earthdiagnostics.data_convention.PrefaceConvention.create_link" + ) def test_create_links_member_not_found(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20010101/mon/ocean/var/r1i1p1') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20010101/mon/ocean/" + "var/r1i1p1", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_not_called() - @mock.patch('earthdiagnostics.data_convention.PrefaceConvention.create_link') + @mock.patch( + "earthdiagnostics.data_convention.PrefaceConvention.create_link" + ) def test_create_links_with_grid(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/experiment_name/S20010101/mon/ocean/var/' - 'r2i1p1/grid') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/experiment_name/S20010101/mon/" + "ocean/var/r2i1p1/grid", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() diff --git a/test/unit/data_convention/test_primavera.py b/test/unit/data_convention/test_primavera.py index 52a06faa..01afd690 100644 --- a/test/unit/data_convention/test_primavera.py +++ b/test/unit/data_convention/test_primavera.py @@ -18,27 +18,27 @@ class TestPrimaveraConvention(TestCase): def setUp(self): """Prepare tests""" self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.config = Mock() self.config.data_dir = self.tmp_dir - self.config.experiment.experiment_name = 'experiment_name' - self.config.experiment.expid = 'expid' - self.config.experiment.institute = 'institute' - self.config.experiment.model = 'model' + self.config.experiment.experiment_name = "experiment_name" + self.config.experiment.expid = "expid" + self.config.experiment.institute = "institute" + self.config.experiment.model = "model" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" - self.config.cmor.activity = 'activity' + self.config.cmor.activity = "activity" self.config.cmor.append_startdate = False self.config.cmor.append_startdate_year_only = False self.config.cmor.initialization_number = 1 - self.config.cmor.version = 'version' - self.config.cmor.default_ocean_grid = 'ocean_grid' - self.config.cmor.default_atmos_grid = 'atmos_grid' + self.config.cmor.version = "version" + self.config.cmor.default_ocean_grid = "ocean_grid" + self.config.cmor.default_atmos_grid = "atmos_grid" - self.convention = PrimaveraConvention('name', self.config) + self.convention = PrimaveraConvention("name", self.config) def tearDown(self): """Cleanup""" @@ -46,108 +46,191 @@ class TestPrimaveraConvention(TestCase): def test_get_startdate_path(self): """Test get startdate path""" - self.assertEqual(self.convention.get_startdate_path('19900101'), - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name')) + self.assertEqual( + self.convention.get_startdate_path("19900101"), + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name", + ), + ) def test_experiment_name(self): """Test get expriment name""" - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_name') + self.assertEqual( + self.convention.experiment_name("19900101"), "experiment_name" + ) def test_experiment_name_append(self): """Test get expriment name when appending startdate""" self.config.cmor.append_startdate = True - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_name') + self.assertEqual( + self.convention.experiment_name("19900101"), "experiment_name" + ) def test_get_cmor_folder_path(self): """Test get cmor foilder path""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name/' - 'r2i1p1f1/Omon/var/ocean_grid/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/ocean_grid/version", + ), + ) def test_get_cmor_folder_path_no_cmor_var(self): """Test get cmor folder path when not passing cmor_var""" - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name/' - 'r2i1p1f1/Omon/var/ocean_grid/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/ocean_grid/version", + ), + ) def test_get_cmor_folder_path_atmos(self): """Test get cmor foilder path for the atmos""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.atmos, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name/' - 'r2i1p1f1/Omon/var/atmos_grid/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.atmos, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/atmos_grid/version", + ), + ) def test_get_cmor_folder_path_custom_grid(self): """Test get cmor foilder path for a custom grid""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - 'grid', cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name/' - 'r2i1p1f1/Omon/var/grid/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", "grid", cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/grid/version", + ), + ) def test_get_cmor_folder_path_no_cmor(self): """Test get cmor folder path with no cmor_var""" frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'mon' - frequency.frequency = 'mon' - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name/' - 'r2i1p1f1/Omon/var/ocean_grid/version')) + frequency.__str__.return_value = "mon" + frequency.frequency = "mon" + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/ocean_grid/version", + ), + ) def test_get_file_path_no_version_primavera(self): """Test get cmor folder path with no version""" - self.config.cmor.version = '' + self.config.cmor.version = "" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon with self.assertRaises(ValueError): - self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', 'grid', cmor_var) + self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + "mon", + "grid", + cmor_var, + ) def test_get_filename(self): """Test get_filename""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_" + "199001-199001.nc", + ) def test_get_filename_with_startdate(self): """Test get_filename with startdate""" self.config.cmor.append_startdate = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_s19900101-r2i1p1f1_ocean_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_s19900101-r2i1p1f1_ocean_grid_" + "199001-199001.nc", + ) def test_get_filename_with_startdate_year_only(self): """Test get_filename with startdate""" @@ -155,223 +238,398 @@ class TestPrimaveraConvention(TestCase): self.config.cmor.append_startdate_year_only = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_s1990-r2i1p1f1_ocean_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_s1990-r2i1p1f1_ocean_grid_" + "199001-199001.nc", + ) def test_get_filename_no_cmor_var(self): """Test get_filename not passing cmor_var""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', None, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + None, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_" + "199001-199001.nc", + ) def test_get_filename_daily(self): """Test get_filename for daily frequency""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.daily, 1, - None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_19900101-19900131.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.daily, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_" + "19900101-19900131.nc", + ) def test_get_filename_6hourly(self): """Test get_filename for 6hourly files""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.six_hourly, 1, - None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001010000-199001311800.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.six_hourly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_" + "199001010000-199001311800.nc", + ) def test_get_filename_3hourly(self): """Test get_filename for 3hourly files""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.three_hourly, 1, - None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001010000-199001312100.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.three_hourly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_" + "199001010000-199001312100.nc", + ) def test_get_filename_atmos(self): """Test get_filename for atmos""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.atmos, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_atmos_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.atmos, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_atmos_grid_" + "199001-199001.nc", + ) def test_get_filename_grid(self): """Test get_filename for a custom grid""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, 'grid') - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_grid_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + "grid", + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_grid_199001-199001.nc", + ) def test_get_filename_year(self): """Test get_filename for a whole year""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.yearly, None, 1990, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_1990.nc') - - def test_get_filename_date_Str(self): + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.yearly, + None, + 1990, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_1990.nc", + ) + + def test_get_filename_date_str(self): """Test get_filename passing date_Str""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, 'date_str', None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_date_str.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + "date_str", + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_date_str.nc", + ) def test_get_filename_no_date_info(self): """Test get_filename with no date info raises ValueError""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, None, None) - - @mock.patch('os.path.isfile') + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + None, + None, + ) + + @mock.patch("os.path.isfile") def test_is_cmorized(self, mock_is_file): """Test is cmorized""" mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var')) - self.assertTrue(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var", + ) + ) + self.assertTrue( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_not_cmorized(self, mock_is_file): """Test is cmorized false""" mock_is_file.return_value = False cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_false_not_member_folder(self): """Test is cmorized false bacause ther is no member folder""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_false_not_table_folder(self): """Test is cmorized false bacause ther is no table folder""" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_cmorized_not_enough_vars(self, mock_is_file): """Test is cmorized false because thera are not eniouch variables""" mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 - os.makedirs(os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S20000101/' - 'mon/ocean/var')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/ocean/var", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch( + "earthdiagnostics.data_convention.PrimaveraConvention.create_link" + ) def test_create_links_primavera(self, mock_create_link): """Test create links""" - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var/gn') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/gn", + ) os.makedirs(member_path) - self.config.var_manager.tables = {'Omon': Mock()} + self.config.var_manager.tables = {"Omon": Mock()} tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() - @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link') + @mock.patch( + "earthdiagnostics.data_convention.PrimaveraConvention.create_link" + ) def test_create_links_with_version_primavera(self, mock_create_link): """Test create links with version""" - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/' - 'var/gn/version') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/gn/version", + ) os.makedirs(member_path) - self.config.var_manager.tables = {'Omon': Mock()} + self.config.var_manager.tables = {"Omon": Mock()} tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() - @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link') - def test_create_links_with_version_primavera_no_member(self, mock_create_link): + @mock.patch( + "earthdiagnostics.data_convention.PrimaveraConvention.create_link" + ) + def test_create_links_with_version_primavera_no_member( + self, mock_create_link + ): """Test create links with version full startdate""" - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/' - 'var/gn/version') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r2i1p1f1/Omon/var/gn/version", + ) os.makedirs(member_path) - self.config.var_manager.tables = {'Omon': Mock()} + self.config.var_manager.tables = {"Omon": Mock()} tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101') + self.convention.create_links("20010101") mock_create_link.assert_called() - @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link') + @mock.patch( + "earthdiagnostics.data_convention.PrimaveraConvention.create_link" + ) def test_create_links_member_not_found_primavera(self, mock_create_link): """Test create links when the member can not be found""" - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/activity/institute/model/experiment_name/r1i1p1f1/Omon/var/gn') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/activity/institute/model/experiment_name/" + "r1i1p1f1/Omon/var/gn", + ) os.makedirs(member_path) - self.config.var_manager.tables = {'Omon': Mock()} + self.config.var_manager.tables = {"Omon": Mock()} tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_not_called() diff --git a/test/unit/data_convention/test_specs.py b/test/unit/data_convention/test_specs.py index 8e054d10..699a2e2b 100644 --- a/test/unit/data_convention/test_specs.py +++ b/test/unit/data_convention/test_specs.py @@ -12,308 +12,565 @@ from earthdiagnostics.data_convention import SPECSConvention class TestSpecsConvention(TestCase): - def setUp(self): self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.config = Mock() self.config.data_dir = self.tmp_dir - self.config.experiment.experiment_name = 'experiment_name' - self.config.experiment.expid = 'expid' - self.config.experiment.institute = 'institute' - self.config.experiment.model = 'model' + self.config.experiment.experiment_name = "experiment_name" + self.config.experiment.expid = "expid" + self.config.experiment.institute = "institute" + self.config.experiment.model = "model" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" - self.config.cmor.activity = 'activity' + self.config.cmor.activity = "activity" self.config.cmor.append_startdate = False self.config.cmor.initialization_number = 1 - self.config.cmor.version = 'version' + self.config.cmor.version = "version" - self.convention = SPECSConvention('name', self.config) + self.convention = SPECSConvention("name", self.config) def tearDown(self): shutil.rmtree(self.tmp_dir) def test_get_startdate_path(self): - self.assertEqual(self.convention.get_startdate_path('19900101'), - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101')) + self.assertEqual( + self.convention.get_startdate_path("19900101"), + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101", + ), + ) def test_experiment_name(self): - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_name') + self.assertEqual( + self.convention.experiment_name("19900101"), "experiment_name" + ) def test_experiment_name_append(self): self.config.cmor.append_startdate = True - self.assertEqual(self.convention.experiment_name('19900101'), - 'experiment_nameS19900101') + self.assertEqual( + self.convention.experiment_name("19900101"), + "experiment_nameS19900101", + ) def test_get_cmor_folder_path(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101/" + "mon/ocean/var/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.atmos, 'var', 'mon', - None, cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101/mon/' - 'atmos/var/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.atmos, "var", "mon", None, cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101/" + "mon/atmos/var/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_custom_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', - 'grid', cmor_var) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101/mon/' - 'ocean/var/grid/r2i1p1/version')) + file_path = self.convention.get_cmor_folder_path( + "19900101", 1, ModelingRealms.ocean, "var", "mon", "grid", cmor_var + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101/" + "mon/ocean/var/grid/r2i1p1/version", + ), + ) def test_get_cmor_folder_path_no_cmor(self): frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'mon' - frequency.frequency = 'mon' - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1/version')) + frequency.__str__.return_value = "mon" + frequency.frequency = "mon" + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101/" + "mon/ocean/var/r2i1p1/version", + ), + ) def test_get_file_path_no_version(self): - self.config.cmor.version = '' + self.config.cmor.version = "" cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', - Frequencies.monthly, None, None) - self.assertEqual(file_path, - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S19900101/mon/' - 'ocean/var/r2i1p1')) + file_path = self.convention.get_cmor_folder_path( + "19900101", + 1, + ModelingRealms.ocean, + "var", + Frequencies.monthly, + None, + None, + ) + self.assertEqual( + file_path, + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S19900101/" + "mon/ocean/var/r2i1p1", + ), + ) def test_get_filename(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc", + ) def test_get_filename_no_cmor_var(self): - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', None, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + None, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc", + ) def test_get_filename_daily(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.daily, 1, - None, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.daily, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc", + ) def test_get_filename_atmos(self): cmor_var = Mock() omon = Mock() - omon.name = 'Amon' + omon.name = "Amon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.atmos, 'var', cmor_var, - Frequencies.monthly, 1, None, None, None) - self.assertEqual(file_path, - 'var_Amon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.atmos, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + None, + ) + self.assertEqual( + file_path, + "var_Amon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc", + ) def test_get_filename_grid(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, 1, None, None, 'grid') - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + 1, + None, + None, + "grid", + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_199001-199001.nc", + ) def test_get_filename_year(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.yearly, None, 1990, None, None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_1990.nc') - - def test_get_filename_date_Str(self): + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.yearly, + None, + 1990, + None, + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_1990.nc", + ) + + def test_get_filename_date_str(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - file_path = self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, 'date_str', None) - self.assertEqual(file_path, - 'var_Omon_model_experiment_name_S19900101_r2i1p1_date_str.nc') + file_path = self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + "date_str", + None, + ) + self.assertEqual( + file_path, + "var_Omon_model_experiment_name_S19900101_r2i1p1_date_str.nc", + ) def test_get_filename_no_date_info(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon with self.assertRaises(ValueError): - self.convention.get_file_name('19900101', 1, ModelingRealms.ocean, 'var', cmor_var, - Frequencies.monthly, None, None, None, None) + self.convention.get_file_name( + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + Frequencies.monthly, + None, + None, + None, + None, + ) def test_get_file_path_raise_incompatible_date_info(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - frequency = 'mon' - - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, None, frequency, year='1998') - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, 1, frequency, year='1998') - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, 1, frequency, date_str='1998') - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, None, frequency, year='1998', date_str='1998') - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, 1, frequency, year='1998', date_str='1998') - self.assertRaises(ValueError, self.convention.get_file_path, '19900101', 1, ModelingRealms.ocean, 'var', - cmor_var, None, frequency) - - @mock.patch('os.path.isfile') + frequency = "mon" + + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + None, + frequency, + year="1998", + ) + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + 1, + frequency, + year="1998", + ) + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + 1, + frequency, + date_str="1998", + ) + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + None, + frequency, + year="1998", + date_str="1998", + ) + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + 1, + frequency, + year="1998", + date_str="1998", + ) + self.assertRaises( + ValueError, + self.convention.get_file_path, + "19900101", + 1, + ModelingRealms.ocean, + "var", + cmor_var, + None, + frequency, + ) + + @mock.patch("os.path.isfile") def test_is_cmorized(self, mock_is_file): mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20000101/mon/ocean/var/' - 'r2i1p1/version')) - self.assertTrue(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/ocean/var/r2i1p1/version", + ) + ) + self.assertTrue( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_not_cmorized(self, mock_is_file): mock_is_file.return_value = False cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20000101/mon/ocean/var/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/ocean/var/r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_cmorized_no_startdate_path(self, mock_is_file): mock_is_file.return_value = False cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, "expid/cmorfiles/institute/model/experiment_name" + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_false_not_member_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 1 - os.makedirs(os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20000101/mon/ocean/var/')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('os.path.isfile') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/ocean/var/", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("os.path.isfile") def test_is_cmorized_not_enough_vars(self, mock_is_file): mock_is_file.return_value = True cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 os.makedirs( - os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S20000101/mon/ocean/var/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/ocean/var/" + "r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_not_domain_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 - os.makedirs(os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S20000101/mon/' - 'r2i1p1/version')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101/" + "mon/r2i1p1/version", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_is_cmorized_not_freq_folder(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon self.config.var_manager.get_variable.return_value = cmor_var self.config.cmor.min_cmorized_vars = 2 - os.makedirs(os.path.join(self.tmp_dir, 'expid/cmorfiles/institute/model/experiment_name/S20000101')) - self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) - - @mock.patch('earthdiagnostics.data_convention.SPECSConvention.create_link') + os.makedirs( + os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20000101", + ) + ) + self.assertFalse( + self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) + + @mock.patch("earthdiagnostics.data_convention.SPECSConvention.create_link") def test_create_links(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20010101/mon/ocean/var/r2i1p1') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20010101/mon/" + "ocean/var/r2i1p1", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() - @mock.patch('earthdiagnostics.data_convention.SPECSConvention.create_link') + @mock.patch("earthdiagnostics.data_convention.SPECSConvention.create_link") def test_create_links_member_not_found(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20010101/mon/ocean/var/r1i1p1') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20010101/mon/" + "ocean/var/r1i1p1", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_not_called() - @mock.patch('earthdiagnostics.data_convention.SPECSConvention.create_link') + @mock.patch("earthdiagnostics.data_convention.SPECSConvention.create_link") def test_create_links_with_grid(self, mock_create_link): - member_path = os.path.join(self.tmp_dir, - 'expid/cmorfiles/institute/model/experiment_name/S20010101/mon/ocean/var/' - 'r2i1p1/grid') + member_path = os.path.join( + self.tmp_dir, + "expid/cmorfiles/institute/model/experiment_name/S20010101/" + "mon/ocean/var/" + "r2i1p1/grid", + ) os.makedirs(member_path) tempfile.mkstemp(dir=member_path) - self.convention.create_links('20010101', 1) + self.convention.create_links("20010101", 1) mock_create_link.assert_called() diff --git a/test/unit/general/test_attribute.py b/test/unit/general/test_attribute.py index 7ee95e18..fe11113b 100644 --- a/test/unit/general/test_attribute.py +++ b/test/unit/general/test_attribute.py @@ -14,18 +14,20 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestAttribute(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.var_file): @@ -34,43 +36,97 @@ class TestAttribute(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Attribute.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'att', 'value']) + jobs = Attribute.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "att", "value"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', '', - 'att', 'value')) - self.assertEqual(jobs[1], Attribute(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', '', - 'att', 'value')) - - jobs = Attribute.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', 'att', 'value', 'grid']) + self.assertEqual( + jobs[0], + Attribute( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "", + "att", + "value", + ), + ) + self.assertEqual( + jobs[1], + Attribute( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "", + "att", + "value", + ), + ) + + jobs = Attribute.generate_jobs( + self.diags, ["diagnostic", "seaice", "var", "att", "value", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, 'var', 'grid', - 'att', 'value')) - self.assertEqual(jobs[1], Attribute(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, 'var', 'grid', - 'att', 'value')) + self.assertEqual( + jobs[0], + Attribute( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + "grid", + "att", + "value", + ), + ) + self.assertEqual( + jobs[1], + Attribute( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + "grid", + "att", + "value", + ), + ) with self.assertRaises(Exception): - Attribute.generate_jobs(self.diags, ['diagnostic']) + Attribute.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - Attribute.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + Attribute.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - mixed = Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', 'att', 'value') - self.assertEqual(str(mixed), - 'Write attributte output Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var ' - 'Attributte: att:value Grid: grid') - - # def test_compute(self): - # dummydata.model3.Model3(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant', - # constant=1) - - # diag = Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'ta', 'grid', 'att', 'value') - - # diag.variable_file = Mock() - # diag.variable_file.local_file = self.var_file - # diag.corrected = Mock() - # diag.compute() - # diag.corrected.set_local_file.assert_called_once() + mixed = Attribute( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "att", + "value", + ) + self.assertEqual( + str(mixed), + "Write attributte output Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: atmos:var Attributte: att:value Grid: grid", + ) diff --git a/test/unit/general/test_dailymean.py b/test/unit/general/test_dailymean.py index 2aae86f6..a0cd3521 100644 --- a/test/unit/general/test_dailymean.py +++ b/test/unit/general/test_dailymean.py @@ -13,16 +13,27 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestDailyMean(TestCase): - def setUp(self): - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - - self.daymean = DailyMean(self.data_manager, '20000101', 1, 1, ModelingRealms.ocean, 'var', 'freq', '') + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + + self.daymean = DailyMean( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.ocean, + "var", + "freq", + "", + ) def tearDown(self): if os.path.exists(self.var_file): @@ -31,39 +42,82 @@ class TestDailyMean(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = DailyMean.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '6hr']) + jobs = DailyMean.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "6hr"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], DailyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - Frequencies.six_hourly, '')) - self.assertEqual(jobs[1], DailyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - Frequencies.six_hourly, '')) - - jobs = DailyMean.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', '3h', 'grid']) + self.assertEqual( + jobs[0], + DailyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + Frequencies.six_hourly, + "", + ), + ) + self.assertEqual( + jobs[1], + DailyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + Frequencies.six_hourly, + "", + ), + ) + + jobs = DailyMean.generate_jobs( + self.diags, ["diagnostic", "seaice", "var", "3h", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], DailyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, 'var', - Frequencies.three_hourly, 'grid')) - self.assertEqual(jobs[1], DailyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, 'var', - Frequencies.three_hourly, 'grid')) + self.assertEqual( + jobs[0], + DailyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + Frequencies.three_hourly, + "grid", + ), + ) + self.assertEqual( + jobs[1], + DailyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + Frequencies.three_hourly, + "grid", + ), + ) with self.assertRaises(Exception): - DailyMean.generate_jobs(self.diags, ['diagnostic']) + DailyMean.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - DailyMean.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + DailyMean.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.daymean), - 'Calculate daily mean Startdate: 20000101 Member: 1 Chunk: 1 ' - 'Variable: ocean:var Original frequency: freq Grid: ') - - # def test_compute(self): - # dummydata.model3.Model3(oname=self.var_file, start_year=2000, stop_year=2000, method='constant', constant=1) - - # self.daymean.variable_file = Mock() - # self.daymean.variable_file.local_file = self.var_file - # self.daymean.mean_file = Mock() - # self.daymean.compute() - # self.daymean.mean_file.set_local_file.assert_called_once() + self.assertEqual( + str(self.daymean), + "Calculate daily mean Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: ocean:var Original frequency: freq Grid: ", + ) diff --git a/test/unit/general/test_module.py b/test/unit/general/test_module.py index 3fa4f684..d6eed3c3 100644 --- a/test/unit/general/test_module.py +++ b/test/unit/general/test_module.py @@ -14,19 +14,21 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestModule(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.varu_file = mktemp('.nc') - self.varv_file = mktemp('.nc') + self.varu_file = mktemp(".nc") + self.varv_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.varu_file): @@ -37,49 +39,98 @@ class TestModule(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Module.generate_jobs(self.diags, ['diagnostic', 'atmos', 'varu', 'varv', 'varmodule']) + jobs = Module.generate_jobs( + self.diags, ["diagnostic", "atmos", "varu", "varv", "varmodule"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, - 'varu', 'varv', 'varmodule', '')) - self.assertEqual(jobs[1], Module(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, - 'varu', 'varv', 'varmodule', '')) - - jobs = Module.generate_jobs(self.diags, ['diagnostic', 'seaIce', 'varu', 'varv', 'varmodule', 'grid']) + self.assertEqual( + jobs[0], + Module( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "varu", + "varv", + "varmodule", + "", + ), + ) + self.assertEqual( + jobs[1], + Module( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "varu", + "varv", + "varmodule", + "", + ), + ) + + jobs = Module.generate_jobs( + self.diags, + ["diagnostic", "seaIce", "varu", "varv", "varmodule", "grid"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Module(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, - 'varu', 'varv', 'varmodule', 'grid')) - self.assertEqual(jobs[1], Module(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, - 'varu', 'varv', 'varmodule', 'grid')) + self.assertEqual( + jobs[0], + Module( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "varu", + "varv", + "varmodule", + "grid", + ), + ) + self.assertEqual( + jobs[1], + Module( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "varu", + "varv", + "varmodule", + "grid", + ), + ) with self.assertRaises(Exception): - Module.generate_jobs(self.diags, ['diagnostic']) + Module.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - Module.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + Module.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - mixed = Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'varu', 'varv', 'varmodule', 'grid') - self.assertEqual(str(mixed), - 'Calculate module Startdate: 20010101 Member: 0 Chunk: 0 ' - 'Variables: atmos:varu,varv,varmodule Grid: grid') - - # def test_compute(self): - # dummydata.model2.Model2(oname=self.varu_file, var='ua', start_year=2000, stop_year=2000, method='constant', - # constant=1) - - # dummydata.model2.Model2(oname=self.varv_file, var='va', start_year=2000, stop_year=2000, method='constant', - # constant=1) - - # diag = Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'ua', 'va', 'varmodule', 'grid') - - # diag.component_u_file = Mock() - # diag.component_u_file.local_file = self.varu_file - - # diag.component_v_file = Mock() - # diag.component_v_file.local_file = self.varv_file - # diag.module_file = Mock() - # diag.compute() - # diag.module_file.set_local_file.assert_called_once() + mixed = Module( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "varu", + "varv", + "varmodule", + "grid", + ) + self.assertEqual( + str(mixed), + "Calculate module Startdate: 20010101 Member: 0 Chunk: 0 " + "Variables: atmos:varu,varv,varmodule Grid: grid", + ) diff --git a/test/unit/general/test_monthlymean.py b/test/unit/general/test_monthlymean.py index 4eda55aa..f1f94d14 100644 --- a/test/unit/general/test_monthlymean.py +++ b/test/unit/general/test_monthlymean.py @@ -4,6 +4,7 @@ from unittest import TestCase from tempfile import mktemp from mock import Mock, patch + # import dummydata.model3 from earthdiagnostics.diagnostic import DiagnosticVariableOption @@ -13,16 +14,27 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestMonthlyMean(TestCase): - def setUp(self): - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - - self.monmean = MonthlyMean(self.data_manager, '20000101', 1, 1, ModelingRealms.atmos, 'ta', 'freq', '') + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + + self.monmean = MonthlyMean( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.atmos, + "ta", + "freq", + "", + ) def tearDown(self): if os.path.exists(self.var_file): @@ -31,38 +43,81 @@ class TestMonthlyMean(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = MonthlyMean.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'monthly']) + jobs = MonthlyMean.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "monthly"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MonthlyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - Frequencies.monthly, '')) - self.assertEqual(jobs[1], MonthlyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - Frequencies.monthly, '')) - - jobs = MonthlyMean.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', 'mon', 'grid']) + self.assertEqual( + jobs[0], + MonthlyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + Frequencies.monthly, + "", + ), + ) + self.assertEqual( + jobs[1], + MonthlyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + Frequencies.monthly, + "", + ), + ) + + jobs = MonthlyMean.generate_jobs( + self.diags, ["diagnostic", "seaice", "var", "mon", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MonthlyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, 'var', - Frequencies.monthly, 'grid')) - self.assertEqual(jobs[1], MonthlyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, 'var', - Frequencies.monthly, 'grid')) + self.assertEqual( + jobs[0], + MonthlyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + Frequencies.monthly, + "grid", + ), + ) + self.assertEqual( + jobs[1], + MonthlyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + Frequencies.monthly, + "grid", + ), + ) with self.assertRaises(Exception): - MonthlyMean.generate_jobs(self.diags, ['diagnostic']) + MonthlyMean.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - MonthlyMean.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + MonthlyMean.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.monmean), - 'Calculate monthly mean Startdate: 20000101 Member: 1 Chunk: 1 ' - 'Variable: atmos:ta Original frequency: freq Grid: ') - - # def test_compute(self): - # dummydata.model3.Model3(oname=self.var_file, start_year=2000, stop_year=2000, method='constant', constant=1) - - # self.monmean.variable_file = Mock() - # self.monmean.variable_file.local_file = self.var_file - # self.monmean.mean_file = Mock() - # self.monmean.compute() - # self.monmean.mean_file.set_local_file.assert_called_once() + self.assertEqual( + str(self.monmean), + "Calculate monthly mean Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: atmos:ta Original frequency: freq Grid: ", + ) diff --git a/test/unit/general/test_relink.py b/test/unit/general/test_relink.py index d2810bd1..f3c87dca 100644 --- a/test/unit/general/test_relink.py +++ b/test/unit/general/test_relink.py @@ -10,13 +10,15 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestRelink(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 @@ -25,38 +27,121 @@ class TestRelink(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Relink.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var']) + jobs = Relink.generate_jobs(self.diags, ["diagnostic", "atmos", "var"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Relink(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, - 'var', True, '')) - self.assertEqual(jobs[1], Relink(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, - 'var', True, '')) + self.assertEqual( + jobs[0], + Relink( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + True, + "", + ), + ) + self.assertEqual( + jobs[1], + Relink( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + True, + "", + ), + ) - jobs = Relink.generate_jobs(self.diags, ['diagnostic', 'seaIce', 'var', 'False']) + jobs = Relink.generate_jobs( + self.diags, ["diagnostic", "seaIce", "var", "False"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Relink(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, - 'var', False, '')) - self.assertEqual(jobs[1], Relink(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, - 'var', False, '')) + self.assertEqual( + jobs[0], + Relink( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + False, + "", + ), + ) + self.assertEqual( + jobs[1], + Relink( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + False, + "", + ), + ) - jobs = Relink.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', 'True', 'grid']) + jobs = Relink.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "True", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Relink(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, - 'var', True, 'grid')) - self.assertEqual(jobs[1], Relink(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, - 'var', True, 'grid')) + self.assertEqual( + jobs[0], + Relink( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + True, + "grid", + ), + ) + self.assertEqual( + jobs[1], + Relink( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + True, + "grid", + ), + ) with self.assertRaises(Exception): - Relink.generate_jobs(self.diags, ['diagnostic']) + Relink.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - Relink.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + Relink.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - mixed = Relink(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', True, 'grid') - self.assertEqual(str(mixed), - 'Relink output Startdate: 20010101 Member: 0 Chunk: 0 Move old: True ' - 'Variable: ocean:var Grid: grid') + mixed = Relink( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + True, + "grid", + ) + self.assertEqual( + str(mixed), + "Relink output Startdate: 20010101 Member: 0 Chunk: 0 " + "Move old: True Variable: ocean:var Grid: grid", + ) diff --git a/test/unit/general/test_relinkall.py b/test/unit/general/test_relinkall.py index 22bd7612..b49595e0 100644 --- a/test/unit/general/test_relinkall.py +++ b/test/unit/general/test_relinkall.py @@ -8,13 +8,17 @@ from mock import Mock, patch class TestRelinkAll(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.startdates = ['20010101', ] + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.startdates = [ + "20010101", + ] self.box = Box() self.box.min_depth = 0 @@ -23,15 +27,15 @@ class TestRelinkAll(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = RelinkAll.generate_jobs(self.diags, ['diagnostic']) + jobs = RelinkAll.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 1) - self.assertEqual(jobs[0], RelinkAll(self.data_manager, '20010101')) + self.assertEqual(jobs[0], RelinkAll(self.data_manager, "20010101")) with self.assertRaises(Exception): - RelinkAll.generate_jobs(self.diags, ['diagnostic', '0']) + RelinkAll.generate_jobs(self.diags, ["diagnostic", "0"]) def test_str(self): - mixed = RelinkAll(self.data_manager, '20010101') - self.assertEqual(str(mixed), 'Relink all output Startdate: 20010101') + mixed = RelinkAll(self.data_manager, "20010101") + self.assertEqual(str(mixed), "Relink all output Startdate: 20010101") diff --git a/test/unit/general/test_rewrite.py b/test/unit/general/test_rewrite.py index d2675d80..70e0fc5a 100644 --- a/test/unit/general/test_rewrite.py +++ b/test/unit/general/test_rewrite.py @@ -14,20 +14,30 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestRewrite(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.diag = Rewrite(self.data_manager, '20000101', 1, 1, ModelingRealms.atmos, 'ta', 'grid') - self.var_file = mktemp('.nc') + self.diag = Rewrite( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.atmos, + "ta", + "grid", + ) + self.var_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.var_file): @@ -36,35 +46,78 @@ class TestRewrite(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Rewrite.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var']) + jobs = Rewrite.generate_jobs( + self.diags, ["diagnostic", "atmos", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Rewrite(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'original')) - self.assertEqual(jobs[1], Rewrite(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', 'original')) - - jobs = Rewrite.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', 'grid']) + self.assertEqual( + jobs[0], + Rewrite( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "original", + ), + ) + self.assertEqual( + jobs[1], + Rewrite( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "original", + ), + ) + + jobs = Rewrite.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Rewrite(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', 'grid')) - self.assertEqual(jobs[1], Rewrite(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', 'grid')) + self.assertEqual( + jobs[0], + Rewrite( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "grid", + ), + ) + self.assertEqual( + jobs[1], + Rewrite( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "grid", + ), + ) with self.assertRaises(Exception): - Rewrite.generate_jobs(self.diags, ['diagnostic']) + Rewrite.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - Rewrite.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + Rewrite.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.diag), - 'Rewrite output Startdate: 20000101 Member: 1 Chunk: 1 Variable: atmos:ta Grid: grid') - - # def test_compute(self): - # dummydata.model2.Model2(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant', - # constant=1) - - # self.diag.variable_file = Mock() - # self.diag.variable_file.local_file = self.var_file - # self.diag.corrected = Mock() - # self.diag.compute() - # self.diag.corrected.set_local_file.assert_called_once() + self.assertEqual( + str(self.diag), + "Rewrite output Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: atmos:ta Grid: grid", + ) diff --git a/test/unit/general/test_scale.py b/test/unit/general/test_scale.py index 2def45d9..e91a4090 100644 --- a/test/unit/general/test_scale.py +++ b/test/unit/general/test_scale.py @@ -5,31 +5,35 @@ from mock import Mock, patch import os from tempfile import mktemp -# import dummydata -import iris - from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import DiagnosticVariableOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableOption, + DiagnosticOptionError, +) from earthdiagnostics.frequency import Frequencies from earthdiagnostics.general.scale import Scale from earthdiagnostics.modelingrealm import ModelingRealms class TestScale(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.startdates = ['20010101', ] + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.startdates = [ + "20010101", + ] self.diags.config.frequency = Frequencies.monthly self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.var_file): @@ -38,91 +42,277 @@ class TestScale(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0']) + jobs = Scale.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "0", "0"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', '', - float('nan'), float('nan'), Frequencies.monthly, False)) - self.assertEqual(jobs[1], Scale(self.data_manager, '20010101', 0, 1, 0, 0, ModelingRealms.atmos, 'var', '', - float('nan'), float('nan'), Frequencies.monthly, False)) + self.assertEqual( + jobs[0], + Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "", + float("nan"), + float("nan"), + Frequencies.monthly, + False, + ), + ) + self.assertEqual( + jobs[1], + Scale( + self.data_manager, + "20010101", + 0, + 1, + 0, + 0, + ModelingRealms.atmos, + "var", + "", + float("nan"), + float("nan"), + Frequencies.monthly, + False, + ), + ) - jobs = Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0', 'grid']) + jobs = Scale.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "0", "0", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', 'grid', - float('nan'), float('nan'), Frequencies.monthly, False)) - self.assertEqual(jobs[1], Scale(self.data_manager, '20010101', 0, 1, 0, 0, ModelingRealms.atmos, 'var', 'grid', - float('nan'), float('nan'), Frequencies.monthly, False)) + self.assertEqual( + jobs[0], + Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + float("nan"), + float("nan"), + Frequencies.monthly, + False, + ), + ) + self.assertEqual( + jobs[1], + Scale( + self.data_manager, + "20010101", + 0, + 1, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + float("nan"), + float("nan"), + Frequencies.monthly, + False, + ), + ) - jobs = Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0', 'grid', '0', '100']) + jobs = Scale.generate_jobs( + self.diags, + ["diagnostic", "atmos", "var", "0", "0", "grid", "0", "100"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.monthly, False)) - self.assertEqual(jobs[1], Scale(self.data_manager, '20010101', 0, 1, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.monthly, False)) + self.assertEqual( + jobs[0], + Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.monthly, + False, + ), + ) + self.assertEqual( + jobs[1], + Scale( + self.data_manager, + "20010101", + 0, + 1, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.monthly, + False, + ), + ) - jobs = Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0', 'grid', '0', '100', '3hr']) + jobs = Scale.generate_jobs( + self.diags, + [ + "diagnostic", + "atmos", + "var", + "0", + "0", + "grid", + "0", + "100", + "3hr", + ], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.three_hourly, False)) - self.assertEqual(jobs[1], Scale(self.data_manager, '20010101', 0, 1, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.three_hourly, False)) + self.assertEqual( + jobs[0], + Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.three_hourly, + False, + ), + ) + self.assertEqual( + jobs[1], + Scale( + self.data_manager, + "20010101", + 0, + 1, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.three_hourly, + False, + ), + ) - jobs = Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0', 'grid', '0', '100', '3hr', - True]) + jobs = Scale.generate_jobs( + self.diags, + [ + "diagnostic", + "atmos", + "var", + "0", + "0", + "grid", + "0", + "100", + "3hr", + True, + ], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.three_hourly, True)) - self.assertEqual(jobs[1], Scale(self.data_manager, '20010101', 0, 1, 0, 0, ModelingRealms.atmos, 'var', 'grid', - 0, 100, Frequencies.three_hourly, True)) + self.assertEqual( + jobs[0], + Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.three_hourly, + True, + ), + ) + self.assertEqual( + jobs[1], + Scale( + self.data_manager, + "20010101", + 0, + 1, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.three_hourly, + True, + ), + ) with self.assertRaises(DiagnosticOptionError): - Scale.generate_jobs(self.diags, ['diagnostic']) + Scale.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - Scale.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '0', 'grid', '0', '100', '3hr', 'True', - 'extra']) + Scale.generate_jobs( + self.diags, + [ + "diagnostic", + "atmos", + "var", + "0", + "0", + "grid", + "0", + "100", + "3hr", + "True", + "extra", + ], + ) def test_str(self): - scale = Scale(self.data_manager, '20010101', 0, 0, 0, 0, ModelingRealms.atmos, 'var', 'grid', 0, 100, - Frequencies.three_hourly, False) - self.assertEqual(str(scale), - 'Scale output Startdate: 20010101 Member: 0 Chunk: 0 Scale value: 0 Offset: 0 ' - 'Variable: atmos:var Frequency: 3hr Apply mask: False') - - # def test_compute_factor(self): - - # scale = Scale(self.data_manager, '20010101', 0, 0, 10, 0, ModelingRealms.atmos, 'ta', 'grid', 1, 100, - # Frequencies.three_hourly, False) - # cube = self._get_data_and_test(scale) - # self.assertEqual(cube.data.max(), 10) - - # def test_compute_offset(self): - # scale = Scale(self.data_manager, '20010101', 0, 0, 1, 10, ModelingRealms.atmos, 'ta', 'grid', 1, 100, - # Frequencies.three_hourly, False) - # cube = self._get_data_and_test(scale) - # self.assertEqual(cube.data.max(), 11) - - # def test_compute_too_low(self): - # scale = Scale(self.data_manager, '20010101', 0, 0, 0, 10, ModelingRealms.atmos, 'ta', 'grid', 10, 100, - # Frequencies.three_hourly, False) - # cube = self._get_data_and_test(scale) - # self.assertEqual(cube.data.max(), 1) - - # def test_compute_too_high(self): - # scale = Scale(self.data_manager, '20010101', 0, 0, 0, 10, ModelingRealms.atmos, 'ta', 'grid', 0, 0.5, - # Frequencies.three_hourly, False) - # cube = self._get_data_and_test(scale) - # self.assertEqual(cube.data.max(), 1) - - # def _get_data_and_test(self, scale): - # dummydata.model2.Model2(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant', - # constant=1) - # scale.variable_file = Mock() - # scale.variable_file.local_file = self.var_file - # scale.corrected = Mock() - # scale.compute() - # scale.corrected.set_local_file.assert_called_once() - # cube = iris.load_cube(scale.corrected.set_local_file.call_args[0][0]) - # self.assertEqual(cube.data.max(), cube.data.min()) - # return cube + scale = Scale( + self.data_manager, + "20010101", + 0, + 0, + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 100, + Frequencies.three_hourly, + False, + ) + self.assertEqual( + str(scale), + "Scale output Startdate: 20010101 Member: 0 Chunk: 0 " + "Scale value: 0 Offset: 0 " + "Variable: atmos:var Frequency: 3hr Apply mask: False", + ) diff --git a/test/unit/general/test_select_levels.py b/test/unit/general/test_select_levels.py index f20c222b..f05b0a2f 100644 --- a/test/unit/general/test_select_levels.py +++ b/test/unit/general/test_select_levels.py @@ -3,11 +3,10 @@ from unittest import TestCase import os from tempfile import mktemp -import numpy as np -# import dummydata -import iris - -from earthdiagnostics.diagnostic import DiagnosticVariableListOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableListOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.general.select_levels import SelectLevels from earthdiagnostics.frequency import Frequencies @@ -17,76 +16,183 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestSelectLevels(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.startdates = ['20010101', ] + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.startdates = [ + "20010101", + ] self.diags.config.frequency = Frequencies.monthly self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.var_file): os.remove(self.var_file) def fake_parse(self, value): - return value.split('-') + return value.split("-") - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = SelectLevels.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '20']) + jobs = SelectLevels.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "0", "20"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - '', 0, 20)) - self.assertEqual(jobs[1], SelectLevels(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - '', 0, 20)) - - jobs = SelectLevels.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var1-var2', '0', '20']) + self.assertEqual( + jobs[0], + SelectLevels( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "", + 0, + 20, + ), + ) + self.assertEqual( + jobs[1], + SelectLevels( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "", + 0, + 20, + ), + ) + + jobs = SelectLevels.generate_jobs( + self.diags, ["diagnostic", "atmos", "var1-var2", "0", "20"] + ) self.assertEqual(len(jobs), 4) - self.assertEqual(jobs[0], SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var1', - '', 0, 20)) - self.assertEqual(jobs[1], SelectLevels(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var1', - '', 0, 20)) - self.assertEqual(jobs[2], SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var2', - '', 0, 20)) - self.assertEqual(jobs[3], SelectLevels(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var2', - '', 0, 20)) - - jobs = SelectLevels.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '20', 'grid']) + self.assertEqual( + jobs[0], + SelectLevels( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var1", + "", + 0, + 20, + ), + ) + self.assertEqual( + jobs[1], + SelectLevels( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var1", + "", + 0, + 20, + ), + ) + self.assertEqual( + jobs[2], + SelectLevels( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var2", + "", + 0, + 20, + ), + ) + self.assertEqual( + jobs[3], + SelectLevels( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var2", + "", + 0, + 20, + ), + ) + + jobs = SelectLevels.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "0", "20", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - 'grid', 0, 20)) - self.assertEqual(jobs[1], SelectLevels(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - 'grid', 0, 20)) + self.assertEqual( + jobs[0], + SelectLevels( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 20, + ), + ) + self.assertEqual( + jobs[1], + SelectLevels( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "grid", + 0, + 20, + ), + ) with self.assertRaises(DiagnosticOptionError): - SelectLevels.generate_jobs(self.diags, ['diagnostic']) + SelectLevels.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - SelectLevels.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', '0', '20', 'grid', 'extra']) + SelectLevels.generate_jobs( + self.diags, + ["diagnostic", "atmos", "var", "0", "20", "grid", "extra"], + ) def test_str(self): - select = SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', 0, 20) - self.assertEqual(str(select), - 'Select levels Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var ' - 'Levels: 0-20 Grid: grid') - - # def test_compute(self): - # dummydata.model3.Model3(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant', - # constant=1) - # select = SelectLevels(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'ta', 'grid', 0, 3) - # select.variable_file = Mock() - # select.variable_file.local_file = self.var_file - - # select.result = Mock() - # select.compute() - # select.result.set_local_file.assert_called_once() - # cube = iris.load_cube(select.result.set_local_file.call_args[0][0]) - # original_cube = iris.load_cube(select.result.set_local_file.call_args[0][0]) - # self.assertTrue(np.all(cube.coord('air_pressure').points == original_cube.coord('air_pressure').points[0:4])) + select = SelectLevels( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + 0, + 20, + ) + self.assertEqual( + str(select), + "Select levels Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: atmos:var Levels: 0-20 Grid: grid", + ) diff --git a/test/unit/general/test_simplify_dimensions.py b/test/unit/general/test_simplify_dimensions.py index 348f7053..7981493b 100644 --- a/test/unit/general/test_simplify_dimensions.py +++ b/test/unit/general/test_simplify_dimensions.py @@ -1,7 +1,10 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.diagnostic import DiagnosticVariableListOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableListOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.general.simplify_dimensions import SimplifyDimensions from earthdiagnostics.frequency import Frequencies @@ -11,48 +14,112 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestSimplifyDimensions(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.startdates = ['20010101', ] + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.startdates = [ + "20010101", + ] self.diags.config.frequency = Frequencies.monthly - self.diags.config.data_convention = 'convention' + self.diags.config.data_convention = "convention" self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 def fake_parse(self, value): - return value.split('-') + return value.split("-") - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = SimplifyDimensions.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var']) + jobs = SimplifyDimensions.generate_jobs( + self.diags, ["diagnostic", "atmos", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], SimplifyDimensions(self.data_manager, '20010101', 0, 0, - ModelingRealms.atmos, 'var', '', 'convention')) - self.assertEqual(jobs[1], SimplifyDimensions(self.data_manager, '20010101', 0, 1, - ModelingRealms.atmos, 'var', '', 'convention')) + self.assertEqual( + jobs[0], + SimplifyDimensions( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "", + "convention", + ), + ) + self.assertEqual( + jobs[1], + SimplifyDimensions( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "", + "convention", + ), + ) - jobs = SimplifyDimensions.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'grid']) + jobs = SimplifyDimensions.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], SimplifyDimensions(self.data_manager, '20010101', 0, 0, - ModelingRealms.atmos, 'var', 'grid', 'convention')) - self.assertEqual(jobs[1], SimplifyDimensions(self.data_manager, '20010101', 0, 1, - ModelingRealms.atmos, 'var', 'grid', 'convention')) + self.assertEqual( + jobs[0], + SimplifyDimensions( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "convention", + ), + ) + self.assertEqual( + jobs[1], + SimplifyDimensions( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "grid", + "convention", + ), + ) with self.assertRaises(DiagnosticOptionError): - SimplifyDimensions.generate_jobs(self.diags, ['diagnostic']) + SimplifyDimensions.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - SimplifyDimensions.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'grid', 'extra']) + SimplifyDimensions.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "grid", "extra"] + ) def test_str(self): - mixed = SimplifyDimensions(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - 'grid', 'convention') - self.assertEqual(str(mixed), - 'Simplify dimension Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var ' - 'Grid: grid') + mixed = SimplifyDimensions( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "convention", + ) + self.assertEqual( + str(mixed), + "Simplify dimension Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: atmos:var Grid: grid", + ) diff --git a/test/unit/general/test_verticalmeanmetersiris.py b/test/unit/general/test_verticalmeanmetersiris.py index 82bd0865..b1bba32e 100644 --- a/test/unit/general/test_verticalmeanmetersiris.py +++ b/test/unit/general/test_verticalmeanmetersiris.py @@ -4,9 +4,14 @@ import os from tempfile import mktemp # import dummydata -from earthdiagnostics.diagnostic import DiagnosticVariableListOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableListOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box -from earthdiagnostics.general.verticalmeanmetersiris import VerticalMeanMetersIris +from earthdiagnostics.general.verticalmeanmetersiris import ( + VerticalMeanMetersIris, +) from earthdiagnostics.frequency import Frequencies from mock import Mock, patch @@ -14,19 +19,23 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestVerticalMeanMetersIris(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.startdates = ['20010101', ] + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.startdates = [ + "20010101", + ] self.diags.config.frequency = Frequencies.monthly self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") def tearDown(self): if os.path.exists(self.var_file): @@ -37,64 +46,126 @@ class TestVerticalMeanMetersIris(TestCase): raise DiagnosticOptionError return [value] - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) def test_generate_jobs(self): box = Box(True) - jobs = VerticalMeanMetersIris.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var']) + jobs = VerticalMeanMetersIris.generate_jobs( + self.diags, ["diagnostic", "ocean", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 0, - ModelingRealms.ocean, 'var', box)) - self.assertEqual(jobs[1], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 1, - ModelingRealms.ocean, 'var', box)) + self.assertEqual( + jobs[0], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + ), + ) box = Box(True) box.min_depth = 0 box.max_depth = 100 - jobs = VerticalMeanMetersIris.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', '0', '100']) + jobs = VerticalMeanMetersIris.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "0", "100"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 0, - ModelingRealms.ocean, 'var', box)) - self.assertEqual(jobs[1], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 1, - ModelingRealms.ocean, 'var', box)) - - jobs = VerticalMeanMetersIris.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', '0', '100']) + self.assertEqual( + jobs[0], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + ), + ) + + jobs = VerticalMeanMetersIris.generate_jobs( + self.diags, ["diagnostic", "seaice", "var", "0", "100"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 0, - ModelingRealms.seaIce, 'var', box)) - self.assertEqual(jobs[1], VerticalMeanMetersIris(self.data_manager, '20010101', 0, 1, - ModelingRealms.seaIce, 'var', box)) + self.assertEqual( + jobs[0], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + box, + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + box, + ), + ) with self.assertRaises(DiagnosticOptionError): - VerticalMeanMetersIris.generate_jobs(self.diags, ['diagnostic']) + VerticalMeanMetersIris.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - VerticalMeanMetersIris.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', '0', '100', 'seaIce', - 'extra']) + VerticalMeanMetersIris.generate_jobs( + self.diags, + ["diagnostic", "ocean", "var", "0", "100", "seaIce", "extra"], + ) def test_str(self): box = Box(True) box.min_depth = 0 box.max_depth = 100 - mixed = VerticalMeanMetersIris(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', box) - self.assertEqual(str(mixed), - 'Vertical mean meters Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var ' - 'Box: 0-100m') - - # def test_compute(self): - # box = Box() - # box.min_depth = 1 - # box.max_depth = 50000 - - # dummydata.model3.Model3(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant', - # constant=1) - - # diag = VerticalMeanMetersIris(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', box) - - # diag.variable_file = Mock() - # diag.variable_file.local_file = self.var_file - # diag.results = Mock() - # diag.compute() - # diag.results.set_local_file.assert_called_once() + mixed = VerticalMeanMetersIris( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + box, + ) + self.assertEqual( + str(mixed), + "Vertical mean meters Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: atmos:var Box: 0-100m", + ) diff --git a/test/unit/general/test_yearlymean.py b/test/unit/general/test_yearlymean.py index 10f17962..d2d38608 100644 --- a/test/unit/general/test_yearlymean.py +++ b/test/unit/general/test_yearlymean.py @@ -3,6 +3,7 @@ from unittest import TestCase from tempfile import mktemp from mock import Mock, patch import os + # import dummydata from earthdiagnostics.diagnostic import DiagnosticVariableOption @@ -12,16 +13,27 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestYearlyMean(TestCase): - def setUp(self): - self.var_file = mktemp('.nc') + self.var_file = mktemp(".nc") self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.diags.config.frequency = Frequencies.monthly - self.yearly_mean = YearlyMean(self.data_manager, '20000101', 1, 1, ModelingRealms.ocean, 'var', 'freq', '') + self.yearly_mean = YearlyMean( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.ocean, + "var", + "freq", + "", + ) def tearDown(self): if os.path.exists(self.var_file): @@ -30,38 +42,81 @@ class TestYearlyMean(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = YearlyMean.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'day']) + jobs = YearlyMean.generate_jobs( + self.diags, ["diagnostic", "atmos", "var", "day"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], YearlyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - Frequencies.daily, '')) - self.assertEqual(jobs[1], YearlyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - Frequencies.daily, '')) + self.assertEqual( + jobs[0], + YearlyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + Frequencies.daily, + "", + ), + ) + self.assertEqual( + jobs[1], + YearlyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + Frequencies.daily, + "", + ), + ) - jobs = YearlyMean.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', 'mon', 'grid']) + jobs = YearlyMean.generate_jobs( + self.diags, ["diagnostic", "seaice", "var", "mon", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], YearlyMean(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, 'var', - Frequencies.monthly, 'grid')) - self.assertEqual(jobs[1], YearlyMean(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, 'var', - Frequencies.monthly, 'grid')) + self.assertEqual( + jobs[0], + YearlyMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.seaIce, + "var", + Frequencies.monthly, + "grid", + ), + ) + self.assertEqual( + jobs[1], + YearlyMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.seaIce, + "var", + Frequencies.monthly, + "grid", + ), + ) with self.assertRaises(Exception): - YearlyMean.generate_jobs(self.diags, ['diagnostic']) + YearlyMean.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - YearlyMean.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + YearlyMean.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.yearly_mean), - 'Calculate yearly mean Startdate: 20000101 Member: 1 Chunk: 1 ' - 'Variable: ocean:var Original frequency: freq Grid: ') - - # def test_compute(self): - # dummydata.model3.Model3(oname=self.var_file, start_year=2000, stop_year=2000, method='constant', constant=1) - - # self.yearly_mean.variable_file = Mock() - # self.yearly_mean.variable_file.local_file = self.var_file - # self.yearly_mean.mean_file = Mock() - # self.yearly_mean.compute() - # self.yearly_mean.mean_file.set_local_file.assert_called_once() + self.assertEqual( + str(self.yearly_mean), + "Calculate yearly mean Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: ocean:var Original frequency: freq Grid: ", + ) diff --git a/test/unit/ocean/test_areamoc.py b/test/unit/ocean/test_areamoc.py index cb7e7e7e..9e9a18b2 100644 --- a/test/unit/ocean/test_areamoc.py +++ b/test/unit/ocean/test_areamoc.py @@ -8,13 +8,12 @@ from mock import Mock, patch class TestAreaMoc(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() self.basins = Mock() - self.basins.Global = Basin('Global') - self.basins.Atlantic = Basin('Atlantic') + self.basins.Global = Basin("Global") + self.basins.Atlantic = Basin("Atlantic") self.box = Box() self.box.min_lat = 0 @@ -22,36 +21,90 @@ class TestAreaMoc(TestCase): self.box.min_depth = 0 self.box.max_depth = 0 - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.psi = AreaMoc(self.data_manager, '20000101', 1, 1, self.basins.Atlantic, self.box) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.psi = AreaMoc( + self.data_manager, "20000101", 1, 1, self.basins.Atlantic, self.box + ) def fake_parse(self, value): if type(value) is Basin: return value - if value == 'atl': - value = 'Atlantic' + if value == "atl": + value = "Atlantic" else: - value = 'Global' + value = "Global" return Basin(value) - @patch.object(Basins, 'parse', fake_parse) + @patch.object(Basins, "parse", fake_parse) def test_generate_jobs(self): - jobs = AreaMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0']) + jobs = AreaMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], AreaMoc(self.data_manager, '20010101', 0, 0, self.basins.Global, self.box)) - self.assertEqual(jobs[1], AreaMoc(self.data_manager, '20010101', 0, 1, self.basins.Global, self.box)) + self.assertEqual( + jobs[0], + AreaMoc( + self.data_manager, + "20010101", + 0, + 0, + self.basins.Global, + self.box, + ), + ) + self.assertEqual( + jobs[1], + AreaMoc( + self.data_manager, + "20010101", + 0, + 1, + self.basins.Global, + self.box, + ), + ) - jobs = AreaMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', 'atl']) + jobs = AreaMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "atl"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], AreaMoc(self.data_manager, '20010101', 0, 0, self.basins.Atlantic, self.box)) - self.assertEqual(jobs[1], AreaMoc(self.data_manager, '20010101', 0, 1, self.basins.Atlantic, self.box)) + self.assertEqual( + jobs[0], + AreaMoc( + self.data_manager, + "20010101", + 0, + 0, + self.basins.Atlantic, + self.box, + ), + ) + self.assertEqual( + jobs[1], + AreaMoc( + self.data_manager, + "20010101", + 0, + 1, + self.basins.Atlantic, + self.box, + ), + ) with self.assertRaises(Exception): - AreaMoc.generate_jobs(self.diags, ['diagnostic']) + AreaMoc.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - AreaMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0']) + AreaMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.psi), - 'Area MOC Startdate: 20000101 Member: 1 Chunk: 1 Box: 0N0 Basin: Atlantic') + self.assertEqual( + str(self.psi), + "Area MOC Startdate: 20000101 Member: 1 Chunk: 1 Box: 0N0 " + "Basin: Atlantic", + ) diff --git a/test/unit/ocean/test_averagesection.py b/test/unit/ocean/test_averagesection.py index 5b00d64b..e74ff2b5 100644 --- a/test/unit/ocean/test_averagesection.py +++ b/test/unit/ocean/test_averagesection.py @@ -10,7 +10,6 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestAverageSection(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() @@ -21,36 +20,111 @@ class TestAverageSection(TestCase): self.box.min_lon = 0 self.box.max_lon = 0 - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) def fake_parse(self, value): return value # noinspection PyUnresolvedReferences - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = AverageSection.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', '0', '0', '0', '0']) + jobs = AverageSection.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "0", "0", "0", "0"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], AverageSection(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - self.box, '')) - self.assertEqual(jobs[1], AverageSection(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - self.box, '')) + self.assertEqual( + jobs[0], + AverageSection( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + self.box, + "", + ), + ) + self.assertEqual( + jobs[1], + AverageSection( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + self.box, + "", + ), + ) - jobs = AverageSection.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', '0', '0', '0', '0', 'grid']) + jobs = AverageSection.generate_jobs( + self.diags, + ["diagnostic", "ocean", "var", "0", "0", "0", "0", "grid"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], AverageSection(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - self.box, 'grid')) - self.assertEqual(jobs[1], AverageSection(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - self.box, 'grid')) + self.assertEqual( + jobs[0], + AverageSection( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + self.box, + "grid", + ), + ) + self.assertEqual( + jobs[1], + AverageSection( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + self.box, + "grid", + ), + ) with self.assertRaises(Exception): - AverageSection.generate_jobs(self.diags, ['diagnostic']) + AverageSection.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - AverageSection.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', '0', '0', '0', '0', 'grid', - 'extra']) + AverageSection.generate_jobs( + self.diags, + [ + "diagnostic", + "ocean", + "var", + "0", + "0", + "0", + "0", + "grid", + "extra", + ], + ) def test_str(self): - diag = AverageSection(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', self.box, 'grid') - self.assertEqual(str(diag), - 'Average section Startdate: 20010101 Member: 0 Chunk: 0 Box: 0N0E ' - 'Variable: ocean:var Grid: grid') + diag = AverageSection( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + self.box, + "grid", + ) + self.assertEqual( + str(diag), + "Average section Startdate: 20010101 Member: 0 Chunk: 0 Box: 0N0E " + "Variable: ocean:var Grid: grid", + ) diff --git a/test/unit/ocean/test_convectionsites.py b/test/unit/ocean/test_convectionsites.py index c9468b08..5ec4a3be 100644 --- a/test/unit/ocean/test_convectionsites.py +++ b/test/unit/ocean/test_convectionsites.py @@ -5,24 +5,43 @@ from mock import Mock class TestConvectionSites(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.psi = ConvectionSites(self.data_manager, '20000101', 1, 1, 'model_version') + self.psi = ConvectionSites( + self.data_manager, "20000101", 1, 1, "model_version" + ) def test_generate_jobs(self): - jobs = ConvectionSites.generate_jobs(self.diags, ['diagnostic']) + jobs = ConvectionSites.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], ConvectionSites(self.data_manager, '20010101', 0, 0, 'model_version')) - self.assertEqual(jobs[1], ConvectionSites(self.data_manager, '20010101', 0, 1, 'model_version')) + self.assertEqual( + jobs[0], + ConvectionSites( + self.data_manager, "20010101", 0, 0, "model_version" + ), + ) + self.assertEqual( + jobs[1], + ConvectionSites( + self.data_manager, "20010101", 0, 1, "model_version" + ), + ) with self.assertRaises(Exception): - ConvectionSites.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + ConvectionSites.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.psi), 'Convection sites Startdate: 20000101 Member: 1 Chunk: 1') + self.assertEqual( + str(self.psi), + "Convection sites Startdate: 20000101 Member: 1 Chunk: 1", + ) diff --git a/test/unit/ocean/test_cutsection.py b/test/unit/ocean/test_cutsection.py index c88bc664..c0cbc5db 100644 --- a/test/unit/ocean/test_cutsection.py +++ b/test/unit/ocean/test_cutsection.py @@ -1,7 +1,10 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.diagnostic import DiagnosticVariableOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.ocean.cutsection import CutSection from mock import Mock, patch @@ -10,7 +13,6 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestCutSection(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() @@ -21,36 +23,100 @@ class TestCutSection(TestCase): self.box.min_lon = 0 self.box.max_lon = 0 - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.psi = CutSection(self.data_manager, '20000101', 1, 1, ModelingRealms.atmos, 'var', True, 0) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.psi = CutSection( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.atmos, + "var", + True, + 0, + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = CutSection.generate_jobs(self.diags, ['diagnostic', 'var', 'true', '10']) + jobs = CutSection.generate_jobs( + self.diags, ["diagnostic", "var", "true", "10"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], CutSection(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - True, 10)) - self.assertEqual(jobs[1], CutSection(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - True, 10)) + self.assertEqual( + jobs[0], + CutSection( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + True, + 10, + ), + ) + self.assertEqual( + jobs[1], + CutSection( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + True, + 10, + ), + ) - jobs = CutSection.generate_jobs(self.diags, ['diagnostic', 'var', 'false', '0', 'atmos']) + jobs = CutSection.generate_jobs( + self.diags, ["diagnostic", "var", "false", "0", "atmos"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], CutSection(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', - False, 0)) - self.assertEqual(jobs[1], CutSection(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', - False, 0)) + self.assertEqual( + jobs[0], + CutSection( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + False, + 0, + ), + ) + self.assertEqual( + jobs[1], + CutSection( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + False, + 0, + ), + ) with self.assertRaises(DiagnosticOptionError): - CutSection.generate_jobs(self.diags, ['diagnostic']) + CutSection.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - CutSection.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + CutSection.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.psi), - 'Cut section Startdate: 20000101 Member: 1 Chunk: 1 Variable: atmos:var ' - 'Zonal: True Value: 0') + self.assertEqual( + str(self.psi), + "Cut section Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: atmos:var Zonal: True Value: 0", + ) diff --git a/test/unit/ocean/test_gyres.py b/test/unit/ocean/test_gyres.py index 9d3a13d3..f5c56e07 100644 --- a/test/unit/ocean/test_gyres.py +++ b/test/unit/ocean/test_gyres.py @@ -6,24 +6,40 @@ from mock import Mock class TestGyres(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.gyres = Gyres(self.data_manager, '20000101', 1, 1, 'model_version') + self.gyres = Gyres( + self.data_manager, "20000101", 1, 1, "model_version" + ) def test_generate_jobs(self): - jobs = Gyres.generate_jobs(self.diags, ['diagnostic']) + jobs = Gyres.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Gyres(self.data_manager, '20010101', 0, 0, 'model_version')) - self.assertEqual(jobs[1], Gyres(self.data_manager, '20010101', 0, 1, 'model_version')) + self.assertEqual( + jobs[0], + Gyres(self.data_manager, "20010101", 0, 0, "model_version"), + ) + self.assertEqual( + jobs[1], + Gyres(self.data_manager, "20010101", 0, 1, "model_version"), + ) with self.assertRaises(Exception): - Gyres.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + Gyres.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.gyres), 'Gyres Startdate: 20000101 Member: 1 Chunk: 1 Model version: model_version') + self.assertEqual( + str(self.gyres), + "Gyres Startdate: 20000101 Member: 1 Chunk: 1 " + "Model version: model_version", + ) diff --git a/test/unit/ocean/test_heatcontent.py b/test/unit/ocean/test_heatcontent.py index 517f1d30..657f2f65 100644 --- a/test/unit/ocean/test_heatcontent.py +++ b/test/unit/ocean/test_heatcontent.py @@ -14,35 +14,81 @@ def _get_levels_from_meters_mock(cls, box): class TestHeatContent(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - @patch('earthdiagnostics.ocean.heatcontent.HeatContent._get_levels_from_meters') + @patch( + "earthdiagnostics.ocean.heatcontent.HeatContent." + "_get_levels_from_meters" + ) def test_generate_jobs(self, levels_mock): levels_mock.return_value = (1, 20) - jobs = HeatContent.generate_jobs(self.diags, ['diagnostic', 'Global', '-1', '0', '100']) + jobs = HeatContent.generate_jobs( + self.diags, ["diagnostic", "Global", "-1", "0", "100"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], HeatContent(self.data_manager, '20010101', 0, 0, Basins().Global, -1, - self.box, 0, 0)) - self.assertEqual(jobs[1], HeatContent(self.data_manager, '20010101', 0, 1, Basins().Global, -1, - self.box, 0, 0)) + self.assertEqual( + jobs[0], + HeatContent( + self.data_manager, + "20010101", + 0, + 0, + Basins().Global, + -1, + self.box, + 0, + 0, + ), + ) + self.assertEqual( + jobs[1], + HeatContent( + self.data_manager, + "20010101", + 0, + 1, + Basins().Global, + -1, + self.box, + 0, + 0, + ), + ) with self.assertRaises(Exception): - HeatContent.generate_jobs(self.diags, ['diagnostic']) + HeatContent.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - HeatContent.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + HeatContent.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - diag = HeatContent(self.data_manager, '20010101', 0, 0, Basins().Global, -1, self.box, 1, 20) - self.assertEqual(str(diag), - 'Heat content Startdate: 20010101 Member: 0 Chunk: 0 Mixed layer: -1 Box: 0-100 Basin: Global') + diag = HeatContent( + self.data_manager, + "20010101", + 0, + 0, + Basins().Global, + -1, + self.box, + 1, + 20, + ) + self.assertEqual( + str(diag), + "Heat content Startdate: 20010101 Member: 0 Chunk: 0 " + "Mixed layer: -1 Box: 0-100 Basin: Global", + ) diff --git a/test/unit/ocean/test_heatcontentlayer.py b/test/unit/ocean/test_heatcontentlayer.py index 6000a732..3afb212d 100644 --- a/test/unit/ocean/test_heatcontentlayer.py +++ b/test/unit/ocean/test_heatcontentlayer.py @@ -6,14 +6,14 @@ from mock import Mock class TestHeatContentLayer(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' + self.diags.model_version = "model_version" self.diags.config.experiment.get_chunk_list.return_value = ( - ('20010101', 0, 0), ('20010101', 0, 1) + ("20010101", 0, 0), + ("20010101", 0, 1), ) self.weight = Mock() @@ -24,9 +24,21 @@ class TestHeatContentLayer(TestCase): def test_str(self): diag = HeatContentLayer( - self.data_manager, '20000101', 1, 1, self.box, self.weight, 0, 10, None, Mock(), 0, 0 + self.data_manager, + "20000101", + 1, + 1, + self.box, + self.weight, + 0, + 10, + None, + Mock(), + 0, + 0, ) self.assertEqual( str(diag), - 'Heat content layer Startdate: 20000101 Member: 1 Chunk: 1 Box: 0-100m' + "Heat content layer Startdate: 20000101 Member: 1 Chunk: 1 " + "Box: 0-100m", ) diff --git a/test/unit/ocean/test_interpolate.py b/test/unit/ocean/test_interpolate.py index 7a348c25..1ad0782f 100644 --- a/test/unit/ocean/test_interpolate.py +++ b/test/unit/ocean/test_interpolate.py @@ -9,69 +9,260 @@ from earthdiagnostics.diagnostic import DiagnosticVariableListOption class TestInterpolate(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.model_version = 'model_version' + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.model_version = "model_version" def fake_parse(self, value): - return value.split('-') + return value.split("-") - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var']) + jobs = Interpolate.generate_jobs(self.diags, ["interp", "grid", "var"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', 'grid', - 'model_version', False, '')) - self.assertEqual(jobs[1], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', 'grid', - 'model_version', False, '')) + self.assertEqual( + jobs[0], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "grid", + "model_version", + False, + "", + ), + ) + self.assertEqual( + jobs[1], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "grid", + "model_version", + False, + "", + ), + ) - jobs = Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var1-var2']) + jobs = Interpolate.generate_jobs( + self.diags, ["interp", "grid", "var1-var2"] + ) self.assertEqual(len(jobs), 4) - self.assertEqual(jobs[0], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var1', 'grid', - 'model_version', False, '')) - self.assertEqual(jobs[1], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var1', 'grid', - 'model_version', False, '')) - self.assertEqual(jobs[2], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var2', 'grid', - 'model_version', False, '')) - self.assertEqual(jobs[3], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var2', 'grid', - 'model_version', False, '')) + self.assertEqual( + jobs[0], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var1", + "grid", + "model_version", + False, + "", + ), + ) + self.assertEqual( + jobs[1], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var1", + "grid", + "model_version", + False, + "", + ), + ) + self.assertEqual( + jobs[2], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var2", + "grid", + "model_version", + False, + "", + ), + ) + self.assertEqual( + jobs[3], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var2", + "grid", + "model_version", + False, + "", + ), + ) - jobs = Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var', 'atmos']) + jobs = Interpolate.generate_jobs( + self.diags, ["interp", "grid", "var", "atmos"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', - 'model_version', False, '')) - self.assertEqual(jobs[1], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', 'grid', - 'model_version', False, '')) + self.assertEqual( + jobs[0], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + False, + "", + ), + ) + self.assertEqual( + jobs[1], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + False, + "", + ), + ) - jobs = Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var', 'atmos', 'true']) + jobs = Interpolate.generate_jobs( + self.diags, ["interp", "grid", "var", "atmos", "true"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', - 'model_version', True, '')) - self.assertEqual(jobs[1], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', 'grid', - 'model_version', True, '')) + self.assertEqual( + jobs[0], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + True, + "", + ), + ) + self.assertEqual( + jobs[1], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + True, + "", + ), + ) - jobs = Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var', 'atmos', 'true', 'original_grid']) + jobs = Interpolate.generate_jobs( + self.diags, + ["interp", "grid", "var", "atmos", "true", "original_grid"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', - 'model_version', True, 'original_grid')) - self.assertEqual(jobs[1], Interpolate(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', 'grid', - 'model_version', True, 'original_grid')) + self.assertEqual( + jobs[0], + Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + True, + "original_grid", + ), + ) + self.assertEqual( + jobs[1], + Interpolate( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + True, + "original_grid", + ), + ) with self.assertRaises(Exception): - Interpolate.generate_jobs(self.diags, ['interp']) + Interpolate.generate_jobs(self.diags, ["interp"]) with self.assertRaises(Exception): - Interpolate.generate_jobs(self.diags, ['interp', 'grid', 'var', 'atmos', 'true', 'original_grid', 'extra']) + Interpolate.generate_jobs( + self.diags, + [ + "interp", + "grid", + "var", + "atmos", + "true", + "original_grid", + "extra", + ], + ) def test_str(self): - diag = Interpolate(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', - 'model_version', True, 'original_grid') - self.assertEqual(str(diag), - 'Interpolate Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var ' - 'Target grid: grid Invert lat: True Model: model_version ' - 'Original grid: original_grid') + diag = Interpolate( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.atmos, + "var", + "grid", + "model_version", + True, + "original_grid", + ) + self.assertEqual( + str(diag), + "Interpolate Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: atmos:var " + "Target grid: grid Invert lat: True Model: model_version " + "Original grid: original_grid", + ) diff --git a/test/unit/ocean/test_interpolatecdo.py b/test/unit/ocean/test_interpolatecdo.py index 73c6d39e..cc2d86ae 100644 --- a/test/unit/ocean/test_interpolatecdo.py +++ b/test/unit/ocean/test_interpolatecdo.py @@ -5,90 +5,324 @@ from earthdiagnostics.ocean.interpolatecdo import InterpolateCDO from mock import Mock, patch from earthdiagnostics.modelingrealm import ModelingRealms -from earthdiagnostics.diagnostic import DiagnosticVariableListOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableListOption, + DiagnosticOptionError, +) class TestInterpolate(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diags.config.experiment.model_version = 'model_version' - self.diags.config.experiment.atmos_grid = 'atmos_grid' + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diags.config.experiment.model_version = "model_version" + self.diags.config.experiment.atmos_grid = "atmos_grid" def fake_parse(self, value): if not value: raise DiagnosticOptionError - return value.split('-') - - @patch('earthdiagnostics.ocean.interpolatecdo.InterpolateCDO.compute_weights') - @patch('earthdiagnostics.ocean.interpolatecdo.InterpolateCDO.get_sample_grid_file') - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) - @patch('os.remove') - @patch('earthdiagnostics.utils.TempFile.get') - def test_generate_jobs(self, mock_weights, mock_grid_file, mock_remove, mock_get): + return value.split("-") + + @patch( + "earthdiagnostics.ocean.interpolatecdo.InterpolateCDO.compute_weights" + ) + @patch( + "earthdiagnostics.ocean.interpolatecdo.InterpolateCDO." + "get_sample_grid_file" + ) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) + @patch("os.remove") + @patch("earthdiagnostics.utils.TempFile.get") + def test_generate_jobs( + self, mock_weights, mock_grid_file, mock_remove, mock_get + ): mock_weights.return_value = None - mock_get.return_value = 'path_to_weights' + mock_get.return_value = "path_to_weights" - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var']) + jobs = InterpolateCDO.generate_jobs( + self.diags, ["interpcdo", "ocean", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'atmos_grid', 'model_version', True, '', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'atmos_grid', 'model_version', True, '', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "atmos_grid", + "model_version", + True, + "", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "atmos_grid", + "model_version", + True, + "", + None, + ), + ) - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'target_grid']) + jobs = InterpolateCDO.generate_jobs( + self.diags, ["interpcdo", "ocean", "var", "target_grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', True, '', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', True, '', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + True, + "", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + True, + "", + None, + ), + ) - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'target_grid', 'bicubic']) + jobs = InterpolateCDO.generate_jobs( + self.diags, ["interpcdo", "ocean", "var", "target_grid", "bicubic"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', True, '', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', True, '', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + True, + "", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + True, + "", + None, + ), + ) - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'target_grid', 'bicubic', - 'false']) + jobs = InterpolateCDO.generate_jobs( + self.diags, + ["interpcdo", "ocean", "var", "target_grid", "bicubic", "false"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, '', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, '', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "", + None, + ), + ) - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'target_grid', 'bicubic', - 'false', 'orig']) + jobs = InterpolateCDO.generate_jobs( + self.diags, + [ + "interpcdo", + "ocean", + "var", + "target_grid", + "bicubic", + "false", + "orig", + ], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, 'orig', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, 'orig', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "orig", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "orig", + None, + ), + ) - jobs = InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'target_grid', 'bicubic', - 'false', 'orig', 'false']) + jobs = InterpolateCDO.generate_jobs( + self.diags, + [ + "interpcdo", + "ocean", + "var", + "target_grid", + "bicubic", + "false", + "orig", + "false", + ], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, 'orig', None)) - self.assertEqual(jobs[1], InterpolateCDO(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 'target_grid', 'model_version', False, 'orig', None)) + self.assertEqual( + jobs[0], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "orig", + None, + ), + ) + self.assertEqual( + jobs[1], + InterpolateCDO( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "target_grid", + "model_version", + False, + "orig", + None, + ), + ) with self.assertRaises(DiagnosticOptionError): - InterpolateCDO.generate_jobs(self.diags, ['interp']) + InterpolateCDO.generate_jobs(self.diags, ["interp"]) with self.assertRaises(DiagnosticOptionError): - InterpolateCDO.generate_jobs(self.diags, ['interpcdo', 'ocean', 'var', 'bicubic', 'false', 'orig', 'false', - 'extra']) + InterpolateCDO.generate_jobs( + self.diags, + [ + "interpcdo", + "ocean", + "var", + "bicubic", + "false", + "orig", + "false", + "extra", + ], + ) def test_str(self): - diag = InterpolateCDO(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 'atmos_grid', 'model_version', False, 'orig', None) - self.assertEqual(str(diag), - 'Interpolate with CDO Startdate: 20010101 Member: 0 Chunk: 0 Variable: ocean:var ' - 'Target grid: atmos_grid Original grid: orig Mask ocean: False ' - 'Model: model_version') + diag = InterpolateCDO( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "atmos_grid", + "model_version", + False, + "orig", + None, + ) + self.assertEqual( + str(diag), + "Interpolate with CDO Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: ocean:var " + "Target grid: atmos_grid Original grid: orig Mask ocean: False " + "Model: model_version", + ) diff --git a/test/unit/ocean/test_maskland.py b/test/unit/ocean/test_maskland.py index 22f4994b..eec42276 100644 --- a/test/unit/ocean/test_maskland.py +++ b/test/unit/ocean/test_maskland.py @@ -1,7 +1,10 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.diagnostic import DiagnosticVariableListOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableListOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.ocean.mask_land import MaskLand from mock import Mock, patch @@ -10,7 +13,6 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestMaskLand(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() @@ -21,45 +23,140 @@ class TestMaskLand(TestCase): self.box.min_lon = 0 self.box.max_lon = 0 - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError - return value.split('-') + return value.split("-") - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) - @patch('earthdiagnostics.ocean.mask_land.MaskLand._get_mask') + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) + @patch("earthdiagnostics.ocean.mask_land.MaskLand._get_mask") def test_generate_jobs(self, get_mask_mock): get_mask_mock.return_value = None - jobs = MaskLand.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var']) + jobs = MaskLand.generate_jobs( + self.diags, ["diagnostic", "ocean", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MaskLand(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', 't', '')) - self.assertEqual(jobs[1], MaskLand(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', 't', '')) + self.assertEqual( + jobs[0], + MaskLand( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "t", + "", + ), + ) + self.assertEqual( + jobs[1], + MaskLand( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "t", + "", + ), + ) - for mask in ('t', 'u', 'v', 'f', 'w'): - jobs = MaskLand.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', mask]) + for mask in ("t", "u", "v", "f", "w"): + jobs = MaskLand.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", mask] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MaskLand(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - mask, '')) - self.assertEqual(jobs[1], MaskLand(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - mask, '')) + self.assertEqual( + jobs[0], + MaskLand( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + mask, + "", + ), + ) + self.assertEqual( + jobs[1], + MaskLand( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + mask, + "", + ), + ) with self.assertRaises(DiagnosticOptionError): - MaskLand.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', 'BAD']) + MaskLand.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "BAD"] + ) - jobs = MaskLand.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', 't', 'grid']) + jobs = MaskLand.generate_jobs( + self.diags, ["diagnostic", "ocean", "var", "t", "grid"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MaskLand(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - 't', 'grid')) - self.assertEqual(jobs[1], MaskLand(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - 't', 'grid')) + self.assertEqual( + jobs[0], + MaskLand( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "t", + "grid", + ), + ) + self.assertEqual( + jobs[1], + MaskLand( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + "t", + "grid", + ), + ) with self.assertRaises(DiagnosticOptionError): - MaskLand.generate_jobs(self.diags, ['diagnostic']) + MaskLand.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - MaskLand.generate_jobs(self.diags, ['diagnostic', 'ocean', 'var', 't', 'grid', 'extra']) + MaskLand.generate_jobs( + self.diags, + ["diagnostic", "ocean", "var", "t", "grid", "extra"], + ) def test_str(self): - diag = MaskLand(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', 't', 'grid') - self.assertEqual(str(diag), 'Land mask Startdate: 20010101 Member: 0 Chunk: 0 Variable: ocean:var Grid: grid') + diag = MaskLand( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + "t", + "grid", + ) + self.assertEqual( + str(diag), + "Land mask Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: ocean:var Grid: grid", + ) diff --git a/test/unit/ocean/test_maxmoc.py b/test/unit/ocean/test_maxmoc.py index 2605e0cc..2f8decbf 100644 --- a/test/unit/ocean/test_maxmoc.py +++ b/test/unit/ocean/test_maxmoc.py @@ -9,12 +9,11 @@ from earthdiagnostics.ocean.maxmoc import MaxMoc class TestMaxMoc(TestCase): - def setUp(self): self.data_manager = Mock() self.basins = Mock() - self.basins.Global = Basin('Global') - self.basins.Atlantic = Basin('Atlantic') + self.basins.Global = Basin("Global") + self.basins.Atlantic = Basin("Atlantic") self.box = Box(True) self.box.min_lat = 0.0 @@ -22,46 +21,104 @@ class TestMaxMoc(TestCase): self.box.min_depth = 0.0 self.box.max_depth = 0.0 - self.maxmoc = MaxMoc(self.data_manager, '20000101', 1, 2000, self.basins.Global, self.box) + self.maxmoc = MaxMoc( + self.data_manager, + "20000101", + 1, + 2000, + self.basins.Global, + self.box, + ) def fake_parse(self, value): if type(value) is Basin: return value - if value == 'atl': - value = 'Atlantic' + if value == "atl": + value = "Atlantic" else: - value = 'Global' + value = "Global" return Basin(value) - @patch.object(Basins, 'parse', fake_parse) + @patch.object(Basins, "parse", fake_parse) def test_generate_jobs(self): self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.startdates = ('20010101',) + self.diags.model_version = "model_version" + self.diags.config.experiment.startdates = ("20010101",) self.diags.config.experiment.members = (0,) self.diags.config.experiment.get_full_years.return_value = (2000, 2001) - jobs = MaxMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0']) + jobs = MaxMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MaxMoc(self.data_manager, '20010101', 0, 2000, self.basins.Global, self.box)) - self.assertEqual(jobs[1], MaxMoc(self.data_manager, '20010101', 0, 2001, self.basins.Global, self.box)) - - jobs = MaxMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', 'atl']) + self.assertEqual( + jobs[0], + MaxMoc( + self.data_manager, + "20010101", + 0, + 2000, + self.basins.Global, + self.box, + ), + ) + self.assertEqual( + jobs[1], + MaxMoc( + self.data_manager, + "20010101", + 0, + 2001, + self.basins.Global, + self.box, + ), + ) + + jobs = MaxMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "atl"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MaxMoc(self.data_manager, '20010101', 0, 2000, self.basins.Atlantic, self.box)) - self.assertEqual(jobs[1], MaxMoc(self.data_manager, '20010101', 0, 2001, self.basins.Atlantic, self.box)) + self.assertEqual( + jobs[0], + MaxMoc( + self.data_manager, + "20010101", + 0, + 2000, + self.basins.Atlantic, + self.box, + ), + ) + self.assertEqual( + jobs[1], + MaxMoc( + self.data_manager, + "20010101", + 0, + 2001, + self.basins.Atlantic, + self.box, + ), + ) self.diags.config.experiment.get_full_years.return_value = list() - jobs = MaxMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0']) + jobs = MaxMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0"] + ) self.assertEqual(len(jobs), 0) with self.assertRaises(Exception): - MaxMoc.generate_jobs(self.diags, ['diagnostic']) + MaxMoc.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(Exception): - MaxMoc.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + MaxMoc.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.maxmoc), - 'Max moc Startdate: 20000101 Member: 1 Year: 2000 Box: 0.0N0m Basin: Global') + self.assertEqual( + str(self.maxmoc), + "Max moc Startdate: 20000101 Member: 1 Year: 2000 Box: 0.0N0m " + "Basin: Global", + ) diff --git a/test/unit/ocean/test_mixedlayerheatcontent.py b/test/unit/ocean/test_mixedlayerheatcontent.py index 8d94bd81..e79bb635 100644 --- a/test/unit/ocean/test_mixedlayerheatcontent.py +++ b/test/unit/ocean/test_mixedlayerheatcontent.py @@ -6,24 +6,35 @@ from mock import Mock class TestMixedLayerHeatContent(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.mixed = MixedLayerHeatContent(self.data_manager, '20000101', 1, 1) + self.mixed = MixedLayerHeatContent(self.data_manager, "20000101", 1, 1) def test_generate_jobs(self): - jobs = MixedLayerHeatContent.generate_jobs(self.diags, ['diagnostic']) + jobs = MixedLayerHeatContent.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MixedLayerHeatContent(self.data_manager, '20010101', 0, 0)) - self.assertEqual(jobs[1], MixedLayerHeatContent(self.data_manager, '20010101', 0, 1)) + self.assertEqual( + jobs[0], MixedLayerHeatContent(self.data_manager, "20010101", 0, 0) + ) + self.assertEqual( + jobs[1], MixedLayerHeatContent(self.data_manager, "20010101", 0, 1) + ) with self.assertRaises(Exception): - MixedLayerHeatContent.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + MixedLayerHeatContent.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.mixed), 'Mixed layer heat content Startdate: 20000101 Member: 1 Chunk: 1') + self.assertEqual( + str(self.mixed), + "Mixed layer heat content Startdate: 20000101 Member: 1 Chunk: 1", + ) diff --git a/test/unit/ocean/test_mixedlayersaltcontent.py b/test/unit/ocean/test_mixedlayersaltcontent.py index 9832e3f8..54a20a54 100644 --- a/test/unit/ocean/test_mixedlayersaltcontent.py +++ b/test/unit/ocean/test_mixedlayersaltcontent.py @@ -6,24 +6,35 @@ from mock import Mock class TestMixedLayerSaltContent(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.mixed = MixedLayerSaltContent(self.data_manager, '20000101', 1, 1) + self.mixed = MixedLayerSaltContent(self.data_manager, "20000101", 1, 1) def test_generate_jobs(self): - jobs = MixedLayerSaltContent.generate_jobs(self.diags, ['diagnostic']) + jobs = MixedLayerSaltContent.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MixedLayerSaltContent(self.data_manager, '20010101', 0, 0)) - self.assertEqual(jobs[1], MixedLayerSaltContent(self.data_manager, '20010101', 0, 1)) + self.assertEqual( + jobs[0], MixedLayerSaltContent(self.data_manager, "20010101", 0, 0) + ) + self.assertEqual( + jobs[1], MixedLayerSaltContent(self.data_manager, "20010101", 0, 1) + ) with self.assertRaises(Exception): - MixedLayerSaltContent.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + MixedLayerSaltContent.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.mixed), 'Mixed layer salt content Startdate: 20000101 Member: 1 Chunk: 1') + self.assertEqual( + str(self.mixed), + "Mixed layer salt content Startdate: 20000101 Member: 1 Chunk: 1", + ) diff --git a/test/unit/ocean/test_moc.py b/test/unit/ocean/test_moc.py index beffc57b..b72f6164 100644 --- a/test/unit/ocean/test_moc.py +++ b/test/unit/ocean/test_moc.py @@ -2,20 +2,24 @@ from unittest import TestCase from mock import Mock -from earthdiagnostics.constants import Basins from earthdiagnostics.ocean.moc import Moc class TestMoc(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.mixed = Moc(self.data_manager, '20000101', 1, 1, {'global': None}) + self.mixed = Moc(self.data_manager, "20000101", 1, 1, {"global": None}) def test_str(self): - self.assertEqual(str(self.mixed), "MOC Startdate: 20000101 Member: 1 Chunk: 1 Basins: ['global']") + self.assertEqual( + str(self.mixed), + "MOC Startdate: 20000101 Member: 1 Chunk: 1 Basins: ['global']", + ) diff --git a/test/unit/ocean/test_mxl.py b/test/unit/ocean/test_mxl.py index 5b63aeea..2b7d5518 100644 --- a/test/unit/ocean/test_mxl.py +++ b/test/unit/ocean/test_mxl.py @@ -6,23 +6,27 @@ from mock import Mock class TestMxl(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) def test_generate_jobs(self): - jobs = Mxl.generate_jobs(self.diags, ['diagnostic']) + jobs = Mxl.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Mxl(self.data_manager, '20010101', 0, 0)) - self.assertEqual(jobs[1], Mxl(self.data_manager, '20010101', 0, 1)) + self.assertEqual(jobs[0], Mxl(self.data_manager, "20010101", 0, 0)) + self.assertEqual(jobs[1], Mxl(self.data_manager, "20010101", 0, 1)) with self.assertRaises(Exception): - Mxl.generate_jobs(self.diags, ['diagnostic', 'extra']) + Mxl.generate_jobs(self.diags, ["diagnostic", "extra"]) def test_str(self): - diag = Mxl(self.data_manager, '20010101', 0, 0) - self.assertEqual(str(diag), 'Mixed layer Startdate: 20010101 Member: 0 Chunk: 0') + diag = Mxl(self.data_manager, "20010101", 0, 0) + self.assertEqual( + str(diag), "Mixed layer Startdate: 20010101 Member: 0 Chunk: 0" + ) diff --git a/test/unit/ocean/test_psi.py b/test/unit/ocean/test_psi.py index e80ffae4..a4b5efdb 100644 --- a/test/unit/ocean/test_psi.py +++ b/test/unit/ocean/test_psi.py @@ -8,31 +8,52 @@ from earthdiagnostics.utils import Utils class TestPsi(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) @staticmethod def fake_get(basin): return None - @patch.object(Utils, 'get_mask', fake_get) + @patch.object(Utils, "get_mask", fake_get) def test_generate_jobs(self): - jobs = Psi.generate_jobs(self.diags, ['diagnostic']) + jobs = Psi.generate_jobs(self.diags, ["diagnostic"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Psi(self.data_manager, '20010101', 0, 0, {Basins().Global: None})) - self.assertEqual(jobs[1], Psi(self.data_manager, '20010101', 0, 1, {Basins().Global: None})) - - jobs = Psi.generate_jobs(self.diags, ['diagnostic', 'atl']) + self.assertEqual( + jobs[0], + Psi(self.data_manager, "20010101", 0, 0, {Basins().Global: None}), + ) + self.assertEqual( + jobs[1], + Psi(self.data_manager, "20010101", 0, 1, {Basins().Global: None}), + ) + + jobs = Psi.generate_jobs(self.diags, ["diagnostic", "atl"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Psi(self.data_manager, '20010101', 0, 0, {Basins().Atlantic: None})) - self.assertEqual(jobs[1], Psi(self.data_manager, '20010101', 0, 1, {Basins().Atlantic: None})) + self.assertEqual( + jobs[0], + Psi( + self.data_manager, "20010101", 0, 0, {Basins().Atlantic: None} + ), + ) + self.assertEqual( + jobs[1], + Psi( + self.data_manager, "20010101", 0, 1, {Basins().Atlantic: None} + ), + ) with self.assertRaises(Exception): - Psi.generate_jobs(self.diags, ['diagnostic', 'atl', 'badoption']) + Psi.generate_jobs(self.diags, ["diagnostic", "atl", "badoption"]) def test_str(self): - psi = Psi(self.data_manager, '20000101', 1, 1, {Basins().Global: None}) - self.assertEqual(str(psi), 'PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global') + psi = Psi(self.data_manager, "20000101", 1, 1, {Basins().Global: None}) + self.assertEqual( + str(psi), + "PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global", + ) diff --git a/test/unit/ocean/test_region_mean.py b/test/unit/ocean/test_region_mean.py index fb554abf..e62c5f0d 100644 --- a/test/unit/ocean/test_region_mean.py +++ b/test/unit/ocean/test_region_mean.py @@ -6,19 +6,23 @@ from mock import Mock, patch from earthdiagnostics.box import Box from earthdiagnostics.frequency import Frequencies from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import DiagnosticOptionError, DiagnosticVariableListOption +from earthdiagnostics.diagnostic import ( + DiagnosticOptionError, + DiagnosticVariableListOption, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.ocean.regionmean import RegionMean from earthdiagnostics.utils import TempFile class TestRegionMean(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() self.diags.config.experiment.get_chunk_list.return_value = ( - ('20010101', 0, 0), ('20010101', 0, 1)) + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.diags.config.frequency = Frequencies.monthly def fake_parse(self, value): @@ -28,111 +32,420 @@ class TestRegionMean(TestCase): @staticmethod def fake_get(): - return 'tempfile' + return "tempfile" - @patch.object(DiagnosticVariableListOption, 'parse', fake_parse) - @patch.object(TempFile, 'get', fake_get) + @patch.object(DiagnosticVariableListOption, "parse", fake_parse) + @patch.object(TempFile, "get", fake_get) def test_generate_jobs(self): box = Box() box.min_depth = -1 box.max_depth = -1 jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', ['var']]) + self.diags, ["diagnostic", "ocean", ["var"]] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 't', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 't', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "t", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "t", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', ['var'], '', 'U']) + self.diags, ["diagnostic", "ocean", ["var"], "", "U"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U']) + self.diags, ["diagnostic", "ocean", ["var"], "global", "U"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) box = Box() box.min_depth = 1.0 box.max_depth = 10.0 jobs = RegionMean.generate_jobs( - self.diags, ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10']) + self.diags, + ["diagnostic", "ocean", ["var"], "global", "U", "1", "10"], + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, True, False, Basins().Global, 'u', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + True, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', - '1', '10', 'false'] + [ + "diagnostic", + "ocean", + ["var"], + "global", + "U", + "1", + "10", + "false", + ], ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, False, False, Basins().Global, 'u', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, False, False, Basins().Global, 'u', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + False, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + False, + False, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', - '1', '10', 'false', 'True'] + [ + "diagnostic", + "ocean", + ["var"], + "global", + "U", + "1", + "10", + "false", + "True", + ], ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, False, True, Basins().Global, 'u', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, False, True, Basins().Global, 'u', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + False, + True, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + False, + True, + Basins().Global, + "u", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', - '10', 'false', 'True', 'grid'] + [ + "diagnostic", + "ocean", + ["var"], + "global", + "U", + "1", + "10", + "false", + "True", + "grid", + ], ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, False, Basins().Global, True, 'grid', Frequencies.monthly)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, False, Basins().Global, True, 'grid', Frequencies.monthly)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + False, + Basins().Global, + True, + "grid", + Frequencies.monthly, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + False, + Basins().Global, + True, + "grid", + Frequencies.monthly, + ), + ) jobs = RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', - '10', 'false', 'True', 'grid', 'day'] + [ + "diagnostic", + "ocean", + ["var"], + "global", + "U", + "1", + "10", + "false", + "True", + "grid", + "day", + ], ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, False, Basins().Global, True, 'grid', Frequencies.daily)) - self.assertEqual(jobs[1], RegionMean(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, False, Basins().Global, True, 'grid', Frequencies.daily)) + self.assertEqual( + jobs[0], + RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + False, + Basins().Global, + True, + "grid", + Frequencies.daily, + ), + ) + self.assertEqual( + jobs[1], + RegionMean( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + False, + Basins().Global, + True, + "grid", + Frequencies.daily, + ), + ) with self.assertRaises(DiagnosticOptionError): - RegionMean.generate_jobs(self.diags, ['diagnostic']) + RegionMean.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): RegionMean.generate_jobs( self.diags, - ['diagnostic', 'ocean', ['var'], 'global', 'U', '1', '10', 'false', - 'True', 'grid', 'day', 'extra'] + [ + "diagnostic", + "ocean", + ["var"], + "global", + "U", + "1", + "10", + "false", + "True", + "grid", + "day", + "extra", + ], ) def test_str(self): box = Box() box.min_depth = 1 box.max_depth = 10 - diag = RegionMean(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', box, False, - True, Basins().Global, 'T', Frequencies.monthly) - self.assertEqual(str(diag), - 'Region mean Startdate: 20010101 Member: 0 Chunk: 0 Variable: var Box: 1-10 ' - 'Save 3D: False Save variance: True Grid point: T') + diag = RegionMean( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + False, + True, + Basins().Global, + "T", + Frequencies.monthly, + ) + self.assertEqual( + str(diag), + "Region mean Startdate: 20010101 Member: 0 Chunk: 0 Variable: var " + "Box: 1-10 Save 3D: False Save variance: True Grid point: T", + ) diff --git a/test/unit/ocean/test_rotation.py b/test/unit/ocean/test_rotation.py index 4cc90e47..93268c32 100644 --- a/test/unit/ocean/test_rotation.py +++ b/test/unit/ocean/test_rotation.py @@ -3,45 +3,115 @@ from unittest import TestCase from earthdiagnostics.ocean.rotation import Rotation from earthdiagnostics.modelingrealm import ModelingRealms -from earthdiagnostics.diagnostic import DiagnosticOptionError, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + DiagnosticOptionError, + DiagnosticVariableOption, +) from mock import Mock, patch class TestMixedLayerHeatContent(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) - self.mixed = Rotation(self.data_manager, '20000101', 1, 1, ModelingRealms.ocean, 'varu', 'varv', 'exe') + self.mixed = Rotation( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.ocean, + "varu", + "varv", + "exe", + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Rotation.generate_jobs(self.diags, ['diagnostic', 'ocean', 'varu', 'varv']) + jobs = Rotation.generate_jobs( + self.diags, ["diagnostic", "ocean", "varu", "varv"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Rotation(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'varu', 'varv', - '/home/Earth/jvegas/pyCharm/cfutools/interpolation/rotateUVorca')) - self.assertEqual(jobs[1], Rotation(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'varu', 'varv', - '/home/Earth/jvegas/pyCharm/cfutools/interpolation/rotateUVorca')) + self.assertEqual( + jobs[0], + Rotation( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "varu", + "varv", + "/home/Earth/jvegas/pyCharm/cfutools/interpolation/" + "rotateUVorca", + ), + ) + self.assertEqual( + jobs[1], + Rotation( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "varu", + "varv", + "/home/Earth/jvegas/pyCharm/cfutools/interpolation/" + "rotateUVorca", + ), + ) - jobs = Rotation.generate_jobs(self.diags, ['diagnostic', 'ocean', 'varu', 'varv', 'exe']) + jobs = Rotation.generate_jobs( + self.diags, ["diagnostic", "ocean", "varu", "varv", "exe"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], - Rotation(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'varu', 'varv', 'exe')) - self.assertEqual(jobs[1], - Rotation(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'varu', 'varv', 'exe')) + self.assertEqual( + jobs[0], + Rotation( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "varu", + "varv", + "exe", + ), + ) + self.assertEqual( + jobs[1], + Rotation( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "varu", + "varv", + "exe", + ), + ) with self.assertRaises(Exception): - Rotation.generate_jobs(self.diags, ['diagnostic', 'ocean', 'varu', 'varv', 'exe', 'extra']) + Rotation.generate_jobs( + self.diags, + ["diagnostic", "ocean", "varu", "varv", "exe", "extra"], + ) def test_str(self): - self.assertEqual(str(self.mixed), - 'Rotate variables Startdate: 20000101 Member: 1 Chunk: 1 Variables: ocean:varu , ocean:varv') + self.assertEqual( + str(self.mixed), + "Rotate variables Startdate: 20000101 Member: 1 Chunk: 1 " + "Variables: ocean:varu , ocean:varv", + ) diff --git a/test/unit/ocean/test_siasiesiv.py b/test/unit/ocean/test_siasiesiv.py index b1f6526d..5507560c 100644 --- a/test/unit/ocean/test_siasiesiv.py +++ b/test/unit/ocean/test_siasiesiv.py @@ -12,15 +12,29 @@ class TestSiasiesiv(TestCase): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.mask = Mock() self.var_manager = Mock() self.convention = Mock() - self.psi = Siasiesiv(self.data_manager, '20000101', 1, 1, {Basins().Global: []}, self.var_manager, - self.convention, False) + self.psi = Siasiesiv( + self.data_manager, + "20000101", + 1, + 1, + {Basins().Global: []}, + self.var_manager, + self.convention, + False, + ) def test_str(self): - self.assertEqual(str(self.psi), - 'Siasiesiv Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global Omit volume: False') + self.assertEqual( + str(self.psi), + "Siasiesiv Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global " + "Omit volume: False", + ) diff --git a/test/unit/ocean/test_vertical_gradient.py b/test/unit/ocean/test_vertical_gradient.py index c55aeded..1752f665 100644 --- a/test/unit/ocean/test_vertical_gradient.py +++ b/test/unit/ocean/test_vertical_gradient.py @@ -2,62 +2,99 @@ from unittest import TestCase from earthdiagnostics.ocean.verticalgradient import VerticalGradient from earthdiagnostics.box import Box -from earthdiagnostics.diagnostic import DiagnosticOptionError, DiagnosticVariableOption +from earthdiagnostics.diagnostic import ( + DiagnosticOptionError, + DiagnosticVariableOption, +) from mock import Mock, patch class TestVerticalGradient(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): box = Box() box.min_depth = 1 box.max_depth = 2 - jobs = VerticalGradient.generate_jobs(self.diags, ['diagnostic', 'var']) + jobs = VerticalGradient.generate_jobs( + self.diags, ["diagnostic", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalGradient(self.data_manager, '20010101', 0, 0, 'var', box)) - self.assertEqual(jobs[1], VerticalGradient(self.data_manager, '20010101', 0, 1, 'var', box)) + self.assertEqual( + jobs[0], + VerticalGradient(self.data_manager, "20010101", 0, 0, "var", box), + ) + self.assertEqual( + jobs[1], + VerticalGradient(self.data_manager, "20010101", 0, 1, "var", box), + ) box = Box() box.min_depth = 2 box.max_depth = 2 - jobs = VerticalGradient.generate_jobs(self.diags, ['diagnostic', 'var', '2']) + jobs = VerticalGradient.generate_jobs( + self.diags, ["diagnostic", "var", "2"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalGradient(self.data_manager, '20010101', 0, 0, 'var', box)) - self.assertEqual(jobs[1], VerticalGradient(self.data_manager, '20010101', 0, 1, 'var', box)) + self.assertEqual( + jobs[0], + VerticalGradient(self.data_manager, "20010101", 0, 0, "var", box), + ) + self.assertEqual( + jobs[1], + VerticalGradient(self.data_manager, "20010101", 0, 1, "var", box), + ) box = Box() box.min_depth = 1 box.max_depth = 10 - jobs = VerticalGradient.generate_jobs(self.diags, ['diagnostic', 'var', '1', 10]) + jobs = VerticalGradient.generate_jobs( + self.diags, ["diagnostic", "var", "1", 10] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalGradient(self.data_manager, '20010101', 0, 0, 'var', box)) - self.assertEqual(jobs[1], VerticalGradient(self.data_manager, '20010101', 0, 1, 'var', box)) + self.assertEqual( + jobs[0], + VerticalGradient(self.data_manager, "20010101", 0, 0, "var", box), + ) + self.assertEqual( + jobs[1], + VerticalGradient(self.data_manager, "20010101", 0, 1, "var", box), + ) with self.assertRaises(DiagnosticOptionError): - VerticalGradient.generate_jobs(self.diags, ['diagnostic']) + VerticalGradient.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - VerticalGradient.generate_jobs(self.diags, ['diagnostic', 'var', '1', '10', 'extra']) + VerticalGradient.generate_jobs( + self.diags, ["diagnostic", "var", "1", "10", "extra"] + ) def test_str(self): box = Box() box.min_depth = 1 box.max_depth = 10 - diag = VerticalGradient(self.data_manager, '20010101', 0, 0, 'var', box) - self.assertEqual(str(diag), 'Vertical gradient Startdate: 20010101 Member: 0 Chunk: 0 Variable: var Box: 1-10') + diag = VerticalGradient( + self.data_manager, "20010101", 0, 0, "var", box + ) + self.assertEqual( + str(diag), + "Vertical gradient Startdate: 20010101 Member: 0 Chunk: 0 " + "Variable: var Box: 1-10", + ) diff --git a/test/unit/ocean/test_verticalmean.py b/test/unit/ocean/test_verticalmean.py index 3b25f50f..f0f89d9f 100644 --- a/test/unit/ocean/test_verticalmean.py +++ b/test/unit/ocean/test_verticalmean.py @@ -1,57 +1,92 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.diagnostic import DiagnosticVariableOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.ocean.verticalmean import VerticalMean from mock import Mock, patch class TestVerticalMean(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box() self.box.min_depth = 0 self.box.max_depth = 100 - self.mixed = VerticalMean(self.data_manager, '20000101', 1, 1, 'var', self.box) + self.mixed = VerticalMean( + self.data_manager, "20000101", 1, 1, "var", self.box + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = VerticalMean.generate_jobs(self.diags, ['diagnostic', 'var', '0', '100']) + jobs = VerticalMean.generate_jobs( + self.diags, ["diagnostic", "var", "0", "100"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMean(self.data_manager, '20010101', 0, 0, 'var', self.box)) - self.assertEqual(jobs[1], VerticalMean(self.data_manager, '20010101', 0, 1, 'var', self.box)) + self.assertEqual( + jobs[0], + VerticalMean(self.data_manager, "20010101", 0, 0, "var", self.box), + ) + self.assertEqual( + jobs[1], + VerticalMean(self.data_manager, "20010101", 0, 1, "var", self.box), + ) - jobs = VerticalMean.generate_jobs(self.diags, ['diagnostic', 'var', '0']) + jobs = VerticalMean.generate_jobs( + self.diags, ["diagnostic", "var", "0"] + ) box = Box() box.min_depth = 0 self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMean(self.data_manager, '20010101', 0, 0, 'var', box)) - self.assertEqual(jobs[1], VerticalMean(self.data_manager, '20010101', 0, 1, 'var', box)) + self.assertEqual( + jobs[0], + VerticalMean(self.data_manager, "20010101", 0, 0, "var", box), + ) + self.assertEqual( + jobs[1], + VerticalMean(self.data_manager, "20010101", 0, 1, "var", box), + ) - jobs = VerticalMean.generate_jobs(self.diags, ['diagnostic', 'var']) + jobs = VerticalMean.generate_jobs(self.diags, ["diagnostic", "var"]) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMean(self.data_manager, '20010101', 0, 0, 'var', Box())) - self.assertEqual(jobs[1], VerticalMean(self.data_manager, '20010101', 0, 1, 'var', Box())) + self.assertEqual( + jobs[0], + VerticalMean(self.data_manager, "20010101", 0, 0, "var", Box()), + ) + self.assertEqual( + jobs[1], + VerticalMean(self.data_manager, "20010101", 0, 1, "var", Box()), + ) with self.assertRaises(DiagnosticOptionError): - VerticalMean.generate_jobs(self.diags, ['diagnostic']) + VerticalMean.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - VerticalMean.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0', '0']) + VerticalMean.generate_jobs( + self.diags, + ["diagnostic", "0", "0", "0", "0", "0", "0", "0", "0"], + ) def test_str(self): - self.assertEqual(str(self.mixed), - 'Vertical mean Startdate: 20000101 Member: 1 Chunk: 1 Variable: var Box: 0-100') + self.assertEqual( + str(self.mixed), + "Vertical mean Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: var Box: 0-100", + ) diff --git a/test/unit/ocean/test_verticalmeanmeters.py b/test/unit/ocean/test_verticalmeanmeters.py index 315e6a36..873a7f13 100644 --- a/test/unit/ocean/test_verticalmeanmeters.py +++ b/test/unit/ocean/test_verticalmeanmeters.py @@ -1,7 +1,10 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.diagnostic import DiagnosticVariableOption, DiagnosticOptionError +from earthdiagnostics.diagnostic import ( + DiagnosticVariableOption, + DiagnosticOptionError, +) from earthdiagnostics.box import Box from earthdiagnostics.ocean.verticalmeanmeters import VerticalMeanMeters from earthdiagnostics.modelingrealm import ModelingRealms @@ -9,56 +12,144 @@ from mock import Mock, patch class TestVerticalMeanMeters(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.model_version = 'model_version' - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) + self.diags.model_version = "model_version" + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) self.box = Box(True) self.box.min_depth = 0 self.box.max_depth = 100 - self.mixed = VerticalMeanMeters(self.data_manager, '20000101', 1, 1, ModelingRealms.ocean, 'var', self.box, 'T') + self.mixed = VerticalMeanMeters( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.ocean, + "var", + self.box, + "T", + ) def fake_parse(self, value): if not value: raise DiagnosticOptionError return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = VerticalMeanMeters.generate_jobs(self.diags, ['diagnostic', 'var', '0', '100']) + jobs = VerticalMeanMeters.generate_jobs( + self.diags, ["diagnostic", "var", "0", "100"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMeters(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - self.box, 'T')) - self.assertEqual(jobs[1], VerticalMeanMeters(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - self.box, 'T')) + self.assertEqual( + jobs[0], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + self.box, + "T", + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + self.box, + "T", + ), + ) - jobs = VerticalMeanMeters.generate_jobs(self.diags, ['diagnostic', 'var', '0']) + jobs = VerticalMeanMeters.generate_jobs( + self.diags, ["diagnostic", "var", "0"] + ) box = Box(True) box.min_depth = 0 self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMeters(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - box, 'T')) - self.assertEqual(jobs[1], VerticalMeanMeters(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - box, 'T')) + self.assertEqual( + jobs[0], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + box, + "T", + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + box, + "T", + ), + ) - jobs = VerticalMeanMeters.generate_jobs(self.diags, ['diagnostic', 'var']) + jobs = VerticalMeanMeters.generate_jobs( + self.diags, ["diagnostic", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], VerticalMeanMeters(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - Box(True), 'T')) - self.assertEqual(jobs[1], VerticalMeanMeters(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - Box(True), 'T')) + self.assertEqual( + jobs[0], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + Box(True), + "T", + ), + ) + self.assertEqual( + jobs[1], + VerticalMeanMeters( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + Box(True), + "T", + ), + ) with self.assertRaises(DiagnosticOptionError): - VerticalMeanMeters.generate_jobs(self.diags, ['diagnostic']) + VerticalMeanMeters.generate_jobs(self.diags, ["diagnostic"]) with self.assertRaises(DiagnosticOptionError): - VerticalMeanMeters.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0']) + VerticalMeanMeters.generate_jobs( + self.diags, ["diagnostic", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.mixed), - 'Vertical mean meters Startdate: 20000101 Member: 1 Chunk: 1 Variable: ocean:var Box: 0-100m') + self.assertEqual( + str(self.mixed), + "Vertical mean meters Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: ocean:var Box: 0-100m", + ) diff --git a/test/unit/statistics/test_climatologicalpercentile.py b/test/unit/statistics/test_climatologicalpercentile.py index 385174b6..6d57ea0a 100644 --- a/test/unit/statistics/test_climatologicalpercentile.py +++ b/test/unit/statistics/test_climatologicalpercentile.py @@ -1,7 +1,9 @@ # coding=utf-8 from unittest import TestCase -from earthdiagnostics.statistics.climatologicalpercentile import ClimatologicalPercentile +from earthdiagnostics.statistics.climatologicalpercentile import ( + ClimatologicalPercentile, +) from earthdiagnostics.diagnostic import DiagnosticVariableOption from mock import Mock, patch @@ -9,7 +11,6 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestClimatologicalPercentile(TestCase): - def setUp(self): self.data_manager = Mock() self.data_manager.variable_list.get_variable.return_value = None @@ -20,23 +21,46 @@ class TestClimatologicalPercentile(TestCase): def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = ClimatologicalPercentile.generate_jobs(self.diags, ['climpercent', 'ocean', 'var', '2000', '2001', '11']) + jobs = ClimatologicalPercentile.generate_jobs( + self.diags, ["climpercent", "ocean", "var", "2000", "2001", "11"] + ) self.assertEqual(len(jobs), 1) - self.assertEqual(jobs[0], ClimatologicalPercentile(self.data_manager, ModelingRealms.ocean, 'var', - 2000, 2001, 11, - self.diags.config.experiment)) + self.assertEqual( + jobs[0], + ClimatologicalPercentile( + self.data_manager, + ModelingRealms.ocean, + "var", + 2000, + 2001, + 11, + self.diags.config.experiment, + ), + ) with self.assertRaises(Exception): - ClimatologicalPercentile.generate_jobs(self.diags, ['climpercent']) + ClimatologicalPercentile.generate_jobs(self.diags, ["climpercent"]) with self.assertRaises(Exception): - ClimatologicalPercentile.generate_jobs(self.diags, ['climpercent', 'ocean', 'var', '2000', '2001', '11', - 'extra']) + ClimatologicalPercentile.generate_jobs( + self.diags, + ["climpercent", "ocean", "var", "2000", "2001", "11", "extra"], + ) def test_str(self): - diagnostic = ClimatologicalPercentile(self.data_manager, ModelingRealms.ocean, 'var', - 2000, 2001, 11, self.diags.config.experiment) + diagnostic = ClimatologicalPercentile( + self.data_manager, + ModelingRealms.ocean, + "var", + 2000, + 2001, + 11, + self.diags.config.experiment, + ) - self.assertEqual(str(diagnostic), - 'Climatological percentile Variable: ocean:var Period: 2000-2001 Forecast month: 11') + self.assertEqual( + str(diagnostic), + "Climatological percentile Variable: ocean:var Period: 2000-2001 " + "Forecast month: 11", + ) diff --git a/test/unit/statistics/test_daysoverpercentile.py b/test/unit/statistics/test_daysoverpercentile.py index caf59553..b03cfc6a 100644 --- a/test/unit/statistics/test_daysoverpercentile.py +++ b/test/unit/statistics/test_daysoverpercentile.py @@ -8,27 +8,67 @@ from earthdiagnostics.statistics.daysoverpercentile import DaysOverPercentile class TestDaysOverPercentile(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() - self.diags.config.experiment.get_chunk_list.return_value = (('20011101', 0, 0), ('20011101', 0, 1)) - self.diags.config.experiment.startdates = ('20001101', '20011101') + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20011101", 0, 0), + ("20011101", 0, 1), + ) + self.diags.config.experiment.startdates = ("20001101", "20011101") def test_generate_jobs(self): - jobs = DaysOverPercentile.generate_jobs(self.diags, ['monpercent', 'ocean', 'var', '2000', '2001', '11']) + jobs = DaysOverPercentile.generate_jobs( + self.diags, ["monpercent", "ocean", "var", "2000", "2001", "11"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], DaysOverPercentile(self.data_manager, ModelingRealms.ocean, 'var', 2000, 2001, - '20001101', 11)) - self.assertEqual(jobs[1], DaysOverPercentile(self.data_manager, ModelingRealms.ocean, 'var', 2000, 2001, - '20011101', 11)) + self.assertEqual( + jobs[0], + DaysOverPercentile( + self.data_manager, + ModelingRealms.ocean, + "var", + 2000, + 2001, + "20001101", + 11, + ), + ) + self.assertEqual( + jobs[1], + DaysOverPercentile( + self.data_manager, + ModelingRealms.ocean, + "var", + 2000, + 2001, + "20011101", + 11, + ), + ) with self.assertRaises(Exception): - DaysOverPercentile.generate_jobs(self.diags, ['monpercent', 'ocean', 'var', '2000', '2001']) + DaysOverPercentile.generate_jobs( + self.diags, ["monpercent", "ocean", "var", "2000", "2001"] + ) with self.assertRaises(Exception): - DaysOverPercentile.generate_jobs(self.diags, ['monpercent', 'ocean', 'var', '2000', '2001', '11', 'extra']) + DaysOverPercentile.generate_jobs( + self.diags, + ["monpercent", "ocean", "var", "2000", "2001", "11", "extra"], + ) def test_str(self): - diagnostic = DaysOverPercentile(self.data_manager, ModelingRealms.ocean, 'var', 2000, 2001, '20001101', 11) - self.assertEqual(str(diagnostic), - 'Days over percentile Startdate: 20001101 Variable: ocean:var Climatology: 2000-2001') + diagnostic = DaysOverPercentile( + self.data_manager, + ModelingRealms.ocean, + "var", + 2000, + 2001, + "20001101", + 11, + ) + self.assertEqual( + str(diagnostic), + "Days over percentile Startdate: 20001101 Variable: ocean:var " + "Climatology: 2000-2001", + ) diff --git a/test/unit/statistics/test_discretize.py b/test/unit/statistics/test_discretize.py index 78a71479..6e03092b 100644 --- a/test/unit/statistics/test_discretize.py +++ b/test/unit/statistics/test_discretize.py @@ -9,52 +9,141 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestClimatologicalPercentile(TestCase): - def setUp(self): self.data_manager = Mock() self.data_manager.variable_list.get_variable.return_value = None self.diags = Mock() self.diags.data_manager = self.data_manager - self.diags.config.experiment.startdates = ('20000101', '20010101') + self.diags.config.experiment.startdates = ("20000101", "20010101") def fake_parse(self, value): return value - @patch.object(DiagnosticVariableOption, 'parse', fake_parse) + @patch.object(DiagnosticVariableOption, "parse", fake_parse) def test_generate_jobs(self): - jobs = Discretize.generate_jobs(self.diags, ['discretize', 'ocean', 'var']) + jobs = Discretize.generate_jobs( + self.diags, ["discretize", "ocean", "var"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], Discretize(self.data_manager, '20000101', ModelingRealms.ocean, 'var', - 2000, float('nan'), float('nan'))) - self.assertEqual(jobs[1], Discretize(self.data_manager, '20010101', ModelingRealms.ocean, 'var', - 2000, float('nan'), float('nan'))) - - jobs = Discretize.generate_jobs(self.diags, ['discretize', 'ocean', 'var', '500']) - self.assertEqual(jobs[0], Discretize(self.data_manager, '20000101', ModelingRealms.ocean, 'var', - 500, float('nan'), float('nan'))) - self.assertEqual(jobs[1], Discretize(self.data_manager, '20010101', ModelingRealms.ocean, 'var', - 500, float('nan'), float('nan'))) - - jobs = Discretize.generate_jobs(self.diags, ['discretize', 'ocean', 'var', '500', '0', '40']) - self.assertEqual(jobs[0], Discretize(self.data_manager, '20000101', ModelingRealms.ocean, 'var', - 500, 0, 40)) - self.assertEqual(jobs[1], Discretize(self.data_manager, '20010101', ModelingRealms.ocean, 'var', - 500, 0, 40)) + self.assertEqual( + jobs[0], + Discretize( + self.data_manager, + "20000101", + ModelingRealms.ocean, + "var", + 2000, + float("nan"), + float("nan"), + ), + ) + self.assertEqual( + jobs[1], + Discretize( + self.data_manager, + "20010101", + ModelingRealms.ocean, + "var", + 2000, + float("nan"), + float("nan"), + ), + ) + + jobs = Discretize.generate_jobs( + self.diags, ["discretize", "ocean", "var", "500"] + ) + self.assertEqual( + jobs[0], + Discretize( + self.data_manager, + "20000101", + ModelingRealms.ocean, + "var", + 500, + float("nan"), + float("nan"), + ), + ) + self.assertEqual( + jobs[1], + Discretize( + self.data_manager, + "20010101", + ModelingRealms.ocean, + "var", + 500, + float("nan"), + float("nan"), + ), + ) + + jobs = Discretize.generate_jobs( + self.diags, ["discretize", "ocean", "var", "500", "0", "40"] + ) + self.assertEqual( + jobs[0], + Discretize( + self.data_manager, + "20000101", + ModelingRealms.ocean, + "var", + 500, + 0, + 40, + ), + ) + self.assertEqual( + jobs[1], + Discretize( + self.data_manager, + "20010101", + ModelingRealms.ocean, + "var", + 500, + 0, + 40, + ), + ) with self.assertRaises(Exception): - Discretize.generate_jobs(self.diags, ['discretize']) + Discretize.generate_jobs(self.diags, ["discretize"]) with self.assertRaises(Exception): - Discretize.generate_jobs(self.diags, ['discretize', 'ocean', 'var', '500', '0', '40', 'extra']) + Discretize.generate_jobs( + self.diags, + ["discretize", "ocean", "var", "500", "0", "40", "extra"], + ) def test_str(self): - diagnostic = Discretize(self.data_manager, '20000101', ModelingRealms.ocean, 'var', 2000, 10, 40) + diagnostic = Discretize( + self.data_manager, + "20000101", + ModelingRealms.ocean, + "var", + 2000, + 10, + 40, + ) - self.assertEqual(str(diagnostic), - 'Discretizing variable: ocean:var Startdate: 20000101 Bins: 2000 Range: [10, 40]') + self.assertEqual( + str(diagnostic), + "Discretizing variable: ocean:var Startdate: 20000101 Bins: 2000 " + "Range: [10, 40]", + ) - diagnostic = Discretize(self.data_manager, '20000101', ModelingRealms.ocean, 'var', 2000, - float('nan'), float('nan')) + diagnostic = Discretize( + self.data_manager, + "20000101", + ModelingRealms.ocean, + "var", + 2000, + float("nan"), + float("nan"), + ) - self.assertEqual(str(diagnostic), - 'Discretizing variable: ocean:var Startdate: 20000101 Bins: 2000 Range: [None, None]') + self.assertEqual( + str(diagnostic), + "Discretizing variable: ocean:var Startdate: 20000101 Bins: 2000 " + "Range: [None, None]", + ) diff --git a/test/unit/statistics/test_monthlypercentile.py b/test/unit/statistics/test_monthlypercentile.py index d0e4da48..5c7ddbc8 100644 --- a/test/unit/statistics/test_monthlypercentile.py +++ b/test/unit/statistics/test_monthlypercentile.py @@ -9,7 +9,6 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestMonthlyPercentile(TestCase): - def setUp(self): self.data_manager = Mock() self.diags = Mock() @@ -20,22 +19,60 @@ class TestMonthlyPercentile(TestCase): self.box.min_lon = 0 self.box.max_lon = 0 - self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1)) - self.diagnostic = MonthlyPercentile(self.data_manager, '20000101', 1, 1, ModelingRealms.ocean, 'var', [10, 90]) + self.diags.config.experiment.get_chunk_list.return_value = ( + ("20010101", 0, 0), + ("20010101", 0, 1), + ) + self.diagnostic = MonthlyPercentile( + self.data_manager, + "20000101", + 1, + 1, + ModelingRealms.ocean, + "var", + [10, 90], + ) def test_generate_jobs(self): - jobs = MonthlyPercentile.generate_jobs(self.diags, ['monpercent', 'ocean', 'var', '10-90']) + jobs = MonthlyPercentile.generate_jobs( + self.diags, ["monpercent", "ocean", "var", "10-90"] + ) self.assertEqual(len(jobs), 2) - self.assertEqual(jobs[0], MonthlyPercentile(self.data_manager, '20010101', 0, 0, ModelingRealms.ocean, 'var', - [10, 90])) - self.assertEqual(jobs[1], MonthlyPercentile(self.data_manager, '20010101', 0, 1, ModelingRealms.ocean, 'var', - [10, 90])) + self.assertEqual( + jobs[0], + MonthlyPercentile( + self.data_manager, + "20010101", + 0, + 0, + ModelingRealms.ocean, + "var", + [10, 90], + ), + ) + self.assertEqual( + jobs[1], + MonthlyPercentile( + self.data_manager, + "20010101", + 0, + 1, + ModelingRealms.ocean, + "var", + [10, 90], + ), + ) with self.assertRaises(Exception): - MonthlyPercentile.generate_jobs(self.diags, ['monpercent']) + MonthlyPercentile.generate_jobs(self.diags, ["monpercent"]) with self.assertRaises(Exception): - MonthlyPercentile.generate_jobs(self.diags, ['monpercent', '0', '0', '0', '0', '0', '0', '0']) + MonthlyPercentile.generate_jobs( + self.diags, ["monpercent", "0", "0", "0", "0", "0", "0", "0"] + ) def test_str(self): - self.assertEqual(str(self.diagnostic), 'Monthly percentile Startdate: 20000101 Member: 1 Chunk: 1 ' - 'Variable: ocean:var Percentiles: 10, 90') + self.assertEqual( + str(self.diagnostic), + "Monthly percentile Startdate: 20000101 Member: 1 Chunk: 1 " + "Variable: ocean:var Percentiles: 10, 90", + ) diff --git a/test/unit/test_box.py b/test/unit/test_box.py index 3134ef6f..ff89a667 100644 --- a/test/unit/test_box.py +++ b/test/unit/test_box.py @@ -87,25 +87,25 @@ class TestBox(TestCase): Box().min_lon = 80 def test_get_lat_str(self): - self.assertEqual('20S0N', self.box1.get_lat_str()) - self.assertEqual('20N', self.box2.get_lat_str()) - self.assertEqual('', self.box3.get_lat_str()) - self.assertEqual('20S10S', self.box4.get_lat_str()) + self.assertEqual("20S0N", self.box1.get_lat_str()) + self.assertEqual("20N", self.box2.get_lat_str()) + self.assertEqual("", self.box3.get_lat_str()) + self.assertEqual("20S10S", self.box4.get_lat_str()) def test_get_lon_str(self): - self.assertEqual('20W0E', self.box1.get_lon_str()) - self.assertEqual('20E', self.box2.get_lon_str()) - self.assertEqual('', self.box3.get_lon_str()) - self.assertEqual('20W10W', self.box4.get_lon_str()) + self.assertEqual("20W0E", self.box1.get_lon_str()) + self.assertEqual("20E", self.box2.get_lon_str()) + self.assertEqual("", self.box3.get_lon_str()) + self.assertEqual("20W10W", self.box4.get_lon_str()) def test_get_depth_str(self): - self.assertEqual('0-20', self.box1.get_depth_str()) - self.assertEqual('20m', self.box2.get_depth_str()) - self.assertEqual('', self.box3.get_depth_str()) - self.assertEqual('0-20', self.box4.get_depth_str()) + self.assertEqual("0-20", self.box1.get_depth_str()) + self.assertEqual("20m", self.box2.get_depth_str()) + self.assertEqual("", self.box3.get_depth_str()) + self.assertEqual("0-20", self.box4.get_depth_str()) def test__str__(self): - self.assertEqual('20S0N20W0E0-20', str(self.box1)) - self.assertEqual('20N20E20m', str(self.box2)) - self.assertEqual('', str(self.box3)) - self.assertEqual('20S10S20W10W0-20', str(self.box4)) + self.assertEqual("20S0N20W0E0-20", str(self.box1)) + self.assertEqual("20N20E20m", str(self.box2)) + self.assertEqual("", str(self.box3)) + self.assertEqual("20S10S20W10W0-20", str(self.box4)) diff --git a/test/unit/test_cdftools.py b/test/unit/test_cdftools.py index 26150c2f..f354186a 100644 --- a/test/unit/test_cdftools.py +++ b/test/unit/test_cdftools.py @@ -9,58 +9,112 @@ from earthdiagnostics.cdftools import CDFTools # noinspection PyUnusedLocal def bad_file(path, access=None): - return not os.path.basename(path).startswith('bad') + return not os.path.basename(path).startswith("bad") class TestCDFTools(TestCase): # noinspection PyUnusedLocal - @mock.patch('os.path.isfile', side_effect=bad_file) - @mock.patch('os.access', side_effect=bad_file) - @mock.patch('earthdiagnostics.utils.Utils.execute_shell_command') + @mock.patch("os.path.isfile", side_effect=bad_file) + @mock.patch("os.access", side_effect=bad_file) + @mock.patch("earthdiagnostics.utils.Utils.execute_shell_command") def test_run(self, mock_path, mock_exists, execute_mock): - self.cdftools = CDFTools('') - execute_mock.return_value = ['Command output'] + self.cdftools = CDFTools("") + execute_mock.return_value = ["Command output"] with self.assertRaises(ValueError): - self.cdftools.run('badcommand', input_file='input_file', output_file='output_file') + self.cdftools.run( + "badcommand", + input_file="input_file", + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file='badinput_file', output_file='output_file') + self.cdftools.run( + "command", + input_file="badinput_file", + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file=['input_file', 'badinput_file'], output_file='output_file') + self.cdftools.run( + "command", + input_file=["input_file", "badinput_file"], + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file='input_file', output_file='input_file') + self.cdftools.run( + "command", input_file="input_file", output_file="input_file" + ) with self.assertRaises(Exception): - self.cdftools.run('command', input_file='input_file', output_file='badoutput_file') + self.cdftools.run( + "command", + input_file="input_file", + output_file="badoutput_file", + ) - self.cdftools.run('command', input_file='input_file', output_file='output_file') - self.cdftools.run('command', input_option='-i', input_file='input_file', output_file='output_file') - self.cdftools.run('command', input_file='input_file') - self.cdftools.run('command', input_file=None) - self.cdftools.run('command', input_file=['input_file', 'input_file2']) - self.cdftools.run('command', input_file='input_file', options='-o -p') - self.cdftools.run('command', input_file='input_file', options=('-o', '-p')) + self.cdftools.run( + "command", input_file="input_file", output_file="output_file" + ) + self.cdftools.run( + "command", + input_option="-i", + input_file="input_file", + output_file="output_file", + ) + self.cdftools.run("command", input_file="input_file") + self.cdftools.run("command", input_file=None) + self.cdftools.run("command", input_file=["input_file", "input_file2"]) + self.cdftools.run("command", input_file="input_file", options="-o -p") + self.cdftools.run( + "command", input_file="input_file", options=("-o", "-p") + ) # noinspection PyUnusedLocal - @mock.patch('os.path.isfile', side_effect=bad_file) - @mock.patch('os.access', side_effect=bad_file) - @mock.patch('earthdiagnostics.utils.Utils.execute_shell_command') + @mock.patch("os.path.isfile", side_effect=bad_file) + @mock.patch("os.access", side_effect=bad_file) + @mock.patch("earthdiagnostics.utils.Utils.execute_shell_command") def test_run_with_path(self, mock_path, mock_exists, execute_mock): - self.cdftools = CDFTools('/some/path') - execute_mock.return_value = ['Command output'] + self.cdftools = CDFTools("/some/path") + execute_mock.return_value = ["Command output"] with self.assertRaises(ValueError): - self.cdftools.run('badcommand', input_file='input_file', output_file='output_file') + self.cdftools.run( + "badcommand", + input_file="input_file", + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file='badinput_file', output_file='output_file') + self.cdftools.run( + "command", + input_file="badinput_file", + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file=['input_file', 'badinput_file'], output_file='output_file') + self.cdftools.run( + "command", + input_file=["input_file", "badinput_file"], + output_file="output_file", + ) with self.assertRaises(ValueError): - self.cdftools.run('command', input_file='input_file', output_file='input_file') + self.cdftools.run( + "command", input_file="input_file", output_file="input_file" + ) with self.assertRaises(Exception): - self.cdftools.run('command', input_file='input_file', output_file='badoutput_file') + self.cdftools.run( + "command", + input_file="input_file", + output_file="badoutput_file", + ) - self.cdftools.run('command', input_file='input_file', output_file='output_file') - self.cdftools.run('command', input_option='-i', input_file='input_file', output_file='output_file') - self.cdftools.run('command', input_file='input_file') - self.cdftools.run('command', input_file=None) - self.cdftools.run('command', input_file=['input_file', 'input_file2']) - self.cdftools.run('command', input_file='input_file', options='-o -p') - self.cdftools.run('command', input_file='input_file', options=('-o', '-p')) + self.cdftools.run( + "command", input_file="input_file", output_file="output_file" + ) + self.cdftools.run( + "command", + input_option="-i", + input_file="input_file", + output_file="output_file", + ) + self.cdftools.run("command", input_file="input_file") + self.cdftools.run("command", input_file=None) + self.cdftools.run("command", input_file=["input_file", "input_file2"]) + self.cdftools.run("command", input_file="input_file", options="-o -p") + self.cdftools.run( + "command", input_file="input_file", options=("-o", "-p") + ) diff --git a/test/unit/test_cmormanager.py b/test/unit/test_cmormanager.py index 84660a49..d9b2811f 100644 --- a/test/unit/test_cmormanager.py +++ b/test/unit/test_cmormanager.py @@ -13,45 +13,49 @@ from earthdiagnostics.data_convention import DataConvention class TestCMORManager(TestCase): - def setUp(self): self.tmp_dir = tempfile.mkdtemp() - os.mkdir(os.path.join(self.tmp_dir, 'expid')) + os.mkdir(os.path.join(self.tmp_dir, "expid")) self.convention = create_autospec(DataConvention) - self.convention.name = 'specs' - self.convention.lat_name = 'lat' - self.convention.lon_name = 'lon' - self.convention.get_startdate_path.return_value = os.path.join(self.tmp_dir, 'expid', 'startdate') - self.convention.get_cmor_folder_path.return_value = os.path.join(self.tmp_dir, 'expid', 'startdate', 'cmor') - self.convention.get_file_name.return_value = 'filename' - self.convention.get_file_path.return_value = os.path.join(self.tmp_dir, 'expid', 'startdate', 'cmor', - 'filename') + self.convention.name = "specs" + self.convention.lat_name = "lat" + self.convention.lon_name = "lon" + self.convention.get_startdate_path.return_value = os.path.join( + self.tmp_dir, "expid", "startdate" + ) + self.convention.get_cmor_folder_path.return_value = os.path.join( + self.tmp_dir, "expid", "startdate", "cmor" + ) + self.convention.get_file_name.return_value = "filename" + self.convention.get_file_path.return_value = os.path.join( + self.tmp_dir, "expid", "startdate", "cmor", "filename" + ) self.convention.is_cmorized.return_value = True self.config = Mock() self.config.data_dir = self.tmp_dir - self.config.scratch_dir = os.path.join(self.tmp_dir, 'scratch') + self.config.scratch_dir = os.path.join(self.tmp_dir, "scratch") self.config.data_convention = self.convention - self.config.frequency = '6hr' - self.config.data_type = 'exp' - self.config.experiment.expid = 'expid' - self.config.experiment.model = 'model' - self.config.experiment.experiment_name = 'expname' - self.config.experiment.institute = 'institute' + self.config.frequency = "6hr" + self.config.data_type = "exp" + self.config.experiment.expid = "expid" + self.config.experiment.model = "model" + self.config.experiment.experiment_name = "expname" + self.config.experiment.institute = "institute" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 12 self.config.experiment.num_chunks = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" self.config.experiment.atmos_timestep = 3 self.config.experiment.ocean_timestep = 6 - self.config.experiment.get_member_str.return_value = 'r1i1p1' + self.config.experiment.get_member_str.return_value = "r1i1p1" self.config.cmor.initialization_number = 1 - self.config.cmor.version = '' - self.config.cmor.default_ocean_grid = 'ocean_grid' - self.config.cmor.default_atmos_grid = 'atmos_grid' - self.config.cmor.activity = 'activity' + self.config.cmor.version = "" + self.config.cmor.default_ocean_grid = "ocean_grid" + self.config.cmor.default_atmos_grid = "atmos_grid" + self.config.cmor.activity = "activity" self.config.cmor.force = False self.config.cmor.force_untar = False self.config.cmor.cmorize_atmosphere = True @@ -64,7 +68,9 @@ class TestCMORManager(TestCase): def test_find_data(self): cmor_manager = CMORManager(self.config) - self.assertEqual(cmor_manager.cmor_path, os.path.join(self.tmp_dir, 'expid')) + self.assertEqual( + cmor_manager.cmor_path, os.path.join(self.tmp_dir, "expid") + ) def test_find_data_fail(self): os.rmdir(os.path.join(self.tmp_dir, self.config.experiment.expid)) @@ -72,156 +78,235 @@ class TestCMORManager(TestCase): CMORManager(self.config) def test_find_data_with_model(self): - os.makedirs(os.path.join(self.tmp_dir, 'model', self.config.experiment.expid)) + os.makedirs( + os.path.join(self.tmp_dir, "model", self.config.experiment.expid) + ) os.rmdir(os.path.join(self.tmp_dir, self.config.experiment.expid)) cmor_manager = CMORManager(self.config) - self.assertEqual(cmor_manager.cmor_path, os.path.join(self.tmp_dir, 'model', 'expid')) + self.assertEqual( + cmor_manager.cmor_path, + os.path.join(self.tmp_dir, "model", "expid"), + ) def test_find_data_with_ecearth_fix(self): - self.config.experiment.model = 'EC-Earth' - os.makedirs(os.path.join(self.tmp_dir, 'ecearth', self.config.experiment.expid)) + self.config.experiment.model = "EC-Earth" + os.makedirs( + os.path.join(self.tmp_dir, "ecearth", self.config.experiment.expid) + ) os.rmdir(os.path.join(self.tmp_dir, self.config.experiment.expid)) cmor_manager = CMORManager(self.config) - self.assertEqual(cmor_manager.cmor_path, os.path.join(self.tmp_dir, 'ecearth', 'expid')) + self.assertEqual( + cmor_manager.cmor_path, + os.path.join(self.tmp_dir, "ecearth", "expid"), + ) def test_find_data_with_type_and_model(self): - os.makedirs(os.path.join(self.tmp_dir, 'exp', 'model', self.config.experiment.expid)) + os.makedirs( + os.path.join( + self.tmp_dir, "exp", "model", self.config.experiment.expid + ) + ) os.rmdir(os.path.join(self.tmp_dir, self.config.experiment.expid)) cmor_manager = CMORManager(self.config) - self.assertEqual(cmor_manager.cmor_path, os.path.join(self.tmp_dir, 'exp', 'model', 'expid')) + self.assertEqual( + cmor_manager.cmor_path, + os.path.join(self.tmp_dir, "exp", "model", "expid"), + ) def test_file_exists(self): - with mock.patch('os.path.isfile') as isfile: + with mock.patch("os.path.isfile") as isfile: cmor_manager = CMORManager(self.config) isfile.return_value = True - self.assertTrue(cmor_manager.file_exists(ModelingRealms.ocean, 'var', '20011101', 1, 1)) + self.assertTrue( + cmor_manager.file_exists( + ModelingRealms.ocean, "var", "20011101", 1, 1 + ) + ) isfile.return_value = False - self.assertFalse(cmor_manager.file_exists(ModelingRealms.ocean, 'var', '20011101', 1, 1)) + self.assertFalse( + cmor_manager.file_exists( + ModelingRealms.ocean, "var", "20011101", 1, 1 + ) + ) def test_file_exists_multiple_versions(self): - with mock.patch('os.path.isfile') as isfile: + with mock.patch("os.path.isfile") as isfile: cmor_manager = CMORManager(self.config) isfile.return_value = True - self.assertTrue(cmor_manager.file_exists(ModelingRealms.ocean, 'var', '20011101', 1, 1, - possible_versions=('version1', 'version2'))) + self.assertTrue( + cmor_manager.file_exists( + ModelingRealms.ocean, + "var", + "20011101", + 1, + 1, + possible_versions=("version1", "version2"), + ) + ) isfile.return_value = False - self.assertFalse(cmor_manager.file_exists(ModelingRealms.ocean, 'var', '20011101', 1, 1, - possible_versions=('version1', 'version2'))) - - @mock.patch('earthdiagnostics.cmormanager.Cmorizer', autospec=True) + self.assertFalse( + cmor_manager.file_exists( + ModelingRealms.ocean, + "var", + "20011101", + 1, + 1, + possible_versions=("version1", "version2"), + ) + ) + + @mock.patch("earthdiagnostics.cmormanager.Cmorizer", autospec=True) def test_prepare_cmorize(self, mock_cmor): mock_instance = mock_cmor.return_value self.convention.is_cmorized.return_value = False cmor_manager = CMORManager(self.config) - self.config.experiment.get_member_list.return_value = (('20000101', 2),) + self.config.experiment.get_member_list.return_value = ( + ("20000101", 2), + ) cmor_manager.prepare() mock_instance.cmorize_ocean.assert_called_once() mock_instance.cmorize_atmos.assert_called_once() - @mock.patch('earthdiagnostics.cmormanager.Cmorizer', autospec=True) + @mock.patch("earthdiagnostics.cmormanager.Cmorizer", autospec=True) def test_prepare_cmorize_force(self, mock_cmor): mock_instance = mock_cmor.return_value self.config.cmor.force = True cmor_manager = CMORManager(self.config) - self.config.experiment.get_member_list.return_value = (('20000101', 2),) + self.config.experiment.get_member_list.return_value = ( + ("20000101", 2), + ) cmor_manager.prepare() mock_instance.cmorize_ocean.assert_called_once() mock_instance.cmorize_atmos.assert_called_once() - @mock.patch('earthdiagnostics.cmormanager.CMORManager.is_cmorized') - @mock.patch('earthdiagnostics.utils.Utils.unzip') - @mock.patch('earthdiagnostics.utils.Utils.untar') - def test_prepare_cmorize_force_untar(self, mock_untar, mock_unzip, mock_cmorized): - original_cmor_path = os.path.join(self.config.data_dir, self.config.experiment.expid, - 'original_files', 'cmorfiles') + @mock.patch("earthdiagnostics.cmormanager.CMORManager.is_cmorized") + @mock.patch("earthdiagnostics.utils.Utils.unzip") + @mock.patch("earthdiagnostics.utils.Utils.untar") + def test_prepare_cmorize_force_untar( + self, mock_untar, mock_unzip, mock_cmorized + ): + original_cmor_path = os.path.join( + self.config.data_dir, + self.config.experiment.expid, + "original_files", + "cmorfiles", + ) os.makedirs(original_cmor_path) - self.config.experiment.get_chunk_start_str.return_value = '20000101' - cmor_prefix = 'CMORT_{0}_{1}_{2}_{3}-'.format(self.config.experiment.expid, '20000101', 'r1i1p1', - '20000101') - tempfile.mkstemp('.tar.gz', cmor_prefix, original_cmor_path) - tempfile.mkstemp('.tar', cmor_prefix, original_cmor_path) + self.config.experiment.get_chunk_start_str.return_value = "20000101" + cmor_prefix = "CMORT_{0}_{1}_{2}_{3}-".format( + self.config.experiment.expid, "20000101", "r1i1p1", "20000101" + ) + tempfile.mkstemp(".tar.gz", cmor_prefix, original_cmor_path) + tempfile.mkstemp(".tar", cmor_prefix, original_cmor_path) mock_cmorized.return_value = True self.config.cmor.force_untar = True cmor_manager = CMORManager(self.config) - self.config.experiment.get_member_list.return_value = (('20000101', 2),) + self.config.experiment.get_member_list.return_value = ( + ("20000101", 2), + ) cmor_manager.prepare() self.convention.create_links.assert_called_once() mock_unzip.assert_called_once() mock_untar.assert_called_once() - @mock.patch('earthdiagnostics.cmormanager.CMORManager.is_cmorized') + @mock.patch("earthdiagnostics.cmormanager.CMORManager.is_cmorized") def test_prepare_already_cmorized(self, mock_cmorized): mock_cmorized.return_value = True self.config.cmor.force_untar = False cmor_manager = CMORManager(self.config) - self.config.experiment.get_member_list.return_value = (('20000101', 2),) + self.config.experiment.get_member_list.return_value = ( + ("20000101", 2), + ) cmor_manager.prepare() - @mock.patch('earthdiagnostics.cmormanager.CMORManager.is_cmorized') + @mock.patch("earthdiagnostics.cmormanager.CMORManager.is_cmorized") def test_prepare_cmorization_not_requested(self, mock_cmorized): mock_cmorized.return_value = False self.config.cmor.chunk_cmorization_requested.return_value = False cmor_manager = CMORManager(self.config) - self.config.experiment.get_member_list.return_value = (('20000101', 2),) + self.config.experiment.get_member_list.return_value = ( + ("20000101", 2), + ) cmor_manager.prepare() def test_prepare_meteofrance(self): - self.convention.name = 'meteofrance' + self.convention.name = "meteofrance" cmor_manager = CMORManager(self.config) cmor_manager.prepare() def test_is_not_cmorized(self): self.convention.is_cmorized.return_value = False cmor_manager = CMORManager(self.config) - self.assertFalse(cmor_manager.is_cmorized('20000101', 1, 1, ModelingRealms.ocean)) + self.assertFalse( + cmor_manager.is_cmorized("20000101", 1, 1, ModelingRealms.ocean) + ) def test_link_file(self): cmor_manager = CMORManager(self.config) cmor_var = Mock() - cmor_manager.link_file(ModelingRealms.ocean, 'var', cmor_var, '20010101', 1, 1) + cmor_manager.link_file( + ModelingRealms.ocean, "var", cmor_var, "20010101", 1, 1 + ) self.convention.create_link.assert_called_once() def test_request_chunk(self): cmor_manager = CMORManager(self.config) - datafile = cmor_manager.request_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) + datafile = cmor_manager.request_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) def _mock_filepath(self): - return os.path.join(self.tmp_dir, 'expid/startdate/cmor/filename') + return os.path.join(self.tmp_dir, "expid/startdate/cmor/filename") def test_declare_chunk(self): cmor_manager = CMORManager(self.config) - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) mock_frequency = Mock() - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1, - frequency=mock_frequency) + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, + "var", + "20010101", + 1, + 1, + frequency=mock_frequency, + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) self.config.var_manager.get_variable.return_value = None - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) def test_request_year(self): self.config.experiment.get_year_chunks.return_value = (1, 2) mock_diagnostic = Mock() cmor_manager = CMORManager(self.config) - datafile = cmor_manager.request_year(mock_diagnostic, ModelingRealms.ocean, 'var', '20010101', 1, 2000) + datafile = cmor_manager.request_year( + mock_diagnostic, ModelingRealms.ocean, "var", "20010101", 1, 2000 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) def test_declare_year(self): cmor_manager = CMORManager(self.config) - datafile = cmor_manager.declare_year(ModelingRealms.ocean, 'var', '20010101', 1, 2001) + datafile = cmor_manager.declare_year( + ModelingRealms.ocean, "var", "20010101", 1, 2001 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) self.config.var_manager.get_variable.return_value = None - datafile = cmor_manager.declare_year(ModelingRealms.ocean, 'var', '20010101', 1, 2001) + datafile = cmor_manager.declare_year( + ModelingRealms.ocean, "var", "20010101", 1, 2001 + ) self.assertEqual(datafile.remote_file, self._mock_filepath()) class TestMergeYear(TestCase): - def setUp(self): self.data_manager = Mock() @@ -229,7 +314,41 @@ class TestMergeYear(TestCase): self.assertIsNone(MergeYear.generate_jobs(None, None)) def test_eq(self): - self.assertEqual(MergeYear(self.data_manager, ModelingRealms.ocean, 'var', 'startdate', 1, 1998), - MergeYear(self.data_manager, ModelingRealms.ocean, 'var', 'startdate', 1, 1998)) - self.assertEqual(MergeYear(self.data_manager, ModelingRealms.ocean, 'var', 'startdate', 1, 1998, 'grid'), - MergeYear(self.data_manager, ModelingRealms.ocean, 'var', 'startdate', 1, 1998, 'grid')) + self.assertEqual( + MergeYear( + self.data_manager, + ModelingRealms.ocean, + "var", + "startdate", + 1, + 1998, + ), + MergeYear( + self.data_manager, + ModelingRealms.ocean, + "var", + "startdate", + 1, + 1998, + ), + ) + self.assertEqual( + MergeYear( + self.data_manager, + ModelingRealms.ocean, + "var", + "startdate", + 1, + 1998, + "grid", + ), + MergeYear( + self.data_manager, + ModelingRealms.ocean, + "var", + "startdate", + 1, + 1998, + "grid", + ), + ) diff --git a/test/unit/test_config.py b/test/unit/test_config.py index b7ead20d..73102e22 100644 --- a/test/unit/test_config.py +++ b/test/unit/test_config.py @@ -5,11 +5,23 @@ import datetime import mock import os -from earthdiagnostics.config import CMORConfig, ConfigException, THREDDSConfig, ReportConfig, ExperimentConfig, Config +from earthdiagnostics.config import ( + CMORConfig, + ConfigException, + THREDDSConfig, + ReportConfig, + ExperimentConfig, + Config, +) from earthdiagnostics.frequency import Frequencies from earthdiagnostics.modelingrealm import ModelingRealms -from earthdiagnostics.data_convention import SPECSConvention, PrimaveraConvention, PrefaceConvention, \ - MeteoFranceConvention, CMIP6Convention +from earthdiagnostics.data_convention import ( + SPECSConvention, + PrimaveraConvention, + PrefaceConvention, + MeteoFranceConvention, + CMIP6Convention, +) class VariableMock(object): @@ -17,10 +29,12 @@ class VariableMock(object): def __init__(self): self.domain = ModelingRealms.ocean - self.short_name = 'tos' + self.short_name = "tos" def __eq__(self, other): - return self.domain == other.domain and self.short_name == other.short_name + return ( + self.domain == other.domain and self.short_name == other.short_name + ) class VariableManagerMock(object): @@ -28,7 +42,7 @@ class VariableManagerMock(object): def get_variable(self, alias, silent=False): """Return a VariableMock given an alias""" - if alias == 'bad': + if alias == "bad": return None var = VariableMock() var.short_name = alias @@ -48,7 +62,7 @@ class ParserMock(mock.Mock): def get_var_string(self, section, var): """Get var string""" - return '{0}:{1}'.format(section, var) + return "{0}:{1}".format(section, var) def get_value(self, section, var, default): """Get value from mock parser""" @@ -69,21 +83,33 @@ class ParserMock(mock.Mock): """Get integer option""" return self.get_value(section, var, default) - def get_choice_option(self, section, var, choices, default, ignore_case=True): + def get_choice_option( + self, section, var, choices, default, ignore_case=True + ): """Get choice option""" return self.get_value(section, var, default) - def get_int_list_option(self, section, var, default=list(), separator=' '): + def get_int_list_option(self, section, var, default=list(), separator=" "): """Get integer list option""" try: - return [int(val) for val in self._values[self.get_var_string(section, var)].split(separator)] + return [ + int(val) + for val in self._values[ + self.get_var_string(section, var) + ].split(separator) + ] except KeyError: return default - def get_list_option(self, section, var, default=list(), separator=' '): + def get_list_option(self, section, var, default=list(), separator=" "): """Get list option""" try: - return [val for val in self._values[self.get_var_string(section, var)].split(separator)] + return [ + val + for val in self._values[ + self.get_var_string(section, var) + ].split(separator) + ] except KeyError: return default @@ -93,12 +119,12 @@ class ParserMock(mock.Mock): def has_section(self, section): """Check if section exists""" - start = '{0}:'.format(section) + start = "{0}:".format(section) return any(x.startswith(start) for x in self._values) def options(self, section): """Return all options for a given section""" - start = '{0}:'.format(section) + start = "{0}:".format(section) return [x[len(start):] for x in self._values if x.startswith(start)] @@ -118,19 +144,19 @@ class TestCMORConfig(TestCase): self.assertEqual(config.force, False) self.assertEqual(config.force_untar, False) self.assertEqual(config.use_grib, True) - self.assertEqual(config.activity, 'CMIP') - self.assertEqual(config.associated_experiment, 'to be filled') - self.assertEqual(config.associated_model, 'to be filled') - self.assertEqual(config.initialization_description, 'to be filled') - self.assertEqual(config.initialization_method, '1') + self.assertEqual(config.activity, "CMIP") + self.assertEqual(config.associated_experiment, "to be filled") + self.assertEqual(config.associated_model, "to be filled") + self.assertEqual(config.initialization_description, "to be filled") + self.assertEqual(config.initialization_method, "1") self.assertEqual(config.initialization_number, 1) - self.assertEqual(config.source, 'to be filled') - self.assertEqual(config.version, '') - self.assertEqual(config.physics_version, '1') - self.assertEqual(config.physics_description, 'to be filled') + self.assertEqual(config.source, "to be filled") + self.assertEqual(config.version, "") + self.assertEqual(config.physics_version, "1") + self.assertEqual(config.physics_description, "to be filled") self.assertEqual(config.filter_files, []) - self.assertEqual(config.default_atmos_grid, 'gr') - self.assertEqual(config.default_ocean_grid, 'gn') + self.assertEqual(config.default_atmos_grid, "gr") + self.assertEqual(config.default_ocean_grid, "gn") self.assertEqual(config.min_cmorized_vars, 10) self.assertEqual(config.append_startdate, False) @@ -142,33 +168,35 @@ class TestCMORConfig(TestCase): def test_cmorize_list(self): """Test cmorize return true for variables in the given list""" - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'ocean:thetao ocean:tos') + self.mock_parser.add_value( + "CMOR", "VARIABLE_LIST", "ocean:thetao ocean:tos" + ) config = CMORConfig(self.mock_parser, self.var_manager) self.assertTrue(config.cmorize(VariableMock())) thetao_mock = VariableMock() thetao_mock.domain = ModelingRealms.ocean - thetao_mock.short_name = 'thetao' + thetao_mock.short_name = "thetao" self.assertTrue(config.cmorize(thetao_mock)) def test_bad_list(self): """Test cmorize return false for malformed variables""" - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', '#ocean:tos') + self.mock_parser.add_value("CMOR", "VARIABLE_LIST", "#ocean:tos") with self.assertRaises(ConfigException): CMORConfig(self.mock_parser, self.var_manager) - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'atmos:tos') + self.mock_parser.add_value("CMOR", "VARIABLE_LIST", "atmos:tos") with self.assertRaises(ConfigException): CMORConfig(self.mock_parser, self.var_manager) - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'ocean:bad') + self.mock_parser.add_value("CMOR", "VARIABLE_LIST", "ocean:bad") with self.assertRaises(ConfigException): CMORConfig(self.mock_parser, self.var_manager) def test_not_cmorize(self): """Test cmorize return false for variables not in the given list""" - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'ocean:tos') + self.mock_parser.add_value("CMOR", "VARIABLE_LIST", "ocean:tos") config = CMORConfig(self.mock_parser, self.var_manager) self.assertTrue(config.cmorize(VariableMock())) @@ -177,38 +205,42 @@ class TestCMORConfig(TestCase): tas_mock = VariableMock() tas_mock.domain = ModelingRealms.atmos - tas_mock.short_name = 'tas' + tas_mock.short_name = "tas" self.assertFalse(config.cmorize(tas_mock)) thetao_mock = VariableMock() thetao_mock.domain = ModelingRealms.ocean - thetao_mock.short_name = 'thetao' + thetao_mock.short_name = "thetao" self.assertFalse(config.cmorize(thetao_mock)) def test_comment(self): """Test cmorize return false for variables after the comment mark""" - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'ocean:tos #ocean:thetao ') + self.mock_parser.add_value( + "CMOR", "VARIABLE_LIST", "ocean:tos #ocean:thetao " + ) config = CMORConfig(self.mock_parser, self.var_manager) self.assertTrue(config.cmorize(VariableMock())) thetao_mock = VariableMock() thetao_mock.domain = ModelingRealms.ocean - thetao_mock.short_name = 'thetao' + thetao_mock.short_name = "thetao" self.assertFalse(config.cmorize(thetao_mock)) - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', '#ocean:tos ocean:thetao ') + self.mock_parser.add_value( + "CMOR", "VARIABLE_LIST", "#ocean:tos ocean:thetao " + ) with self.assertRaises(ConfigException): CMORConfig(self.mock_parser, self.var_manager) def test_cmorization_chunk(self): - """Test chunk cmorization requested return always true if not given a list""" + """Test chunk cmorization requested returns true if not given a list""" config = CMORConfig(self.mock_parser, self.var_manager) self.assertTrue(config.chunk_cmorization_requested(1)) def test_cmorize_only_some_chunks(self): - """Test chunk cmorization requested return true only if the value is in the list""" - self.mock_parser.add_value('CMOR', 'CHUNKS', '3 5') + """Test chunk cmor requested returns true only if it is in the list""" + self.mock_parser.add_value("CMOR", "CHUNKS", "3 5") config = CMORConfig(self.mock_parser, self.var_manager) self.assertTrue(config.chunk_cmorization_requested(3)) self.assertTrue(config.chunk_cmorization_requested(5)) @@ -219,72 +251,106 @@ class TestCMORConfig(TestCase): def test_any_required(self): """Test any_required method""" config = CMORConfig(self.mock_parser, self.var_manager) - self.assertTrue(config.any_required(['tos'])) + self.assertTrue(config.any_required(["tos"])) - self.mock_parser.add_value('CMOR', 'VARIABLE_LIST', 'ocean:tos ocean:thetao') + self.mock_parser.add_value( + "CMOR", "VARIABLE_LIST", "ocean:tos ocean:thetao" + ) config = CMORConfig(self.mock_parser, self.var_manager) - self.assertTrue(config.any_required(['tos', 'thetao', 'tas'])) - self.assertTrue(config.any_required(['tos', 'tas'])) - self.assertTrue(config.any_required(['thetao'])) + self.assertTrue(config.any_required(["tos", "thetao", "tas"])) + self.assertTrue(config.any_required(["tos", "tas"])) + self.assertTrue(config.any_required(["thetao"])) - self.assertFalse(config.any_required(['tas'])) + self.assertFalse(config.any_required(["tas"])) def test_hourly_vars(self): """Test hourly vars specification""" config = CMORConfig(self.mock_parser, self.var_manager) self.assertEqual(config.get_variables(Frequencies.six_hourly), {}) - self.mock_parser.add_value('CMOR', 'ATMOS_HOURLY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') + self.mock_parser.add_value( + "CMOR", + "ATMOS_HOURLY_VARS", + "128,129:1,130:1-2,131:1:10,132:0:10:5", + ) config = CMORConfig(self.mock_parser, self.var_manager) - self.assertEqual(config.get_variables(Frequencies.six_hourly), {128: None, - 129: '1', - 130: '1,2', - 131: '1,2,3,4,5,6,7,8,9', - 132: '0,5'}) + self.assertEqual( + config.get_variables(Frequencies.six_hourly), + { + 128: None, + 129: "1", + 130: "1,2", + 131: "1,2,3,4,5,6,7,8,9", + 132: "0,5", + }, + ) self.assertEqual(config.get_levels(Frequencies.six_hourly, 128), None) - self.assertEqual(config.get_levels(Frequencies.six_hourly, 129), '1') - self.assertEqual(config.get_levels(Frequencies.six_hourly, 130), '1,2') - self.assertEqual(config.get_levels(Frequencies.six_hourly, 131), '1,2,3,4,5,6,7,8,9',) - self.assertEqual(config.get_levels(Frequencies.six_hourly, 132), '0,5') + self.assertEqual(config.get_levels(Frequencies.six_hourly, 129), "1") + self.assertEqual(config.get_levels(Frequencies.six_hourly, 130), "1,2") + self.assertEqual( + config.get_levels(Frequencies.six_hourly, 131), + "1,2,3,4,5,6,7,8,9", + ) + self.assertEqual(config.get_levels(Frequencies.six_hourly, 132), "0,5") def test_daily_vars(self): """Test daily vars specification""" config = CMORConfig(self.mock_parser, self.var_manager) self.assertEqual(config.get_variables(Frequencies.daily), {}) - self.mock_parser.add_value('CMOR', 'ATMOS_DAILY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') + self.mock_parser.add_value( + "CMOR", "ATMOS_DAILY_VARS", "128,129:1,130:1-2,131:1:10,132:0:10:5" + ) config = CMORConfig(self.mock_parser, self.var_manager) - self.assertEqual(config.get_variables(Frequencies.daily), {128: None, - 129: '1', - 130: '1,2', - 131: '1,2,3,4,5,6,7,8,9', - 132: '0,5'}) + self.assertEqual( + config.get_variables(Frequencies.daily), + { + 128: None, + 129: "1", + 130: "1,2", + 131: "1,2,3,4,5,6,7,8,9", + 132: "0,5", + }, + ) self.assertEqual(config.get_levels(Frequencies.daily, 128), None) - self.assertEqual(config.get_levels(Frequencies.daily, 129), '1') - self.assertEqual(config.get_levels(Frequencies.daily, 130), '1,2') - self.assertEqual(config.get_levels(Frequencies.daily, 131), '1,2,3,4,5,6,7,8,9',) - self.assertEqual(config.get_levels(Frequencies.daily, 132), '0,5') + self.assertEqual(config.get_levels(Frequencies.daily, 129), "1") + self.assertEqual(config.get_levels(Frequencies.daily, 130), "1,2") + self.assertEqual( + config.get_levels(Frequencies.daily, 131), "1,2,3,4,5,6,7,8,9", + ) + self.assertEqual(config.get_levels(Frequencies.daily, 132), "0,5") def test_monthly_vars(self): """Test monthly vars specification""" config = CMORConfig(self.mock_parser, self.var_manager) self.assertEqual(config.get_variables(Frequencies.monthly), {}) - self.mock_parser.add_value('CMOR', 'ATMOS_MONTHLY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') + self.mock_parser.add_value( + "CMOR", + "ATMOS_MONTHLY_VARS", + "128,129:1,130:1-2,131:1:10,132:0:10:5", + ) config = CMORConfig(self.mock_parser, self.var_manager) - self.assertEqual(config.get_variables(Frequencies.monthly), {128: None, - 129: '1', - 130: '1,2', - 131: '1,2,3,4,5,6,7,8,9', - 132: '0,5'}) + self.assertEqual( + config.get_variables(Frequencies.monthly), + { + 128: None, + 129: "1", + 130: "1,2", + 131: "1,2,3,4,5,6,7,8,9", + 132: "0,5", + }, + ) self.assertEqual(config.get_levels(Frequencies.monthly, 128), None) - self.assertEqual(config.get_levels(Frequencies.monthly, 129), '1') - self.assertEqual(config.get_levels(Frequencies.monthly, 130), '1,2') - self.assertEqual(config.get_levels(Frequencies.monthly, 131), '1,2,3,4,5,6,7,8,9',) - self.assertEqual(config.get_levels(Frequencies.monthly, 132), '0,5') + self.assertEqual(config.get_levels(Frequencies.monthly, 129), "1") + self.assertEqual(config.get_levels(Frequencies.monthly, 130), "1,2") + self.assertEqual( + config.get_levels(Frequencies.monthly, 131), "1,2,3,4,5,6,7,8,9", + ) + self.assertEqual(config.get_levels(Frequencies.monthly, 132), "0,5") def test_bad_frequency_vars(self): """Test get variables fails if a bada frequency is specified""" @@ -293,13 +359,25 @@ class TestCMORConfig(TestCase): config.get_variables(Frequencies.climatology) def test_requested_codes(self): - """Test requested codes returns a set of the codes requested at all frequencies""" - self.mock_parser.add_value('CMOR', 'ATMOS_HOURLY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') - self.mock_parser.add_value('CMOR', 'ATMOS_DAILY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') - self.mock_parser.add_value('CMOR', 'ATMOS_MONTHLY_VARS', '128,129:1,130:1-2,131:1:10,132:0:10:5') + """Test requested codes returns a set of codes at all frequencies""" + self.mock_parser.add_value( + "CMOR", + "ATMOS_HOURLY_VARS", + "128,129:1,130:1-2,131:1:10,132:0:10:5", + ) + self.mock_parser.add_value( + "CMOR", "ATMOS_DAILY_VARS", "128,129:1,130:1-2,131:1:10,132:0:10:5" + ) + self.mock_parser.add_value( + "CMOR", + "ATMOS_MONTHLY_VARS", + "128,129:1,130:1-2,131:1:10,132:0:10:5", + ) config = CMORConfig(self.mock_parser, self.var_manager) - self.assertEqual(config.get_requested_codes(), {128, 129, 130, 131, 132}) + self.assertEqual( + config.get_requested_codes(), {128, 129, 130, 131, 132} + ) class TestTHREDDSConfig(TestCase): @@ -312,13 +390,13 @@ class TestTHREDDSConfig(TestCase): def test_basic_config(self): """Test basic configuration""" config = THREDDSConfig(self.mock_parser) - self.assertEqual(config.server_url, '') + self.assertEqual(config.server_url, "") def test_url(self): """Test SERVER_URL parameter""" - self.mock_parser.add_value('THREDDS', 'SERVER_URL', 'test_url') + self.mock_parser.add_value("THREDDS", "SERVER_URL", "test_url") config = THREDDSConfig(self.mock_parser) - self.assertEqual(config.server_url, 'test_url') + self.assertEqual(config.server_url, "test_url") class TestReportConfig(TestCase): @@ -331,18 +409,18 @@ class TestReportConfig(TestCase): def test_basic_config(self): """test default config""" config = ReportConfig(self.mock_parser) - self.assertEqual(config.path, '') + self.assertEqual(config.path, "") self.assertEqual(config.maximum_priority, 10) def test_path(self): """Test path configuration""" - self.mock_parser.add_value('REPORT', 'PATH', 'new_path') + self.mock_parser.add_value("REPORT", "PATH", "new_path") config = ReportConfig(self.mock_parser) - self.assertEqual(config.path, 'new_path') + self.assertEqual(config.path, "new_path") def test_priority(self): """Test maximum priority configuration""" - self.mock_parser.add_value('REPORT', 'MAXIMUM_PRIORITY', 3) + self.mock_parser.add_value("REPORT", "MAXIMUM_PRIORITY", 3) config = ReportConfig(self.mock_parser) self.assertEqual(config.maximum_priority, 3) @@ -364,7 +442,7 @@ class TestExperimentConfig(TestCase): self.assertEqual(config.chunk_size, 0) self.assertEqual(config.num_chunks, 0) - self.assertEqual(config.atmos_grid, '') + self.assertEqual(config.atmos_grid, "") self.assertEqual(config.atmos_timestep, 6) self.assertEqual(config.ocean_timestep, 6) @@ -375,24 +453,25 @@ class TestExperimentConfig(TestCase): Including: - Adding the prefix string or not - Providing them as a space separated list - - Providing them as a range with start and end separated by - (both extremes included) + - Providing them as a range with start and end separated by - + (both extremes included) """ - self.mock_parser.add_value('EXPERIMENT', 'MEMBERS', 'fc0 1') + self.mock_parser.add_value("EXPERIMENT", "MEMBERS", "fc0 1") config = ExperimentConfig() config.parse_ini(self.mock_parser) self.assertEqual(config.members, [0, 1]) - self.mock_parser.add_value('EXPERIMENT', 'MEMBERS', 'fc00 fc01') + self.mock_parser.add_value("EXPERIMENT", "MEMBERS", "fc00 fc01") config = ExperimentConfig() config.parse_ini(self.mock_parser) self.assertEqual(config.members, [0, 1]) - self.mock_parser.add_value('EXPERIMENT', 'MEMBERS', 'fc1-fc3') + self.mock_parser.add_value("EXPERIMENT", "MEMBERS", "fc1-fc3") config = ExperimentConfig() config.parse_ini(self.mock_parser) self.assertEqual(config.members, [1, 2, 3]) - self.mock_parser.add_value('EXPERIMENT', 'MEMBERS', '1-3') + self.mock_parser.add_value("EXPERIMENT", "MEMBERS", "1-3") config = ExperimentConfig() config.parse_ini(self.mock_parser) self.assertEqual(config.members, [1, 2, 3]) @@ -406,63 +485,107 @@ class TestExperimentConfig(TestCase): - Using a regex """ - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '20001101 20011101') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "20001101 20011101" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20011101']) + self.assertEqual(config.startdates, ["20001101", "20011101"]) - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '200(0|1)1101') + self.mock_parser.add_value("EXPERIMENT", "STARTDATES", "200(0|1)1101") config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20011101']) + self.assertEqual(config.startdates, ["20001101", "20011101"]) - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '200[0-2](02|05|08|11)01') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "200[0-2](02|05|08|11)01" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, [u'20000201', u'20000501', u'20000801', u'20001101', u'20010201', - u'20010501', u'20010801', u'20011101', u'20020201', u'20020501', - u'20020801', u'20021101']) + self.assertEqual( + config.startdates, + [ + u"20000201", + u"20000501", + u"20000801", + u"20001101", + u"20010201", + u"20010501", + u"20010801", + u"20011101", + u"20020201", + u"20020501", + u"20020801", + u"20021101", + ], + ) def test_auto_startdates(self): """ Test parsing startdates using the automatic generation - Reference syntax: {START,STOP,INTERVAL} where start and stop are dates and the interval is given with the code - NI where N is the number of intervals (default 1) and I is the base interval (Y,M,W,D) + Reference syntax: {START,STOP,INTERVAL} where start and stop are + dates and the interval is given with the code NI where N is the number + of intervals (default 1) and I is the base interval (Y,M,W,D) """ - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{2000,2001,1Y}') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{2000,2001,1Y}" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20000101', '20010101']) + self.assertEqual(config.startdates, ["20000101", "20010101"]) - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20011101,1Y}') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20011101,1Y}" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20011101']) + self.assertEqual(config.startdates, ["20001101", "20011101"]) - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20011101,6M} ') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20011101,6M} " + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20010501', '20011101']) + self.assertEqual( + config.startdates, ["20001101", "20010501", "20011101"] + ) - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20001201,1W}') + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20001201,1W}" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20001108', '20001115', '20001122', '20001129']) - - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20001201,W}') + self.assertEqual( + config.startdates, + ["20001101", "20001108", "20001115", "20001122", "20001129"], + ) + + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20001201,W}" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20001108', '20001115', '20001122', '20001129']) - - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20001201,7D}') + self.assertEqual( + config.startdates, + ["20001101", "20001108", "20001115", "20001122", "20001129"], + ) + + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20001201,7D}" + ) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.startdates, ['20001101', '20001108', '20001115', '20001122', '20001129']) - - self.mock_parser.add_value('EXPERIMENT', 'STARTDATES', '{20001101,20001201,7F}') + self.assertEqual( + config.startdates, + ["20001101", "20001108", "20001115", "20001122", "20001129"], + ) + + self.mock_parser.add_value( + "EXPERIMENT", "STARTDATES", "{20001101,20001201,7F}" + ) with self.assertRaises(ConfigException): config = ExperimentConfig() config.parse_ini(self.mock_parser) @@ -471,76 +594,101 @@ class TestExperimentConfig(TestCase): """Test get member str""" config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.get_member_str(1), 'fc1') + self.assertEqual(config.get_member_str(1), "fc1") def test_get_full_years(self): """Test get full years method""" - self.mock_parser.add_value('EXPERIMENT', 'CHUNK_SIZE', 3) - self.mock_parser.add_value('EXPERIMENT', 'CHUNKS', 15) + self.mock_parser.add_value("EXPERIMENT", "CHUNK_SIZE", 3) + self.mock_parser.add_value("EXPERIMENT", "CHUNKS", 15) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.get_full_years('20000601'), [2001, 2002, 2003, 2004]) - self.assertEqual(config.get_full_years('20000101'), [2000, 2001, 2002, 2003]) + self.assertEqual( + config.get_full_years("20000601"), [2001, 2002, 2003, 2004] + ) + self.assertEqual( + config.get_full_years("20000101"), [2000, 2001, 2002, 2003] + ) def test_get_year_chunks(self): """Test get chunk years""" - self.mock_parser.add_value('EXPERIMENT', 'CHUNK_SIZE', 3) - self.mock_parser.add_value('EXPERIMENT', 'CHUNKS', 13) + self.mock_parser.add_value("EXPERIMENT", "CHUNK_SIZE", 3) + self.mock_parser.add_value("EXPERIMENT", "CHUNKS", 13) config = ExperimentConfig() config.parse_ini(self.mock_parser) - self.assertEqual(config.get_year_chunks('20000601', 2003), [11, 12, 13]) - self.assertEqual(config.get_year_chunks('20000601', 2001), [3, 4, 5, 6, 7]) - self.assertEqual(config.get_year_chunks('20000101', 2000), [1, 2, 3, 4]) - - self.assertEqual(config.get_year_chunks('20000601', 2000), [1, 2, 3]) - self.assertEqual(config.get_year_chunks('20000601', 1999), []) + self.assertEqual( + config.get_year_chunks("20000601", 2003), [11, 12, 13] + ) + self.assertEqual( + config.get_year_chunks("20000601", 2001), [3, 4, 5, 6, 7] + ) + self.assertEqual( + config.get_year_chunks("20000101", 2000), [1, 2, 3, 4] + ) + + self.assertEqual(config.get_year_chunks("20000601", 2000), [1, 2, 3]) + self.assertEqual(config.get_year_chunks("20000601", 1999), []) def test_get_chunk_list(self): """Test get complete chunk list""" config = ExperimentConfig() - config.startdates = ('20010101', ) + config.startdates = ("20010101",) config.members = (0, 1, 2) config.chunk_list = [0] config.num_chunks = 2 - self.assertEqual(config.get_chunk_list(), [('20010101', 0, 0), ('20010101', 1, 0), ('20010101', 2, 0)]) + self.assertEqual( + config.get_chunk_list(), + [("20010101", 0, 0), ("20010101", 1, 0), ("20010101", 2, 0)], + ) config.chunk_list = [] - self.assertEqual(config.get_chunk_list(), [('20010101', 0, 1), ('20010101', 0, 2), ('20010101', 1, 1), - ('20010101', 1, 2), ('20010101', 2, 1), ('20010101', 2, 2)]) + self.assertEqual( + config.get_chunk_list(), + [ + ("20010101", 0, 1), + ("20010101", 0, 2), + ("20010101", 1, 1), + ("20010101", 1, 2), + ("20010101", 2, 1), + ("20010101", 2, 2), + ], + ) def test_get_member_list(self): """Test get member list""" config = ExperimentConfig() - config.startdates = ('20010101', ) + config.startdates = ("20010101",) config.members = (0, 1, 2) - self.assertEqual(config.get_member_list(), [('20010101', 0), ('20010101', 1), ('20010101', 2)]) + self.assertEqual( + config.get_member_list(), + [("20010101", 0), ("20010101", 1), ("20010101", 2)], + ) def test_get_chunk_start_str(self): """Test get_chunk_start_str""" config = ExperimentConfig() - self.mock_parser.add_value('EXPERIMENT', 'CHUNK_SIZE', 12) - self.mock_parser.add_value('EXPERIMENT', 'CHUNKS', 3) + self.mock_parser.add_value("EXPERIMENT", "CHUNK_SIZE", 12) + self.mock_parser.add_value("EXPERIMENT", "CHUNKS", 3) config.parse_ini(self.mock_parser) - self.assertEqual(config.get_chunk_start_str('20001101', 3), '20021101') + self.assertEqual(config.get_chunk_start_str("20001101", 3), "20021101") def test_get_chunk_start_str_datetime(self): """Test get_chunk_start_str when receiving a date object""" config = ExperimentConfig() - self.mock_parser.add_value('EXPERIMENT', 'CHUNK_SIZE', 12) - self.mock_parser.add_value('EXPERIMENT', 'CHUNKS', 3) + self.mock_parser.add_value("EXPERIMENT", "CHUNK_SIZE", 12) + self.mock_parser.add_value("EXPERIMENT", "CHUNKS", 3) date = datetime.datetime(year=2000, month=11, day=1) config.parse_ini(self.mock_parser) - self.assertEqual(config.get_chunk_start_str(date, 3), '20021101') + self.assertEqual(config.get_chunk_start_str(date, 3), "20021101") def test_get_chunk_end_str(self): """Test get_chunk_end_str""" config = ExperimentConfig() - self.mock_parser.add_value('EXPERIMENT', 'CHUNK_SIZE', 12) - self.mock_parser.add_value('EXPERIMENT', 'CHUNKS', 3) + self.mock_parser.add_value("EXPERIMENT", "CHUNK_SIZE", 12) + self.mock_parser.add_value("EXPERIMENT", "CHUNKS", 3) config.parse_ini(self.mock_parser) - self.assertEqual(config.get_chunk_end_str('20001101', 3), '20031101') + self.assertEqual(config.get_chunk_end_str("20001101", 3), "20031101") class TestConfig(TestCase): @@ -549,9 +697,9 @@ class TestConfig(TestCase): def setUp(self): """Prepare tests""" self.mock_parser = ParserMock() - self.mock_parser.add_value('DIAGNOSTICS', 'FREQUENCY', 'mon') - self.mock_parser.add_value('DIAGNOSTICS', 'DIAGS', 'diag1 diag2') - self.mock_parser.add_value('DIAGNOSTICS', 'SCRATCH_DIR', 'scratch') + self.mock_parser.add_value("DIAGNOSTICS", "FREQUENCY", "mon") + self.mock_parser.add_value("DIAGNOSTICS", "DIAGS", "diag1 diag2") + self.mock_parser.add_value("DIAGNOSTICS", "SCRATCH_DIR", "scratch") self._environ = dict(os.environ) def tearDown(self): @@ -565,28 +713,36 @@ class TestConfig(TestCase): def mock_new_exp(): mock_exp = mock.Mock() - mock_exp.expid = 'expid' + mock_exp.expid = "expid" return mock_exp - with mock.patch('earthdiagnostics.config.ConfigParser', new=mock_new): - with mock.patch('earthdiagnostics.config.VariableManager'): - with mock.patch('earthdiagnostics.config.ExperimentConfig', new=mock_new_exp): - with mock.patch('earthdiagnostics.config.CMORConfig'): - with mock.patch('earthdiagnostics.config.THREDDSConfig'): - with mock.patch('os.path.isfile'): - config.parse('path') + + with mock.patch("earthdiagnostics.config.ConfigParser", new=mock_new): + with mock.patch("earthdiagnostics.config.VariableManager"): + with mock.patch( + "earthdiagnostics.config.ExperimentConfig", + new=mock_new_exp, + ): + with mock.patch("earthdiagnostics.config.CMORConfig"): + with mock.patch( + "earthdiagnostics.config.THREDDSConfig" + ): + with mock.patch("os.path.isfile"): + config.parse("path") def test_diags(self): """Test diags parsing""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'DIAGS', 'diag1 diag2,opt1,opt2 # Commented diag') + self.mock_parser.add_value( + "DIAGNOSTICS", "DIAGS", "diag1 diag2,opt1,opt2 # Commented diag" + ) self._parse(config) - self.assertEqual(config.get_commands(), (['diag1', 'diag2,opt1,opt2'])) + self.assertEqual(config.get_commands(), (["diag1", "diag2,opt1,opt2"])) def test_file_not_found(self): """Test value error is raised if config file is not found""" config = Config() with self.assertRaises(ValueError): - config.parse('path') + config.parse("path") def test_parse(self): """Test parse minimal config""" @@ -594,28 +750,34 @@ class TestConfig(TestCase): self._parse(config) self.assertEqual(config.frequency, Frequencies.monthly) self.assertEqual(config.auto_clean, True) - self.assertEqual(config.cdftools_path, '') - self.assertEqual(config.con_files, '') - self.assertEqual(config.data_adaptor, 'CMOR') - self.assertEqual(config.get_commands(), (['diag1', 'diag2'])) + self.assertEqual(config.cdftools_path, "") + self.assertEqual(config.con_files, "") + self.assertEqual(config.data_adaptor, "CMOR") + self.assertEqual(config.get_commands(), (["diag1", "diag2"])) self.assertIsInstance(config.data_convention, SPECSConvention) - self.assertEqual(config.scratch_masks, '/scratch/Earth/ocean_masks') - namelist = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', - 'earthdiagnostics/CDFTOOLS_specs.namlist')) - self.assertEqual(os.environ['NAM_CDF_NAMES'], namelist) + self.assertEqual(config.scratch_masks, "/scratch/Earth/ocean_masks") + namelist = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "earthdiagnostics/CDFTOOLS_specs.namlist", + ) + ) + self.assertEqual(os.environ["NAM_CDF_NAMES"], namelist) def test_alias(self): """Test alias parsing""" config = Config() - self.mock_parser.add_value('ALIAS', 'diag1', 'diag3') + self.mock_parser.add_value("ALIAS", "diag1", "diag3") self._parse(config) - self.assertEqual(config.get_commands(), ['diag3', 'diag2']) + self.assertEqual(config.get_commands(), ["diag3", "diag2"]) def test_auto_clean_ram_disk(self): """Test that USE_RAMDISK forces AUTO_CLEAN to true""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'AUTO_CLEAN', False) - self.mock_parser.add_value('DIAGNOSTICS', 'USE_RAMDISK', True) + self.mock_parser.add_value("DIAGNOSTICS", "AUTO_CLEAN", False) + self.mock_parser.add_value("DIAGNOSTICS", "USE_RAMDISK", True) self._parse(config) self.assertEqual(config.auto_clean, True) self.assertEqual(config.use_ramdisk, True) @@ -623,43 +785,75 @@ class TestConfig(TestCase): def test_data_convention_primavera(self): """Test parsing data convention PRIMAVERA""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'DATA_CONVENTION', 'primavera') + self.mock_parser.add_value( + "DIAGNOSTICS", "DATA_CONVENTION", "primavera" + ) self._parse(config) self.assertIsInstance(config.data_convention, PrimaveraConvention) - self.assertEqual(config.scratch_masks, '/scratch/Earth/ocean_masks/primavera') - namelist = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', - 'earthdiagnostics/CDFTOOLS_primavera.namlist')) - self.assertEqual(os.environ['NAM_CDF_NAMES'], namelist) + self.assertEqual( + config.scratch_masks, "/scratch/Earth/ocean_masks/primavera" + ) + namelist = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "earthdiagnostics/CDFTOOLS_primavera.namlist", + ) + ) + self.assertEqual(os.environ["NAM_CDF_NAMES"], namelist) def test_data_convention_cmip6(self): """Test parsing data convention CMIP6""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'DATA_CONVENTION', 'cmip6') + self.mock_parser.add_value("DIAGNOSTICS", "DATA_CONVENTION", "cmip6") self._parse(config) self.assertIsInstance(config.data_convention, CMIP6Convention) - self.assertEqual(config.scratch_masks, '/scratch/Earth/ocean_masks/cmip6') - namelist = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', - 'earthdiagnostics/CDFTOOLS_cmip6.namlist')) - self.assertEqual(os.environ['NAM_CDF_NAMES'], namelist) + self.assertEqual( + config.scratch_masks, "/scratch/Earth/ocean_masks/cmip6" + ) + namelist = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "earthdiagnostics/CDFTOOLS_cmip6.namlist", + ) + ) + self.assertEqual(os.environ["NAM_CDF_NAMES"], namelist) def test_data_convention_meteofrance(self): """Test parsing data convention MeteoFrance""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'DATA_CONVENTION', 'meteofrance') + self.mock_parser.add_value( + "DIAGNOSTICS", "DATA_CONVENTION", "meteofrance" + ) self._parse(config) self.assertIsInstance(config.data_convention, MeteoFranceConvention) - self.assertEqual(config.scratch_masks, '/scratch/Earth/ocean_masks') - namelist = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', - 'earthdiagnostics/CDFTOOLS_meteofrance.namlist')) - self.assertEqual(os.environ['NAM_CDF_NAMES'], namelist) + self.assertEqual(config.scratch_masks, "/scratch/Earth/ocean_masks") + namelist = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "earthdiagnostics/CDFTOOLS_meteofrance.namlist", + ) + ) + self.assertEqual(os.environ["NAM_CDF_NAMES"], namelist) def test_data_convention_preface(self): """Test parsing data convention Preface""" config = Config() - self.mock_parser.add_value('DIAGNOSTICS', 'DATA_CONVENTION', 'preface') + self.mock_parser.add_value("DIAGNOSTICS", "DATA_CONVENTION", "preface") self._parse(config) self.assertIsInstance(config.data_convention, PrefaceConvention) - self.assertEqual(config.scratch_masks, '/scratch/Earth/ocean_masks') - namelist = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', - 'earthdiagnostics/CDFTOOLS_preface.namlist')) - self.assertEqual(os.environ['NAM_CDF_NAMES'], namelist) + self.assertEqual(config.scratch_masks, "/scratch/Earth/ocean_masks") + namelist = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "earthdiagnostics/CDFTOOLS_preface.namlist", + ) + ) + self.assertEqual(os.environ["NAM_CDF_NAMES"], namelist) diff --git a/test/unit/test_constants.py b/test/unit/test_constants.py index 6ad614b4..f78ea16d 100644 --- a/test/unit/test_constants.py +++ b/test/unit/test_constants.py @@ -1,52 +1,50 @@ # coding=utf-8 from unittest import TestCase -from mock import Mock -from earthdiagnostics.constants import Basin, Basins +from earthdiagnostics.constants import Basin class TestBasin(TestCase): - def setUp(self): - self.basin = Basin('Basin') + self.basin = Basin("Basin") def test_name(self): - self.assertEqual('Basin', self.basin.name) + self.assertEqual("Basin", self.basin.name) def test__eq__(self): - self.assertTrue(Basin('Basin') == self.basin) - self.assertFalse(Basin('bas') == self.basin) + self.assertTrue(Basin("Basin") == self.basin) + self.assertFalse(Basin("bas") == self.basin) def test__neq__(self): - self.assertFalse(Basin('Basin') != self.basin) - self.assertTrue(Basin('bas') != self.basin) + self.assertFalse(Basin("Basin") != self.basin) + self.assertTrue(Basin("bas") != self.basin) def test__greater(self): - self.assertFalse(Basin('Aasin') > self.basin) - self.assertFalse(Basin('Basin') > self.basin) - self.assertTrue(Basin('Casin') > self.basin) + self.assertFalse(Basin("Aasin") > self.basin) + self.assertFalse(Basin("Basin") > self.basin) + self.assertTrue(Basin("Casin") > self.basin) def test__lower(self): - self.assertTrue(Basin('Aasin') < self.basin) - self.assertFalse(Basin('Basin') < self.basin) - self.assertFalse(Basin('Casin') < self.basin) + self.assertTrue(Basin("Aasin") < self.basin) + self.assertFalse(Basin("Basin") < self.basin) + self.assertFalse(Basin("Casin") < self.basin) def test_ge(self): - self.assertFalse(Basin('Aasin') >= self.basin) - self.assertTrue(Basin('Basin') >= self.basin) - self.assertTrue(Basin('Casin') >= self.basin) + self.assertFalse(Basin("Aasin") >= self.basin) + self.assertTrue(Basin("Basin") >= self.basin) + self.assertTrue(Basin("Casin") >= self.basin) def test__le(self): - self.assertTrue(Basin('Aasin') <= self.basin) - self.assertTrue(Basin('Basin') <= self.basin) - self.assertFalse(Basin('Casin') <= self.basin) + self.assertTrue(Basin("Aasin") <= self.basin) + self.assertTrue(Basin("Basin") <= self.basin) + self.assertFalse(Basin("Casin") <= self.basin) def test__str__(self): - self.assertEqual(str(self.basin), 'Basin') + self.assertEqual(str(self.basin), "Basin") def test__repr__(self): - self.assertEqual(repr(self.basin), 'Basin') + self.assertEqual(repr(self.basin), "Basin") def test_order(self): - self.assertTrue(self.basin < Basin('Vasin')) - self.assertTrue(self.basin > Basin('Asin')) + self.assertTrue(self.basin < Basin("Vasin")) + self.assertTrue(self.basin > Basin("Asin")) diff --git a/test/unit/test_data_manager.py b/test/unit/test_data_manager.py index 3d5a3d60..66976792 100644 --- a/test/unit/test_data_manager.py +++ b/test/unit/test_data_manager.py @@ -5,11 +5,9 @@ from unittest import TestCase from mock import Mock from earthdiagnostics.datamanager import DataManager -from earthdiagnostics.modelingrealm import ModelingRealms class TestDataManager(TestCase): - def setUp(self): self.config = Mock() self.config.experiment.atmos_timestep = 3 @@ -21,32 +19,37 @@ class TestDataManager(TestCase): self.data_manager.prepare() def test_link_file(self): - self.data_manager.link_file(None, '', None, '', 0, 0) + self.data_manager.link_file(None, "", None, "", 0, 0) def test_file_exists(self): with self.assertRaises(NotImplementedError): - self.data_manager.file_exists(None, '', '', 0, 0) + self.data_manager.file_exists(None, "", "", 0, 0) def test_request_chunk(self): with self.assertRaises(NotImplementedError): - self.data_manager.request_chunk(None, '', '', 0, 0) + self.data_manager.request_chunk(None, "", "", 0, 0) def test_request_year(self): with self.assertRaises(NotImplementedError): - self.data_manager.request_year(None, '', '', '', 0, 0) + self.data_manager.request_year(None, "", "", "", 0, 0) def test_declare_chunk(self): with self.assertRaises(NotImplementedError): - self.data_manager.declare_chunk(None, '', '', 0, 0) + self.data_manager.declare_chunk(None, "", "", 0, 0) def test_declare_year(self): with self.assertRaises(NotImplementedError): - self.data_manager.declare_year(None, '', '', 0, 0) + self.data_manager.declare_year(None, "", "", 0, 0) def test_get_final_varname(self): - self.assertEqual(self.data_manager._get_final_var_name(None, 'var'), 'var') + self.assertEqual( + self.data_manager._get_final_var_name(None, "var"), "var" + ) box = Mock() - box.get_lon_str.return_value = '_lon' - box.get_lat_str.return_value = '_lat' - box.get_depth_str.return_value = '_depth' - self.assertEqual(self.data_manager._get_final_var_name(box, 'var'), 'var_lon_lat_depth') + box.get_lon_str.return_value = "_lon" + box.get_lat_str.return_value = "_lat" + box.get_depth_str.return_value = "_depth" + self.assertEqual( + self.data_manager._get_final_var_name(box, "var"), + "var_lon_lat_depth", + ) diff --git a/test/unit/test_datafile.py b/test/unit/test_datafile.py index 0b47650e..685674c1 100644 --- a/test/unit/test_datafile.py +++ b/test/unit/test_datafile.py @@ -2,43 +2,66 @@ from unittest import TestCase from mock import Mock -from earthdiagnostics.datafile import UnitConversion, DataFile, LocalStatus, StorageStatus +from earthdiagnostics.datafile import ( + UnitConversion, + DataFile, + LocalStatus, + StorageStatus, +) class TestConversion(TestCase): - def test__init(self): - conversion = UnitConversion('km', 'm', 1000, 0) - self.assertEqual(conversion.source, 'km') - self.assertEqual(conversion.destiny, 'm') + conversion = UnitConversion("km", "m", 1000, 0) + self.assertEqual(conversion.source, "km") + self.assertEqual(conversion.destiny, "m") self.assertEqual(conversion.factor, 1000) self.assertEqual(conversion.offset, 0) def test_add_conversion(self): UnitConversion._dict_conversions = dict() - conversion = UnitConversion('km', 'm', 1000, 0) + conversion = UnitConversion("km", "m", 1000, 0) UnitConversion.add_conversion(conversion) - self.assertIs(UnitConversion._dict_conversions[('km', 'm')], conversion) + self.assertIs( + UnitConversion._dict_conversions[("km", "m")], conversion + ) UnitConversion._dict_conversions = dict() def test_get_factor_offset(self): UnitConversion._dict_conversions = dict() - conversion = UnitConversion('km', 'm', 1000, 0) + conversion = UnitConversion("km", "m", 1000, 0) UnitConversion.add_conversion(conversion) - self.assertEqual(UnitConversion.get_conversion_factor_offset('km', 'm'), (1000, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('m', 'km'), (1 / 1000.0, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('1e3 m', 'km'), (1, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('10^3 m', 'km'), (1, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('km', '1e3 m'), (1, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('km', '10^3 m'), (1, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('m', 'm'), (1, 0)) - self.assertEqual(UnitConversion.get_conversion_factor_offset('m²', 'km'), (None, None)) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("km", "m"), (1000, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("m", "km"), + (1 / 1000.0, 0), + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("1e3 m", "km"), (1, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("10^3 m", "km"), (1, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("km", "1e3 m"), (1, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("km", "10^3 m"), (1, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("m", "m"), (1, 0) + ) + self.assertEqual( + UnitConversion.get_conversion_factor_offset("m²", "km"), + (None, None), + ) UnitConversion._dict_conversions = dict() class TestDatafile(TestCase): - def setUp(self): self.data_file = DataFile() diff --git a/test/unit/test_diagnostic.py b/test/unit/test_diagnostic.py index f4baa226..78c2fd33 100644 --- a/test/unit/test_diagnostic.py +++ b/test/unit/test_diagnostic.py @@ -1,9 +1,21 @@ """Test for diagnostic module""" # coding=utf-8 -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticOptionError, DiagnosticFloatOption, \ - DiagnosticDomainOption, DiagnosticIntOption, DiagnosticBoolOption, DiagnosticComplexStrOption, \ - DiagnosticListIntOption, DiagnosticChoiceOption, DiagnosticVariableOption, DiagnosticVariableListOption, \ - DiagnosticBasinOption, DiagnosticStatus +from earthdiagnostics.diagnostic import ( + Diagnostic, + DiagnosticOption, + DiagnosticOptionError, + DiagnosticFloatOption, + DiagnosticDomainOption, + DiagnosticIntOption, + DiagnosticBoolOption, + DiagnosticComplexStrOption, + DiagnosticListIntOption, + DiagnosticChoiceOption, + DiagnosticVariableOption, + DiagnosticVariableListOption, + DiagnosticBasinOption, + DiagnosticStatus, +) from earthdiagnostics.datafile import StorageStatus from earthdiagnostics.constants import Basins @@ -19,79 +31,81 @@ class TestDiagnosticOption(TestCase): def test_good_default_value(self): """Test default_value""" - diag = DiagnosticOption('option', 'default') - self.assertEqual('default', diag.parse('')) + diag = DiagnosticOption("option", "default") + self.assertEqual("default", diag.parse("")) def test_no_default_value(self): """Test raises if not value provided and no default is available""" - diag = DiagnosticOption('option') + diag = DiagnosticOption("option") with self.assertRaises(DiagnosticOptionError): - self.assertEqual('default', diag.parse('')) + self.assertEqual("default", diag.parse("")) def test_parse_value(self): """Test parse simple value""" - diag = DiagnosticOption('option') - self.assertEqual('value', diag.parse('value')) + diag = DiagnosticOption("option") + self.assertEqual("value", diag.parse("value")) class TestDiagnosticFloatOption(TestCase): """Test diagnostic float option""" + def test_float_default_value(self): """Test float default_value""" - diag = DiagnosticFloatOption('option', 3.0) - self.assertEqual(3.0, diag.parse('')) + diag = DiagnosticFloatOption("option", 3.0) + self.assertEqual(3.0, diag.parse("")) def test_str_default_value(self): """Test str default_value""" - diag = DiagnosticFloatOption('option', '3') - self.assertEqual(3.0, diag.parse('')) + diag = DiagnosticFloatOption("option", "3") + self.assertEqual(3.0, diag.parse("")) def test_bad_default_value(self): """Test fails if default_value can not be converted to float""" - diag = DiagnosticFloatOption('option', 'default') + diag = DiagnosticFloatOption("option", "default") with self.assertRaises(ValueError): - self.assertEqual('default', diag.parse('')) + self.assertEqual("default", diag.parse("")) def test_no_default_value(self): """Test raises if not value provided and no default is available""" - diag = DiagnosticFloatOption('option') + diag = DiagnosticFloatOption("option") with self.assertRaises(DiagnosticOptionError): - self.assertEqual('default', diag.parse('')) + self.assertEqual("default", diag.parse("")) def test_parse_value(self): """Test parse simple value""" - diag = DiagnosticFloatOption('option') - self.assertEqual(3.25, diag.parse('3.25')) + diag = DiagnosticFloatOption("option") + self.assertEqual(3.25, diag.parse("3.25")) class TestDiagnosticDomainOption(TestCase): """Test diagnostic domain option""" + def test_domain_default_value(self): """Test domain default_value""" - diag = DiagnosticDomainOption('option', ModelingRealms.ocean) - self.assertEqual(ModelingRealms.ocean, diag.parse('')) + diag = DiagnosticDomainOption("option", ModelingRealms.ocean) + self.assertEqual(ModelingRealms.ocean, diag.parse("")) def test_str_default_value(self): """Test str default_value""" - diag = DiagnosticDomainOption('option', 'atmos') - self.assertEqual(ModelingRealms.atmos, diag.parse('')) + diag = DiagnosticDomainOption("option", "atmos") + self.assertEqual(ModelingRealms.atmos, diag.parse("")) def test_bad_default_value(self): - """Test raises if default_value can not be converted to ModellingRealm""" - diag = DiagnosticDomainOption('option', 'default') + """Test raises if can not convert default_value to ModellingRealm""" + diag = DiagnosticDomainOption("option", "default") with self.assertRaises(ValueError): - diag.parse('') + diag.parse("") def test_no_default_value(self): """Test no value and default_value is provided""" - diag = DiagnosticDomainOption('option') + diag = DiagnosticDomainOption("option") with self.assertRaises(DiagnosticOptionError): - diag.parse('') + diag.parse("") def test_parse_value(self): """Test parse option""" - diag = DiagnosticDomainOption('option') - self.assertEqual(ModelingRealms.seaIce, diag.parse('seaice')) + diag = DiagnosticDomainOption("option") + self.assertEqual(ModelingRealms.seaIce, diag.parse("seaice")) class TestDiagnosticIntOption(TestCase): @@ -99,58 +113,58 @@ class TestDiagnosticIntOption(TestCase): def test_int_default_value(self): """Test int default_value""" - diag = DiagnosticIntOption('option', 3) - self.assertEqual(3, diag.parse('')) + diag = DiagnosticIntOption("option", 3) + self.assertEqual(3, diag.parse("")) def test_str_default_value(self): """Test str default_value""" - diag = DiagnosticIntOption('option', '3') - self.assertEqual(3, diag.parse('')) + diag = DiagnosticIntOption("option", "3") + self.assertEqual(3, diag.parse("")) def test_bad_default_value(self): """Test raises if default_value can not be parsed to int""" - diag = DiagnosticIntOption('option', 'default') + diag = DiagnosticIntOption("option", "default") with self.assertRaises(ValueError): - diag.parse('') + diag.parse("") def test_no_default_value(self): """Test raises if no value and default_value are provided""" - diag = DiagnosticIntOption('option') + diag = DiagnosticIntOption("option") with self.assertRaises(DiagnosticOptionError): - diag.parse('') + diag.parse("") def test_parse_value(self): """Test parse int value""" - diag = DiagnosticIntOption('option') - self.assertEqual(3, diag.parse('3')) + diag = DiagnosticIntOption("option") + self.assertEqual(3, diag.parse("3")) def test_parse_bad_value(self): """Test raises if value can not be parsed to int""" - diag = DiagnosticIntOption('option') + diag = DiagnosticIntOption("option") with self.assertRaises(ValueError): - diag.parse('3.5') + diag.parse("3.5") def test_good_low_limit(self): """Test parse if low limit is configured""" - diag = DiagnosticIntOption('option', None, 0) - self.assertEqual(1, diag.parse('1')) + diag = DiagnosticIntOption("option", None, 0) + self.assertEqual(1, diag.parse("1")) def test_bad_low_limit(self): - """"Test raises if low limit is configured and value does not conform to it""" - diag = DiagnosticIntOption('option', None, 0) + """"Test raises if low limit is configured and value is higher""" + diag = DiagnosticIntOption("option", None, 0) with self.assertRaises(DiagnosticOptionError): - diag.parse('-1') + diag.parse("-1") def test_good_high_limit(self): """Test parse if high limit is configured""" - diag = DiagnosticIntOption('option', None, None, 0) - self.assertEqual(-1, diag.parse('-1')) + diag = DiagnosticIntOption("option", None, None, 0) + self.assertEqual(-1, diag.parse("-1")) def test_bad_high_limit(self): - """Test raises if high limit is configured and value does not conform to it""" - diag = DiagnosticIntOption('option', None, None, 0) + """Test raises if high limit is configured and value is lower""" + diag = DiagnosticIntOption("option", None, None, 0) with self.assertRaises(DiagnosticOptionError): - diag.parse('1') + diag.parse("1") class TestDiagnosticBoolOption(TestCase): @@ -158,44 +172,44 @@ class TestDiagnosticBoolOption(TestCase): def test_bool_default_value(self): """Test parse bool default_value""" - diag = DiagnosticBoolOption('option', True) - self.assertEqual(True, diag.parse('')) + diag = DiagnosticBoolOption("option", True) + self.assertEqual(True, diag.parse("")) def test_str_default_value(self): """Test parse str default_value""" - diag = DiagnosticBoolOption('option', 'False') - self.assertEqual(False, diag.parse('')) + diag = DiagnosticBoolOption("option", "False") + self.assertEqual(False, diag.parse("")) def test_no_default_value(self): """Test raises if neither value or default_value are provided""" - diag = DiagnosticBoolOption('option') + diag = DiagnosticBoolOption("option") with self.assertRaises(DiagnosticOptionError): - diag.parse('') + diag.parse("") - def test_parse_True(self): + def test_parse_caps_true(self): """Test parses 'True' as boolean True""" - diag = DiagnosticBoolOption('option') - self.assertTrue(diag.parse('True')) + diag = DiagnosticBoolOption("option") + self.assertTrue(diag.parse("True")) def test_parse_true(self): """Test parses 'true' as boolean True""" - diag = DiagnosticBoolOption('option') - self.assertTrue(diag.parse('true')) + diag = DiagnosticBoolOption("option") + self.assertTrue(diag.parse("true")) def test_parse_t(self): """Test parses 't' as boolean True""" - diag = DiagnosticBoolOption('option') - self.assertTrue(diag.parse('t')) + diag = DiagnosticBoolOption("option") + self.assertTrue(diag.parse("t")) def test_parse_yes(self): """Test parses 'YES' as boolean True""" - diag = DiagnosticBoolOption('option') - self.assertTrue(diag.parse('YES')) + diag = DiagnosticBoolOption("option") + self.assertTrue(diag.parse("YES")) def test_parse_bad_value(self): """Test parses random str as boolean False""" - diag = DiagnosticBoolOption('option') - self.assertFalse(diag.parse('3.5')) + diag = DiagnosticBoolOption("option") + self.assertFalse(diag.parse("3.5")) class TestDiagnosticComplexStrOption(TestCase): @@ -204,26 +218,32 @@ class TestDiagnosticComplexStrOption(TestCase): It replaces '&.' por ' ' and '&;' por ',' """ + def test_complex_default_value(self): """Test default value""" - diag = DiagnosticComplexStrOption('option', 'default&.str&;&.working') - self.assertEqual('default str, working', diag.parse('')) + diag = DiagnosticComplexStrOption("option", "default&.str&;&.working") + self.assertEqual("default str, working", diag.parse("")) def test_simple_default_value(self): """Test simple default value""" - diag = DiagnosticComplexStrOption('default str, working', 'default str, working') - self.assertEqual('default str, working', diag.parse('')) + diag = DiagnosticComplexStrOption( + "default str, working", "default str, working" + ) + self.assertEqual("default str, working", diag.parse("")) def test_no_default_value(self): """Test raises if neither value or default_value are provided""" - diag = DiagnosticComplexStrOption('option') + diag = DiagnosticComplexStrOption("option") with self.assertRaises(DiagnosticOptionError): - diag.parse('') + diag.parse("") def test_parse_value(self): """Test parse value""" - diag = DiagnosticComplexStrOption('option') - self.assertEqual('complex string, for testing', diag.parse('complex&.string&;&.for&.testing')) + diag = DiagnosticComplexStrOption("option") + self.assertEqual( + "complex string, for testing", + diag.parse("complex&.string&;&.for&.testing"), + ) class TestDiagnosticListIntOption(TestCase): @@ -231,57 +251,57 @@ class TestDiagnosticListIntOption(TestCase): def test_tuple_default_value(self): """Test use tuple default_value if no option is provided""" - diag = DiagnosticListIntOption('option', (3,)) - self.assertEqual((3,), diag.parse('')) + diag = DiagnosticListIntOption("option", (3,)) + self.assertEqual((3,), diag.parse("")) def test_list_default_value(self): """Test use list default_value if no option is provided""" - diag = DiagnosticListIntOption('option', [3]) - self.assertEqual([3], diag.parse('')) + diag = DiagnosticListIntOption("option", [3]) + self.assertEqual([3], diag.parse("")) def test_str_default_value(self): """Test use str default_value if no option is provided""" - diag = DiagnosticListIntOption('option', '3-4') - self.assertEqual([3, 4], diag.parse('')) + diag = DiagnosticListIntOption("option", "3-4") + self.assertEqual([3, 4], diag.parse("")) def test_bad_default_value(self): """Test raises if default value can not be converted to int""" - diag = DiagnosticListIntOption('option', 'default') + diag = DiagnosticListIntOption("option", "default") with self.assertRaises(ValueError): - diag.parse('') + diag.parse("") def test_no_default_value(self): """Test raises if neither value or default_value are provided""" - diag = DiagnosticListIntOption('option') + diag = DiagnosticListIntOption("option") with self.assertRaises(DiagnosticOptionError): - diag.parse('') + diag.parse("") def test_parse_value(self): """Test parse a list of integers""" - diag = DiagnosticListIntOption('option') - self.assertEqual([3, 2], diag.parse('3-2')) + diag = DiagnosticListIntOption("option") + self.assertEqual([3, 2], diag.parse("3-2")) def test_parse_single_value(self): """Test parse a single integer""" - diag = DiagnosticListIntOption('option') - self.assertEqual([3], diag.parse('3')) + diag = DiagnosticListIntOption("option") + self.assertEqual([3], diag.parse("3")) def test_too_low(self): - """Test raises if low limit is configured and value does not conform to it""" - diag = DiagnosticListIntOption('option', min_limit=5) + """Test raises if low limit is configured and value is lower""" + diag = DiagnosticListIntOption("option", min_limit=5) with self.assertRaises(DiagnosticOptionError): - diag.parse('3') + diag.parse("3") def test_too_high(self): - """Test raises if high limit is configured and value does not conform to it""" - diag = DiagnosticListIntOption('option', max_limit=5) + """Test raises if high limit is configured and value is higher""" + diag = DiagnosticListIntOption("option", max_limit=5) with self.assertRaises(DiagnosticOptionError): - diag.parse('8') + diag.parse("8") def test_parse_bad_value(self): - diag = DiagnosticListIntOption('option') + diag = DiagnosticListIntOption("option") with self.assertRaises(ValueError): - diag.parse('3.5') + diag.parse("3.5") class TestDiagnosticChoiceOption(TestCase): @@ -289,29 +309,29 @@ class TestDiagnosticChoiceOption(TestCase): def test_choice_value(self): """Test parsing a choice""" - diag = DiagnosticChoiceOption('option', ('a', 'b')) - self.assertEqual('a', diag.parse('a')) + diag = DiagnosticChoiceOption("option", ("a", "b")) + self.assertEqual("a", diag.parse("a")) def test_choice_default_value(self): """Test use default value""" - diag = DiagnosticChoiceOption('option', ('a', 'b'), default_value='a') - self.assertEqual('a', diag.parse('')) + diag = DiagnosticChoiceOption("option", ("a", "b"), default_value="a") + self.assertEqual("a", diag.parse("")) def test_bad_default_value(self): """Test error is raised if default option is not a valid choice""" with self.assertRaises(DiagnosticOptionError): - DiagnosticChoiceOption('option', ('a', 'b'), default_value='c') + DiagnosticChoiceOption("option", ("a", "b"), default_value="c") def test_ignore_case_value(self): """Test ignore_case option""" - diag = DiagnosticChoiceOption('option', ('a', 'b')) - self.assertEqual('b', diag.parse('b')) - self.assertEqual('b', diag.parse('B')) + diag = DiagnosticChoiceOption("option", ("a", "b")) + self.assertEqual("b", diag.parse("b")) + self.assertEqual("b", diag.parse("B")) - diag = DiagnosticChoiceOption('option', ('a', 'b'), ignore_case=False) - self.assertEqual('b', diag.parse('b')) + diag = DiagnosticChoiceOption("option", ("a", "b"), ignore_case=False) + self.assertEqual("b", diag.parse("b")) with self.assertRaises(DiagnosticOptionError): - self.assertEqual('b', diag.parse('B')) + self.assertEqual("b", diag.parse("B")) class TestDiagnosticVariableOption(TestCase): @@ -325,10 +345,10 @@ class TestDiagnosticVariableOption(TestCase): def test_parse(self): """Test parse""" var_manager_mock = Mock() - var_manager_mock.get_variable.return_value = self._get_var_mock('var1') + var_manager_mock.get_variable.return_value = self._get_var_mock("var1") diag = DiagnosticVariableOption(var_manager_mock) - self.assertEqual('var1', diag.parse('var1')) + self.assertEqual("var1", diag.parse("var1")) def test_not_recognized(self): """Test parsing a not recognized variable""" @@ -336,7 +356,7 @@ class TestDiagnosticVariableOption(TestCase): var_manager_mock.get_variable.return_value = None diag = DiagnosticVariableOption(var_manager_mock) - self.assertEqual('var1', diag.parse('var1')) + self.assertEqual("var1", diag.parse("var1")) class TestDiagnosticVariableListOption(TestCase): @@ -345,23 +365,26 @@ class TestDiagnosticVariableListOption(TestCase): def test_parse_multiple(self): """Test parsing multiple vars""" var_manager_mock = Mock() - var_manager_mock.get_variable.side_effect = (self._get_var_mock('var1'), self._get_var_mock('var2')) - diag = DiagnosticVariableListOption(var_manager_mock, 'variables') - self.assertEqual(['var1', 'var2'], diag.parse('var1:var2')) + var_manager_mock.get_variable.side_effect = ( + self._get_var_mock("var1"), + self._get_var_mock("var2"), + ) + diag = DiagnosticVariableListOption(var_manager_mock, "variables") + self.assertEqual(["var1", "var2"], diag.parse("var1:var2")) def test_parse_one(self): """Test parsing only one var""" var_manager_mock = Mock() - var_manager_mock.get_variable.return_value = self._get_var_mock('var1') - diag = DiagnosticVariableListOption(var_manager_mock, 'variables') - self.assertEqual(['var1'], diag.parse('var1')) + var_manager_mock.get_variable.return_value = self._get_var_mock("var1") + diag = DiagnosticVariableListOption(var_manager_mock, "variables") + self.assertEqual(["var1"], diag.parse("var1")) def test_not_recognized(self): """Test parsing one var that can not be recognized""" var_manager_mock = Mock() var_manager_mock.get_variable.return_value = None - diag = DiagnosticVariableListOption(var_manager_mock, 'variables') - self.assertEqual(['var1'], diag.parse('var1')) + diag = DiagnosticVariableListOption(var_manager_mock, "variables") + self.assertEqual(["var1"], diag.parse("var1")) def _get_var_mock(self, name): mock = Mock() @@ -374,6 +397,7 @@ class TestDiagnostic(TestCase): def setUp(self): """Prepare tests""" + class MockDiag(Diagnostic): @classmethod def generate_jobs(cls, diags, options): @@ -392,7 +416,10 @@ class TestDiagnostic(TestCase): def test_str(self): """Test base __str__ implementation""" - self.assertEqual(str(Diagnostic(None)), 'Developer must override base class __str__ method') + self.assertEqual( + str(Diagnostic(None)), + "Developer must override base class __str__ method", + ) def test_repr(self): """Test base __repr__ implementation""" @@ -413,7 +440,7 @@ class TestDiagnostic(TestCase): with self.assertRaises(NotImplementedError): Diagnostic(None).request_data() - @patch.object(Diagnostic, 'dispatch') + @patch.object(Diagnostic, "dispatch") def test_set_status_call_dispatch(self, dispatch_mock): """Test dispatch is called when setting a different statuts""" diag = Diagnostic(None) @@ -421,12 +448,12 @@ class TestDiagnostic(TestCase): diag.status = DiagnosticStatus.FAILED dispatch_mock.assert_called_once_with(diag, old_status) - @patch.object(Diagnostic, 'dispatch') + @patch.object(Diagnostic, "dispatch") def test_set_status_call_dispatch_not_called(self, dispatch_mock): """Test dispatch is not called when setting again the same status""" diag = Diagnostic(None) diag.status = diag.status - assert not dispatch_mock.called, 'Dispatch should not have been called' + assert not dispatch_mock.called, "Dispatch should not have been called" def test_register(self): """Test register diagnostic""" @@ -436,31 +463,35 @@ class TestDiagnostic(TestCase): with self.assertRaises(ValueError): Diagnostic.register(self.MockDiag) - self.MockDiag.alias = 'mock' + self.MockDiag.alias = "mock" Diagnostic.register(self.MockDiag) def test_get_diagnostic(self): """Test get diagnostic""" - self.assertIsNone(Diagnostic.get_diagnostic('none')) - self.MockDiag.alias = 'mock' + self.assertIsNone(Diagnostic.get_diagnostic("none")) + self.MockDiag.alias = "mock" Diagnostic.register(self.MockDiag) - self.assertIs(self.MockDiag, Diagnostic.get_diagnostic('mock')) + self.assertIs(self.MockDiag, Diagnostic.get_diagnostic("mock")) def test_generate_jobs(self): """Test generate_jobs raises NotImplementedError""" with self.assertRaises(NotImplementedError): - Diagnostic.generate_jobs(None, ['']) + Diagnostic.generate_jobs(None, [""]) def test_empty_process_options(self): """Test process options works even without options""" - self.assertEqual(len(Diagnostic.process_options(('diag_name',), tuple())), 0) + self.assertEqual( + len(Diagnostic.process_options(("diag_name",), tuple())), 0 + ) def test_diagnostic_can_be_skipped(self): """Test diagnostic can be skipped""" diag = self._get_diag_for_skipping() self.assertTrue(diag.can_skip_run()) - def _get_diag_for_skipping(self, modified=False, status=StorageStatus.READY): + def _get_diag_for_skipping( + self, modified=False, status=StorageStatus.READY + ): mock = Mock() datafile_mock = Mock() datafile_mock.storage_status = status @@ -471,7 +502,7 @@ class TestDiagnostic(TestCase): return diag def test_diagnostic_can_not_be_skipped_do_not_generate_files(self): - """Test diagnostic can not be skipped because it does not generate files""" + """Test can not skip diagnostic because it does not generate files""" diag = Diagnostic(None) self.assertFalse(diag.can_skip_run()) @@ -481,7 +512,7 @@ class TestDiagnostic(TestCase): self.assertFalse(diag.can_skip_run()) def test_diagnostic_can_not_be_skipped_data_not_ready(self): - """Test diagnostic can not be skipped because results do not yet exist""" + """Test diagnostic can not be skipped because results do not exist""" diag = self._get_diag_for_skipping(status=StorageStatus.PENDING) self.assertFalse(diag.can_skip_run()) @@ -489,17 +520,17 @@ class TestDiagnostic(TestCase): class TestDiagnosticBasinOption(TestCase): """Test for DiagnosticBasinOPtion""" - @patch.object(Basins, 'parse') + @patch.object(Basins, "parse") def test_parse(self, parse_mock): """Test parse good basin""" - parse_mock.return_value = 'var1' - diag = DiagnosticBasinOption('basin') - self.assertEqual('var1', diag.parse('var1')) + parse_mock.return_value = "var1" + diag = DiagnosticBasinOption("basin") + self.assertEqual("var1", diag.parse("var1")) - @patch.object(Basins, 'parse') + @patch.object(Basins, "parse") def test_not_recognized(self, parse_mock): """Test parsing a bad basin""" parse_mock.return_value = None - diag = DiagnosticBasinOption('basin') + diag = DiagnosticBasinOption("basin") with self.assertRaises(DiagnosticOptionError): - diag.parse('bad') + diag.parse("bad") diff --git a/test/unit/test_earthdiags.py b/test/unit/test_earthdiags.py index 79beeb9d..f0b78a34 100644 --- a/test/unit/test_earthdiags.py +++ b/test/unit/test_earthdiags.py @@ -13,13 +13,14 @@ from mock import Mock class TestEarthDiags(TestCase): - def setUp(self): self.earthdiags = EarthDiags() def test_clean(self): self.earthdiags.config.scratch_dir = tempfile.mkdtemp() - os.makedirs(os.path.join(self.earthdiags.config.scratch_dir, 'extra_folder')) + os.makedirs( + os.path.join(self.earthdiags.config.scratch_dir, "extra_folder") + ) self.assertTrue(os.path.isdir(self.earthdiags.config.scratch_dir)) self.earthdiags.clean() @@ -29,7 +30,7 @@ class TestEarthDiags(TestCase): temp_folder = tempfile.mkdtemp() scratch_dir = tempfile.mktemp() os.symlink(temp_folder, scratch_dir) - os.makedirs(os.path.join(temp_folder, 'extra_folder')) + os.makedirs(os.path.join(temp_folder, "extra_folder")) self.earthdiags.config.scratch_dir = scratch_dir self.assertTrue(os.path.islink(self.earthdiags.config.scratch_dir)) @@ -40,95 +41,118 @@ class TestEarthDiags(TestCase): def test_read_config(self): self.earthdiags.config = Mock() - self.earthdiags.config.scratch_dir = 'scratch_path' - self.earthdiags.config.cdftools_path = 'cdftools_path' + self.earthdiags.config.scratch_dir = "scratch_path" + self.earthdiags.config.cdftools_path = "cdftools_path" self.earthdiags.config.data_convention = Mock() - self.earthdiags.config.data_convention.lat_name = 'lat' - self.earthdiags.config.data_convention.lon_name = 'lon' + self.earthdiags.config.data_convention.lat_name = "lat" + self.earthdiags.config.data_convention.lon_name = "lon" - self.earthdiags.read_config('config/path') + self.earthdiags.read_config("config/path") - self.earthdiags.config.parse.assert_called_once_with('config/path') - self.assertEqual(os.environ['HDF5_USE_FILE_LOCKING'], 'FALSE') - self.assertEqual(self.earthdiags.config.scratch_dir, TempFile.scratch_folder) + self.earthdiags.config.parse.assert_called_once_with("config/path") + self.assertEqual(os.environ["HDF5_USE_FILE_LOCKING"], "FALSE") + self.assertEqual( + self.earthdiags.config.scratch_dir, TempFile.scratch_folder + ) self.assertEqual(self.earthdiags.config.cdftools_path, cdftools.path) - self.assertDictEqual(self.earthdiags.dic_variables, - {'nav_lat': 'lat', - 'nav_lev': 'lev', - 'nav_lon': 'lon', - 't': 'time', - 'time_counter': 'time', - 'x': 'i', - 'y': 'j', - 'z': 'lev'}) + self.assertDictEqual( + self.earthdiags.dic_variables, + { + "nav_lat": "lat", + "nav_lev": "lev", + "nav_lon": "lon", + "t": "time", + "time_counter": "time", + "x": "i", + "y": "j", + "z": "lev", + }, + ) def test_read_config_primavera(self): self.earthdiags.config = Mock() - self.earthdiags.config.scratch_dir = 'scratch_path' - self.earthdiags.config.cdftools_path = 'cdftools_path' + self.earthdiags.config.scratch_dir = "scratch_path" + self.earthdiags.config.cdftools_path = "cdftools_path" self.earthdiags.config.data_convention = Mock() - self.earthdiags.config.data_convention.lat_name = 'latitude' - self.earthdiags.config.data_convention.lon_name = 'longitude' + self.earthdiags.config.data_convention.lat_name = "latitude" + self.earthdiags.config.data_convention.lon_name = "longitude" - self.earthdiags.read_config('config/path') + self.earthdiags.read_config("config/path") - self.earthdiags.config.parse.assert_called_once_with('config/path') - self.assertEqual(os.environ['HDF5_USE_FILE_LOCKING'], 'FALSE') - self.assertEqual(self.earthdiags.config.scratch_dir, TempFile.scratch_folder) + self.earthdiags.config.parse.assert_called_once_with("config/path") + self.assertEqual(os.environ["HDF5_USE_FILE_LOCKING"], "FALSE") + self.assertEqual( + self.earthdiags.config.scratch_dir, TempFile.scratch_folder + ) self.assertEqual(self.earthdiags.config.cdftools_path, cdftools.path) - self.assertDictEqual(self.earthdiags.dic_variables, - {'nav_lat': 'latitude', - 'nav_lev': 'lev', - 'nav_lon': 'longitude', - 't': 'time', - 'time_counter': 'time', - 'x': 'i', - 'y': 'j', - 'z': 'lev'}) - - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.open_documentation') + self.assertDictEqual( + self.earthdiags.dic_variables, + { + "nav_lat": "latitude", + "nav_lev": "lev", + "nav_lon": "longitude", + "t": "time", + "time_counter": "time", + "x": "i", + "y": "j", + "z": "lev", + }, + ) + + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.open_documentation") def test_parse_args_open_doc(self, open_doc): - EarthDiags.parse_args(['--doc']) + EarthDiags.parse_args(["--doc"]) open_doc.assert_called_once() - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.read_config') - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.clean') + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.read_config") + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.clean") def test_parse_args_clean(self, clean, read_config): - EarthDiags.parse_args(['--clean']) + EarthDiags.parse_args(["--clean"]) clean.assert_called_once() - read_config.assert_called_once_with('diags.conf') + read_config.assert_called_once_with("diags.conf") - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.read_config') - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.report') + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.read_config") + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.report") def test_parse_args_report(self, report, read_config): - EarthDiags.parse_args(['--report']) + EarthDiags.parse_args(["--report"]) report.assert_called_once() - read_config.assert_called_once_with('diags.conf') + read_config.assert_called_once_with("diags.conf") - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.read_config') - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.run') + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.read_config") + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.run") def test_parse_args_run(self, run, read_config): EarthDiags.parse_args([]) run.assert_called_once() - read_config.assert_called_once_with('diags.conf') + read_config.assert_called_once_with("diags.conf") - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.read_config') - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.run') + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.read_config") + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.run") def test_log_levels(self, run, read_config): log_file = tempfile.mktemp() - for level in ('EVERYTHING', 'DEBUG', 'INFO', 'RESULT', 'USER_WARNING', - 'WARNING', 'ERROR', 'CRITICAL', 'NO_LOG'): - EarthDiags.parse_args(['-lc', level, '-lf', level, '-log', log_file]) + for level in ( + "EVERYTHING", + "DEBUG", + "INFO", + "RESULT", + "USER_WARNING", + "WARNING", + "ERROR", + "CRITICAL", + "NO_LOG", + ): + EarthDiags.parse_args( + ["-lc", level, "-lf", level, "-log", log_file] + ) self.assertEqual(getattr(Log, level), Log.console_handler.level) self.assertEqual(getattr(Log, level), Log.file_handler.level) if os.path.isfile(log_file): os.remove(log_file) - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.read_config') - @mock.patch('earthdiagnostics.earthdiags.EarthDiags.run') + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.read_config") + @mock.patch("earthdiagnostics.earthdiags.EarthDiags.run") def test_log_cdo_nco(self, run, read_config): - EarthDiags.parse_args(['-lc', 'DEBUG']) + EarthDiags.parse_args(["-lc", "DEBUG"]) run.assert_called_once() - read_config.assert_called_once_with('diags.conf') + read_config.assert_called_once_with("diags.conf") self.assertTrue(Utils.cdo().debug) self.assertTrue(Utils.nco().debug) diff --git a/test/unit/test_frequency.py b/test/unit/test_frequency.py index be424181..515b4b06 100644 --- a/test/unit/test_frequency.py +++ b/test/unit/test_frequency.py @@ -6,37 +6,57 @@ from earthdiagnostics.variable import VariableType class TestFrequency(TestCase): - def test_parse(self): - freq = Frequency('m') - self.assertEqual(Frequency.parse('m'), freq) + freq = Frequency("m") + self.assertEqual(Frequency.parse("m"), freq) self.assertIs(Frequency.parse(freq), freq) def test_not_supported(self): with self.assertRaises(ValueError): - Frequency('badfreq') + Frequency("badfreq") def test_get_monthly_mean(self): - self.assertEqual(Frequency('m').folder_name(VariableType.MEAN), 'monthly_mean') + self.assertEqual( + Frequency("m").folder_name(VariableType.MEAN), "monthly_mean" + ) def test_get_monthly_stats(self): - self.assertEqual(Frequency('m').folder_name(VariableType.STATISTIC), 'monthly_statistics') + self.assertEqual( + Frequency("m").folder_name(VariableType.STATISTIC), + "monthly_statistics", + ) def test_get_daily_mean(self): - self.assertEqual(Frequency('d').folder_name(VariableType.MEAN), 'daily_mean') + self.assertEqual( + Frequency("d").folder_name(VariableType.MEAN), "daily_mean" + ) def test_get_weekly_mean(self): - self.assertEqual(Frequency('w').folder_name(VariableType.MEAN), 'weekly_mean') + self.assertEqual( + Frequency("w").folder_name(VariableType.MEAN), "weekly_mean" + ) def test_get_daily_stats(self): - self.assertEqual(Frequency('d').folder_name(VariableType.STATISTIC), 'daily_statistics') + self.assertEqual( + Frequency("d").folder_name(VariableType.STATISTIC), + "daily_statistics", + ) def test_get_6hourlymean(self): - self.assertEqual(Frequency('6hr').folder_name(VariableType.MEAN), '6hourly') + self.assertEqual( + Frequency("6hr").folder_name(VariableType.MEAN), "6hourly" + ) def test_get_6hourlystatistics(self): - self.assertEqual(Frequency('6hr').folder_name(VariableType.STATISTIC), '6hourly_statistics') + self.assertEqual( + Frequency("6hr").folder_name(VariableType.STATISTIC), + "6hourly_statistics", + ) def test_get_climatology(self): - self.assertEqual(Frequency('clim').folder_name(VariableType.STATISTIC), 'clim') - self.assertEqual(Frequency('clim').folder_name(VariableType.MEAN), 'clim') + self.assertEqual( + Frequency("clim").folder_name(VariableType.STATISTIC), "clim" + ) + self.assertEqual( + Frequency("clim").folder_name(VariableType.MEAN), "clim" + ) diff --git a/test/unit/test_lint.py b/test/unit/test_lint.py index 4c1d1d91..f4d8da73 100644 --- a/test/unit/test_lint.py +++ b/test/unit/test_lint.py @@ -6,22 +6,21 @@ import pycodestyle # formerly known as pep8 class TestLint(unittest.TestCase): - def test_pep8_conformance(self): """Test that we conform to PEP-8.""" check_paths = [ - 'earthdiagnostics', - 'test', - ] - exclude_paths = [ - 'earthdiagnostics/cmor_tables/primavera' + "earthdiagnostics", + "test", ] + exclude_paths = ["earthdiagnostics/cmor_tables/primavera"] - print("PEP8 check of directories: {}\n".format(', '.join(check_paths))) + print("PEP8 check of directories: {}\n".format(", ".join(check_paths))) # Get paths wrt package root - package_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + package_root = os.path.dirname( + os.path.dirname(os.path.dirname(__file__)) + ) for paths in (check_paths, exclude_paths): for i, path in enumerate(paths): paths[i] = os.path.join(package_root, path) diff --git a/test/unit/test_modelling_realm.py b/test/unit/test_modelling_realm.py index 8340ffa7..ca2dd872 100644 --- a/test/unit/test_modelling_realm.py +++ b/test/unit/test_modelling_realm.py @@ -7,60 +7,97 @@ from earthdiagnostics.modelingrealm import ModelingRealm, ModelingRealms class TestModellingRealms(TestCase): - def test_parse(self): - self.assertEqual(ModelingRealms.parse('atmos'), ModelingRealms.atmos) - self.assertEqual(ModelingRealms.parse('atmoschem'), ModelingRealms.atmosChem) - self.assertEqual(ModelingRealms.parse('atmoSChem'), ModelingRealms.atmosChem) - self.assertIsNone(ModelingRealms.parse('')) + self.assertEqual(ModelingRealms.parse("atmos"), ModelingRealms.atmos) + self.assertEqual( + ModelingRealms.parse("atmoschem"), ModelingRealms.atmosChem + ) + self.assertEqual( + ModelingRealms.parse("atmoSChem"), ModelingRealms.atmosChem + ) + self.assertIsNone(ModelingRealms.parse("")) with self.assertRaises(ValueError): - ModelingRealms.parse('badrealm') + ModelingRealms.parse("badrealm") class TestModellingRealm(TestCase): - def setUp(self): - self.realm = ModelingRealm('ocean') + self.realm = ModelingRealm("ocean") self.data_convention = Mock() - self.data_convention.name = 'specs' + self.data_convention.name = "specs" def test__repr__(self): self.assertEqual(str(self.realm), repr(self.realm)) def test_constructor_fail_on_bad_realm(self): with self.assertRaises(ValueError): - ModelingRealm('badrealm') + ModelingRealm("badrealm") def test_comparison(self): - self.assertEqual(ModelingRealm('ocean'), self.realm) - self.assertNotEqual(ModelingRealm('OCEAN'), self.realm) - self.assertNotEqual(ModelingRealm('atmos'), self.realm) + self.assertEqual(ModelingRealm("ocean"), self.realm) + self.assertNotEqual(ModelingRealm("OCEAN"), self.realm) + self.assertNotEqual(ModelingRealm("atmos"), self.realm) def test_get_omon(self): - self.assertEqual(self.realm.get_table_name(Frequencies.monthly, 'specs'), 'Omon') + self.assertEqual( + self.realm.get_table_name(Frequencies.monthly, "specs"), "Omon" + ) def test_get_oimon(self): - self.assertEqual(ModelingRealm('seaIce').get_table_name(Frequencies.monthly, self.data_convention), 'OImon') + self.assertEqual( + ModelingRealm("seaIce").get_table_name( + Frequencies.monthly, self.data_convention + ), + "OImon", + ) def test_get_simon(self): - self.data_convention.name = 'cmip6' - self.assertEqual(ModelingRealm('seaIce').get_table_name(Frequencies.monthly, self.data_convention), 'SImon') + self.data_convention.name = "cmip6" + self.assertEqual( + ModelingRealm("seaIce").get_table_name( + Frequencies.monthly, self.data_convention + ), + "SImon", + ) def test_get_limon(self): - self.assertEqual(ModelingRealm('landIce').get_table_name(Frequencies.monthly, self.data_convention), 'LImon') + self.assertEqual( + ModelingRealm("landIce").get_table_name( + Frequencies.monthly, self.data_convention + ), + "LImon", + ) def test_get_day(self): - self.assertEqual(ModelingRealm('atmos').get_table_name(Frequencies.daily, self.data_convention), 'day') + self.assertEqual( + ModelingRealm("atmos").get_table_name( + Frequencies.daily, self.data_convention + ), + "day", + ) def test_get_6hrplev(self): - self.assertEqual(ModelingRealm('atmos').get_table_name(Frequencies.six_hourly, self.data_convention), '6hrPlev') + self.assertEqual( + ModelingRealm("atmos").get_table_name( + Frequencies.six_hourly, self.data_convention + ), + "6hrPlev", + ) def test_get_varfolder(self): - self.assertEqual(ModelingRealm('ocean').get_varfolder('var', 6, 3), 'var_f6h') - self.assertEqual(ModelingRealm('atmos').get_varfolder('var', 6, 3), 'var_f3h') + self.assertEqual( + ModelingRealm("ocean").get_varfolder("var", 6, 3), "var_f6h" + ) + self.assertEqual( + ModelingRealm("atmos").get_varfolder("var", 6, 3), "var_f3h" + ) frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = '3hr' - frequency.frequency = '3hr' - self.assertEqual(ModelingRealm('atmos').get_varfolder('var', 6, 3, frequency=frequency), - 'var') + frequency.__str__.return_value = "3hr" + frequency.frequency = "3hr" + self.assertEqual( + ModelingRealm("atmos").get_varfolder( + "var", 6, 3, frequency=frequency + ), + "var", + ) diff --git a/test/unit/test_obsreconmanager.py b/test/unit/test_obsreconmanager.py index b04cc568..a77feb6d 100644 --- a/test/unit/test_obsreconmanager.py +++ b/test/unit/test_obsreconmanager.py @@ -11,34 +11,33 @@ from earthdiagnostics.modelingrealm import ModelingRealms class TestObsReconManager(TestCase): - def setUp(self): self.config = Mock() - self.config.data_convention = 'specs' - self.config.data_type = 'exp' - self.config.experiment.expid = 'expid' - self.config.experiment.model = 'model' - self.config.experiment.experiment_name = 'expname' - self.config.experiment.institute = 'institute' + self.config.data_convention = "specs" + self.config.data_type = "exp" + self.config.experiment.expid = "expid" + self.config.experiment.model = "model" + self.config.experiment.experiment_name = "expname" + self.config.experiment.institute = "institute" self.config.experiment.member_count_start = 0 self.config.experiment.chunk_size = 12 self.config.experiment.num_chunks = 1 - self.config.experiment.calendar = 'standard' + self.config.experiment.calendar = "standard" self.config.experiment.atmos_timestep = 3 self.config.experiment.ocean_timestep = 6 frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'mon' - frequency.folder_name.return_value = 'monthly_mean' + frequency.__str__.return_value = "mon" + frequency.folder_name.return_value = "monthly_mean" self.config.frequency = frequency self.config.cmor.initialization_number = 1 - self.config.cmor.version = '' - self.config.cmor.default_ocean_grid = 'ocean_grid' - self.config.cmor.default_atmos_grid = 'atmos_grid' - self.config.cmor.activity = 'activity' + self.config.cmor.version = "" + self.config.cmor.default_ocean_grid = "ocean_grid" + self.config.cmor.default_atmos_grid = "atmos_grid" + self.config.cmor.activity = "activity" self.config.cmor.force = False self.config.cmor.force_untar = False @@ -47,8 +46,15 @@ class TestObsReconManager(TestCase): self.config.data_dir = self.tmp_dir def _make_data_folder(self): - os.makedirs(os.path.join(self.tmp_dir, self.config.data_type, self.config.experiment.institute, - self.config.experiment.model, self.config.experiment.expid)) + os.makedirs( + os.path.join( + self.tmp_dir, + self.config.data_type, + self.config.experiment.institute, + self.config.experiment.model, + self.config.experiment.expid, + ) + ) def tearDown(self): if os.path.exists(self.tmp_dir): @@ -61,7 +67,7 @@ class TestObsReconManager(TestCase): def test_bad_num_chunks(self): self.config.experiment.chunk_size = 3 - self.config.data_type = 'obs' + self.config.data_type = "obs" self._make_data_folder() with self.assertRaises(Exception): ObsReconManager(self.config) @@ -69,70 +75,113 @@ class TestObsReconManager(TestCase): def test_get_file_path(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon frequency = Mock() frequency.__str__ = Mock() - frequency.__str__.return_value = 'frequency' - frequency.folder_name.return_value = 'folder_name' + frequency.__str__.return_value = "frequency" + frequency.folder_name.return_value = "folder_name" manager = ObsReconManager(self.config) - self.assertEqual(manager.get_file_path('19900101', ModelingRealms.ocean, 'var', frequency, None), - os.path.join(self.tmp_dir, 'exp/institute/model/folder_name/var/var_19900101.nc')) + self.assertEqual( + manager.get_file_path( + "19900101", ModelingRealms.ocean, "var", frequency, None + ), + os.path.join( + self.tmp_dir, + "exp/institute/model/folder_name/var/var_19900101.nc", + ), + ) def test_get_file_path_default_freq(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon manager = ObsReconManager(self.config) - self.assertEqual(manager.get_file_path('19900101', ModelingRealms.ocean, 'var', None, None), - os.path.join(self.tmp_dir, 'exp/institute/model/monthly_mean/var/var_19900101.nc')) + self.assertEqual( + manager.get_file_path( + "19900101", ModelingRealms.ocean, "var", None, None + ), + os.path.join( + self.tmp_dir, + "exp/institute/model/monthly_mean/var/var_19900101.nc", + ), + ) def test_get_file_path_no_startdate(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon manager = ObsReconManager(self.config) - self.assertEqual(manager.get_file_path(None, ModelingRealms.ocean, 'var', None, None), - os.path.join(self.tmp_dir, 'exp/institute/model/monthly_mean/var/var.nc')) + self.assertEqual( + manager.get_file_path( + None, ModelingRealms.ocean, "var", None, None + ), + os.path.join( + self.tmp_dir, "exp/institute/model/monthly_mean/var/var.nc" + ), + ) def test_get_file_path_recon_weekly(self): cmor_var = Mock() omon = Mock() - omon.name = 'Omon' + omon.name = "Omon" cmor_var.get_table.return_value = omon - self.config.data_type = 'recon' + self.config.data_type = "recon" self.config.experiment.chunk_size = 1 self._make_data_folder() manager = ObsReconManager(self.config) - self.assertEqual(manager.get_file_path('19900101', ModelingRealms.ocean, 'var', None, None), - os.path.join(self.tmp_dir, 'recon/institute/model/monthly_mean/var/var_199001.nc')) - - @mock.patch('earthdiagnostics.obsreconmanager.ObsReconManager.get_file_path') + self.assertEqual( + manager.get_file_path( + "19900101", ModelingRealms.ocean, "var", None, None + ), + os.path.join( + self.tmp_dir, + "recon/institute/model/monthly_mean/var/var_199001.nc", + ), + ) + + @mock.patch( + "earthdiagnostics.obsreconmanager.ObsReconManager.get_file_path" + ) def test_declare_chunk(self, mock_get_file_path): - mock_get_file_path.return_value = '/path/to/file' + mock_get_file_path.return_value = "/path/to/file" cmor_manager = ObsReconManager(self.config) - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) - self.assertEqual(datafile.remote_file, '/path/to/file') + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) + self.assertEqual(datafile.remote_file, "/path/to/file") mock_frequency = Mock() - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1, - frequency=mock_frequency) - self.assertEqual(datafile.remote_file, '/path/to/file') + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, + "var", + "20010101", + 1, + 1, + frequency=mock_frequency, + ) + self.assertEqual(datafile.remote_file, "/path/to/file") self.config.var_manager.get_variable.return_value = None - datafile = cmor_manager.declare_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) - self.assertEqual(datafile.remote_file, '/path/to/file') - - @mock.patch('earthdiagnostics.obsreconmanager.ObsReconManager.get_file_path') + datafile = cmor_manager.declare_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) + self.assertEqual(datafile.remote_file, "/path/to/file") + + @mock.patch( + "earthdiagnostics.obsreconmanager.ObsReconManager.get_file_path" + ) def test_request_chunk(self, mock_get_file_path): - mock_get_file_path.return_value = '/path/to/file' + mock_get_file_path.return_value = "/path/to/file" cmor_manager = ObsReconManager(self.config) - datafile = cmor_manager.request_chunk(ModelingRealms.ocean, 'var', '20010101', 1, 1) - self.assertEqual(datafile.remote_file, '/path/to/file') + datafile = cmor_manager.request_chunk( + ModelingRealms.ocean, "var", "20010101", 1, 1 + ) + self.assertEqual(datafile.remote_file, "/path/to/file") diff --git a/test/unit/test_publisher.py b/test/unit/test_publisher.py index f884f46a..4a0a8484 100644 --- a/test/unit/test_publisher.py +++ b/test/unit/test_publisher.py @@ -5,7 +5,6 @@ from mock import Mock class TestPublisher(TestCase): - def test_suscribe(self): suscriber = Mock() pub = Publisher() @@ -16,7 +15,7 @@ class TestPublisher(TestCase): suscriber = Mock() pub = Publisher() pub.subscribe(suscriber) - self.assertTrue(hasattr(suscriber, 'update')) + self.assertTrue(hasattr(suscriber, "update")) self.assertIn(suscriber, pub.suscribers) def test_unsuscribe(self): diff --git a/test/unit/test_utils.py b/test/unit/test_utils.py index a277f7ca..a3f64e51 100644 --- a/test/unit/test_utils.py +++ b/test/unit/test_utils.py @@ -7,55 +7,70 @@ from earthdiagnostics.utils import TempFile, Utils class TestTempFile(TestCase): def setUp(self): - TempFile.scratch_folder = '/tmp' - TempFile.prefix = 'prefix' + TempFile.scratch_folder = "/tmp" + TempFile.prefix = "prefix" def test_get(self): - self.assertEqual(TempFile.get('tempfile', clean=False), '/tmp/tempfile') - self.assertEqual(TempFile.get('tempfile2', clean=True), '/tmp/tempfile2') - self.assertNotIn('/tmp/tempfile', TempFile.files) - self.assertIn('/tmp/tempfile2', TempFile.files) + self.assertEqual( + TempFile.get("tempfile", clean=False), "/tmp/tempfile" + ) + self.assertEqual( + TempFile.get("tempfile2", clean=True), "/tmp/tempfile2" + ) + self.assertNotIn("/tmp/tempfile", TempFile.files) + self.assertIn("/tmp/tempfile2", TempFile.files) TempFile.autoclean = True - self.assertEqual(TempFile.get('tempfile3'), '/tmp/tempfile3') - self.assertIn('/tmp/tempfile3', TempFile.files) + self.assertEqual(TempFile.get("tempfile3"), "/tmp/tempfile3") + self.assertIn("/tmp/tempfile3", TempFile.files) TempFile.autoclean = False - self.assertEqual(TempFile.get('tempfile4'), '/tmp/tempfile4') - self.assertNotIn('/tmp/tempfile4', TempFile.files) + self.assertEqual(TempFile.get("tempfile4"), "/tmp/tempfile4") + self.assertNotIn("/tmp/tempfile4", TempFile.files) - with mock.patch('tempfile.mkstemp') as mkstemp_mock: - with mock.patch('os.close') as close_mock: - mkstemp_mock.return_value = (34, 'path_to_tempfile') + with mock.patch("tempfile.mkstemp") as mkstemp_mock: + with mock.patch("os.close") as close_mock: + mkstemp_mock.return_value = (34, "path_to_tempfile") TempFile.get() - TempFile.get(suffix='suffix') + TempFile.get(suffix="suffix") - mkstemp_mock.assert_has_calls((mock.call(dir='/tmp', prefix='prefix', suffix='.nc'), - mock.call(dir='/tmp', prefix='prefix', suffix='suffix'))) + mkstemp_mock.assert_has_calls( + ( + mock.call(dir="/tmp", prefix="prefix", suffix=".nc"), + mock.call( + dir="/tmp", prefix="prefix", suffix="suffix" + ), + ) + ) close_mock.assert_has_calls((mock.call(34), mock.call(34))) def test_clean(self): - with mock.patch('os.path.exists') as exists_mock: - with mock.patch('tempfile.mkstemp'): - with mock.patch('os.close'): - with mock.patch('os.remove'): + with mock.patch("os.path.exists") as exists_mock: + with mock.patch("tempfile.mkstemp"): + with mock.patch("os.close"): + with mock.patch("os.remove"): TempFile.clean() TempFile.clean() exists_mock.side_effect = [True, False] TempFile.autoclean = True - TempFile.get('tempfile') - TempFile.get('tempfile2') + TempFile.get("tempfile") + TempFile.get("tempfile2") TempFile.clean() self.assertEqual(len(TempFile.files), 0) class TestUtils(TestCase): - def test_rename_variable(self): - with mock.patch('earthdiagnostics.utils.Utils.rename_variables') as rename_mock: - Utils.rename_variable('file', 'old', 'new') - Utils.rename_variable('file', 'old', 'new', False) - Utils.rename_variable('file', 'old', 'new', False, False) - rename_mock.assert_has_calls((mock.call('file', {'old': 'new'}, True, True), - mock.call('file', {'old': 'new'}, False, True), - mock.call('file', {'old': 'new'}, False, False))) + with mock.patch( + "earthdiagnostics.utils.Utils.rename_variables" + ) as rename_mock: + Utils.rename_variable("file", "old", "new") + Utils.rename_variable("file", "old", "new", False) + Utils.rename_variable("file", "old", "new", False, False) + rename_mock.assert_has_calls( + ( + mock.call("file", {"old": "new"}, True, True), + mock.call("file", {"old": "new"}, False, True), + mock.call("file", {"old": "new"}, False, False), + ) + ) diff --git a/test/unit/test_variable.py b/test/unit/test_variable.py index affc50ff..90c51709 100644 --- a/test/unit/test_variable.py +++ b/test/unit/test_variable.py @@ -3,7 +3,13 @@ from mock import Mock from unittest import TestCase -from earthdiagnostics.variable import CMORTable, VariableAlias, Variable, VariableJsonException, VariableManager +from earthdiagnostics.variable import ( + CMORTable, + VariableAlias, + Variable, + VariableJsonException, + VariableManager, +) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.constants import Basins from earthdiagnostics.frequency import Frequencies @@ -18,147 +24,157 @@ class TestCMORTable(TestCase): def test_str(self): """Test string representation""" - self.assertEqual(str(CMORTable('name', 'm', 'Month YEAR', 'realm')), 'name') + self.assertEqual( + str(CMORTable("name", "m", "Month YEAR", "realm")), "name" + ) def test_repr(self): """Test string Representation""" - self.assertEqual(repr(CMORTable('name', 'm', 'Month YEAR', 'realm')), 'name (realm m, Month YEAR)') + self.assertEqual( + repr(CMORTable("name", "m", "Month YEAR", "realm")), + "name (realm m, Month YEAR)", + ) def test_lt(self): - self.assertLess(CMORTable('a', 'm', 'Month YEAR', 'realm'), CMORTable('b', 'm', 'Month YEAR', 'realm')) + self.assertLess( + CMORTable("a", "m", "Month YEAR", "realm"), + CMORTable("b", "m", "Month YEAR", "realm"), + ) class TestVariableAlias(TestCase): - def test_str(self): - alias = VariableAlias('alias') - self.assertEqual(str(alias), 'alias') - alias.basin = 'basin' - self.assertEqual(str(alias), 'alias Basin: basin') - alias.grid = 'grid' - self.assertEqual(str(alias), 'alias Basin: basin Grid: grid') + alias = VariableAlias("alias") + self.assertEqual(str(alias), "alias") + alias.basin = "basin" + self.assertEqual(str(alias), "alias Basin: basin") + alias.grid = "grid" + self.assertEqual(str(alias), "alias Basin: basin Grid: grid") def test_eq(self): - alias = VariableAlias('alias') - self.assertEqual(alias, VariableAlias('alias')) - alias.basin = 'basin' - self.assertEqual(alias, VariableAlias('alias', 'basin')) - alias.grid = 'grid' - self.assertEqual(alias, VariableAlias('alias', 'basin', 'grid')) + alias = VariableAlias("alias") + self.assertEqual(alias, VariableAlias("alias")) + alias.basin = "basin" + self.assertEqual(alias, VariableAlias("alias", "basin")) + alias.grid = "grid" + self.assertEqual(alias, VariableAlias("alias", "basin", "grid")) self.assertNotEqual(alias, None) def test_ne(self): - alias = VariableAlias('alias') - self.assertNotEqual(alias, VariableAlias('alias2')) - alias.basin = 'basin' - self.assertNotEqual(alias, VariableAlias('alias', 'basin2')) - alias.grid = 'grid' - self.assertNotEqual(alias, VariableAlias('alias', 'basin', 'grid2')) + alias = VariableAlias("alias") + self.assertNotEqual(alias, VariableAlias("alias2")) + alias.basin = "basin" + self.assertNotEqual(alias, VariableAlias("alias", "basin2")) + alias.grid = "grid" + self.assertNotEqual(alias, VariableAlias("alias", "basin", "grid2")) class TestVariable(TestCase): - def test_parse_json(self): var = Variable() - json = {'out_name': 'out_name', - 'standard_name': 'standard_name', - 'long_name': 'long_name', - 'modeling_realm': 'ocean', - 'valid_min': 'valid_min', - 'valid_max': 'valid_max', - 'units': 'units', - } - var.parse_json(json, 'out_name') - - self.assertEqual(var.short_name, 'out_name') - self.assertEqual(var.standard_name, 'standard_name') - self.assertEqual(var.long_name, 'long_name') - - self.assertEqual(var.valid_min, 'valid_min') - self.assertEqual(var.valid_max, 'valid_max') - self.assertEqual(var.units, 'units') + json = { + "out_name": "out_name", + "standard_name": "standard_name", + "long_name": "long_name", + "modeling_realm": "ocean", + "valid_min": "valid_min", + "valid_max": "valid_max", + "units": "units", + } + var.parse_json(json, "out_name") + + self.assertEqual(var.short_name, "out_name") + self.assertEqual(var.standard_name, "standard_name") + self.assertEqual(var.long_name, "long_name") + + self.assertEqual(var.valid_min, "valid_min") + self.assertEqual(var.valid_max, "valid_max") + self.assertEqual(var.units, "units") self.assertEqual(var.priority, 1) self.assertEqual(var.domain, ModelingRealms.ocean) def test_parse_json_no_out_name(self): var = Variable() - json = {'standard_name': 'standard_name', - 'long_name': 'long_name', - 'modeling_realm': 'ocean', - 'valid_min': 'valid_min', - 'valid_max': 'valid_max', - 'units': 'units', - } + json = { + "standard_name": "standard_name", + "long_name": "long_name", + "modeling_realm": "ocean", + "valid_min": "valid_min", + "valid_max": "valid_max", + "units": "units", + } with self.assertRaises(VariableJsonException): - var.parse_json(json, 'out_name') + var.parse_json(json, "out_name") def test_parse_json_with_priority(self): var = Variable() - json = {'out_name': 'out_name', - 'standard_name': 'standard_name', - 'long_name': 'long_name', - 'modeling_realm': 'ocean', - 'valid_min': 'valid_min', - 'valid_max': 'valid_max', - 'units': 'units', - 'priority': '2', - } - var.parse_json(json, 'out_name') - - self.assertEqual(var.short_name, 'out_name') - self.assertEqual(var.standard_name, 'standard_name') - self.assertEqual(var.long_name, 'long_name') - - self.assertEqual(var.valid_min, 'valid_min') - self.assertEqual(var.valid_max, 'valid_max') - self.assertEqual(var.units, 'units') + json = { + "out_name": "out_name", + "standard_name": "standard_name", + "long_name": "long_name", + "modeling_realm": "ocean", + "valid_min": "valid_min", + "valid_max": "valid_max", + "units": "units", + "priority": "2", + } + var.parse_json(json, "out_name") + + self.assertEqual(var.short_name, "out_name") + self.assertEqual(var.standard_name, "standard_name") + self.assertEqual(var.long_name, "long_name") + + self.assertEqual(var.valid_min, "valid_min") + self.assertEqual(var.valid_max, "valid_max") + self.assertEqual(var.units, "units") self.assertEqual(var.priority, 2) self.assertEqual(var.domain, ModelingRealms.ocean) def test_parse_json_with_primavera_priority(self): var = Variable() - json = {'out_name': 'out_name', - 'standard_name': 'standard_name', - 'long_name': 'long_name', - 'modeling_realm': 'ocean', - 'valid_min': 'valid_min', - 'valid_max': 'valid_max', - 'units': 'units', - 'primavera_priority': '2', - } - var.parse_json(json, 'out_name') - - self.assertEqual(var.short_name, 'out_name') - self.assertEqual(var.standard_name, 'standard_name') - self.assertEqual(var.long_name, 'long_name') - - self.assertEqual(var.valid_min, 'valid_min') - self.assertEqual(var.valid_max, 'valid_max') - self.assertEqual(var.units, 'units') + json = { + "out_name": "out_name", + "standard_name": "standard_name", + "long_name": "long_name", + "modeling_realm": "ocean", + "valid_min": "valid_min", + "valid_max": "valid_max", + "units": "units", + "primavera_priority": "2", + } + var.parse_json(json, "out_name") + + self.assertEqual(var.short_name, "out_name") + self.assertEqual(var.standard_name, "standard_name") + self.assertEqual(var.long_name, "long_name") + + self.assertEqual(var.valid_min, "valid_min") + self.assertEqual(var.valid_max, "valid_max") + self.assertEqual(var.units, "units") self.assertEqual(var.priority, 2) self.assertEqual(var.domain, ModelingRealms.ocean) def test_get_modelling_realm(self): var = Variable() - domain = var.get_modelling_realm(('ocean',)) + domain = var.get_modelling_realm(("ocean",)) self.assertEqual(ModelingRealms.ocean, domain) - domain = var.get_modelling_realm(('ocean', 'atmos')) + domain = var.get_modelling_realm(("ocean", "atmos")) self.assertEqual(ModelingRealms.ocean, domain) - domain = var.get_modelling_realm(('ocean', 'ocnBgchem')) + domain = var.get_modelling_realm(("ocean", "ocnBgchem")) self.assertEqual(ModelingRealms.ocnBgchem, domain) - domain = var.get_modelling_realm(('ocean', 'seaIce')) + domain = var.get_modelling_realm(("ocean", "seaIce")) self.assertEqual(ModelingRealms.seaIce, domain) - domain = var.get_modelling_realm(('atmos', 'atmosChem')) + domain = var.get_modelling_realm(("atmos", "atmosChem")) self.assertEqual(ModelingRealms.atmosChem, domain) - domain = var.get_modelling_realm(('land', 'landIce')) + domain = var.get_modelling_realm(("land", "landIce")) self.assertEqual(ModelingRealms.landIce, domain) domain = var.get_modelling_realm(tuple()) @@ -166,16 +182,29 @@ class TestVariable(TestCase): def test_parse_csv(self): var = Variable() - var.parse_csv(['not_used', 'out_name', 'standard_name', 'long_name', 'ocean', 'global', 'units', - 'valid_min', 'valid_max', 'grid', 'Amon: ']) - self.assertEqual(var.short_name, 'out_name') - self.assertEqual(var.standard_name, 'standard_name') - self.assertEqual(var.long_name, 'long_name') - - self.assertEqual(var.valid_min, 'valid_min') - self.assertEqual(var.valid_max, 'valid_max') - self.assertEqual(var.units, 'units') - self.assertEqual(var.grid, 'grid') + var.parse_csv( + [ + "not_used", + "out_name", + "standard_name", + "long_name", + "ocean", + "global", + "units", + "valid_min", + "valid_max", + "grid", + "Amon: ", + ] + ) + self.assertEqual(var.short_name, "out_name") + self.assertEqual(var.standard_name, "standard_name") + self.assertEqual(var.long_name, "long_name") + + self.assertEqual(var.valid_min, "valid_min") + self.assertEqual(var.valid_max, "valid_max") + self.assertEqual(var.units, "units") + self.assertEqual(var.grid, "grid") self.assertEqual(var.domain, ModelingRealms.ocean) self.assertEqual(var.basin, Basins().Global) @@ -183,48 +212,73 @@ class TestVariable(TestCase): def test_get_table(self): var = Variable() var.domain = ModelingRealms.atmos - table = var.get_table(Frequencies.monthly, 'specs') + table = var.get_table(Frequencies.monthly, "specs") self.assertEqual(table.frequency, Frequencies.monthly) - self.assertEqual(table.name, 'Amon') - self.assertEqual(table.date, 'December 2013') + self.assertEqual(table.name, "Amon") + self.assertEqual(table.date, "December 2013") def test_get_table_added(self): convention = Mock() - convention.name = 'specs' + convention.name = "specs" var = Variable() var.domain = ModelingRealms.atmos - var.add_table(CMORTable('Amon', Frequencies.monthly, 'December 2013', ModelingRealms.atmos)) + var.add_table( + CMORTable( + "Amon", + Frequencies.monthly, + "December 2013", + ModelingRealms.atmos, + ) + ) table = var.get_table(Frequencies.monthly, convention) self.assertEqual(table.frequency, Frequencies.monthly) - self.assertEqual(table.name, 'Amon') - self.assertEqual(table.date, 'December 2013') + self.assertEqual(table.name, "Amon") + self.assertEqual(table.date, "December 2013") def test_get_table_multiple_added(self): convention = Mock() - convention.name = 'specs' + convention.name = "specs" var = Variable() var.domain = ModelingRealms.atmos - var.add_table(CMORTable('day', Frequencies.daily, 'December 2013', ModelingRealms.atmos)) - var.add_table(CMORTable('cfDay', Frequencies.daily, 'December 2013', ModelingRealms.atmos)) + var.add_table( + CMORTable( + "day", Frequencies.daily, "December 2013", ModelingRealms.atmos + ) + ) + var.add_table( + CMORTable( + "cfDay", + Frequencies.daily, + "December 2013", + ModelingRealms.atmos, + ) + ) table = var.get_table(Frequencies.daily, convention) self.assertEqual(table.frequency, Frequencies.daily) - self.assertEqual(table.name, 'day') - self.assertEqual(table.date, 'December 2013') + self.assertEqual(table.name, "day") + self.assertEqual(table.date, "December 2013") def test_get_table_not_added(self): convention = Mock() - convention.name = 'specs' + convention.name = "specs" var = Variable() var.domain = ModelingRealms.atmos - var.add_table(CMORTable('Amon', Frequencies.monthly, 'December 2013', ModelingRealms.atmos)) + var.add_table( + CMORTable( + "Amon", + Frequencies.monthly, + "December 2013", + ModelingRealms.atmos, + ) + ) table = var.get_table(Frequencies.daily, convention) self.assertEqual(table.frequency, Frequencies.daily) - self.assertEqual(table.name, 'day') - self.assertEqual(table.date, 'December 2013') + self.assertEqual(table.name, "day") + self.assertEqual(table.date, "December 2013") def test_get_table_not_matching(self): convention = Mock() - convention.name = 'specs' + convention.name = "specs" var = Variable() with self.assertRaises(ValueError): var.get_table(Frequencies.daily, convention) @@ -243,52 +297,62 @@ class TestVariableManager(TestCase): def test_load_primavera(self): """Test loading SPECS tables""" - self.var_manager.load_variables('primavera') + self.var_manager.load_variables("primavera") self.assertTrue(self.var_manager.get_all_variables()) def test_load_cmip6(self): """Test loading SPECS tables""" - self.var_manager.load_variables('cmip6') + self.var_manager.load_variables("cmip6") self.assertTrue(self.var_manager.get_all_variables()) def test_load_specs(self): """Test loading SPECS tables""" - self.var_manager.load_variables('specs') + self.var_manager.load_variables("specs") self.assertTrue(self.var_manager.get_all_variables()) def test_load_preface(self): """Test loading SPECS tables""" - self.var_manager.load_variables('preface') + self.var_manager.load_variables("preface") self.assertTrue(self.var_manager.get_all_variables()) def test_bad_load(self): """Test loading a bad table raises an exception""" with self.assertRaises(Exception): - self.var_manager.load_variables('badconvention') + self.var_manager.load_variables("badconvention") def test_get_variable(self): """Test get variable""" - var1 = self._get_var_mock('var1', ['var1_alias']) + var1 = self._get_var_mock("var1", ["var1_alias"]) self.var_manager.register_variable(var1) self.var_manager.create_aliases_dict() - self.assertIs(self.var_manager.get_variable('var1'), var1) - self.assertIs(self.var_manager.get_variable('var1_alias'), var1) + self.assertIs(self.var_manager.get_variable("var1"), var1) + self.assertIs(self.var_manager.get_variable("var1_alias"), var1) - self.assertIsNone(self.var_manager.get_variable('var2')) - self.assertIsNone(self.var_manager.get_variable('var2', True)) + self.assertIsNone(self.var_manager.get_variable("var2")) + self.assertIsNone(self.var_manager.get_variable("var2", True)) def test_get_variable_and_alias(self): """Test get variable and alias""" - var1 = self._get_var_mock('var1', ['var1_alias']) + var1 = self._get_var_mock("var1", ["var1_alias"]) self.var_manager.register_variable(var1) self.var_manager.create_aliases_dict() - self.assertEqual(self.var_manager.get_variable_and_alias('var1'), (VariableAlias('var1'), var1)) - self.assertEqual(self.var_manager.get_variable_and_alias('var1_alias'), (VariableAlias('var1_alias'), var1)) - - self.assertEqual(self.var_manager.get_variable_and_alias('var2'), (None, None)) - self.assertEqual(self.var_manager.get_variable_and_alias('var2', True), (None, None)) + self.assertEqual( + self.var_manager.get_variable_and_alias("var1"), + (VariableAlias("var1"), var1), + ) + self.assertEqual( + self.var_manager.get_variable_and_alias("var1_alias"), + (VariableAlias("var1_alias"), var1), + ) + + self.assertEqual( + self.var_manager.get_variable_and_alias("var2"), (None, None) + ) + self.assertEqual( + self.var_manager.get_variable_and_alias("var2", True), (None, None) + ) def _get_var_mock(self, name, aliases): var1 = Mock() @@ -301,13 +365,15 @@ class TestVariableManager(TestCase): alias_mock.basin = None return alias_mock - var1.known_aliases = [get_alias_mock(alias) for alias in [name] + aliases] + var1.known_aliases = [ + get_alias_mock(alias) for alias in [name] + aliases + ] return var1 def test_get_all_variables(self): """Test get all variables""" - var1 = self._get_var_mock('var1', ['var1_alias']) - var2 = self._get_var_mock('var2', ['var2_alias']) + var1 = self._get_var_mock("var1", ["var1_alias"]) + var2 = self._get_var_mock("var2", ["var2_alias"]) self.var_manager.register_variable(var1) self.var_manager.register_variable(var2) diff --git a/test/unit/test_variable_type.py b/test/unit/test_variable_type.py index 52cf76df..cc0af24e 100644 --- a/test/unit/test_variable_type.py +++ b/test/unit/test_variable_type.py @@ -10,13 +10,15 @@ class TestVariableType(TestCase): def test_mean(self): """Test to_str of MEAN""" - self.assertEqual(VariableType.to_str(VariableType.MEAN), 'mean') + self.assertEqual(VariableType.to_str(VariableType.MEAN), "mean") def test_statistics(self): """Test to_str of STATISTIC""" - self.assertEqual(VariableType.to_str(VariableType.STATISTIC), 'statistics') + self.assertEqual( + VariableType.to_str(VariableType.STATISTIC), "statistics" + ) def test_bad_one(self): """Test to_str raises if type not recognized""" with self.assertRaises(ValueError): - VariableType.to_str('bad type') + VariableType.to_str("bad type") diff --git a/test/unit/test_workmanager.py b/test/unit/test_workmanager.py index fa2a71e0..8a1a12c7 100644 --- a/test/unit/test_workmanager.py +++ b/test/unit/test_workmanager.py @@ -41,7 +41,8 @@ class TestDownloader(TestCase): datafile = Mock() def _download(): - raise Exception('Error') + raise Exception("Error") + datafile.download = _download self.downloader.submit(datafile) if six.PY3: @@ -62,9 +63,13 @@ class TestDownloader(TestCase): with self.assertLogs(log.Log.log) as cmd: self.downloader.start() self.downloader.shutdown() - self.assertListEqual(cmd.output, - ['INFO:bscearth.utils:Suscribers: () Size: 1', - 'INFO:bscearth.utils:Suscribers: () Size: 1']) + self.assertListEqual( + cmd.output, + [ + "INFO:bscearth.utils:Suscribers: () Size: 1", + "INFO:bscearth.utils:Suscribers: () Size: 1", + ], + ) else: self.downloader.start() self.downloader.shutdown() @@ -74,21 +79,28 @@ class TestDownloader(TestCase): no_suscribers = self._create_datafile_mock() self.downloader.submit(no_suscribers) suscriber = Mock() - one_suscriber = self._create_datafile_mock(suscribers=(suscriber, )) + one_suscriber = self._create_datafile_mock(suscribers=(suscriber,)) self.downloader.submit(one_suscriber) if six.PY3: with self.assertLogs(log.Log.log) as cmd: self.downloader.start() self.downloader.shutdown() - self.assertListEqual(cmd.output, - ['INFO:bscearth.utils:Suscribers: ({0},) Size: 0'.format(str(suscriber)), - 'INFO:bscearth.utils:Suscribers: () Size: 0']) + self.assertListEqual( + cmd.output, + [ + "INFO:bscearth.utils:Suscribers: ({0},) Size: 0".format( + str(suscriber) + ), + "INFO:bscearth.utils:Suscribers: () Size: 0", + ], + ) else: self.downloader.start() self.downloader.shutdown() def test_download_more_suscribers_waiting_first(self): """Test downloads with more diagnostics waiting go first""" + class StatusDiag(Diagnostic): def __init__(self, data_manager, pending): super(StatusDiag, self).__init__(data_manager) @@ -102,19 +114,27 @@ class TestDownloader(TestCase): suscriber = StatusDiag(Mock(), 1) suscriber.status = DiagnosticStatus.WAITING - one_waiting_suscriber = self._create_datafile_mock(suscribers=(suscriber,)) + one_waiting_suscriber = self._create_datafile_mock( + suscribers=(suscriber,) + ) self.downloader.submit(one_waiting_suscriber) suscriber = StatusDiag(Mock(), 2) suscriber.status = DiagnosticStatus.FAILED - one_failed_suscriber = self._create_datafile_mock(suscribers=(suscriber, )) + one_failed_suscriber = self._create_datafile_mock( + suscribers=(suscriber,) + ) self.downloader.submit(one_failed_suscriber) if six.PY3: with self.assertLogs(log.Log.log) as cmd: self.downloader.start() self.downloader.shutdown() - self.assertListEqual(cmd.output, - ['INFO:bscearth.utils:Suscribers: (1,) Size: 0', - 'INFO:bscearth.utils:Suscribers: (2,) Size: 0']) + self.assertListEqual( + cmd.output, + [ + "INFO:bscearth.utils:Suscribers: (1,) Size: 0", + "INFO:bscearth.utils:Suscribers: (2,) Size: 0", + ], + ) else: self.downloader.start() self.downloader.shutdown() @@ -125,7 +145,10 @@ class TestDownloader(TestCase): data_mock.suscribers = suscribers def _download(): - log.Log.info('Suscribers: {0.suscribers} Size: {0.size}'.format(data_mock)) + log.Log.info( + "Suscribers: {0.suscribers} Size: {0.size}".format(data_mock) + ) + data_mock.download = _download return data_mock @@ -174,8 +197,9 @@ class TestWorkManager(TestCase): def test_prepare_job_list(self): """Test job list preparation""" + class Diag1(Diagnostic): - alias = 'diag1' + alias = "diag1" @classmethod def generate_jobs(cls, diags, options): @@ -184,7 +208,11 @@ class TestWorkManager(TestCase): diag.add_subjob(Diag1(self.data_manager)) return (diag,) - self.config.get_commands.return_value = ['diag1', 'baddiag', 'diag1,badoption'] + self.config.get_commands.return_value = [ + "diag1", + "baddiag", + "diag1,badoption", + ] Diagnostic.register(Diag1) self.work_manager.prepare_job_list() @@ -193,10 +221,10 @@ class TestWorkManager(TestCase): self.config.max_cores = -1 class EmptyDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Null diag' + return "Null diag" def compute(self): pass @@ -219,13 +247,13 @@ class TestWorkManager(TestCase): self.config.skip_diags_done = True class SkippedDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Skipped diag' + return "Skipped diag" def compute(self): - raise Exception('Diagnostic should not be computed') + raise Exception("Diagnostic should not be computed") def request_data(self): pass @@ -252,13 +280,13 @@ class TestWorkManager(TestCase): self.config.skip_diags_done = False class SkippedDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Skipped diag' + return "Skipped diag" def compute(self): - raise Exception('Diagnostic should not be computed') + raise Exception("Diagnostic should not be computed") def request_data(self): pass @@ -278,14 +306,15 @@ class TestWorkManager(TestCase): def test_failed_run(self): """Test run when a diagnostic fails""" + class FailDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Fail diag' + return "Fail diag" def compute(self): - raise Exception('Must fail') + raise Exception("Must fail") def request_data(self): pass @@ -307,13 +336,13 @@ class TestWorkManager(TestCase): req = MockFile() req.storage_status = StorageStatus.READY req.local_status = LocalStatus.PENDING - req.remote_file = 'requested' + req.remote_file = "requested" self.data_manager.requested_files[req.remote_file] = req return req def _declare_chunk(self, *args, **kwargs): req = MockFile() - req.remote_file = 'generated' + req.remote_file = "generated" self.data_manager.requested_files[req.remote_file] = req req.storage_status = StorageStatus.PENDING req.local_status = LocalStatus.NOT_REQUESTED @@ -324,10 +353,10 @@ class TestWorkManager(TestCase): self.data_manager.config.max_cores = -1 class DataDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Data diag' + return "Data diag" def compute(self): self.declared.local_status = LocalStatus.READY @@ -336,7 +365,9 @@ class TestWorkManager(TestCase): self.request = self.request_chunk(None, None, None, None, None) def declare_data_generated(self): - self.declared = self.declare_chunk(None, None, None, None, None) + self.declared = self.declare_chunk( + None, None, None, None, None + ) diag = DataDiag(self.data_manager) self.work_manager.add_job(diag) @@ -350,10 +381,10 @@ class TestWorkManager(TestCase): self.data_manager.config.max_cores = -1 class DataDiag(Diagnostic): - alias = 'diag1' + alias = "diag1" def __str__(self): - return 'Data diag' + return "Data diag" def compute(self): self.declared.local_status = LocalStatus.READY @@ -362,14 +393,18 @@ class TestWorkManager(TestCase): self.request = self.request_chunk(None, None, None, None, None) def declare_data_generated(self): - self.declared = self.declare_chunk(None, None, None, None, None) + self.declared = self.declare_chunk( + None, None, None, None, None + ) self.declared.upload_fails = True data_diag = DataDiag(self.data_manager) self.work_manager.add_job(data_diag) self.work_manager.run() - self.assertIn(data_diag, self.work_manager.jobs[DiagnosticStatus.COMPLETED]) + self.assertIn( + data_diag, self.work_manager.jobs[DiagnosticStatus.COMPLETED] + ) self.assertTrue(self.work_manager.had_errors) def test_run_empty(self): @@ -377,7 +412,12 @@ class TestWorkManager(TestCase): if six.PY3: with self.assertLogs(log.Log.log) as cmd: self.work_manager.run() - self.assertTrue([record for record in cmd.records if - record.message == 'No diagnostics to run']) + self.assertTrue( + [ + record + for record in cmd.records + if record.message == "No diagnostics to run" + ] + ) else: self.work_manager.run() -- GitLab From aa3c51d005cb12e2b424ea4633a177ca8490b9ce Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 31 Jan 2020 11:57:01 +0100 Subject: [PATCH 12/14] Fix region dim names --- earthdiagnostics/datafile.py | 17 +++-- earthdiagnostics/ocean/regionmean.py | 11 ++-- earthdiagnostics/ocean/regionsum.py | 99 +++++++++++++--------------- earthdiagnostics/publisher.py | 2 +- test/unit/ocean/test_region_mean.py | 72 +------------------- test/unit/test_lint.py | 32 --------- 6 files changed, 63 insertions(+), 170 deletions(-) delete mode 100644 test/unit/test_lint.py diff --git a/earthdiagnostics/datafile.py b/earthdiagnostics/datafile.py index e32c5cdf..f43b83b0 100644 --- a/earthdiagnostics/datafile.py +++ b/earthdiagnostics/datafile.py @@ -462,12 +462,19 @@ class DataFile(Publisher): final_cube = cube_list.merge_cube() temp = TempFile.get() iris.save(final_cube, temp, zlib=True) + handler = Utils.open_cdf(temp) + renames = {} + for dim in handler.dimensions: + if dim.startswith('dim'): + renames[dim] = 'region' + if dim.startswith('string'): + renames[dim] = 'region_length' if '-' in final_cube.var_name: - Utils.rename_variable( - temp, - final_cube.var_name.replace('-', '_'), - final_cube.var_name, must_exist=False - ) + renames[final_cube.var_name.replace('-', '_')] = \ + final_cube.var_name + + Utils.rename_variables( + temp, renames, must_exist=False, rename_dimension=True) Utils.move_file(temp, self.local_file) self._correct_metadata() diff --git a/earthdiagnostics/ocean/regionmean.py b/earthdiagnostics/ocean/regionmean.py index 2eb41416..fe9cc813 100644 --- a/earthdiagnostics/ocean/regionmean.py +++ b/earthdiagnostics/ocean/regionmean.py @@ -68,7 +68,6 @@ class RegionMean(Diagnostic): variable, box, save3d, - variance, basins, grid_point, frequency, @@ -81,7 +80,6 @@ class RegionMean(Diagnostic): self.variable = variable self.box = box self.save3d = save3d - self.variance = variance self.basins = basins self.grid_point = grid_point self.frequency = frequency @@ -106,7 +104,7 @@ class RegionMean(Diagnostic): return ( "Region mean Startdate: {0.startdate} Member: {0.member} " "Chunk: {0.chunk} Variable: {0.variable} " - "Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance} " + "Box: {0.box} Save 3D: {0.save3d} " "Grid point: {0.grid_point}".format(self) ) @@ -134,7 +132,6 @@ class RegionMean(Diagnostic): DiagnosticIntOption("min_depth", -1), DiagnosticIntOption("max_depth", -1), DiagnosticBoolOption("save3D", True), - DiagnosticBoolOption("variance", False), DiagnosticOption("grid", ""), DiagnosticFrequencyOption("frequency", diags.config.frequency), ) @@ -165,7 +162,6 @@ class RegionMean(Diagnostic): var, box, options["save3D"], - options["variance"], options["basins"], options["grid_point"].lower(), options["frequency"], @@ -206,7 +202,7 @@ class RegionMean(Diagnostic): def _mean_2d_var(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) - mean = regmean.compute_regmean_2D(data.data, masks, areacello) + mean = regmean.compute_regmean_2d(data.data, masks, areacello) self._save_result_2d("mean", mean, data) def _meand_3d_variable(self, data, mesh, masks): @@ -221,7 +217,7 @@ class RegionMean(Diagnostic): e3 = e3.extract(depth_constraint) data = data.extract(depth_constraint) volcello = areacello * e3.data.astype(np.float32) - mean = regmean.compute_regmean_3D(data.data, masks, volcello) + mean = regmean.compute_regmean_3d(data.data, masks, volcello) self._save_result_2d("mean", mean, data) if self.save3d: mean3d = regmean.compute_regmean_levels(data.data, masks, volcello) @@ -328,6 +324,7 @@ class RegionMean(Diagnostic): self.chunk, box=box_save, region=self.basins, + frequency=self.frequency, ) def _save_result_2d(self, var, result, data): diff --git a/earthdiagnostics/ocean/regionsum.py b/earthdiagnostics/ocean/regionsum.py index adadf303..ded849f6 100644 --- a/earthdiagnostics/ocean/regionsum.py +++ b/earthdiagnostics/ocean/regionsum.py @@ -19,7 +19,8 @@ from earthdiagnostics.diagnostic import ( DiagnosticDomainOption, DiagnosticBoolOption, DiagnosticBasinListOption, - DiagnosticVariableOption, + DiagnosticVariableListOption, + DiagnosticFrequencyOption ) from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -70,6 +71,7 @@ class RegionSum(Diagnostic): save3d, basins, grid, + frequency, ): Diagnostic.__init__(self, data_manager) self.startdate = startdate @@ -82,11 +84,15 @@ class RegionSum(Diagnostic): self.save3d = save3d self.basins = basins self.grid = grid + self.frequency = frequency + self.declared = {} self.lat_name = "lat" self.lon_name = "lon" + self._hash = None + def __eq__(self, other): if self._different_type(other): return False @@ -99,14 +105,14 @@ class RegionSum(Diagnostic): and self.variable == other.variable and self.grid_point == other.grid_point and self.grid == other.grid - and self.basin == other.basin + and self.basins == other.basins ) def __str__(self): return ( - "Region sum Startdate: {0.startdate} Member: {0.member}" + "Region sum Startdate: {0.startdate} Member: {0.member} " "Chunk: {0.chunk} Variable: {0.variable} " - "Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d}" + "Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d} " "Original grid: {0.grid} Basin: {0.basins}".format(self) ) @@ -126,54 +132,52 @@ class RegionSum(Diagnostic): """ options_available = ( DiagnosticDomainOption(), - DiagnosticVariableOption(diags.data_manager.config.var_manager), - DiagnosticOption("grid_point", "T"), + DiagnosticVariableListOption( + diags.data_manager.config.var_manager, "variable" + ), DiagnosticBasinListOption("basins", "global"), + DiagnosticOption("grid_point", "T"), DiagnosticIntOption("min_depth", -1), DiagnosticIntOption("max_depth", -1), - DiagnosticIntOption("min_lat", -1), - DiagnosticIntOption("max_lat", -1), - DiagnosticIntOption("min_lon", -1), - DiagnosticIntOption("max_lon", -1), DiagnosticBoolOption("save3D", True), DiagnosticOption("grid", ""), + DiagnosticFrequencyOption("frequency", diags.config.frequency), ) options = cls.process_options(options, options_available) box = Box() box.min_depth = options["min_depth"] box.max_depth = options["max_depth"] - box.min_lat = options["min_lat"] - box.max_lat = options["max_lat"] - box.min_lon = options["min_lon"] - box.max_lon = options["max_lon"] basins = options["basins"] if not basins: Log.error("Basins not recognized") return () + basins.sort() job_list = list() - for ( - startdate, - member, - chunk, - ) in diags.config.experiment.get_chunk_list(): - job_list.append( - RegionSum( - diags.data_manager, - startdate, - member, - chunk, - options["domain"], - options["variable"], - options["grid_point"].lower(), - box, - options["save3D"], - options["basins"], - options["grid"], + for var in options["variable"]: + for ( + startdate, + member, + chunk, + ) in diags.config.experiment.get_chunk_list(): + job_list.append( + RegionSum( + diags.data_manager, + startdate, + member, + chunk, + options["domain"], + var, + options["grid_point"].lower(), + box, + options["save3D"], + options["basins"], + options["grid"], + options["frequency"], + ) ) - ) return job_list def request_data(self): @@ -185,6 +189,7 @@ class RegionSum(Diagnostic): self.member, self.chunk, grid=self.grid, + frequency=self.frequency ) def declare_data_generated(self): @@ -202,27 +207,10 @@ class RegionSum(Diagnostic): has_levels = self._fix_file_metadata() data = self._load_data() masks = {} - self.basins.sort() for basin in self.basins: masks[basin] = Utils.get_mask(basin) mesh = Nemo("mesh_hgr.nc", "mask_regions.nc") - if ( - self.box.min_lat is not -1 - and self.box.max_lat is not -1 - and self.box.min_lon is not -1 - and self.box.max_lat is not -1 - ): - name = "{0}_{1}".format( - Box.get_lat_str(self.box), Box.get_lon_str(self.box) - ) - - masks[name] = mesh.get_region_mask( - self.box.min_lat, - self.box.max_lat, - self.box.min_lon, - self.box.max_lon, - ) if has_levels: self._sum_3d_variable(data, mesh, masks) else: @@ -230,7 +218,7 @@ class RegionSum(Diagnostic): def _sum_2d_var(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) - varsum = regsum.compute_regsum_2D(data.data, masks, areacello) + varsum = regsum.compute_regsum_2d(data.data, masks, areacello) self._save_result_2d("sum", varsum, data) def _sum_3d_variable(self, data, mesh, masks): @@ -249,13 +237,13 @@ class RegionSum(Diagnostic): data = data.extract(depth_constraint) tmask = tmask.extract(depth_constraint) volcello = areacello * e3.data.astype(np.float32) - varsum = regsum.compute_regsum_3D( + varsum = regsum.compute_regsum_3d( data.data, masks, volcello, tmask.data ) self._save_result_2d("sum", varsum, data) if self.save3d: varsum3d = regsum.compute_regsum_levels( - data.data, masks, volcello, tmask + data.data, masks, volcello, tmask.data ) self._save_result_3d("sum", varsum3d, data) @@ -350,7 +338,7 @@ class RegionSum(Diagnostic): var = "sum" if threed: if not self.save3d: - return False + return original_name = "{0}_{1}".format(var, self.variable) final_name = "{1}3d{0}".format(var, self.variable) levels = ",lev" @@ -386,7 +374,7 @@ class RegionSum(Diagnostic): def _declare_var(self, var, threed, box_save): if threed: if not self.save3d: - return False + return final_name = "{1}3d{0}".format(var, self.variable) else: final_name = "{1}{0}".format(var, self.variable) @@ -400,6 +388,7 @@ class RegionSum(Diagnostic): box=box_save, region=self.basins, grid=self.grid, + frequency=self.frequency, ) def _save_result_2d(self, var, result, data): diff --git a/earthdiagnostics/publisher.py b/earthdiagnostics/publisher.py index 4f808731..695c9ca2 100644 --- a/earthdiagnostics/publisher.py +++ b/earthdiagnostics/publisher.py @@ -30,7 +30,7 @@ class Publisher(object): :param who: suscriber to remove :type who: object """ - del self._subscribers[who] + self._subscribers.pop(who) def dispatch(self, *args): """ diff --git a/test/unit/ocean/test_region_mean.py b/test/unit/ocean/test_region_mean.py index e62c5f0d..44a5ad01 100644 --- a/test/unit/ocean/test_region_mean.py +++ b/test/unit/ocean/test_region_mean.py @@ -56,7 +56,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "t", Frequencies.monthly, @@ -73,7 +72,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "t", Frequencies.monthly, @@ -95,7 +93,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -112,7 +109,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -134,7 +130,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -151,7 +146,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -178,7 +172,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -195,7 +188,6 @@ class TestRegionMean(TestCase): "var", box, True, - False, Basins().Global, "u", Frequencies.monthly, @@ -227,7 +219,6 @@ class TestRegionMean(TestCase): "var", box, False, - False, Basins().Global, "u", Frequencies.monthly, @@ -244,7 +235,6 @@ class TestRegionMean(TestCase): "var", box, False, - False, Basins().Global, "u", Frequencies.monthly, @@ -262,57 +252,6 @@ class TestRegionMean(TestCase): "1", "10", "false", - "True", - ], - ) - self.assertEqual(len(jobs), 2) - self.assertEqual( - jobs[0], - RegionMean( - self.data_manager, - "20010101", - 0, - 0, - ModelingRealms.ocean, - "var", - box, - False, - True, - Basins().Global, - "u", - Frequencies.monthly, - ), - ) - self.assertEqual( - jobs[1], - RegionMean( - self.data_manager, - "20010101", - 0, - 1, - ModelingRealms.ocean, - "var", - box, - False, - True, - Basins().Global, - "u", - Frequencies.monthly, - ), - ) - - jobs = RegionMean.generate_jobs( - self.diags, - [ - "diagnostic", - "ocean", - ["var"], - "global", - "U", - "1", - "10", - "false", - "True", "grid", ], ) @@ -329,7 +268,6 @@ class TestRegionMean(TestCase): box, False, Basins().Global, - True, "grid", Frequencies.monthly, ), @@ -346,7 +284,6 @@ class TestRegionMean(TestCase): box, False, Basins().Global, - True, "grid", Frequencies.monthly, ), @@ -363,7 +300,6 @@ class TestRegionMean(TestCase): "1", "10", "false", - "True", "grid", "day", ], @@ -381,7 +317,6 @@ class TestRegionMean(TestCase): box, False, Basins().Global, - True, "grid", Frequencies.daily, ), @@ -398,7 +333,6 @@ class TestRegionMean(TestCase): box, False, Basins().Global, - True, "grid", Frequencies.daily, ), @@ -419,7 +353,6 @@ class TestRegionMean(TestCase): "1", "10", "false", - "True", "grid", "day", "extra", @@ -439,13 +372,12 @@ class TestRegionMean(TestCase): "var", box, False, - True, - Basins().Global, + [Basins().Global], "T", Frequencies.monthly, ) self.assertEqual( str(diag), "Region mean Startdate: 20010101 Member: 0 Chunk: 0 Variable: var " - "Box: 1-10 Save 3D: False Save variance: True Grid point: T", + "Box: 1-10 Save 3D: False Grid point: T", ) diff --git a/test/unit/test_lint.py b/test/unit/test_lint.py deleted file mode 100644 index f4d8da73..00000000 --- a/test/unit/test_lint.py +++ /dev/null @@ -1,32 +0,0 @@ -""" Lint tests """ -import os -import unittest - -import pycodestyle # formerly known as pep8 - - -class TestLint(unittest.TestCase): - def test_pep8_conformance(self): - """Test that we conform to PEP-8.""" - - check_paths = [ - "earthdiagnostics", - "test", - ] - exclude_paths = ["earthdiagnostics/cmor_tables/primavera"] - - print("PEP8 check of directories: {}\n".format(", ".join(check_paths))) - - # Get paths wrt package root - package_root = os.path.dirname( - os.path.dirname(os.path.dirname(__file__)) - ) - for paths in (check_paths, exclude_paths): - for i, path in enumerate(paths): - paths[i] = os.path.join(package_root, path) - - style = pycodestyle.StyleGuide() - style.options.exclude.extend(exclude_paths) - style.options.max_line_length = 120 - - self.assertEqual(style.check_files(check_paths).total_errors, 0) -- GitLab From ec0c8e80421a68ec851054fb3715464607322cec Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Fri, 31 Jan 2020 15:30:06 +0100 Subject: [PATCH 13/14] Remove region from psi --- earthdiagnostics/ocean/psi.py | 26 +++++++------------------- test/unit/ocean/test_psi.py | 19 ++----------------- 2 files changed, 9 insertions(+), 36 deletions(-) diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index 24d8b322..e1624b6a 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -1,15 +1,13 @@ # coding=utf-8 import numpy as np -import netCDF4 import iris import iris.analysis import iris.coords import iris.util -from bscearth.utils.log import Log -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinListOption +from earthdiagnostics.diagnostic import Diagnostic from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile @@ -60,7 +58,7 @@ class Psi(Diagnostic): ) def __str__(self): - return "PSI Startdate: {0} Member: {1} Chunk: {2} Basins: {3}".format( + return "PSI Startdate: {0} Member: {1} Chunk: {2}".format( self.startdate, self.member, self.chunk, @@ -81,17 +79,11 @@ class Psi(Diagnostic): :type options: list[str] :return: """ - options_available = (DiagnosticBasinListOption("basins", "global"),) + options_available = () options = cls.process_options(options, options_available) - basins = options["basins"] - if not basins: - Log.error("Basins not recognized") - return () - masks = {} - basins.sort() - for basin in basins: + for basin in ('global', ): masks[basin] = Utils.get_mask(basin) job_list = list() @@ -158,14 +150,11 @@ class Psi(Diagnostic): Utils.copy_variable(handler_source, handler_temp, lon_name, True, True) handler_temp.createDimension("region", len(result)) handler_temp.createDimension("region_length", 50) - var_region = handler_temp.createVariable( - "region", "S1", ("region", "region_length") - ) var = handler_temp.createVariable( - "vsftbarot", float, ("time", "j", "i", "region") + "vsftbarot", float, ("time", "j", "i",) ) var.units = "m3/s" - var.coordinates = " ".join((lat_name, lon_name)) + var.coordinates = " ".join((lat_name, lon_name,)) var.missing_value = 1e20 var.fill_value = 1e20 var.valid_min = -300e6 @@ -173,9 +162,8 @@ class Psi(Diagnostic): var.long_name = "Barotropic_Stream_Function" for i, basin in enumerate(result): - var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) result[basin].mask = self.masks[basin] < 1 - var[..., i] = result[basin] + var[...] = result[basin] handler_temp.close() self.psi.set_local_file(temp, diagnostic=self) diff --git a/test/unit/ocean/test_psi.py b/test/unit/ocean/test_psi.py index a4b5efdb..c94ff3ea 100644 --- a/test/unit/ocean/test_psi.py +++ b/test/unit/ocean/test_psi.py @@ -33,27 +33,12 @@ class TestPsi(TestCase): Psi(self.data_manager, "20010101", 0, 1, {Basins().Global: None}), ) - jobs = Psi.generate_jobs(self.diags, ["diagnostic", "atl"]) - self.assertEqual(len(jobs), 2) - self.assertEqual( - jobs[0], - Psi( - self.data_manager, "20010101", 0, 0, {Basins().Atlantic: None} - ), - ) - self.assertEqual( - jobs[1], - Psi( - self.data_manager, "20010101", 0, 1, {Basins().Atlantic: None} - ), - ) - with self.assertRaises(Exception): - Psi.generate_jobs(self.diags, ["diagnostic", "atl", "badoption"]) + Psi.generate_jobs(self.diags, ["diagnostic", "extra"]) def test_str(self): psi = Psi(self.data_manager, "20000101", 1, 1, {Basins().Global: None}) self.assertEqual( str(psi), - "PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins: Global", + "PSI Startdate: 20000101 Member: 1 Chunk: 1", ) -- GitLab From 09228cbaba0052a8022093b7c2263e3bf944892e Mon Sep 17 00:00:00 2001 From: Javier Vegas-Regidor Date: Mon, 3 Feb 2020 15:00:26 +0100 Subject: [PATCH 14/14] Precipitation now get directly from the GRIB --- earthdiagnostics/cmorizer.py | 11 ----------- earthdiagnostics/cmormanager.py | 2 +- test/integration/test_cmorizer.py | 18 ++++++++---------- 3 files changed, 9 insertions(+), 22 deletions(-) diff --git a/earthdiagnostics/cmorizer.py b/earthdiagnostics/cmorizer.py index ae200776..cc254ead 100644 --- a/earthdiagnostics/cmorizer.py +++ b/earthdiagnostics/cmorizer.py @@ -507,8 +507,6 @@ class Cmorizer(object): # remap on regular Gauss grid codes = self.cmor.get_requested_codes() - if 228 in codes: - codes.update((142, 143)) codes_str = ",".join([str(code) for code in codes]) try: if grid == "SH": @@ -526,15 +524,6 @@ class Cmorizer(object): output=gribfile + "_", options="-R -f nc4 -t ecmwf", ) - # total precipitation (remove negative values) - if 228 in codes: - Utils.cdo().setcode( - 228, - input="-chname,LSP,TP -setmisstoc,0 -setvrange,0,Inf " - "-add {0}_142.128.nc {0}_143.128.nc".format(gribfile), - output="{0}_228.128.nc".format(gribfile), - options="-f nc4", - ) return True except CDOException: Log.info("No requested codes found in {0} file".format(grid)) diff --git a/earthdiagnostics/cmormanager.py b/earthdiagnostics/cmormanager.py index 9e9acaa3..e9a0fc90 100644 --- a/earthdiagnostics/cmormanager.py +++ b/earthdiagnostics/cmormanager.py @@ -511,7 +511,7 @@ class CMORManager(DataManager): else: raise ( Exception( - "Error appeared while cmorizing startdate {0}" + "Error appeared while cmorizing startdate {0} " "member {1}!".format(startdate, member_str) ) ) diff --git a/test/integration/test_cmorizer.py b/test/integration/test_cmorizer.py index 74a04024..14ee5db2 100644 --- a/test/integration/test_cmorizer.py +++ b/test/integration/test_cmorizer.py @@ -570,7 +570,7 @@ class TestCmorizer(TestCase): coord_data, folder_path, filename.replace("??", "GG"), - [142, 143, 129, 169, 180], + [142, 143, 129, 169, 180, 228], month, ) self._create_file_for_grib( @@ -682,17 +682,14 @@ class TestCmorizer(TestCase): daily_offsets[0] = 0.5 hourly_offsets = np.arange(0.25, 59, 0.25) - factor = 1.0 if code == 129: factor = 9.81 elif code in (180, 169): factor = 6 * 3600.0 - elif code == 228: - base_data = np.ones((2, 2, 2), np.float) * (142 + 143) - month_offsets *= 2 - daily_offsets *= 2 - hourly_offsets *= 2 - factor = 6 * 3600.0 / 1000 + elif code in (228, ): + factor = 6 * 3.6 + else: + factor = 1.0 base_data /= factor month_offsets /= factor @@ -719,7 +716,8 @@ class TestCmorizer(TestCase): for x, offset in enumerate(offsets): self.assertTrue( np.allclose(data.data[x, ...], base_data + offset), - "{} {} data wrong for {}: {}".format( - freq, x, var, data.data[x, ...] - base_data + "{} {} data wrong for {}: {} {} {}".format( + freq, x, var, data.data[x, ...] - base_data, + data.data[x, ...], base_data ), ) -- GitLab