geometry.py

""" Sum factorized geometry evaluations """

from dune.codegen.generation import (backend,
                                     class_member,
                                     domain,
                                     geometry_mixin,
                                     get_counted_variable,
                                     iname,
                                     include_file,
                                     instruction,
                                     kernel_cached,
                                     preamble,
                                     silenced_warning,
                                     temporary_variable,
                                     get_global_context_value,
                                     globalarg,
                                     valuearg,
                                     )
from dune.codegen.options import get_option
from dune.codegen.pdelab.geometry import (local_dimension,
                                          world_dimension,
                                          name_cell_geometry,
                                          name_geometry,
                                          GenericPDELabGeometryMixin,
                                          AxiparallelGeometryMixin,
                                          EquidistantGeometryMixin,
                                          )
from dune.codegen.pdelab.localoperator import (name_ansatz_gfs_constructor_param,
                                               lop_template_ansatz_gfs,
                                               lop_template_range_field,
                                               )
from dune.codegen.pdelab.restriction import restricted_name
from dune.codegen.sumfact.accumulation import basis_sf_kernels
from dune.codegen.sumfact.basis import construct_basis_matrix_sequence
from dune.codegen.sumfact.quadrature import (additional_inames,
                                             default_quadrature_inames)
from dune.codegen.sumfact.switch import get_facedir, get_facemod
from dune.codegen.sumfact.symbolic import SumfactKernelInterfaceBase, SumfactKernel
from dune.codegen.tools import get_pymbolic_basename
from dune.codegen.options import get_form_option, option_switch
from dune.codegen.ufl.modified_terminals import Restriction

from pytools import ImmutableRecord

from loopy.match import Writes

import pymbolic.primitives as prim
import numpy as np


@geometry_mixin("sumfact_multilinear")
class SumfactMultiLinearGeometryMixin(GenericPDELabGeometryMixin):
    def pymbolic_unit_outer_normal(self):
        return pymbolic_unit_outer_normal(self)

    def pymbolic_jacobian_inverse(self, i, j, restriction):
        return pymbolic_jacobian_inverse(i, j, restriction, self.visitor)


@geometry_mixin("sumfact_axiparallel")
class SumfactAxiParallelGeometryMixin(AxiparallelGeometryMixin):
    def facet_normal(self, o):
        i, = self.indices
        self.indices = None
        assert isinstance(i, int)

        # Use facemod_s and facedir_s
        if i == get_facedir(Restriction.POSITIVE):
            if get_facemod(Restriction.POSITIVE):
                return 1
            else:
                return -1
        else:
            return 0


@geometry_mixin("sumfact_equidistant")
class SumfactEqudistantGeometryMixin(EquidistantGeometryMixin, SumfactAxiParallelGeometryMixin):
    def facet_jacobian_determinant(self, o):
        name = "fdetjac"
        self.define_facet_jacobian_determinant(name)
        facedir = get_facedir(Restriction.POSITIVE)
        globalarg(name, shape=(world_dimension(),))
        return prim.Subscript(prim.Variable(name), (facedir,))

    @class_member(classtag="operator")
    def define_facet_jacobian_determinant(self, name):
        self._define_facet_jacobian_determinant_eval(name)
        gfst = lop_template_ansatz_gfs()
        return "std::array<typename {}::Traits::GridView::template Codim<0>::Geometry::ctype, {}> {};".format(gfst, world_dimension(), name)

    @kernel_cached(kernel="operator")
    def _define_facet_jacobian_determinant_eval(self, name):
        gfs = name_ansatz_gfs_constructor_param()
        rft = lop_template_range_field()
        code = ["{",
                "  auto e = *({}.gridView().template begin<0>());".format(gfs),
                "  int dir=0;",
                "  for(auto is = {0}.gridView().ibegin(e); is != {0}.gridView().iend(e); ++(++is), ++dir)".format(gfs),
                "    {}[dir] = is->geometry().integrationElement(Dune::FieldVector<{}, {}>());".format(name, rft, world_dimension() - 1),
                "}",
                ]
        instruction(code="\n".join(code),
                    kernel="operator",
                    )

    def spatial_coordinate(self, o):
        index, = self.indices

        # Urgh: *SOMEHOW* construct a face direction. This is not breaking in the unstructured
        # case, because
        from dune.codegen.pdelab.restriction import Restriction
        restriction = Restriction.NONE
        if get_global_context_value("integral_type") == "interior_facet":
            restriction = Restriction.POSITIVE
        from dune.codegen.sumfact.switch import get_facedir
        face = get_facedir(restriction)

        lowcorner = name_lowerleft_corner()
        meshwidth = name_meshwidth()

        # If we have to decide which boundary condition to take for this
        # intersection we always take the boundary condition of the center
        # of the intersection. We assume that there are no cells with more
        # than one boundary condition.
        if self.do_predicates:
            x = 0.5
        elif index == face:
            x = 0
        else:
            iindex = index
            if face is not None and index > face:
                iindex = iindex - 1
            x = self.quadrature_position(iindex)

        self.indices = None
        return prim.Subscript(prim.Variable(lowcorner), (index,)) + x * prim.Subscript(prim.Variable(meshwidth), (index,))

@iname
def global_corner_iname(restriction):
    name = get_counted_variable(restricted_name("global_corneriname", restriction))
    domain(name, 2 ** world_dimension())
    return name


class GeoCornersInput(SumfactKernelInterfaceBase, ImmutableRecord):
    def __init__(self, direction, restriction):
        """Base class for sum-factorized evaluation of geometry mappings

        At the moment we only do this for cells and not faces. For
        intersections we do this corresponding reference elements of the
        neigboring cells.

        Each spatial component needs a seperate sum factorization kernel.  The
        argument 'direction' specifies the component (x-component: 0,
        y-component: 1, z-component: 2).
        """
        ImmutableRecord.__init__(self, direction=direction, restriction=restriction)

    def __repr__(self):
        return ImmutableRecord.__repr__(self)

    def __str__(self):
        return repr(self)

    @property
    def stage(self):
        return 1

    @property
    def direct_is_possible(self):
        return False

    def realize(self, sf, insn_dep, index=0):
        # Note: world_dimension, since we only do evaluation of cell geometry mappings
        from dune.codegen.sumfact.realization import name_buffer_storage
        name = "input_{}".format(sf.buffer)
        temporary_variable(name,
                           shape=(2 ** world_dimension(), sf.vector_width),
                           custom_base_storage=name_buffer_storage(sf.buffer, 0),
                           managed=True,
                           )
        ciname = global_corner_iname(self.restriction)

        if self.restriction == 0:
            geo = name_geometry()
        else:
            geo = name_cell_geometry(self.restriction)

        # NB: We need to realize this as a C instruction, because the corner
        #     method does return a non-scalar, which does not fit into the current
        #     loopy philosophy for function calls. This problem will be solved once
        #     #11 is resolved. Admittedly, the code looks *really* ugly until that happens.
        code = "{}[{}*{}+{}] = {}.corner({})[{}];".format(name,
                                                          sf.vector_width,
                                                          ciname,
                                                          index,
                                                          geo,
                                                          ciname,
                                                          self.direction,
                                                          )

        insn = instruction(code=code,
                           within_inames=frozenset({ciname}),
                           assignees=(name,),
                           tags=frozenset({"sumfact_stage{}".format(sf.stage)}),
                           )

        return insn_dep.union(frozenset({insn}))


@backend(interface="spatial_coordinate", name="default")
def pymbolic_spatial_coordinate_multilinear(do_predicates, visitor):
    """Calculate global coordinates of quadrature points for multilinear geometry mappings

    The evalualation is done using sum factorization techniques.

    On facets we use the geometry mapping of the self cell to get a global
    evaluation of the quadrature points. Avoiding the geometry mapping of the
    face itself greatly simplifies the generation of code for unstructured
    grids. Instead of looking at all the possible embeddings of the reference
    element of the face into the reference element of the cell we can make the
    grid edge consistent and choose a suitable convention for the order of
    directions in our sum factorization kernels.
    """
    assert len(visitor.indices) == 1

    # Adjust restriction for boundary facets
    restriction = visitor.restriction
    if get_global_context_value("integral_type") == "exterior_facet":
        restriction = 1

    # Generate sum factorization kernel and add vectorization info
    matrix_sequence = construct_basis_matrix_sequence(facedir=get_facedir(restriction),
                                                      facemod=get_facemod(restriction),
                                                      basis_size=(2,) * world_dimension())
    inp = GeoCornersInput(visitor.indices[0], restriction)
    sf = SumfactKernel(matrix_sequence=matrix_sequence,
                       interface=inp,
                       )

    from dune.codegen.sumfact.vectorization import attach_vectorization_info
    vsf = attach_vectorization_info(sf)

    # If this sum factorization kernel was not used in the dry run we
    # just return 0
    if vsf == 0:
        visitor.indices = None
        return 0

    # Add a sum factorization kernel that implements the evaluation of
    # the basis functions at quadrature points (stage 1)
    from dune.codegen.sumfact.realization import realize_sum_factorization_kernel
    var, _ = realize_sum_factorization_kernel(vsf)

    # The results of this function is already the right component of the
    # spatial coordinate and we don't need to index this in the visitor
    visitor.indices = None

    return prim.Subscript(var, vsf.quadrature_index(sf, visitor))


@preamble
def define_corner(name, low):
    geo = name_geometry()
    return "auto {} = {}.corner({});".format(name,
                                             geo,
                                             0 if low else 2 ** local_dimension() - 1)


@class_member(classtag="operator")
def define_mesh_width(name):
    from dune.codegen.pdelab.localoperator import lop_template_range_field
    rft = lop_template_range_field()
    define_mesh_width_eval(name)
    return "Dune::FieldVector<{}, {}> {};".format(rft, world_dimension(), name)


def define_mesh_width_eval(name):
    from dune.codegen.pdelab.localoperator import name_ansatz_gfs_constructor_param
    gfs = name_ansatz_gfs_constructor_param()
    code = ["{",
            "  auto e = *({}.gridView().template begin<0>());".format(gfs),
            "  {} = e.geometry().corner((1<<{}) - 1);".format(name, world_dimension()),
            "  {} -= e.geometry().corner(0);".format(name),
            "}",
            ]
    instruction(code="\n".join(code),
                kernel="operator")


def name_lowerleft_corner():
    name = "lowerleft_corner"
    globalarg(name, shape=(world_dimension(),))
    define_corner(name, True)
    return name


def name_meshwidth():
    name = "meshwidth"
    globalarg(name, shape=(world_dimension(),))
    define_mesh_width(name)
    return name


@backend(interface="spatial_coordinate", name="diagonal_transformation_matrix")
def pymbolic_spatial_coordinate_axiparallel(do_predicates, visitor):
    assert len(visitor.indices) == 1
    index, = visitor.indices

    # Urgh: *SOMEHOW* construct a face direction
    from dune.codegen.pdelab.restriction import Restriction
    restriction = Restriction.NONE
    if get_global_context_value("integral_type") == "interior_facet":
        restriction = Restriction.POSITIVE
    from dune.codegen.sumfact.switch import get_facedir
    face = get_facedir(restriction)

    lowcorner = name_lowerleft_corner()
    meshwidth = name_meshwidth()

    # If we have to decide which boundary condition to take for this
    # intersection we always take the boundary condition of the center
    # of the intersection. We assume that there are no cells with more
    # than one boundary condition.
    if do_predicates:
        x = 0.5
    elif index == face:
        x = 0
    else:
        iindex = index
        if face is not None and index > face:
            iindex = iindex - 1
        from dune.codegen.sumfact.quadrature import pymbolic_indexed_quadrature_position
        x = pymbolic_indexed_quadrature_position(iindex, visitor)

    visitor.indices = None
    return prim.Subscript(prim.Variable(lowcorner), (index,)) + x * prim.Subscript(prim.Variable(meshwidth), (index,))


def define_outer_normal(name, visitor):
    """Calculate outer normal (length not necessarily 1)

    This will be used to calculate the unit outer normal and the determinant of
    the Jacobian of the geometry mapping of the face.

    The outer normal is calculated by multiplying the Jacobian inverse
    transposed of the geometry mapping of the 'self'-cell with the unit outer
    normal \hat{n_s} of the face corresponding to the intersetcion in the
    reference element of the 'self'-cell:

    n = (\nabla \mu_{T_s})^{-T} \hat{n_s}

    Instead of setting up \hat{n_s} we just look at facedir_s and facemod_s.
    """
    facedir_s = get_facedir(Restriction.POSITIVE)
    facemod_s = get_facemod(Restriction.POSITIVE)
    temporary_variable(name, shape=(world_dimension(),))
    for i in range(world_dimension()):
        assignee = prim.Subscript(prim.Variable(name), i)
        ji = pymbolic_jacobian_inverse(facedir_s, i, Restriction.POSITIVE, visitor)

        # Note: 2*facemod_s-1 because of
        # -1 if facemod_s = 0
        # +1 if facemod_s = 1
        expression = ji * (2 * facemod_s - 1)

        inames = default_quadrature_inames(visitor) + additional_inames(visitor)
        instruction(assignee=assignee,
                    expression=expression,
                    within_inames=frozenset(inames),
                    )


def name_outer_normal(visitor):
    name = "outer_normal"
    define_outer_normal(name, visitor)
    return name


def pymbolic_outer_normal(visitor):
    name = name_outer_normal(visitor)
    return prim.Variable(name)


def define_norm_of_outer_normal(name, visitor):
    outer_normal = pymbolic_outer_normal(visitor)

    # Calculate the norm of outer_normal
    temporary_variable(name, shape=())
    expression = 0
    for i in range(world_dimension()):
        expression = prim.Sum((expression, prim.Subscript(outer_normal, i) * prim.Subscript(outer_normal, i)))
    expression = prim.Call(prim.Variable("sqrt"), (expression,))
    inames = default_quadrature_inames(visitor) + additional_inames(visitor)

    # Apparently the call of sqrt shadows the dependency on outer normal. Add manually.
    instruction(assignee=prim.Variable(name),
                expression=expression,
                depends_on=frozenset({Writes(get_pymbolic_basename(outer_normal))}),
                within_inames=frozenset(inames),
                )


def name_norm_of_outer_normal(visitor):
    name = "norm_of_outer_normal"
    define_norm_of_outer_normal(name, visitor)
    return name


def pymbolic_norm_of_outer_normal(visitor):
    name = name_norm_of_outer_normal(visitor)
    return prim.Variable(name)


def pymbolic_unit_outer_normal(visitor):
    assert (len(visitor.indices) is 1)
    index = visitor.indices[0]
    assert isinstance(index, int)
    if get_form_option("diagonal_transformation_matrix"):
        # In this case we just return -1, 0 or 1 depending on the current
        # facemod and facedir. We don't need to (and cannot) index these ints
        # so we set the indices of the visitor to None.
        visitor.indices = None

        # Use facemod_s and facedir_s
        from dune.codegen.sumfact.switch import get_facedir, get_facemod
        if index == get_facedir(Restriction.POSITIVE):
            if get_facemod(Restriction.POSITIVE):
                return 1
            else:
                return -1
        else:
            return 0
    else:
        # In pymbolic_outer_normal we need to create the
        # jacobian_inverse_transposed which doesn't play nicely with the
        # visitor having indices. Maybe this should be improved. For now let's
        # just set them to None and restore them afterwards.
        visitor.indices = None
        outer_normal = pymbolic_outer_normal(visitor)
        norm = pymbolic_norm_of_outer_normal(visitor)
        visitor.indices = (index,)

        # Create unit_outer_normal by scaling outer_normal to length 1
        name = "unit_outer_normal"
        temporary_variable(name, shape=(world_dimension(),))
        inames = default_quadrature_inames(visitor) + additional_inames(visitor)
        for i in range(world_dimension()):
            instruction(assignee=prim.Subscript(prim.Variable(name), i),
                        expression=prim.Subscript(outer_normal, i) / norm,
                        within_inames=frozenset(inames),
                        )

        return prim.Variable(name)


def pymbolic_unit_inner_normal(visitor_indices):
    # Note: The implementation of this was adjusted for unstructured grids but
    # it was not properly tested. If you need the unit inner normal just remove
    # the assert(False) statement but make sure to verify that this does the
    # right thing.
    assert(False)

    assert (len(visitor.indices) is 1)
    index = visitor.indices[0]
    assert isinstance(index, int)
    if get_form_option("diagonal_transformation_matrix"):
        # In this case we just return -1, 0 or 1 depending on the current
        # facemod and facedir. We don't need to (and cannot) index these ints
        # so we set the indices of the visitor to None.
        visitor.indices = None
        if index == get_facedir(Restriction.POSITIVE):
            if get_facemod(Restriction.POSITIVE):
                return -1
            else:
                return 1
        else:
            return 0
    else:
        # In pymbolic_outer_normal we need to create the
        # jacobian_inverse_transposed which doesn't play nicely with the
        # visitor having indices. Maybe this should be improved. For now let's
        # just set them to None and restore them afterwards.
        visitor.indices = None
        outer_normal = pymbolic_outer_normal(visitor)
        norm = pymbolic_norm_of_outer_normal(visitor)
        visitor.indices = (index,)

        # Create unit_outer_normal by scaling outer_normal to length 1
        name = "unit_outer_normal"
        temporary_variable(name, shape=(world_dimension(),))

        inames = default_quadrature_inames(visitor) + additional_inames(visitor)
        for i in range(world_dimension()):
            instruction(assignee=prim.Subscript(prim.Variable(name), i),
                        expression=-1 * prim.Subscript(outer_normal, i) / norm,
                        within_inames=frozenset(inames),
                        )

        return prim.Variable(name)


def define_facet_jacobian_determinant(name, visitor):
    """Absolute value of determinant of jacobian of facet geometry mapping

    Calculate the absolute value of the determinant of the jacobian of the
    geometry mapping from the reference cell of the intersection to the
    intersection:

    |\det(\nabla \mu_F)|

    This is done using the surface metric tensor lemma from the lecture notes
    of Jean-Luc Guermond:

    |\det(\nabla \mu_F)| = \| \tilde{n} \| |\det(\nabla \mu_{T_s})|

    Here \tilde{n} is the outer normal defined by:

    \tilde{n} = (\nabla \mu_{T_s})^{-T} \hat{n}_s

    where \hat{n}_s is the unit outer normal of the corresponding face of the
    reference cell.
    """
    # Calculate norm of outer normal and determinant of jacobian of geometry
    # mapping of the inside cell -> set restriction to 1 for boundary facets
    if get_global_context_value("integral_type") == "exterior_facet":
        assert visitor.restriction is 0
        visitor.restriction = 1
    norm_of_outer_normal = pymbolic_norm_of_outer_normal(visitor)
    detjac = pymbolic_jacobian_determinant(visitor)
    if get_global_context_value("integral_type") == "exterior_facet":
        visitor.restriction = 0

    # Multiply detjac and norm_of_outer_normal
    temporary_variable(name, shape=())
    inames = default_quadrature_inames(visitor) + additional_inames(visitor)
    instruction(assignee=prim.Variable(name),
                expression=detjac * norm_of_outer_normal,
                within_inames=frozenset(inames),
                )


def name_facet_jacobian_determinant(visitor):
    name = "fdetjac"
    define_facet_jacobian_determinant(name, visitor)
    return name


def pymbolic_facet_jacobian_determinant(visitor):
    name = name_facet_jacobian_determinant(visitor)
    return prim.Variable(name)


def pymbolic_facet_area(visitor):
    from dune.codegen.pdelab.geometry import pymbolic_facet_area as nonconstant_facet_area
    return nonconstant_facet_area()


def _name_jacobian(i, j, restriction, visitor):
    """Return the (i, j) component of the jacobian of the geometry mapping

    Evaluation of the derivative of the geometry mapping is done using sum
    factorization.

    Note: At the moment this only works for the mappings from reference cells
    to the cell and not for the geometry mappings of intersections.
    """
    do_predicates = visitor.do_predicates

    # Create matrix sequence with derivative in j direction
    matrix_sequence = construct_basis_matrix_sequence(derivative=j,
                                                      facedir=get_facedir(restriction),
                                                      facemod=get_facemod(restriction),
                                                      basis_size=(2,) * world_dimension())

    # Sum factorization input for the i'th component of the geometry mapping
    inp = GeoCornersInput(i, restriction)

    sf = SumfactKernel(matrix_sequence=matrix_sequence,
                       interface=inp,
                       )

    from dune.codegen.sumfact.vectorization import attach_vectorization_info
    vsf = attach_vectorization_info(sf)

    # If this sum factorization kernel was not used in the dry run we
    # just return 0
    if vsf == 0:
        visitor.indices = None
        return 0

    # Add a sum factorization kernel that implements the evaluation of
    # the basis functions at quadrature points (stage 1)
    from dune.codegen.sumfact.realization import realize_sum_factorization_kernel
    var, _ = realize_sum_factorization_kernel(vsf)

    assert(visitor.indices is None)
    return prim.Subscript(var, vsf.quadrature_index(sf, visitor))


def define_jacobian_inverse(name, restriction, visitor):
    """Return jacobian inverse of the geometry mapping (of a cell)

    At the moment this only works for geometry mappings of cells and not for
    intersection. We only consider this case as it greatly simplifies code
    generation for unstructured grids by making the grid edge consistent and
    avoiding the face-twist problem.
    """
    # Calculate the jacobian of the geometry mapping using sum factorization
    dim = world_dimension()
    names_jacobian = []
    for j in range(dim):
        for i in range(dim):
            name_jacobian = restricted_name("jacobian_geometry_{}{}".format(i, j), restriction)
            temporary_variable(name_jacobian)
            assignee = prim.Variable(name_jacobian)
            expression = _name_jacobian(i, j, restriction, visitor)
            inames = default_quadrature_inames(visitor) + additional_inames(visitor)
            instruction(assignee=assignee,
                        expression=expression,
                        within_inames=frozenset(inames),
                        )
            names_jacobian.append(name_jacobian)

            # The sum factorization kernels from the geometry evaluation of the
            # jacobians will never appear in the expression for the input of
            # stage 3. This way the SumfactCollectMapper will never see them
            # and they will be marked as inactive. Here we explicitly mark the
            # as used.
            basis_sf_kernels(expression.aggregate)

    # Calculate the inverse of the jacobian of the geometry mapping and the
    # determinant by calling a c++ function. Note: The result will be column
    # major -> fortran style.
    name_detjac = name_jacobian_determinant(visitor)
    temporary_variable(name_detjac, shape=())
    ftags = ",".join(["f"] * 2)
    temporary_variable(name, shape=(dim, dim), dim_tags=ftags, managed=True)
    include_file('dune/codegen/sumfact/invertgeometry.hh', filetag='operatorfile')
    code = "{} = invert_and_return_determinant({}, {});".format(name_detjac,
                                                                ", ".join(names_jacobian),
                                                                name)
    silenced_warning("read_no_write({})".format(name_detjac))
    inames = default_quadrature_inames(visitor) + additional_inames(visitor)
    return instruction(code=code,
                       assignees=frozenset([name, name_detjac]),
                       read_variables=frozenset(names_jacobian),
                       inames=frozenset(inames),
                       )


def name_jacobian_inverse(restriction, visitor):
    name = restricted_name("jacobian_inverse", restriction)
    define_jacobian_inverse(name, restriction, visitor)
    return name


def pymbolic_jacobian_inverse(i, j, restriction, visitor):
    name = name_jacobian_inverse(restriction, visitor)
    return prim.Subscript(prim.Variable(name), (i, j))


def name_jacobian_determinant(visitor):
    name = "detjac"
    return name


def pymbolic_jacobian_determinant(visitor):
    # The calculation of the jacobian determinant happens in this function
    if visitor.measure == 'cell':
        restriction = 0
    else:
        restriction = 1
    name_jacobian_inverse(restriction, visitor)

    return prim.Variable(name_jacobian_determinant(visitor))