quadraturegenerator_utils.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 os
from sets import Set


# 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 *

# Should be in dictionary!!
index_names = {0: lambda i: "f%s" %(i), 1: lambda i: "p%s" %(i), 2: lambda i: "s%s" %(i),\
               4: lambda i: "fu%s" %(i), 5: lambda i: "pj%s" %(i), 6: lambda i: "c%s" %(i), 7: lambda i: "a%s" %(i)}


def compute_macro_idims(monomial, idims, irank):
    "Compute macro dimensions in case of restricted basisfunctions"

    # copy dims
    macro_idims = [dim for dim in idims]
    for i in range(irank):
        for v in monomial.basisfunctions:
            if v.index.type == Index.PRIMARY and v.index.index == i:
                if v.restriction != None:
                    macro_idims[i] = idims[i] * 2
                    break
    return macro_idims

def generate_psi_declaration(tensor_number, psi_indices, vindex, aindices, bindices,\
                                   num_quadrature_points, num_dofs, format):

    # Prefetch formats to speed up code generation
    format_secondary_index  = format["secondary index"]

    multi_index = []

    indices = ""
    for index in psi_indices:
        if index == vindex:
            indices = format_secondary_index(index_names[index.type](index.index)) + indices
        else:
            indices += format_secondary_index(index_names[index.type](index([], aindices, bindices, [])))
            multi_index += [index([], aindices, bindices, [])]

    name = format["table declaration"] + format["psis"] + format_secondary_index("t%d" %tensor_number)\
              + indices + format["matrix access"](num_quadrature_points, num_dofs)

    return (name, multi_index)

def generate_psi_entry(tensor_number, aindices, bindices, psi_indices, vindices, name_map, format):

    # Prefetch formats to speed up code generation
    format_secondary_index  = format["secondary index"]

    primary_indices = [format["first free index"], format["second free index"]]

    indices = ""
    for index in psi_indices:
        if index in vindices:
            indices = format_secondary_index(index_names[index.type](index.index)) + indices
            dof_num = index(primary_indices, aindices, bindices, [])
        else:
            indices += format_secondary_index(index_names[index.type](index([], aindices, bindices, [])))

    name = format["psis"] + format_secondary_index("t%d" %tensor_number) + indices
    if name in name_map:
        entry = name_map[name] + format["matrix access"](format["integration points"], dof_num)
    else:
        entry = name + format["matrix access"](format["integration points"], dof_num)

    return entry

def generate_loop(name, value, loop_vars, Indent, format, connect = None):
    "This function generates a loop over a vector or matrix."

    code = []

    # Prefetch formats to speed up code generation
    format_loop     = format["loop"]

    for ls in loop_vars:

        # Get index and lower and upper bounds
        index, lower, upper = (ls[0], ls[1], ls[2])

        # Loop index
        code += [Indent.indent(format_loop(index, lower, upper))]

        # Increase indentation
        Indent.increase()

        if name:
            # If this is the last loop, write values
            if index == loop_vars[-1][0]:
                if connect:
                    code += [connect(Indent.indent(name), value)]
                else:
                    code += [(Indent.indent(name), value)]

    # Decrease indentation
    if name:
        for i in range(len(loop_vars)):
            Indent.decrease()

    return code

def extract_unique(aa, bb):
    "Remove redundant names and entries in the psi table"

    uaa = []
    ubb = []

    for i in range(len(aa)):
        a = aa[i]
        if not a in uaa:
            uaa += [a]
            ubb += [bb[i]]
    return (uaa, ubb)

def equal_loops(tensors):

    group_tensors = {}
    for i in range(len(tensors)):
        tens = tensors[i]
        num_points = len(tens.quadrature.weights)
        if num_points in group_tensors:
            group_tensors[num_points] += [i]
        else:
            group_tensors[num_points] = [i]

    for g in group_tensors:
        tensor_nums = group_tensors[g]
        prims = {}
        for t in tensor_nums:
            tens = tensors[t]
            idims = tens.i.dims
            if tuple(idims) in prims:
                prims[tuple(idims)] += [t]
            else:
                prims[tuple(idims)] = [t]
        group_tensors[g] = prims

    return group_tensors

# def generate_signs(tensors, format):
#     "Generate list of declarations for computation of signs"
#     code = []
#     for j in range(len(tensors)):
#         monomial = tensors[j].monomial
#         # Inspect each basis function (identified by its index)
#         # and check whether sign changes are relevant.
#         for basisfunction in monomial.basisfunctions:
#             index = basisfunction.index
#             values = []
#             necessary = False
#             element = basisfunction.element
#             declarations = []
#             dof_entities = DofMap(element).dof_entities();

#             # Go through the topological entities associated
#             # with each basis function/dof. If the element is
#             # a piola mapped element and the basis function is
#             # associated with a facet, we calculate the
#             # possible sign change.
#             for no in dof_entities:
#                 (entity, entity_no) = dof_entities[no]
#                 if entity == 1 and element.space_mapping(no) == Mapping.PIOLA:
#                     necessary = True
#                     values += [format["facet sign"](entity_no)]
#                 else:
#                     values += ["1"]
#             num_dofs = str(len(dof_entities))
#             name = format["sign tensor declaration"](format["signs"] + format["secondary index"]\
#                    (str(j))+ format["secondary index"](str(index.index)) + format["array access"](num_dofs))
#             value = format["block"](format["separator"].join([str(val) for val in values]))
 
#             # Add declarations for this basis function to the code
#             code += [(name,value)]
                    
#     if necessary:
#         code.insert(0, format["comment"]("Compute signs"))
#         code.insert(0, format["snippet facet signs"](2))
#         return (code, True)
#     else:
#         return ([], False) # Return [] is the case of no sign changes...)

# def add_sign(value, j, i, format):
#     s_set = Set()
#     if value:
#         value = format["grouping"](value)
#         for k in range(len(i)):
#             value = format["multiply"]([format["signs"] + format["secondary index"]\
#                    (str(j))+ format["secondary index"](str(k)) + format["array access"](i[k]),value])
#             s_set.add(format["signs"] + format["secondary index"](str(j))+ format["secondary index"](str(k)))
#     return (value, s_set)

def generate_factor(tensor, a, bgindices, format):
    "Optimise level 2"

    trans_set = Set()

        # Compute product of factors outside sum
    factors = []
    for j in range(len(tensor.coefficients)):
        c = tensor.coefficients[j]
        if not c.index.type == Index.AUXILIARY_G:
            offset = tensor.coefficient_offsets[c]
            coefficient = format["coefficient"](c.n1.index, c.index([], a, [], [])+offset)
            for l in range(len(c.ops)):
                op = c.ops[len(c.ops) - 1 - l]
                if op == Operators.INVERSE:
                    coefficient = format["inverse"](coefficient)
                elif op == Operators.MODULUS:
                    coefficient = format["MODULUSolute value"](coefficient)
                elif op == Operators.SQRT:
                    coefficient = format["sqrt"](coefficient)
            factors += [coefficient]
    for t in tensor.transforms:
        if not (t.index0.type == Index.AUXILIARY_G or  t.index1.type == Index.AUXILIARY_G):
            trans = format["transform"](t.type, t.index0([], a, [], []), \
                                                    t.index1([], a, [], []), \
                                                    t.restriction)
            factors += [trans]
            trans_set.add(trans)

    monomial = format["multiply"](factors)
    if monomial:
        f_out = [monomial]
    else:
        f_out = []
    
    # Compute sum of monomials inside sum
    terms = []
    for b in bgindices:
        factors = []
        for j in range(len(tensor.coefficients)):
            c = tensor.coefficients[j]
            if c.index.type == Index.AUXILIARY_G:
                offset = tensor.coefficient_offsets[c]
                coefficient = format["coefficient"](c.n1.index, c.index([], a, [], b)+offset)
                for l in range(len(c.ops)):
                    op = c.ops[len(c.ops) - 1 - l]
                    if op == Operators.INVERSE:
                        coefficient = format["inverse"](coefficient)
                    elif op == Operators.MODULUS:
                        coefficient = format["absolute value"](coefficient)
                    elif op == Operators.SQRT:
                        coefficient = format["sqrt"](coefficient)
                factors += [coefficient]
        for t in tensor.transforms:
            if t.index0.type == Index.AUXILIARY_G or t.index1.type == Index.AUXILIARY_G:
                trans = format["transform"](t.type, t.index0([], a, [], b), \
                                                        t.index1([], a, [], b), \
                                                        t.restriction)
                factors += [trans]
                trans_set.add(trans)

        terms += [format["multiply"](factors)]

    f_in = format["add"](terms)
    if f_in: f_in = [format["grouping"](f_in)]
    else: f_in = []

    return (f_out, f_in, trans_set)

def generate_factor_old(tensor, a, b, format):
    "Optimise level 0 and 1"
    # Compute product of factors outside sum
    factors = []
    trans_set = Set()

    for j in range(len(tensor.coefficients)):
        c = tensor.coefficients[j]
        if not c.index.type == Index.AUXILIARY_G:
            offset = tensor.coefficient_offsets[c]
            coefficient = format["coefficient"](c.n1.index, c.index([], a, [], [])+offset)
            for l in range(len(c.ops)):
                op = c.ops[len(c.ops) - 1 - l]
                if op == Operators.INVERSE:
                    coefficient = format["inverse"](coefficient)
                elif op == Operators.MODULUS:
                    coefficient = format["absolute value"](coefficient)
                elif op == Operators.SQRT:
                    coefficient = format["sqrt"](coefficient)
            factors += [coefficient]
    for t in tensor.transforms:
        if not (t.index0.type == Index.AUXILIARY_G or  t.index1.type == Index.AUXILIARY_G):
            trans = format["transform"](t.type, t.index0([], a, [], []), \
                                                    t.index1([], a, [], []), \
                                                    t.restriction)
            factors += [trans]
            trans_set.add(trans)

    # Compute sum of monomials inside sum
    for j in range(len(tensor.coefficients)):