diagnostic.py

# coding=utf-8
from earthdiagnostics.constants import Basins
from earthdiagnostics.frequency import Frequency, Frequencies
from earthdiagnostics.variable_type import VariableType
from earthdiagnostics.modelingrealm import ModelingRealm


class Diagnostic(object):
    """
    Base class for the diagnostics. 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
    :type data_manager: DataManager
    """

    alias = None
    """
    Alias to call the diagnostic. Must be overridden at the derived clases
    """
    _diag_list = dict()

    def __init__(self, data_manager):
        self.data_manager = data_manager
        self.required_vars = []
        self.generated_vars = []
        self.can_run_multiple_instances = True

    def __repr__(self):
        return str(self)

    @staticmethod
    def register(cls):
        """
        Register a new diagnostic using the given alias. It must be call using the derived class.
        :param cls: diagnostic class to register
        :type cls: Diagnostic
        """
        if not issubclass(cls, Diagnostic):
            raise ValueError('Class {0} must be derived from Diagnostic'.format(cls))
        if cls.alias is None:
            raise ValueError('Diagnostic class {0} must have defined an alias'.format(cls))
        Diagnostic._diag_list[cls.alias] = cls

    # noinspection PyProtectedMember
    @staticmethod
    def get_diagnostic(name):
        """
        Return the class for a diagnostic given its name

        :param name: diagnostic alias
        :type name: str
        :return: the selected Diagnostic class, None if name can not be found
        :rtype: Diagnostic
        """
        if name in Diagnostic._diag_list.keys():
            return Diagnostic._diag_list[name]
        return None

    def send_file(self, filetosend, domain, var, startdate, member, chunk=None, grid=None, region=None,
                  box=None, rename_var=None, frequency=None, year=None, date_str=None, move_old=False,
                  vartype=VariableType.MEAN):
        """

        :param filetosend:
        :param domain:
        :type domain: ModelingRealm
        :param var:
        :param startdate:
        :param member:
        :param chunk:
        :param grid:
        :param region:
        :param box:
        :param rename_var:
        :param frequency:
        :type frequency: Frequency
        :param year:
        :param date_str:
        :param move_old:
        :param vartype: Variable type (mean, statistic)
        :type vartype: VariableType
        :return:
        """
        self.data_manager.send_file(filetosend, domain, var, startdate, member, chunk, grid, region,
                                    box, rename_var, frequency, year, date_str, move_old, diagnostic=self,
                                    vartype=vartype)

    def compute(self):
        """
        Calculates the diagnostic and stores the output

        Must be implemented by derived classes
        """
        raise NotImplementedError("Class must override compute method")

    @classmethod
    def generate_jobs(cls, diags, options):
        """
        Generate the instances of the diagnostics that will be run by the manager

        Must be implemented by derived classes.

        :param diags: diagnostics manager
        :type diags: Diags
        :param options: list of strings containing the options passed to the diagnostic
        :type options: list[str]
        :return:
        """
        raise NotImplementedError("Class must override generate_jobs class method")

    @classmethod
    def process_options(cls, options, options_available):
        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))
        for x in range(len(options_available)):
            option_definition = options_available[x]
            if len(options) <= x:
                option_value = ''
            else:
                option_value = options[x]
            processed[option_definition.name] = option_definition.parse(option_value)
        return processed

    def __str__(self):
        """
        Must be implemented by derived classes
        :return:
        """
        return 'Developer must override base class __str__ method'


class DiagnosticOption(object):

    def __init__(self, name, default_value=None):
        self.name = name
        self.default_value = default_value

    def parse(self, option_value):
        option_value = self.check_default(option_value)
        return option_value

    def check_default(self, option_value):
        if option_value == '':
            if self.default_value is None:
                raise DiagnosticOptionError('Option {0} is not optional'.format(self.name))
            else:
                return self.default_value
        return option_value


class DiagnosticFloatOption(DiagnosticOption):
    def parse(self, option_value):
        return float(self.check_default(option_value))


class DiagnosticIntOption(DiagnosticOption):

    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

    def parse(self, option_value):
        value = int(self.check_default(option_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))
        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))
        return value


class DiagnosticListIntOption(DiagnosticOption):
    def parse(self, option_value):
        option_value = self.check_default(option_value)
        if isinstance(option_value, tuple) or isinstance(option_value, list):
            return option_value
        return [int(i) for i in option_value.split('-')]


class DiagnosticDomainOption(DiagnosticOption):
    def parse(self, option_value):
        return ModelingRealm.parse(self.check_default(option_value))


class DiagnosticFrequencyOption(DiagnosticOption):
    def parse(self, option_value):
        return Frequency.parse(self.check_default(option_value))


class DiagnosticBasinOption(DiagnosticOption):
    def parse(self, option_value):
        return Basins.parse(self.check_default(option_value))


class DiagnosticComplexStrOption(DiagnosticOption):
    def parse(self, option_value):
        return self.check_default(option_value).replace('&;', ',').replace('&.', ' ')


class DiagnosticBoolOption(DiagnosticOption):
    def parse(self, option_value):
        option_value = self.check_default(option_value)
        if isinstance(option_value, bool):
            return option_value
        else:
            return option_value.lower() in ('true', 't', 'yes')


class DiagnosticOptionError(Exception):
    pass