main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import copy
import os
import re
import threading
import time

import iris
from iris.experimental.equalise_cubes import equalise_attributes
import netCDF4
import numpy as np
import pkg_resources
import sentry_sdk
import six
import yaml

import climix
from .dask_setup import SCHEDULERS, setup_scheduler
from .index import Index
from .index_functions import SUPPORTED_OPERATORS, SUPPORTED_REDUCERS
from .period import PeriodSpecification

import logging
logging.basicConfig(level=logging.INFO)


sentry_sdk.init("https://d3ac73a62877407b848dfc3f318bed85@sentry.io/1458386")


MISSVAL = 1.0e20


def parse_args():
    parser = argparse.ArgumentParser(
        description=(f'A climate index thing, version {climix.__version__}.'))
    parser.add_argument('-d', '--dask-scheduler', choices=SCHEDULERS.keys(),
                        default='distributed-local-cluster')
    parser.add_argument('-k', '--keep-open', action='store_true',
                        help='keep climix running until key press '
                        '(useful for debugging)')
    parser.add_argument('-s', '--sliced-mode', action='store_true',
                        help='activate calculation per period to avoid memory '
                        'problems')
    parser.add_argument('-o', '--output', dest='output_template',
                        help='output filename')
    parser.add_argument('-x', '--index', action='append',
                        required=True, metavar='INDEX', dest='indices',
                        help='the index to calculcate')
    parser.add_argument('datafiles', nargs='+', metavar="DATAFILE",
                        help='the input data files')
    return parser.parse_args()


def ignore_cb(cube, field, filename):
    cube.attributes.pop('creation_date', None)
    cube.attributes.pop('tracking_id', None)


def load_metadata():
    metadata_filename = os.path.join(os.path.dirname(__file__),
                                     'etc', 'metadata.yml')
    with open(metadata_filename) as md_file:
        metadata = yaml.safe_load(md_file)
    return metadata


def build_parameters(parameters_metadata):
    parameters = {}
    for name, md in parameters_metadata.items():
        kind = md.pop('kind')
        if kind == 'quantity':
            parameter = iris.cube.Cube(**md)
        elif kind == 'operator':
            op = md['operator']
            if op not in SUPPORTED_OPERATORS:
                raise ValueError(f'Unknown operator <{op}>')
            parameter = op
        elif kind == 'reducer':
            red = md['reducer']
            if red not in SUPPORTED_REDUCERS:
                raise ValueError(f'Unknown reducer <{red}>')
            parameter = red
        else:
            raise ValueError(f'Unknown parameter kind <{kind}>')
        parameters[name] = parameter
    return parameters


def build_index_function(spec):
    name = spec['name']
    candidates = list(pkg_resources.iter_entry_points('climix.index_functions',
                                                      name=name))
    if len(candidates) == 0:
        raise ValueError(f'No implementation found for index_function <{name}>')
    elif len(candidates) > 1:
        distributions = [candidate.dist for candidate in candidates]
        raise ValueError(
            f'Found several implementations for index_function <{name}>. '
            f'Please make sure only one is installed at any time. '
            f'The implementations come from the distributions {distributions}')
    candidate = candidates[0]
    logging.info(f'Found implementation for index_function <{name}> '
                 f'from distribution <{candidate.dist}>')
    index_function_factory = candidates[0].load()
    parameters = build_parameters(spec['parameters'])
    index_function = index_function_factory(**parameters)
    return index_function


def build_template_index(index_definitions):
    expr = re.compile(r'{([^}]+)}')
    template_index = {}
    for index in index_definitions.keys():
        split = expr.split(index)
        if len(split) == 1:
            continue
        signature = tuple(split[::2])
        parameter_names = split[1::2]
        template_index[signature] = (index, parameter_names)
    return template_index


def replace_parameters_in_dict(dictionary, parameter_dict):
    for key, value in dictionary.items():
        if isinstance(value, dict):
            k, v = list(value.items())[0]
            if len(value) == 1 and v is None and k in parameter_dict.keys():
                dictionary[key] = parameter_dict[k]
            else:
                replace_parameters_in_dict(value, parameter_dict)
        elif isinstance(value, six.string_types):
            dictionary[key] = dictionary[key].format(**parameter_dict)


def get_index_definition(index_definitions, index):
    try:
        return index_definitions[index]
    except KeyError:
        index_expr = re.compile(r'(\d+)')
        split = index_expr.split(index)
        if len(split) == 1:
            raise
        templates = build_template_index(index_definitions)
        signature = tuple(split[::2])
        template, parameter_names = templates[signature]
        parameter_values = split[1::2]
        parameter_dict = {name: int(value)
                          for (name, value) in
                          zip(parameter_names, parameter_values)}
        index_definition = copy.deepcopy(index_definitions[template])
        replace_parameters_in_dict(index_definition, parameter_dict)
        return index_definition


def prepare_indices(index_definitions, requested_indices):
    def select_period(metadata):
        selected_period = metadata['default']
        period_metadata = metadata['allowed'][selected_period]
        return PeriodSpecification(selected_period, period_metadata)
    indices = []
    for index in requested_indices:
        definition = get_index_definition(index_definitions, index)
        period_spec = select_period(definition['period'])
        index_function = build_index_function(definition['index_function'])
        index = Index(index_function, definition['output'], period_spec)
        indices.append(index)
    return indices


def prepare_input_data(datafiles):
    datacube = iris.load_raw(datafiles, callback=ignore_cb)
    iris.util.unify_time_units(datacube)
    equalise_attributes(datacube)
    cube = datacube.concatenate_cube()
    return cube


def guess_output_template(datafiles):
    output_template = '{var_name}_{frequency}.nc'
    def filename_stripper(path):
        # remove directory part...
        basename = os.path.basename(path)
        # ...and extension
        root, ext = os.path.splitext(basename)
        # split at _...
        parts = root.split('_')
        # and remove
        # first part (usually variable)
        # and last part (usually time)
        base = '_'.join(parts[1:-1])
        try:
            time = [int(t) for t in parts[-1].split('-')]
            if len(time) == 1:
                time *= 2
        except ValueError:
            time = [None, None]
        return (base, time[0], time[1])
    files = [filename_stripper(p) for p in datafiles]
    bases, starts, ends = zip(*files)
    unique_bases = set(bases)
    if len(unique_bases) == 1:
        base = unique_bases.pop()
        start = min(starts)
        end = max(ends)
        output_template = f'{{var_name}}_{base}_{{frequency}}_{start}-{end}.nc'
    return output_template


def build_output_filename(index, datafiles, output_template):
    if output_template is None:
        output_template = guess_output_template(datafiles)
    frequency = index.period.label
    return output_template.format(frequency=frequency, **index.output_metadata)


def save(result, output_filename, sliced_mode=False):
    if sliced_mode:
        logging.info('Performing aggregation in sliced mode')
        data = result.core_data()
        logging.info('creating empty data')
        result.data = np.empty(data.shape, data.dtype)
        result.data
        logging.info('saving empty cube')
        iris.save(result, output_filename, fill_value=MISSVAL,
                  local_keys=['proposed_standard_name'])
        logging.info('opening')
        result.data = data
        with netCDF4.Dataset(output_filename, 'a') as ds:
            var = ds[result.var_name]
            time_dim = result.coord_dims('time')[0]
            no_slices = result.shape[time_dim]
            def store(i, data):
                var[i, ...] = data
            thread = threading.Thread()
            thread.start()
            start = time.time()
            for i, result_cube in enumerate(result.slices_over(time_dim)):
                logging.info('Starting with {}'.format(result_cube.coord('time')))
                result_cube.data
                logging.info('Waiting for previous save to finish')
                thread.join()
                thread = threading.Thread(target=store,
                                          args=(i, result_cube.data))
                thread.start()
                end = time.time()
                partial = end - start
                total_estimate = partial/(i+1)*no_slices
                logging.info(f'Finished {i+1}/{no_slices} in {partial:4f}s. '
                             f'Estimated total time is {total_estimate:4f}s.')
            thread.join()
    else:
        logging.info('Performing aggregation in normal mode')
        iris.save(result, output_filename, fill_value=MISSVAL,
                  local_keys=['proposed_standard_name'])


def do_main(requested_indices, datafiles, output_template, sliced_mode):
    logging.info('Loading metadata')
    metadata = load_metadata()
    logging.info('Preparing indices')
    indices = prepare_indices(metadata['indices'], requested_indices)
    for index in indices:
        logging.info(f'Starting calculations for index {index}')
        logging.info('Building output filename')
        output_filename = build_output_filename(index, datafiles, output_template)
        logging.info('Preparing input data')
        input_data = prepare_input_data(datafiles)
        logging.info('Calculating index')
        result = index(input_data)
        logging.info('Saving result')
        save(result, output_filename, sliced_mode=sliced_mode)


def main():
    args = parse_args()
    with setup_scheduler(args) as scheduler:
        logging.info('Scheduler ready; starting main program.')
        start = time.time()
        try:
            do_main(args.indices, args.datafiles,
                    args.output_template, args.sliced_mode)
        finally:
            end = time.time()
            logging.info(f'Calculation took {end-start:.4f} seconds.')
        if args.keep_open:
            input('Press enter to close the cluster ')


if __name__ == "__main__":
    main()