from parser import Parser from autosubmit.config.log import Log from earthdiagnostics.ocean import Salinity, Convection, Circulation, Heat from utils import Utils from earthdiagnostics import cdftools from autosubmit.date.chunk_date_lib import * import shutil import os class Diags: def __init__(self, config_file): self._read_config(config_file) Utils.scratch_folder = os.path.join(self.scratch_dir) cdftools.path = self.cdftools_path def run(self): if not os.path.exists(self.scratch_dir): os.makedirs(self.scratch_dir) os.chdir(self.scratch_dir) self._prepare_mesh_files() # Check if cmorized and convert if not # Run diagnostics Log.info('Running diagnostics') for diag in self.diags.split(): Log.info("Running {0}", diag) if diag == 'vert_mean_sal': min_depth = 0 max_depth = 300 for [input_file, output_file] in self._get_file_names('ocean', 'so', 'vertmeansal'.format(min_depth, max_depth)): Salinity.vertical_mean(str(input_file), str(output_file), min_depth, max_depth) elif diag == 'convection': for [input_file, output_file] in self._get_file_names('ocean', 'mlotst', 'mlotst_sites'): Convection.main_sites(str(input_file), self.nemo_version, str(output_file)) elif diag == 'psi': for [input_file, output_file] in self._get_file_names('ocean', 'sobarstf', 'psi'): Circulation.psi(str(input_file), self.nemo_version, str(output_file)) elif diag == 'gyres': for [input_file, output_file] in self._get_file_names('ocean', 'psi', 'msftbarot'): Circulation.gyres(str(input_file), self.nemo_version, str(output_file)) elif diag == 'ohc_specified_layer': for [input_file, output_file] in self._get_file_names('ocean', 'ohc', 'ohc_2d_avg_0-300m'): Heat.layer(input_file, output_file, 0, 300) Heat.layer(input_file, output_file.replace('ohc_2d_avg_0-300m', 'ohc_2d_avg_300-800m'), 300, 800) Log.result('Finished {0}', diag) Utils.clean() def _prepare_mesh_files(self): dic_variables = dict() dic_variables['x'] = 'i' dic_variables['y'] = 'j' dic_variables['z'] = 'lev' dic_variables['nav_lon'] = 'lon' dic_variables['nav_lat'] = 'lat' dic_variables['nav_lev'] = 'lev' dic_variables['time_counter'] = 'time' dic_variables['t'] = 'time' Log.info('Copying mesh files') if not os.path.exists('mesh_hgr.nc'): shutil.copy(os.path.join(self.con_files, 'mesh_mask_nemo.{0}.nc'.format(self.nemo_version)), 'mesh_hgr.nc') Utils.rename_variables('mesh_hgr.nc', dic_variables, False, True) if not os.path.exists('mesh_zgr.nc'): os.link('mesh_hgr.nc', 'mesh_zgr.nc') if not os.path.exists('mask.nc'): os.link('mesh_hgr.nc', 'mask.nc') if not os.path.exists('new_maskglo.nc'): shutil.copy(os.path.join(self.con_files, 'new_maskglo.{0}.nc'.format(self.nemo_version)), 'new_maskglo.nc') if os.path.exists(os.path.join(self.con_files, 'mask.regions.{0}.nc'.format(self.nemo_version))) and \ os.path.exists('mask_regions.nc'): shutil.copy(os.path.join(self.con_files, 'mask.regions.{0}.nc'.format(self.nemo_version)), 'mask_regions.nc') Utils.rename_variables('mask_regions.nc', dic_variables, False, True) def _read_config(self, config_file): self.parser = Parser() self.parser.optionxform = str self.parser.read(config_file) # Read diags config self.scratch_dir = self.parser.get_option('DIAGNOSTICS', 'SCRATCH_DIR') self.data_dir = self.parser.get_option('DIAGNOSTICS', 'DATA_DIR') self.con_files = self.parser.get_option('DIAGNOSTICS', 'CON_FILES') self.diags = self.parser.get_option('DIAGNOSTICS', 'DIAGS') self.frequency = self.parser.get_option('DIAGNOSTICS', 'FREQUENCY') self.cdftools_path = self.parser.get_option('DIAGNOSTICS', 'CDFTOOLS_PATH') # Read experiment config self.institute = self.parser.get_option('EXPERIMENT', 'INSTITUTE') self.expid = self.parser.get_option('EXPERIMENT', 'EXPID') self.members = self.parser.get_option('EXPERIMENT', 'MEMBERS') self.startdates = self.parser.get_option('EXPERIMENT', 'STARTDATES') self.chunk_size = self.parser.get_int_option('EXPERIMENT', 'CHUNK_SIZE') self.chunks = self.parser.get_int_option('EXPERIMENT', 'CHUNKS') self.model = self.parser.get_option('EXPERIMENT', 'MODEL') self.nemo_version = self.parser.get_option('EXPERIMENT', 'NEMO_VERSION') self.scratch_dir = os.path.join(self.scratch_dir, 'diags', self.expid) def _get_file_names(self, domain, input_var, output_var): file_names = list() for startdate in self.startdates.split(): start = parse_date(startdate) for member in self.members: member_plus = str(int(member)+1) member_path = os.path.join(self.data_dir,self.expid, startdate, 'fc'+member, 'outputs', 'output', self.institute, self.model, self.model, 'S'+startdate, self.frequency, domain) for chunk in range(1, self.chunks + 1): chunk_start = chunk_start_date(start, chunk, self.chunk_size, 'month', 'standard') chunk_end = chunk_end_date(chunk_start, self.chunk_size, 'month', 'standard') chunk_end = previous_day(chunk_end, 'standard') input_file = os.path.join(member_path, input_var, 'r{0}i1p1'.format(member_plus), '{0}_{1[0]}{2}_{3}_output_r{4}i1p1_' '{5}-{6}.nc'.format(input_var, domain.upper(), self.frequency, self.model, member_plus, "{0:04}{1:02}".format(chunk_start.year, chunk_start.month), "{0:04}{1:02}".format(chunk_end.year, chunk_end.month))) output_file = os.path.join(member_path, input_var, 'r{0}i1p1'.format(member_plus), '{0}_{1[0]}{2}_{3}_output_r{4}i1p1_' '{5}-{6}.nc'.format(output_var, domain.upper(), self.frequency, self.model, member_plus, "{0:04}{1:02}".format(chunk_start.year, chunk_start.month), "{0:04}{1:02}".format(chunk_end.year, chunk_end.month))) file_names.append([input_file, output_file]) return file_names def main(): Log.set_console_level(Log.DEBUG) diags = Diags('/home/Earth/jvegas/pyCharm/ocean_diagnostics/earthdiagnostics/diags.conf') diags.run() # Utils.clean() if __name__ == "__main__": main()