quadraturegenerator.py

finite element library for mathematic majored research

"Code generator for quadrature representation"

__author__ = "Kristian B. Oelgaard ([e-mail])"
__date__ = "2007-03-16 -- 2007-06-19"
__copyright__ = "Copyright (C) 2007 Kristian B. Oelgaard"
__license__  = "GNU GPL version 3 or any later version"

# Modified by Anders Logg 2007

# Python modules
import numpy
from sets import Set

# FFC common modules
from ffc.common.constants import *
from ffc.common.utils import *

# FFC language modules
from ffc.compiler.language.index import *
from ffc.compiler.language.restriction import *

# FFC code generation modules
from ffc.compiler.codegeneration.common.codegenerator import *
from ffc.compiler.codegeneration.common.utils import *

# FFC tensor representation modules
from ffc.compiler.representation.tensor.multiindex import *

# Utility functions for quadraturegenerator
from quadraturegenerator_utils import *

# FFC format modules
from ffc.compiler.format.removeunused import *

class QuadratureGenerator(CodeGenerator):
    "Code generator for for tensor representation"

    def __init__(self):
        "Constructor"

        # Initialize common code generator
        CodeGenerator.__init__(self)
        self.optimise_level = 3
        self.save_tables = False
        self.unique_tables = True

    def generate_cell_integral(self, form_data, form_representation, sub_domain, format):
        """Generate dictionary of code for cell integral from the given
        form representation according to the given format"""

        code = []

        # Object to control the code indentation
        Indent = IndentControl()

        # Extract tensors
        tensors = form_representation.cell_tensor
        if len(tensors) == 0:
            return None


        # Generate element code + set of used geometry terms
        element_code, members_code, trans_set = self.__generate_element_tensor\
                                                     (tensors, None, None, Indent, format)

        # Get Jacobian snippet
        jacobi_code = [format["generate jacobian"](form_data.cell_dimension, Integral.CELL)]

        # Remove unused declarations
        code = self.__remove_unused(jacobi_code, trans_set, format)

        # Add element code
        code += [""] + [format["comment"]("Compute element tensor")]
        code += element_code

        return {"tabulate_tensor": code, "members":members_code}

    def generate_exterior_facet_integral(self, form_data, form_representation, sub_domain, format):
        """Generate dictionary of code for exterior facet integral from the given
        form representation according to the given format"""

        # Object to control the code indentation
        Indent = IndentControl()

        # Prefetch formats to speed up code generation
        format_block_begin  = format["block begin"]
        format_block_end    = format["block end"]

        # Extract tensors
        tensors = form_representation.exterior_facet_tensors
        if len(tensors) == 0:
            return None

        num_facets = len(tensors)
        cases = [None for i in range(num_facets)]
        trans_set = Set()
        for i in range(num_facets):
            case = [format_block_begin]

            # Assuming all tables have same dimensions for all facets (members_code)
            c, members_code, t_set = self.__generate_element_tensor(tensors[i], i, None, Indent, format)
            case += c
            case += [format_block_end]
            cases[i] = case
            trans_set = trans_set | t_set

        # Get Jacobian snippet
        jacobi_code = [format["generate jacobian"](form_data.cell_dimension, Integral.EXTERIOR_FACET)]

        # Remove unused declarations
        common = self.__remove_unused(jacobi_code, trans_set, format)

        # Add element code
        common += [""] + [format["comment"]("Compute element tensor for all facets")]

        return {"tabulate_tensor": (common, cases), "constructor":"// Do nothing", "members":members_code}
    
    def generate_interior_facet_integral(self, form_data, form_representation, sub_domain, format):
        """Generate dictionary of code for interior facet integral from the given
        form representation according to the given format"""

        # Object to control the code indentation
        Indent = IndentControl()

        # Prefetch formats to speed up code generation
        format_block_begin  = format["block begin"]
        format_block_end    = format["block end"]

        # Extract tensors
        tensors = form_representation.interior_facet_tensors
        if len(tensors) == 0:
            return None

        num_facets = len(tensors)
        cases = [[None for j in range(num_facets)] for i in range(num_facets)]
        trans_set = Set()
        for i in range(num_facets):
            for j in range(num_facets):
                case = [format_block_begin]

                # Assuming all tables have same dimensions for all facet-facet combinations (members_code)
                c, members_code, t_set = self.__generate_element_tensor(tensors[i][j], i, j, Indent, format)
                case += c
                case += [format_block_end]
                cases[i][j] = case
                trans_set = trans_set | t_set

        # Get Jacobian snippet
        jacobi_code = [format["generate jacobian"](form_data.cell_dimension, Integral.INTERIOR_FACET)]

        # Remove unused declarations
        common = self.__remove_unused(jacobi_code, trans_set, format)

        # Add element code
        common += [""] + [format["comment"]("Compute element tensor for all facets")]

        return {"tabulate_tensor": (common, cases), "constructor":"// Do nothing", "members":members_code}

    def __generate_element_tensor(self, tensors, facet0, facet1, Indent, format):
        "Construct quadrature code for element tensors"

        # Initialize code segments
        tabulate_code = []
        element_code = []
        trans_set = Set()

        # Prefetch formats to speed up code generation
        format_comment      = format["comment"]
        format_loop         = format["loop"]
        format_ip           = format["integration points"]
        format_block_begin  = format["block begin"]
        format_block_end    = format["block end"]

        # Group tensors after number of quadrature points, and dimension of primary indices
        # to reduce number of loops
        group_tensors = equal_loops(tensors)
        tables = None
        name_map = {}
        # Get dictionary of unique tables, and the name_map
        if self.unique_tables:
            name_map, tables = unique_tables(tensors, format)

        # Generate load_table.h if tables should be saved.
        members_code = ""
        if self.save_tables:
            members_code = generate_load_table(tensors)

        for i in range(len(tensors)):
            # Tabulate the quadrature weights
            tabulate_code += self.__tabulate_weights(tensors[i].quadrature.weights, i, Indent, format)

        if self.save_tables:
            # Save psi tables instead of tabulating
            tabulate_code += save_psis(tensors, facet0, facet1, Indent, format, tables)
        else:
            # Tabulate values of basis functions and their derivatives at quadrature points
            tabulate_code += self.__tabulate_psis(tensors, Indent, format, tables)

        # Reset values of the element tensor (assuming same dimensions for all tensors)
        tabulate_code += self.__reset_element_tensor(tensors[0], Indent, format)

        for points in group_tensors:
            # Loop all quadrature points
            # Create list of tensors for comment
            ts = [group_tensors[points][idims][i] for idims in group_tensors[points]\
                                                  for i in range(len(group_tensors[points][idims]))]
            element_code += [Indent.indent(format_comment\
            ("Loop quadrature points (tensor/monomial terms %s)" %(str(tuple(ts)))))]
            element_code += [Indent.indent(format_loop(format_ip, 0, points))]
            element_code += [Indent.indent(format_block_begin)]

            # Increase indentation
            Indent.increase()

            # Get dictionary of primary indices
            prim_dic = group_tensors[points]

            # Generate loop over primary indices
            for idims in prim_dic:
                tensor_numbers = prim_dic[idims]
                indices = [format["first free index"], format["second free index"]]
                # Create loop variables
                loop_vars = [[indices[i], 0, idims[i]] for i in range(len(idims))]

                # Generate loop
                element_code += generate_loop("", "", loop_vars, Indent, format, "")
                Indent.decrease()
                element_code += [Indent.indent(format_block_begin)]
                Indent.increase()

                # Generate element tensors for all tensors with the current number of quadrature points
                for i in tensor_numbers:
                    e_code, t_set = self.__element_tensor(tensors[i], i, Indent, format, name_map)
#                    element_code += self.__element_tensor(tensors[i], i, sign_changes, Indent, format, name_map)
                    element_code += e_code
                    trans_set = trans_set | t_set

                # End loop primary indices
                # Decrease indentation
                Indent.decrease()
                element_code += [Indent.indent(format_block_end)]
                # Decrease indentation
                for i in range(len(idims) - 1):
                    Indent.decrease()

            # End the quadrature loop
            # Decrease indentation
            Indent.decrease()
            element_code += [Indent.indent(format_block_end)]

            if i + 1 < len(tensors):
                element_code += [""]
        return (tabulate_code + element_code, members_code, trans_set)