localoperator.py

from __future__ import absolute_import
from os.path import splitext

import logging

import numpy as np

from dune.codegen.options import (get_form_option,
                                  get_option,
                                  form_option_context,
                                  )
from dune.codegen.generation import (accumulation_mixin,
                                     base_class,
                                     class_basename,
                                     class_member,
                                     construct_from_mixins,
                                     constructor_parameter,
                                     domain,
                                     dump_accumulate_timer,
                                     register_liwkid_timer,
                                     end_of_file,
                                     function_mangler,
                                     generator_factory,
                                     get_global_context_value,
                                     global_context,
                                     iname,
                                     include_file,
                                     initializer_list,
                                     kernel_cached,
                                     post_include,
                                     retrieve_cache_functions,
                                     retrieve_cache_items,
                                     ReturnArg,
                                     run_hook,
                                     template_parameter,
                                     dump_ssc_marks
                                     )
from dune.codegen.cgen.clazz import (AccessModifier,
                                     BaseClass,
                                     ClassMember,
                                     )
from dune.codegen.loopy.target import type_floatingpoint
from dune.codegen.ufl.modified_terminals import Restriction

import dune.codegen.loopy.mangler

from pymbolic.primitives import Variable
import pymbolic.primitives as prim
from pytools import Record, ImmutableRecord

import ufl.classes as uc
import loopy as lp
import cgen


@template_parameter(classtag="operator")
def lop_template_ansatz_gfs():
    name = "GFSU"
    constructor_parameter("const {}&".format(name), name_ansatz_gfs_constructor_param(), classtag="operator")
    return name


def name_ansatz_gfs_constructor_param():
    return "gfsu"


@template_parameter(classtag="operator")
def lop_template_test_gfs():
    name = "GFSV"
    constructor_parameter("const {}&".format(name), name_test_gfs_constructor_param(), classtag="operator")
    return name


def name_test_gfs_constructor_param():
    return "gfsv"


@class_member(classtag="operator")
def lop_domain_field(name):
    # TODO: Rethink for not Galerkin Method
    gfs = lop_template_ansatz_gfs()
    return "using {} = typename {}::Traits::GridView::ctype;".format(name, gfs)


def name_domain_field():
    name = "DF"
    lop_domain_field(name)
    return name


def lop_template_gfs(ma):
    from ufl.classes import Argument, Coefficient
    if isinstance(ma.argexpr, Argument):
        if ma.argexpr.number() == 0:
            return lop_template_test_gfs()
        if ma.argexpr.number() == 1:
            return lop_template_ansatz_gfs()
    if isinstance(ma.argexpr, Coefficient):
        # Index 0 is reserved for trialfunction, index 1 is reserved for jacobian apply function
        assert ma.argexpr.count() < 2
        return lop_template_ansatz_gfs()
    assert False


def name_initree_constructor_param():
    return "iniParams"


@class_member(classtag="operator")
def define_initree(name):
    param_name = name_initree_constructor_param()
    include_file('dune/common/parametertree.hh', filetag="operatorfile")
    constructor_parameter("const Dune::ParameterTree&", param_name, classtag="operator")
    initializer_list(name, [param_name], classtag="operator")

    return "const Dune::ParameterTree& {};".format(name)


@class_member(classtag="operator")
def _enum_pattern(which):
    return "enum {{ doPattern{} = true }};".format(which)


def enum_pattern():
    from dune.codegen.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    from dune.codegen.pdelab.signatures import ufl_measure_to_pdelab_measure
    return _enum_pattern(ufl_measure_to_pdelab_measure(integral_type))


def _pattern_baseclass(measure):
    return base_class('Dune::PDELab::Full{}Pattern'.format(measure), classtag="operator")


def pattern_baseclass():
    from dune.codegen.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    from dune.codegen.pdelab.signatures import ufl_measure_to_pdelab_measure
    return _pattern_baseclass(ufl_measure_to_pdelab_measure(integral_type))


@class_member(classtag="operator")
def _enum_alpha(which):
    return "enum {{ doAlpha{} = true }};".format(which)


def enum_alpha():
    from dune.codegen.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    from dune.codegen.pdelab.signatures import ufl_measure_to_pdelab_measure
    return _enum_alpha(ufl_measure_to_pdelab_measure(integral_type))


@class_member(classtag="operator")
def enum_skeleton_twosided():
    return "enum { doSkeletonTwoSided = true };"


def name_initree_member():
    define_initree("_iniParams")
    return "_iniParams"


@class_basename(classtag="operator")
def localoperator_basename(form_ident):
    return get_form_option("classname", form_ident)


def name_gridfunction_member(coeff, restriction, diffOrder=0):
    # We reuse the grid function for volume integrals in skeleton integrals
    if restriction == Restriction.POSITIVE:
        restriction = Restriction.NONE
    restr = "_n" if restriction == Restriction.NEGATIVE else ""
    name = "local_gridfunction_coeff{}_diff{}{}".format(coeff.count(), diffOrder, restr)
    define_gridfunction_member(name, coeff, restriction, diffOrder)
    return name


def name_gridfunction_constructor_argument(coeff):
    _type = type_gridfunction_template_parameter(coeff)
    name = "gridfunction_coeff{}_".format(coeff.count())
    constructor_parameter("const {}&".format(_type), name, classtag="operator")
    return name


@class_member(classtag="operator")
def define_gridfunction_member(name, coeff, restriction, diffOrder):
    _type = type_gridfunction_template_parameter(coeff)
    param = name_gridfunction_constructor_argument(coeff)
    if diffOrder > 0:
        other = name_gridfunction_member(coeff, restriction, diffOrder - 1)
        init = "derivative({})".format(other)
        initializer_list(name, [init], classtag="operator")
        return "mutable decltype({}) {};".format(init, name)
    else:
        init = "localFunction({})".format(param)
        initializer_list(name, [init], classtag="operator")
        return "mutable typename {}::LocalFunction {};".format(_type, name)