evaluatebasis.py

finite element library for mathematic majored research

"""Code generation for evaluation of finite element basis values. This module generates
code which is more or less a C++ representation of FIAT code. More specifically the
functions from the modules expansion.py and jacobi.py are translated into C++"""

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

# Modified by Anders Logg 2007

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

# FFC fem modules
from ffc.fem.finiteelement import *
from ffc.fem.mixedelement import *

# FFC language modules
from ffc.compiler.language.tokens import *

# FFC code generation common modules
from utils import *

# Python modules
import math
import numpy

class IndentControl:
    "Class to control the indentation of code"

    def __init__(self):
        "Constructor"
        self.size = 0
        self.increment = 2

    def increase(self):
        "Increase indentation by increment"
        self.size += self.increment

    def decrease(self):
        "Decrease indentation by increment"
        self.size -= self.increment

    def indent(self, a):
        "Indent string input string by size"
        return indent(a, self.size)

def evaluate_basis(element, format):
    """Evaluate an element basisfunction at a point. The value(s) of the basisfunction is/are
    computed as in FIAT as the dot product of the coefficients (computed at compile time)
    and basisvalues which are dependent on the coordinate and thus have to be computed at
    run time.

    Currently the following elements are supported in 2D and 3D:

    Lagrange                + mixed/vector valued
    Discontinuous Lagrange  + mixed/vector valued
    Crouzeix-Raviart        + mixed/vector valued
    Brezzi-Douglas-Marini   + mixed/vector valued

    Not supported in 2D or 3D:

    Raviart-Thomas ? (not tested since it is broken in FFC, but should work)
    Nedelec (broken?)
    
    Tensor valued elements!"""

    code = []

    Indent = IndentControl()

    # Get coordinates and generate map
    code += generate_map(element, Indent, format)

    # Check if we have just one element
    if (element.num_sub_elements() == 1):

        # Reset values, change for tensor valued elements
        code += reset_values(element.value_dimension(0), False, Indent, format)

        # Map degree of freedom to local degree of freedom for current element
        code += dof_map(0, Indent, format)

        # Generate element code
        code += generate_element_code(element, 0, False, Indent, format)

    # If the element is of type VectorElement or MixedElement
    else:

        # Reset values, change for tensor valued elements
        code += reset_values(element.value_dimension(0), True, Indent, format)

        # Generate element code, for all sub-elements
        code += mixed_elements(element, Indent, format)

    return code

def generate_map(element, Indent, format):
    """Generates map from physical element to the UFC reference element, and from this element
    to reference square/cube. An affine map is assumed for the first mapping and for the second
    map this function implements the UFC version of the FIAT functions, eta_triangle( xi )
    and eta_tetrahedron( xi ) from reference.py"""

    code = []

    # Prefetch formats to speed up code generation
    format_comment        = format["comment"]
    format_floating_point = format["floating point"]
    format_epsilon        = format["epsilon"]

    # Get coordinates and map to the UFC reference element from codesnippets.py
    code += [Indent.indent(format["coordinate map"](element.cell_shape()))] + [""]

    if (element.cell_shape() == 2):

        # Debug code
#        code += [Indent.indent("""// Debug code
#                 std::cout << "coordinates : " << coordinates[0] << " " << coordinates[1] << std::endl;
#                 std::cout << " mapped coordinates : " << x << " " << y << std::endl;""")]

        # Map coordinates to the reference square
        code += [Indent.indent(format_comment("Map coordinates to the reference square"))]
 
        # Code snippet reproduced from FIAT: reference.py: eta_triangle(xi) & eta_tetrahedron(xi) 
        code += [Indent.indent(format["snippet eta_triangle"]) %(format_floating_point(format_epsilon))]

    elif (element.cell_shape() == 3):

        # Debug code
#        code += [Indent.indent("""// Debug code
#        std::cout << "coordinates : " << coordinates[0] << " " << coordinates[1] << " " << coordinates[2] << std::endl;
#        std::cout << " mapped coordinates : " << x << " " << y << " " << z << std::endl;""")]

        # Map coordinates to the reference cube
        code += [Indent.indent(format_comment("Map coordinates to the reference cube"))]

        # Code snippet reproduced from FIAT: reference.py: eta_triangle(xi) & eta_tetrahedron(xi)
        code += [Indent.indent(format["snippet eta_tetrahedron"]) %(format_floating_point(format_epsilon),\
                       format_floating_point(format_epsilon))]
    else:
        raise RuntimeError, "Cannot generate map for shape: %d" %(element.cell_shape())
 
    return code + [""]

def reset_values(num_components, vector, Indent, format):
    "Reset all components of the 'values' array as it is a pointer to an array."

    code = []

    if (vector or num_components != 1):
        # Reset values as it is a pointer
        code += [Indent.indent(format["comment"]("Reset values"))]
        code += [(Indent.indent(format["argument values"] + format["array access"](i)),\
                                format["floating point"](0.0)) for i in range(num_components)]
    else:
        code += [Indent.indent(format["comment"]("Reset values"))]
        code += [(Indent.indent(format["pointer"] + format["argument values"]), format["floating point"](0.0))]

    return code + [""]

def dof_map(sum_space_dim, Indent, format):
    """This function creates code to map a basis function to a local basis function.
    Example, the following mixed element:

    element = VectorElement("Lagrange", "triangle", 2)

    has the element list, elements = [Lagrange order 2, Lagrange order 2] and 12 dofs (6 each).

    The evaluation of basis function 8 is then mapped to 2 (8-6) for local element no. 2."""

    # In case of only one element or the first element in a series then we don't subtract anything
    if sum_space_dim == 0:
        code = [Indent.indent(format["comment"]("Map degree of freedom to element degree of freedom"))]
        code += [(Indent.indent(format["const uint declaration"] + format["local dof"]), format["argument basis num"])]
    else:
        code = [Indent.indent(format["comment"]("Map degree of freedom to element degree of freedom"))]
        code += [(Indent.indent(format["const uint declaration"] + format["local dof"]),\
                format["subtract"]([format["argument basis num"], "%d" %sum_space_dim]))]

    return code + [""]

def mixed_elements(element, Indent, format):
    "Generate code for each sub-element in the event of vector valued elements or mixed elements"

    code = []

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

    # Extract basis elements, and determine number of elements
#    elements = extract_elements(element)
    elements = element.basis_elements()
    num_elements = len(elements)

    sum_value_dim = 0
    sum_space_dim = 0

    # Generate code for each element
    for i in range(num_elements):

        # Get sub element
        basis_element = elements[i]

        # FIXME: This must most likely change for tensor valued elements
        value_dim = basis_element.value_dimension(0)
        space_dim = basis_element.space_dimension()

        # Determine if the element has a value, for the given dof
        code += [Indent.indent(format_dof_map_if(sum_space_dim, sum_space_dim + space_dim -1))]
        code += [Indent.indent(format_block_begin)]
        # Increase indentation
        Indent.increase()

        # Generate map from global to local dof
        code += dof_map(sum_space_dim, Indent, format)

        # Generate code for basis element
        code += generate_element_code(basis_element, sum_value_dim, True, Indent, format)

        # Decrease indentation, finish block - end element code
        Indent.decrease()
        code += [Indent.indent(format_block_end)] + [""]

        # Increase sum of value dimension, and space dimension
        sum_value_dim += value_dim
        sum_space_dim += space_dim

    return code

def generate_element_code(element, sum_value_dim, vector, Indent, format):
    "Generate code for a single basis element"

    code = []

    # Generate basisvalues
    code += generate_basisvalues(element, Indent, format)

    # Tabulate coefficients
    code += tabulate_coefficients(element, Indent, format)

    # Extract relevant coefficients
    code += relevant_coefficients(element, Indent, format)

    # Compute the value of the basisfunction as the dot product of the coefficients
    # and basisvalues
    code += compute_values(element, sum_value_dim, vector, Indent, format)

    return code

def generate_basisvalues(element, Indent, format):
    "Generate code to compute basis values"

    code = []

    # Compute scaling of y and z 1/2(1-y)^n and 1/2(1-z)^n
    code += compute_scaling(element, Indent, format)

    # Compute auxilliary functions currently only 2D and 3D is supported
    if (element.cell_shape() == 2):
        code += compute_psitilde_a(element, Indent, format)
        code += compute_psitilde_b(element, Indent, format)
    elif (element.cell_shape() == 3):
        code += compute_psitilde_a(element, Indent, format)
        code += compute_psitilde_b(element, Indent, format)
        code += compute_psitilde_c(element, Indent, format)
    else:
        raise RuntimeError, "Cannot compute auxilliary functions for shape: %d" %(element.cell_shape())

    # Compute the basisvalues
    code += compute_basisvalues(element, Indent, format)

    return code

def tabulate_coefficients(element, Indent, format):
    """This function tabulates the element coefficients that are generated by FIAT at
    compile time."""

    code = []

    # Prefetch formats to speed up code generation
    format_comment            = format["comment"]
    format_table_declaration  = format["table declaration"]
    format_coefficients       = format["coefficients table"]
    format_matrix_access      = format["matrix access"]
    format_const_float        = format["const float declaration"]

    # Get coefficients from basis functions, computed by FIAT at compile time
    coefficients = element.basis().coeffs

    # Scalar valued basis element [Lagrange, Discontinuous Lagrange, Crouzeix-Raviart]
    if (element.value_rank() == 0):
        coefficients = [coefficients]

    # Vector valued basis element [Raviart-Thomas, Brezzi-Douglas-Marini (BDM)]
    elif (element.value_rank() == 1):
        coefficients = numpy.transpose(coefficients, [1,0,2])

    else:
        raise RuntimeError, "Tensor elements not supported!"

    # Get number of components, must change for tensor valued elements
    num_components = element.value_dimension(0)

    # Get polynomial dimension of basis
    poly_dim = len(element.basis().base.bs)

    # Get the number of dofs from element
    num_dofs = element.space_dimension()

    code += [Indent.indent(format_comment("Table(s) of coefficients"))]

    # Generate tables for each component
    for i in range(num_components):

        # Extract coefficients for current component
        coeffs = coefficients[i]

        # Declare varable name for coefficients
        name = format_table_declaration + format_coefficients(i) + format_matrix_access(num_dofs, poly_dim)
        value = tabulate_matrix(coeffs, format)

        # Generate array of values
        code += [(Indent.indent(name), Indent.indent(value))] + [""]

    return code

def relevant_coefficients(element, Indent, format):
    "Declare relevant coefficients as const floats"

    code = []

    # Prefetch formats to speed up code generation
    format_comment            = format["comment"]
    format_const_float        = format["const float declaration"]
    format_coeff              = format["coefficient scalar"]
    format_secondary_index    = format["secondary index"]
    format_coefficients       = format["coefficients table"]
    format_matrix_access      = format["matrix access"]
    format_local_dof          = format["local dof"]

    # Get number of components, must change for tensor valued elements
    num_components = element.value_dimension(0)

    # Get polynomial dimension of basis
    poly_dim = len(element.basis().base.bs)

    # Extract relevant coefficients and declare as floats
    code += [Indent.indent(format_comment("Extract relevant coefficients"))]
    for i in range(num_components):
        for j in range(poly_dim):
            name = format_const_float + format_coeff(i) + format_secondary_index(j)
            value = format_coefficients(i) + format_matrix_access(format_local_dof, j)
            # Generate values
            code += [(Indent.indent(name), Indent.indent(value))]

    return code + [""]

def compute_values(element, sum_value_dim, vector, Indent, format):
    """This function computes the value of the basisfunction as the dot product of the
    coefficients and basisvalues """

    code = []

    # Prefetch formats to speed up code generation
    format_values           = format["argument values"]
    format_array_access     = format["array access"]
    format_add              = format["add"]
    format_multiply         = format["multiply"]
    format_coefficients     = format["coefficient scalar"]
    format_secondary_index  = format["secondary index"]
    format_basisvalue       = format["basisvalue"]
    format_pointer          = format["pointer"]
    format_det              = format["determinant"]
    format_inv              = format["inverse"]
    format_add              = format["add"]
    format_mult             = format["multiply"]
    format_group            = format["grouping"]
    format_tmp              = format["tmp declaration"]
    format_tmp_access       = format["tmp access"]

    code += [Indent.indent(format["comment"]("Compute value(s)"))]

    # Get number of components, change for tensor valued elements