diags.py

from earthdiagnostics.ocean import Salinity, Circulation, Heat, General, Ice
from utils import Utils
from earthdiagnostics import cdftools, TempFile
from autosubmit.date.chunk_date_lib import *
from parser import Parser
from autosubmit.config.log import Log
import shutil
import os


class Diags:
    def __init__(self, config_file):
        Log.debug('Initialising Diags')
        self._read_config(config_file)
        TempFile.scratch_folder = self.scratch_dir
        cdftools.path = self.cdftools_path
        Log.debug('Diags ready')

    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 fulldiag in self._get_commands():
            diag_options = fulldiag.split(',')
            diag = diag_options[0]
            Log.info("Running {0}", diag)
            Log.debug('Full command: {0}', fulldiag)
            if diag == 'vertmeansal':
                depth_min = int(diag_options[1])
                depth_max = int(diag_options[2])
                depth = '{0}-{1}'.format(depth_min, depth_max)
                for [input_file, output_file] in self._get_file_names('ocean', 'so',
                                                                      'someanm{0}'.format(depth)):
                    Salinity.vertical_mean(input_file, output_file, depth_min, depth_max)
            elif diag == 'vertmean':
                variable = diag_options[1]
                lev_min = int(diag_options[2])
                lev_max = int(diag_options[3])
                lev = '{0}-{1}'.format(lev_min, lev_max)
                for [input_file, output_file] in self._get_file_names('ocean', variable,
                                                                      '{0}mean{1}'.format(variable, lev)):
                    General.vertical_mean(input_file, output_file, variable, lev_min, lev_max)
            elif diag == 'convection':
                for [input_file, output_file] in self._get_file_names('ocean', 'mlotst',
                                                                      'mlotstsites'):
                    Circulation.convection_sites(input_file, self.nemo_version, 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':
                basin = diag_options[1]
                mixed_layer = int(diag_options[2])
                depth_min = int(diag_options[3])
                depth_max = int(diag_options[4])

                if mixed_layer == 1:
                    mxl = 'mlotst'
                    depth = ''
                elif mixed_layer == 0:
                    mxl = ''
                    depth = '{0}-{1}'.format(depth_min, depth_max)
                else:
                    mxl = 'nomlotst'
                    depth = ''

                for [input_file, mlotst_file, output_file] in self._get_file_names('ocean', 'thetao', 'mlotst',
                                                                      'ohcsum{0}{1}'.format(mxl, depth)):
                    Heat.total(input_file, mlotst_file, output_file, basin, mixed_layer, depth_min, depth_max)
            elif diag == 'ohclayer':
                depth_min = int(diag_options[1])
                depth_max = int(diag_options[2])

                depth = '{0}-{1}'.format(depth_min, depth_max)

                for [input_file, output_file] in self._get_file_names('ocean', 'thetao',
                                                                      'ohc{0}'.format(depth)):
                    Heat.layer(input_file, output_file,  depth_min, depth_max)
            elif diag in ['moc', 'msftmyz']:
                for [input_file, output_file] in self._get_file_names('ocean', 'vo',
                                                                      'msftmyz'):
                    Circulation.moc(input_file, output_file)
            elif diag in ['mocmax', 'msftmyzmax']:
                if len(diag_options) != 5:
                    Log.warning('Moc max requires 4 arguments. Skipping!')
                    continue

                lat_min = int(diag_options[1])
                lat_max = int(diag_options[2])
                lat = '{0}-{1}'.format(lat_min, lat_max)

                depth_min = int(diag_options[3])
                depth_max = int(diag_options[4])
                depth = '{0}-{1}'.format(depth_min, depth_max)

                for [input_file, output_file] in self._get_file_names('ocean', 'msftmyz',
                                                                      'msftmyzmax'):

                    Circulation.max_moc(input_file, lat_min, lat_max,
                                        output_file.replace('msftmyzmax', 'msftmyzmax{0}-{1}'.format(lat, depth)),
                                        depth_min, depth_max)
            elif diag in ['mocarea', 'msftmyzarea']:
                if len(diag_options) != 6:
                    Log.warning('area_moc requires between  5 arguments. Skipping!')
                    continue

                lat_min = int(diag_options[1])
                lat_max = int(diag_options[2])
                lat = '{0}-{1}'.format(lat_min, lat_max)

                depth_min = int(diag_options[3])
                depth_max = int(diag_options[4])
                depth = '{0}-{1}'.format(depth_min, depth_max)

                basin = diag_options[5]

                for [input_file, output_file] in self._get_file_names('ocean', 'msftmyz',
                                                                      'msftmyz{0}-{1}'.format(lat, depth)):
                    Circulation.area_moc(input_file, lat_min, lat_max, output_file, depth_min, depth_max, basin)
            elif diag == 'mlotsthc':
                for [input_file, mlotst_file, output_file] in self._get_file_names('ocean', 'thetao',
                                                                                   'mlotst', 'mlotsthc'):
                    Heat.mixed_layer_content(input_file, mlotst_file, output_file)
            elif diag == 'mlotstsc':
                for [input_file, mlotst_file, output_file] in self._get_file_names('ocean', 'so', 'mlotst', 'mlotstsc'):
                    Salinity.mixed_layer_content(input_file, mlotst_file, output_file)
            elif diag == 'interp3d':
                variable = diag_options[1]
                for [input_file, output_file] in self._get_file_names('ocean', variable,
                                                                      '{0}interp'.format(variable)):
                    General.interpolate3d(input_file, output_file, variable, self.nemo_version)
            elif diag == 'cutsection':
                variable = diag_options[1]
                meridional = diag_options[2] == 'm'
                value = int(diag_options[3])
                for [input_file, output_file] in self._get_file_names('ocean', variable,
                                                                      '{0}interp'.format(variable)):
                    General.cut_section(input_file, output_file, variable, meridional, value)
            elif diag == 'siasiesiv':
                for [thickness_file, fraction_file, output_file] in self._get_file_names('seaIce', 'sit', 'sic',
                                                                                         'siasiesiv'):
                    Ice.siasiesiv(thickness_file, fraction_file, output_file, 'global')
            else:
                Log.warning('Diagnostic {0} not available', diag)
                continue

            Log.result('Finished {0}', diag)

        TempFile.clean()

    def _get_commands(self):
        Log.debug('Preparing command list')
        commands = self.diags.split()

        for alias, added_commands in self._aliases.items():
            if alias in commands:
                Log.debug('Changing alias {0} for {1}', alias, ' '.join(added_commands))
                commands.remove(alias)
                for add_command in added_commands:
                    commands.append(add_command)
        Log.debug('Command list ready ')
        return commands

    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'):
            Log.info('Copying 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'):
            Log.info('Linking mesh_zgr.nc')
            os.link('mesh_hgr.nc', 'mesh_zgr.nc')

        if not os.path.exists('mask.nc'):
            Log.info('Linking mask.nc')
            os.link('mesh_hgr.nc', 'mask.nc')

        if not os.path.exists('new_maskglo.nc'):
            Log.info('Copying 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 \
                not os.path.exists('mask_regions.nc'):
            Log.info('Copying 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)

        if os.path.exists(os.path.join(self.con_files, 'mask.regions.3d.{0}.nc'.format(self.nemo_version))) and \
                not os.path.exists('mask_regions.3d.nc'):
            Log.info('Copying mask_regions.3d.nc')
            shutil.copy(os.path.join(self.con_files, 'mask.regions.3d.{0}.nc'.format(self.nemo_version)),
                        'mask_regions.3d.nc')
            Utils.rename_variables('mask_regions.3d.nc', dic_variables, False, True)

        if not os.path.exists('toto_N.nc'):
            Log.info('Copying toto_N.nc')
            shutil.copy(os.path.join(self.con_files, 'ice_template.nc'), 'toto_N.nc')
            Utils.rename_variables('toto_N.nc', dic_variables, False, True)

        if not os.path.exists('toto_S.nc'):
            Log.info('Linking toto_S.nc')
            os.link('toto_N.nc', 'toto_S.nc')
            Utils.rename_variables('toto_S.nc', dic_variables, False, True)

        Log.result('Mesh files ready!')

    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').lower()
        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')

        # Read aliases
        self._aliases = dict()
        if self.parser.has_section('ALIAS'):
            for option in self.parser.options('ALIAS'):
                self._aliases[option.lower()] = self.parser.get_option('ALIAS', option).lower().split()

        self.scratch_dir = os.path.join(self.scratch_dir, 'diags', self.expid)

    def _get_file_names(self, domain, *variables):
        file_names = list()

        if domain == 'seaIce':
            domain_abreviattion = 'OI'
        else:
            domain_abreviattion = domain[0].upper()
        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')

                    var_file = list()
                    for var in variables:
                        var_file.append(os.path.join(member_path, var, 'r{0}i1p1'.format(member_plus),
                                                     '{0}_{1}{2}_{3}_output_r{4}i1p1_'
                                                     '{5}-{6}.nc'.format(var, domain_abreviattion,
                                                                         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(var_file)

        return file_names


def main():
    TempFile.scratch_folder = '/scratch/Earth/jvegas'
    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()