options.py

""" Manage the command line options to the form compiler executable """

from argparse import ArgumentParser
from os.path import abspath
from pytools import ImmutableRecord, memoize
import cerberus
import yaml
import pkg_resources
from six.moves import configparser
from six import StringIO
from contextlib import contextmanager


class CodegenOptionsValidator(cerberus.Validator):
    # A validator that accepts the helpstr field in the scheme
    def _validate_helpstr(self, helpstr, field, value):
        """ Describe the option

        The rule's arguments are validated against this schema:
        {'type': 'string'}
        """
        return True


def _load_scheme(form=False):
    resource_package = __name__
    if form:
        resource_path = 'options_form.yaml'
    else:
        resource_path = 'options_global.yaml'
    yaml_stream = pkg_resources.resource_string(resource_package, resource_path)
    try:
        scheme = yaml.safe_load(yaml_stream)
    except Exception as e:
        raise e
    return scheme


class CodegenGlobalOptionsArray(ImmutableRecord):
    """ A collection of form compiler arguments """
    def __init__(self, **kwargs):
        # Set the default values from the yaml scheme as defaults
        scheme = _load_scheme()
        opts = {k: v['default'] for k, v in scheme.items()}
        opts.update(**kwargs)
        ImmutableRecord.__init__(self, **opts)


class CodegenFormOptionsArray(ImmutableRecord):
    """ A collection of form-specific form compiler arguments """
    def __init__(self, **kwargs):
        # Set the default values from the yaml scheme as defaults
        scheme = _load_scheme(form=True)
        opts = {k: v['default'] for k, v in scheme.items()}
        opts.update(**kwargs)
        ImmutableRecord.__init__(self, **opts)


# Until more sophisticated logic is needed, we keep the actual option data in this module
_global_options = CodegenGlobalOptionsArray()
_form_options = {}


def show_options():
    # TODO: This needs to be adjusted to options-validation
    def subopt(arr):
        for k, v in arr.__dict__.items():
            if isinstance(v, CodegenOption) and v.helpstr is not None:
                print("{}\n    {}".format(k, v.helpstr))

    print("This is a summary of options available for the code generation process:\n")
    print("The following options can be given in the [formcompiler] section:")
    subopt(CodegenGlobalOptionsArray)

    print("\nThefollowing options can be given in a form-specific subsection of [formcompiler]:")
    subopt(CodegenFormOptionsArray)


def initialize_options():
    """ Initialize the options from the command line """
    global _global_options
    _global_options = update_options_from_commandline(_global_options)
    _global_options = update_options_from_inifile(_global_options)

    # Validate global options
    scheme_global = _load_scheme()
    validator_global = CodegenOptionsValidator(scheme_global, require_all=True)
    if not validator_global.validate(_global_options.__dict__):
        raise RuntimeError("Global options validation failed: {}".format(validator_global.errors))

    # Validate form options
    scheme_form = _load_scheme(form=True)
    validator_form = CodegenOptionsValidator(scheme_form, require_all=True)
    for form in [i.strip() for i in _global_options.operators.split(",")]:
        if not validator_form.validate(_form_options[form].__dict__):
            raise RuntimeError("Form options validation failed: {}".format(validator_form.errors))


def _scheme_type_to_type(scheme_type):
    assert isinstance(scheme_type, str)
    if scheme_type == 'string':
        return str
    if scheme_type == 'boolean':
        return bool
    if scheme_type == 'integer':
        return int
    if scheme_type == 'float':
        return float


def _transform_type(scheme_type, a):
    if scheme_type == 'boolean':
        return bool(int(a))
    else:
        return _scheme_type_to_type(scheme_type)(a)


def update_options_from_commandline(opt):
    """ Return an options array object with updated values from the commandline """
    assert isinstance(opt, CodegenGlobalOptionsArray)
    parser = ArgumentParser(description="Compile UFL files to PDELab C++ code",
                            epilog="Please report bugs to dominic.kempf@iwr.uni-heidelberg.de",
                            )
    parser.add_argument('--version', action='version', version='%(prog)s 0.1')

    # Load global options scheme
    scheme = _load_scheme()

    # Add all options that have a helpstr to the command line parser
    for k, v in scheme.items():
        if v['helpstr'] is not None:
            cmdopt = "--{}".format(k.replace('_', '-'))
            parser.add_argument(cmdopt, help=v['helpstr'], type=_scheme_type_to_type(v['type']))
    parsedargs = {k: v for k, v in vars(parser.parse_args()).items() if v is not None}
    return opt.copy(**parsedargs)


def update_options_from_inifile(opt):
    """ Return an options array object with updated values from an inifile """
    if opt.ini_file:
        config = configparser.ConfigParser()

        # Read ini file
        try:
            config.read(opt.ini_file)
        except configparser.MissingSectionHeaderError:
            # Config parser doesn't like ini files where without section. For
            # this case we introduce a [root] section on top.
            ini_str = '[root]\n' + open(opt.ini_file, 'r').read()
            ini_fp = StringIO(ini_str)
            config = configparser.RawConfigParser()
            config.readfp(ini_fp)

        # Parse global options
        scheme = _load_scheme()
        options = dict(config.items('formcompiler'))
        for k, v in options.items():
            assert k in scheme
            options[k] = _transform_type(scheme[k]['type'], v)
        opt = opt.copy(**options)

        # Parse form options
        scheme = _load_scheme(form=True)
        for form in [i.strip() for i in opt.operators.split(",")]:
            section = 'formcompiler.{}'.format(form)
            options = {}
            if config.has_section(section):
                options = dict(config.items('formcompiler.{}'.format(form)))
                for k, v in options.items():
                    assert k in scheme
                    options[k] = _transform_type(scheme[k]['type'], v)
            _form_options[form] = CodegenFormOptionsArray(**options)

    return opt


@memoize
def process_global_options(opt):
    """ Make sure that the options have been fully processed """
    opt = expand_architecture_options(opt)

    if opt.overlapping:
        opt = opt.copy(parallel=True)

    return opt


@memoize
def process_form_options(opt, form):
    if opt.sumfact:
        opt = opt.copy(unroll_dimension_loops=True,
                       quadrature_mixins="sumfact",
                       basis_mixins="sumfact",
                       accumulation_mixins="sumfact",
                       )

    if opt.blockstructured:
        opt = opt.copy(accumulation_mixins="blockstructured",
                       quadrature_mixins="blockstructured",
                       basis_mixins="blockstructured"
                       )

    if opt.control:
        opt = opt.copy(accumulation_mixins="control")

    if opt.numerical_jacobian:
        opt = opt.copy(generate_jacobians=False, generate_jacobian_apply=False)

    if opt.form is None:
        opt = opt.copy(form=form)

    if opt.classname is None:
        opt = opt.copy(classname="{}Operator".format(form))

    if opt.filename is None:
        opt = opt.copy(filename="{}_{}_file.hh".format(get_option("target_name"), opt.classname))

    if opt.block_preconditioner_pointdiagonal:
        opt = opt.copy(generate_jacobians=False,
                       basis_mixins="sumfact_pointdiagonal",
                       accumulation_mixins="sumfact_pointdiagonal",
                       )

    if opt.block_preconditioner_diagonal or opt.block_preconditioner_offdiagonal:
        assert opt.numerical_jacobian is False
        opt = opt.copy(generate_residuals=False,
                       generate_jacobians=True,
                       matrix_free=True,
                       )

    if opt.matrix_free:
        opt = opt.copy(generate_jacobian_apply=True)

    return opt


def expand_architecture_options(opt):
    if opt.architecture == "haswell":
        return opt.copy(max_vector_width=256)
    elif opt.architecture == "knl":
        return opt.copy(max_vector_width=512)
    elif opt.architecture == "skylake":
        return opt.copy(max_vector_width=512)
    else:
        raise NotImplementedError("Architecture {} not known!".format(opt.architecture))


def set_option(key, value):
    """Add the key value pair to the options.

    If the key is already in the options dictionary its value will be
    overwritten.  Form compiler arguments will always be set before
    any other options.
    """
    global _global_options
    _global_options = process_global_options(_global_options).copy(**{key: value})


def set_form_option(key, value, form=None):
    if form is None:
        from dune.codegen.generation import get_global_context_value
        form = get_global_context_value("form_identifier", 0)
    if isinstance(form, int):
        form = get_option("operators").split(",")[form].strip()
    _form_options[form] = _form_options[form].copy(**{key: value})


def get_option(key):
    processed_global_opts = process_global_options(_global_options)
    return getattr(processed_global_opts, key)


def get_form_option(key, form=None):
    if form is None:
        from dune.codegen.generation import get_global_context_value
        form = get_global_context_value("form_identifier", 0)
    if isinstance(form, int):
        form = get_option("operators").split(",")[form].strip()
    processed_form_opts = process_form_options(_form_options[form], form)
    return getattr(processed_form_opts, key)


@contextmanager
def option_context(conditional=True, **opts):
    """ A context manager that sets a given option and restores it on exit. """
    # Backup old values and set to new ones
    if conditional:
        backup = {}
        for k, v in opts.items():
            backup[k] = get_option(k)
            set_option(k, v)

    yield

    if conditional:
        # Restore old values
        for k in opts.keys():
            set_option(k, backup[k])


@contextmanager
def form_option_context(conditional=True, **opts):
    """ A context manager that sets a given form option and restores it on exit """
    if conditional:
        form = opts.pop("form", None)

        # Backup old values and set to new ones
        backup = {}
        for k, v in opts.items():
            backup[k] = get_form_option(k, form=form)
            set_form_option(k, v, form=form)

    yield

    # Restore old values
    if conditional:
        for k in opts.keys():
            set_form_option(k, backup[k], form=form)