config.py 8.4 KB
Newer Older
# coding=utf-8
import os

from autosubmit.config.log import Log
from autosubmit.date.chunk_date_lib import parse_date, chunk_start_date, chunk_end_date

from earthdiagnostics.parser import Parser


class Config(object):   
    
    def __init__(self, path):
        parser = Parser()
        parser.optionxform = str
        parser.read(path)

        # Read diags config
        self.scratch_dir = parser.get_option('DIAGNOSTICS', 'SCRATCH_DIR')
        self.data_dir = parser.get_option('DIAGNOSTICS', 'DATA_DIR')
        self.con_files = parser.get_option('DIAGNOSTICS', 'CON_FILES')
        self.diags = parser.get_option('DIAGNOSTICS', 'DIAGS').lower()
        self.frequency = parser.get_option('DIAGNOSTICS', 'FREQUENCY')
        self.cdftools_path = parser.get_option('DIAGNOSTICS', 'CDFTOOLS_PATH')
        self.max_cores = parser.get_int_option('DIAGNOSTICS', 'MAX_CORES', 100000)
        self.restore_meshes = parser.get_bool_option('DIAGNOSTICS', 'RESTORE_MESHES', False)

        # Read experiment config
        self.experiment = ExperimentConfig(parser)

        # Read aliases
        self._aliases = dict()
        if parser.has_section('ALIAS'):
            for option in parser.options('ALIAS'):
                self._aliases[option.lower()] = parser.get_option('ALIAS', option).lower().split()
        Log.debug('Preparing command list')
        commands = self.diags.split()
        self._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:
                    self._real_commands.append(add_command)
            else:
                self._real_commands.append(command)
        Log.debug('Command list ready ')

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

        self.cmor = CMORConfig(parser)

    def get_commands(self):
        return self._real_commands
        

class CMORConfig(object):
    
    def __init__(self, parser):
        self.force = parser.get_bool_option('CMOR', 'FORCE', False)
        self.ocean = parser.get_bool_option('CMOR', 'OCEAN_FILES', True)
        self.atmosphere = parser.get_bool_option('CMOR', 'ATMOSPHERE_FILES', True)
        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.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.add_name = parser.get_bool_option('CMOR', 'ADD_NAME')
        self.add_startdate = parser.get_bool_option('CMOR', 'ADD_STARTDATE')

        self._var_hourly = self._parse_variables(parser.get_option('CMOR', 'ATMOS_HOURLY_VARS', ''))
        self._var_daily = self._parse_variables(parser.get_option('CMOR', 'ATMOS_DAILY_VARS', ''))
        self._var_monthly = self._parse_variables(parser.get_option('CMOR', 'ATMOS_MONTHLY_VARS', ''))

    def _parse_variables(self, raw_string):
        variables = dict()
        if raw_string:
            splitted = raw_string.split(',')
            for var_section in splitted:
                splitted_var = var_section.split(':')
                if len(splitted_var) == 1:
                    levels = None
                else:
                    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('-'))
        start = int(levels_splitted[0])
        end = int(levels_splitted[1])
        if len(levels_splitted) == 3:
            step = int(levels_splitted[2])
        else:
            step = 1
        return range(start, end, step)


class ExperimentConfig(object):
    """
    Encapsulates all chunk related tasks

    :param parser: parser for the config file
    :type parser: Parser
    """

    def __init__(self, parser):
        self.institute = parser.get_option('EXPERIMENT', 'INSTITUTE')
        self.expid = parser.get_option('EXPERIMENT', 'EXPID')
        self.experiment_name = parser.get_option('EXPERIMENT', 'NAME', self.expid)

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

        member_digits = parser.get_int_option('EXPERIMENT', 'MEMBER_DIGITS', 1)
        startdates = parser.get_option('EXPERIMENT', 'STARTDATES').split()
        chunk_size = parser.get_int_option('EXPERIMENT', 'CHUNK_SIZE')
        chunks = parser.get_int_option('EXPERIMENT', 'CHUNKS')
        calendar = parser.get_option('EXPERIMENT', 'CALENDAR', 'standard')
        self.model = parser.get_option('EXPERIMENT', 'MODEL')
        self.nfrp = parser.get_int_option('EXPERIMENT', 'NFRP')
        self.model_version = parser.get_option('EXPERIMENT', 'MODEL_VERSION')

        self.startdates = startdates
        self.members = members
        self.num_chunks = chunks
        self.chunk_size = chunk_size
        self.member_digits = member_digits
        self.calendar = calendar

    def get_chunk_list(self):
        """
        Return a list with all the chunks
        :return: List containing tuples of startdate, member and chunk
        :rtype: tuple[str, int, int]
        """
        chunk_list = list()
        for startdate in self.startdates:
            for member in self.members:
                for chunk in range(1, self.num_chunks + 1):
                    chunk_list.append((startdate, member, chunk))
        return chunk_list

    def get_member_list(self):
        """
        Return a list with all the members
        :return: List containing tuples of startdate and member
        :rtype: tuple[str, int, int]
        """
        member_list = list()
        for startdate in self.startdates:
            for member in self.members:
                    member_list.append((startdate, member))
        return member_list

    def get_year_chunks(self, startdate, year):
        """
        Get the list of chunks containing timesteps from the given year
        :param startdate: startdate to use
        :type startdate: str
        :param year: reference year
        :type year: int
        :return: list of chunks containing data from the given year
        :rtype: list[int]
        """
        date = parse_date(startdate)
        chunks = list()
        for chunk in range(1, self.num_chunks + 1):
            chunk_start = chunk_start_date(date, chunk, self.chunk_size, 'month', self.calendar)
            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:
                chunks.append(chunk)

        return chunks

    def get_full_years(self, startdate):
        """
        Returns the list of full years that are in the given startdate
        :param startdate: startdate to use
        :type startdate: str
        :return: list of full years
        :rtype: list[int]
        """
        chunks_per_year = 12 / self.chunk_size
        date = parse_date(startdate)
        first_january = 0
        first_year = date.year
        if date.month != 1:
            month = date.month
            first_year += 1
            while month + self.chunk_size < 12:
                month += self.chunk_size
                first_january += 1

        years = list()
        for chunk in range(first_january, self.num_chunks - chunks_per_year, chunks_per_year):
            years.append(first_year)
            first_year += 1
        return years

    def get_member_str(self, member):
        """
        Returns the member name for a given member number.
        :param member: member's number
        :type member: int
        :return: member's name
        :rtype: str
        """
        return 'fc{0}'.format(str(member).zfill(self.member_digits))