diags.py

#!/usr/bin/env python
# coding=utf-8

import Queue
import argparse
import shutil
import threading

import netCDF4
import operator
import os
from autosubmit.date.chunk_date_lib import *

from datamanager import DataManager
from earthdiagnostics import cdftools
from earthdiagnostics.utils import TempFile
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.ocean import *
from earthdiagnostics.general import *
from experimentmanager import ExperimentManager
from ocean import ConvectionSites, Gyres, Psi, MaxMoc, AreaMoc, Moc, VerticalMean, VerticalMeanMeters, Interpolate, \
    AverageSection, CutSection, MixedLayerSaltContent, Siasiesiv
from parser import Parser
from utils import Utils


class Diags(object):
    """
    Launcher class for the diagnostics

    :param config_file: path to the configuration file
    :type config_file: str
    """
    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
        self._create_dic_variables()
        self.time = dict()
        Log.debug('Diags ready')

    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'] = 'lon'
        self.dic_variables['nav_lat'] = 'lat'
        self.dic_variables['nav_lev'] = 'lev'
        self.dic_variables['time_counter'] = 'time'
        self.dic_variables['t'] = 'time'

    def run(self):
        """
        Run the diagnostics
        """
        Log.debug('Using netCDF version {0}', netCDF4.getlibversion())
        if not os.path.exists(self.scratch_dir):
            os.makedirs(self.scratch_dir)
        os.chdir(self.scratch_dir)

        self._prepare_mesh_files()

        Diagnostic.register(MixedLayerSaltContent)
        Diagnostic.register(Siasiesiv)
        Diagnostic.register(VerticalMean)
        Diagnostic.register(VerticalMeanMeters)
        Diagnostic.register(Interpolate)
        Diagnostic.register(Moc)
        Diagnostic.register(AreaMoc)
        Diagnostic.register(MaxMoc)
        Diagnostic.register(Psi)
        Diagnostic.register(Gyres)
        Diagnostic.register(ConvectionSites)
        Diagnostic.register(CutSection)
        Diagnostic.register(AverageSection)
        Diagnostic.register(MixedLayerHeatContent)
        Diagnostic.register(HeatContentLayer)
        Diagnostic.register(HeatContent)
        Diagnostic.register(MonthlyMean)
        Diagnostic.register(Rewrite)

        parse_date('20000101')

        self.data_manager.prepare_CMOR_files(self.CMOR_force, self.CMOR_ocean, self.CMOR_atmosphere)

        # Run diagnostics
        Log.info('Running diagnostics')
        list_jobs = Queue.Queue()
        for fulldiag in self._get_commands():
            Log.info("Adding {0} to diagnostic list", fulldiag)
            diag_options = fulldiag.split(',')

            diag_class = Diagnostic.get_diagnostic(diag_options[0])
            if diag_class:
                for job in diag_class.generate_jobs(self, diag_options):
                    list_jobs.put(job)
                continue
            else:
                Log.error('{0} is not an available diagnostic', diag_options[0])

        time = datetime.datetime.now()
        Log.info("Starting to compute at {0}", time)
        num_threads = min(Utils.available_cpu_count(), self.max_cores)
        Log.info('Using {0} threads', num_threads)
        threads = list()
        for num_thread in range(0, num_threads):
            self.time[num_thread] = dict()
            t = threading.Thread(target=Diags._run_jobs, args=(self, list_jobs, num_thread))
            threads.append(t)
            t.start()

        list_jobs.join()

        TempFile.clean()
        finsih_time = datetime.datetime.now()
        Log.result("Diagnostics finished at {0}", finsih_time)
        Log.result("Time ellapsed: {0}\n", finsih_time - time)

        Log.info('Time consumed by each diagnostic class')
        Log.info('--------------------------------------')
        total = dict()
        for num_thread in range(0, num_threads):
            for key, value in self.time[num_thread].items():
                if key in total:
                    total[key] += value
                else:
                    total[key] = value

        for diag, time in sorted(total.items(), key=operator.itemgetter(1)):
            Log.info('{0:23} {1:}', diag.__name__, time)

    def _run_jobs(self, queue, numthread):
        def _run_job(current_job, retrials=1):
            while retrials >= 0:
                try:
                    Log.info('Starting {0}', current_job)
                    time = datetime.datetime.now()
                    current_job.compute()
                    time = datetime.datetime.now() - time
                    if type(current_job) in self.time[numthread]:
                        self.time[numthread][type(current_job)] += time
                    else:
                        self.time[numthread][type(current_job)] = time
                    Log.result('Finished {0}', current_job)
                    return True
                except Exception as ex:
                    retrials -= 1
                    Log.error('Job {0} failed: {1}', job, ex)
            return False
        count = 0
        failed_jobs = list()

        while not queue.empty():
            try:
                job = queue.get(timeout=1)
                if _run_job(job):
                    count += 1
                else:
                    failed_jobs.append(str(job))
                queue.task_done()
            except Queue.Empty:
                continue

        if len(failed_jobs) == 0:
            Log.result('Thread {0} finished after taking care of {1} tasks', numthread, count)
        else:
            Log.result('Thread {0} finished after running successfully {1} of {2} tasks', numthread, count,
                       count + len(failed_jobs))
            for job in failed_jobs:
                Log.error('Job {0} could not be run', job)
        return

    def _get_commands(self):
        Log.debug('Preparing command list')
        commands = self.diags.split()
        real_commands = list()
        for command in commands:
            if command in self._aliases:
                added_commands = self._aliases[command]
                Log.info('Changing alias {0} for {1}', command, ' '.join(added_commands))
                for add_command in added_commands:
                    real_commands.append(add_command)
            else:
                real_commands.append(command)
        Log.debug('Command list ready ')
        return real_commands

    def _prepare_mesh_files(self):
        Log.info('Copying mesh files')
        self._copy_file(os.path.join(self.con_files, 'mesh_mask_nemo.{0}.nc'.format(self.model_version)), 'mesh_hgr.nc',
                        self.restore_meshes)
        self._link_file('mesh_hgr.nc', 'mesh_zgr.nc')
        self._link_file('mesh_hgr.nc', 'mask.nc')

        self._copy_file(os.path.join(self.con_files, 'new_maskglo.{0}.nc'.format(self.model_version)), 'new_maskglo.nc',
                        self.restore_meshes)
        self._copy_file(os.path.join(self.con_files, 'mask.regions.{0}.nc'.format(self.model_version)),
                        'mask_regions.nc', self.restore_meshes)
        self._copy_file(os.path.join(self.con_files, 'mask.regions.3d.{0}.nc'.format(self.model_version)),
                        'mask_regions.3d.nc', self.restore_meshes)

        Log.result('Mesh files ready!')

    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.model_version)
            return

        if not force and os.path.exists(destiny):
            if os.stat(source).st_size == os.stat(destiny).st_size:
                Log.info('File {0} already exists', destiny)
                return

        Log.info('Creating file {0}', destiny)
        shutil.copy(source, destiny)
        Log.info('File {0} ready', destiny)
        Utils.rename_variables('mesh_hgr.nc', self.dic_variables, False, 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.model_version)
            return

        if os.path.lexists(destiny):
            os.remove(destiny)

        os.symlink(source, destiny)
        Log.info('File {0} ready', destiny)

    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')
        self.max_cores = self.parser.get_int_option('DIAGNOSTICS', 'MAX_CORES', 100000)
        self.restore_meshes = self.parser.get_bool_option('DIAGNOSTICS', 'RESTORE_MESHES', False)

        # Read experiment config
        self.institute = self.parser.get_option('EXPERIMENT', 'INSTITUTE')
        self.expid = self.parser.get_option('EXPERIMENT', 'EXPID')
        self.experiment_name = self.parser.get_option('EXPERIMENT', 'NAME', self.expid)

        members = list()
        for member in self.parser.get_option('EXPERIMENT', 'MEMBERS').split():
            members.append(int(member))

        member_digits = self.parser.get_int_option('EXPERIMENT', 'MEMBER_DIGITS', 1)
        startdates = self.parser.get_option('EXPERIMENT', 'STARTDATES').split()
        chunk_size = self.parser.get_int_option('EXPERIMENT', 'CHUNK_SIZE')
        chunks = self.parser.get_int_option('EXPERIMENT', 'CHUNKS')
        cal = self.parser.get_option('EXPERIMENT', 'CALENDAR', 'standard')

        self.exp_manager = ExperimentManager(startdates, members, chunks, chunk_size, member_digits, cal)

        self.model = self.parser.get_option('EXPERIMENT', 'MODEL')
        self.nfrp = self.parser.get_int_option('EXPERIMENT', 'NFRP')
        self.model_version = self.parser.get_option('EXPERIMENT', 'MODEL_VERSION')

        self.add_name = self.parser.get_bool_option('CMOR', 'ADD_NAME')
        self.add_startdate = self.parser.get_bool_option('CMOR', 'ADD_STARTDATE')

        # 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)
        if not os.path.exists(self.scratch_dir):
            os.makedirs(self.scratch_dir)
        os.chdir(self.scratch_dir)

        self.data_manager = DataManager(self.exp_manager, self.institute, self.model, self.expid, self.data_dir,
                                        self.frequency, self.experiment_name, self.scratch_dir, self.nfrp)

        self.data_manager.add_startdate = self.add_startdate
        self.data_manager.add_name = self.add_name

        self.CMOR_force = self.parser.get_bool_option('CMOR', 'FORCE', False)
        self.CMOR_ocean = self.parser.get_bool_option('CMOR', 'OCEAN_FILES', True)
        self.CMOR_atmosphere = self.parser.get_bool_option('CMOR', 'ATMOSPHERE_FILES', True)
        self.data_manager.associated_experiment = self.parser.get_option('CMOR', 'ASSOCIATED_EXPERIMENT',
                                                                         'to be filled')
        self.data_manager.associated_model = self.parser.get_option('CMOR', 'ASSOCIATED_MODEL', 'to be filled')
        self.data_manager.initialization_description = self.parser.get_option('CMOR', 'INITIALIZATION_DESCRIPTION',
                                                                              'to be filled')
        self.data_manager.initialization_method = self.parser.get_option('CMOR', 'INITIALIZATION_METHOD',
                                                                         '1')
        self.data_manager.physics_description = self.parser.get_option('CMOR', 'PHYSICS_DESCRIPTION', 'to be filled')
        self.data_manager.physics_version = self.parser.get_option('CMOR', 'PHYSICS_VERSION', '1')
        self.data_manager.source = self.parser.get_option('CMOR', 'SOURCE', 'to be filled')


def main():
    """
    Entry point for the Earth Diagnostics. For more detailed documentation, use -h option
    """
    parser = argparse.ArgumentParser(description='Main executable for Earth Diagnostics.')
    parser.add_argument('-v', '--version', action='version', version='3.0.0b3',
                        help="returns Earth Diagnostics's version number and exit")
    parser.add_argument('--doc', action='store_true',
                        help="opens documentation and exits")
    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()
    if args.doc:
        Log.info('Opening documentation...')
        Utils.execute_shell_command(('xdg-open', os.path.join(os.path.dirname(os.path.realpath(__file__)), '..',
                                                              'EarthDiagnostics.pdf')))
        Log.result('Documentation opened!')
        return
    Log.set_console_level(args.logconsole)
    Log.set_file_level(args.logfile)
    if args.logfilepath:
        Log.set_file(Utils.expand_path(args.logfilepath))

    diags = Diags(Utils.expand_path(args.configfile))
    diags.run()
    TempFile.clean()


if __name__ == "__main__":
    main()