tensorgenerator.py

finite element library for mathematic majored research

        code += [format["comment"]("Compute geometry tensors")]

        # Iterate over all terms
        j = 0
        coeff_set = Set()
        trans_set = Set()
        for i in range(len(terms)):

            term = terms[i]

            # Get list of secondary indices (should be the same so pick first)
            aindices = terms[i].G[0].a.indices

            # Iterate over secondary indices
            for a in aindices:

                # Skip code generation if term is not used
                if not format["geometry tensor access"](i,a) in geo_set:
                    continue

                # Compute factorized values
                values = []
                jj = j
                for G in term.G:
                    val, c_set, t_set = self.__generate_entry(G, a, jj, format)
                    values += [val]
                    coeff_set = coeff_set | c_set
                    trans_set = trans_set | t_set
                    jj += 1

                # Sum factorized values
                name = format["geometry tensor declaration"](i, a)
                value = format["add"](values)

                # Multiply with determinant factor
                # FIXME: dets = pick_first([G.determinants for G in term.G])
                dets = term.G[0].determinants
                value = self.__multiply_value_by_det(value, dets, format, len(values) > 1)

                # Add determinant to transformation set
                if dets:
                    d0 = [format["power"](format["determinant"](det.restriction),
                                          det.power) for det in dets]
                    trans_set.add(format["multiply"](d0))

                # Add declaration
                code += [(name, value)]

            j += len(term.G)

        # Add scale factor
        trans_set.add(format["scale factor"])

        return (code, coeff_set, trans_set)

    def __generate_element_tensor(self, terms, format):
        "Generate list of declaration for computation of element tensor"

        # Generate code as a list of declarations
        code = []    
    
        # Get list of primary indices (should be the same so pick first)
        iindices = terms[0].A0.i.indices

        # Prefetch formats to speed up code generation
        format_element_tensor  = format["element tensor"]
        format_geometry_tensor = format["geometry tensor access"]
        format_add             = format["add"]
        format_subtract        = format["subtract"]
        format_multiply        = format["multiply"]
        format_floating_point  = format["floating point"]
        format_epsilon         = format["epsilon"]

        # Generate code for geometry tensor entries
        gk_tensor = [ ( [(format_geometry_tensor(j, a), a) for a in terms[j].A0.a.indices], j) for j in range(len(terms)) ]

        # Generate code for computing the element tensor
        k = 0
        num_dropped = 0
        num_ops = 0
        zero = format_floating_point(0.0)
        geo_set = Set()
        for i in iindices:
            name = format_element_tensor(i, k)
            value = None
            for (gka, j) in gk_tensor:
                A0 = terms[j].A0
                for (gk, a) in gka:
                    a0 = A0.A0[tuple(i + a)]
                    if abs(a0) > format_epsilon:
                        if value and a0 < 0.0:
                            value = format_subtract([value, format_multiply([format_floating_point(-a0), gk])])
                            geo_set.add(gk)
                        elif value:
                            value = format_add([value, format_multiply([format_floating_point(a0), gk])])
                            geo_set.add(gk)
                        else:
                            value = format_multiply([format_floating_point(a0), gk])
                            geo_set.add(gk)
                        num_ops += 1
                    else:
                        num_dropped += 1
            value = value or zero
            code += [(name, value)]
            k += 1

        return (code, geo_set)

    def __generate_entry(self, G, a, i, format):
        "Generate code for the value of entry a of geometry tensor G"

        coeff_set = Set()
        trans_set = Set()

        # Compute product of factors outside sum
        factors = []
        for j in range(len(G.coefficients)):
            c = G.coefficients[j]
            if not c.index.type == Index.AUXILIARY_G:
                coefficient = format["modified coefficient access"](c.n1.index, i, j, c.index([], a, [], []))
                coeff_set.add(coefficient)
                factors += [coefficient]
            
        for t in G.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: f0 = [monomial]
        else: f0 = []
    
        # Compute sum of monomials inside sum
        terms = []
        for b in G.b.indices:
            factors = []
            for j in range(len(G.coefficients)):
                c = G.coefficients[j]
                if c.index.type == Index.AUXILIARY_G:
                    coefficient = format["modified coefficient access"](c.n1.index, i, j, c.index([], a, [], b))
                    coeff_set.add(coefficient)
                    factors += [coefficient]
            for t in G.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)]

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

        fs = f0 + f1
        if not fs: fs = ["1.0"]

        # Compute product of all factors
        return (format["multiply"](fs), coeff_set, trans_set)

    def __multiply_value_by_det(self, value, dets, format, is_sum):
        if dets:
            d0 = [format["power"](format["determinant"](det.restriction),
                                  det.power) for det in dets]
        else:
            d0 = []
        if value == "1.0":
            v = []
        elif is_sum:
            v = [format["grouping"](value)]
        else:
            v = [value]
        return format["multiply"](d0 + [format["scale factor"]] + v)

    def __remove_unused(self, code, set, format):
        "Remove unused variables so that the compiler will not complain"

        # Normally, the removal of unused variables should happen at the
        # formatting stage, but since the code for the tensor contraction
        # may grow to considerable size, we make an exception and remove
        # unused variables here when we know the names of the used
        # variables. No searching necessary and much, much, much faster.
        
        if code:
            # Generate body of code, using the format
            lines = format["generate body"](code)

            # Generate auxiliary code line that uses all members of the set (to trick remove_unused)
            line_set = format["add equal"]("A", format["multiply"](set))
            lines += "\n" + line_set

            # Remove unused Jacobi declarations
            code = remove_unused(lines)

            # Delete auxiliary line
            code = code.replace("\n" + line_set, "")

            return [code]
        else:
            return code