evaluatebasis.py

finite element library for mathematic majored research

    num_components = element.value_dimension(0)

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

    # Check which transform we should use to map the basis functions
    mapping = pick_first([element.value_mapping(dim) for dim in range(element.value_dimension(0))])
    if mapping == Mapping.PIOLA:
        code += [Indent.indent(format["comment"]("Using Piola transform to map values back to the physical element"))]

    if (vector or num_components != 1):

        # Loop number of components
        for i in range(num_components):
            name = format_values + format_array_access(i + sum_value_dim)
            value = format_add([format_multiply([format_coefficients(i) + format_secondary_index(j),\
                    format_basisvalue(j)]) for j in range(poly_dim)])

            # Use Piola transform to map basisfunctions back to physical element if needed
            if mapping == Mapping.PIOLA:
                code.insert(i+1,(Indent.indent(format_tmp(0, i)), value))
                basis_col = [format_tmp_access(0, j) for j in range(element.cell_dimension())]
                jacobian_row = [format["transform"](Transform.J, j, i, None) for j in range(element.cell_dimension())]
                inner = [format_mult([jacobian_row[j], basis_col[j]]) for j in range(element.cell_dimension())]
                sum = format_group(format_add(inner))
                value = format_mult([format_inv(format_det(None)), sum])

            code += [(Indent.indent(name), value)]
    else:
         name = format_pointer + format_values
         value = format_add([format_multiply([format_coefficients(0) + format_secondary_index(j),\
                 format_basisvalue(j)]) for j in range(poly_dim)])

         code += [(Indent.indent(name), value)]

    return code

def compute_scaling(element, Indent, format):
    """Generate the scalings of y and z coordinates. This function is an implementation of
    the FIAT function make_scalings( etas ) from expansions.py"""

    code = []

    # Prefetch formats to speed up code generation
    format_y            = format["y coordinate"]
    format_z            = format["z coordinate"]
    format_grouping     = format["grouping"]
    format_subtract     = format["subtract"]
    format_multiply     = format["multiply"]
    format_const_float  = format["const float declaration"]
    format_scalings     = format["scalings"]

    # Get the element degree
    degree = element.degree()

    # Get the element shape
    element_shape = element.cell_shape()

    # Currently only 2D and 3D is supported
    if (element_shape == 2):
        scalings = [format_y]
        # Scale factor, for triangles 1/2*(1-y)^i i being the order of the element
#        factors = ["(0.5 - 0.5 * y)"]
        factors = [format_grouping(format_subtract(["0.5", format_multiply(["0.5", format_y])]))]
    elif (element_shape == 3):
        scalings = [format_y, format_z]
#        factors = ["(0.5 - 0.5 * y)", "(0.5 - 0.5 * z)"]
        factors = [format_grouping(format_subtract(["0.5", format_multiply(["0.5", format_y])])),\
                   format_grouping(format_subtract(["0.5", format_multiply(["0.5", format_z])]))]
#        factors = ["(0.5 - 0.5 * y)", "(0.5 - 0.5 * z)"]
    else:
        raise RuntimeError, "Cannot compute scaling for shape: %d" %(elemet_shape)

    code += [Indent.indent(format["comment"]("Generate scalings"))]

    # Can be optimized by leaving out the 1.0 variable
    for i in range(len(scalings)):

      name = format_const_float + format_scalings(scalings[i], 0)
      value = format["floating point"](1.0)
      code += [(Indent.indent(name), value)]

      for j in range(1, degree+1):
          name = format_const_float + format_scalings(scalings[i], j)
          value = format_multiply([format_scalings(scalings[i],j-1), factors[i]])
          code += [(Indent.indent(name), value)]

    return code + [""]

def compute_psitilde_a(element, Indent, format):
    """Compute Legendre functions in x-direction. The function relies on
    eval_jacobi_batch(a,b,n) to compute the coefficients.

    The format is:
    psitilde_a[0] = 1.0
    psitilde_a[1] = a + b * x
    psitilde_a[n] = a * psitilde_a[n-1] + b * psitilde_a[n-1] * x + c * psitilde_a[n-2]
    where a, b and c are coefficients computed by eval_jacobi_batch(0,0,n)
    and n is the element degree"""

    code = []

    # Get the element degree
    degree = element.degree()

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

    # Create list of variable names
    variables = [format["x coordinate"], format["psitilde_a"]]

    for n in range(degree+1):
        # Declare variable
        name = format["const float declaration"] + variables[1] + format["secondary index"](n)

        # Compute value
        value = eval_jacobi_batch_scalar(0, 0, n, variables, format)

        code += [(Indent.indent(name), value)]

    return code + [""]

def compute_psitilde_b(element, Indent, format):
    """Compute Legendre functions in y-direction. The function relies on
    eval_jacobi_batch(a,b,n) to compute the coefficients.

    The format is:
    psitilde_bs_0[0] = 1.0
    psitilde_bs_0[1] = a + b * y
    psitilde_bs_0[n] = a * psitilde_bs_0[n-1] + b * psitilde_bs_0[n-1] * x + c * psitilde_bs_0[n-2]
    psitilde_bs_(n-1)[0] = 1.0
    psitilde_bs_(n-1)[1] = a + b * y
    psitilde_bs_n[0] = 1.0
    where a, b and c are coefficients computed by eval_jacobi_batch(2*i+1,0,n-i) with i in range(0,n+1)
    and n is the element degree + 1"""

    code = []

    # Prefetch formats to speed up code generation
    format_y                = format["y coordinate"]
    format_psitilde_bs      = format["psitilde_bs"]
    format_const_float      = format["const float declaration"]
    format_secondary_index  = format["secondary index"]

    # Get the element degree
    degree = element.degree()

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

    for i in range(0, degree + 1):

        # Compute constants for jacobi function
        a = 2*i+1
        b = 0
        n = degree - i
            
        # Create list of variable names
        variables = [format_y, format_psitilde_bs(i)]

        for j in range(n+1):
            # Declare variable
            name = format_const_float + variables[1] + format_secondary_index(j)

            # Compute values
            value = eval_jacobi_batch_scalar(a, b, j, variables, format)

            code += [(Indent.indent(name), value)]

    return code + [""]

def compute_psitilde_c(element, Indent, format):
    """Compute Legendre functions in y-direction. The function relies on
    eval_jacobi_batch(a,b,n) to compute the coefficients.

    The format is:
    psitilde_cs_0[0] = 1.0
    psitilde_cs_0[1] = a + b * y
    psitilde_cs_0[n] = a * psitilde_cs_0[n-1] + b * psitilde_cs_0[n-1] * x + c * psitilde_cs_0[n-2]
    psitilde_cs_(n-1)[0] = 1.0
    psitilde_cs_(n-1)[1] = a + b * y
    psitilde_cs_n[0] = 1.0
    where a, b and c are coefficients computed by 

    [[jacobi.eval_jacobi_batch(2*(i+j+1),0, n-i-j) for j in range(0,n+1-i)] for i in range(0,n+1)]"""

    code = []

    # Prefetch formats to speed up code generation
    format_z                = format["z coordinate"]
    format_psitilde_cs      = format["psitilde_cs"]
    format_const_float      = format["const float declaration"]
    format_secondary_index  = format["secondary index"]

    # Get the element degree
    degree = element.degree()

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

    for i in range(0, degree + 1):
        for j in range(0, degree + 1 - i):

            # Compute constants for jacobi function
            a = 2*(i+j+1)
            b = 0
            n = degree - i - j
            
            # Create list of variable names
            variables = [format_z, format_psitilde_cs(i,j)]

            for k in range(n+1):
                # Declare variable
                name = format_const_float + variables[1] + format_secondary_index(k)

                # Compute values
                value = eval_jacobi_batch_scalar(a, b, k, variables, format)

                code += [(Indent.indent(name), value)]

    return code + [""]

def compute_basisvalues(element, Indent, format):
    """This function is an implementation of the loops inside the FIAT functions
    tabulate_phis_triangle( n , xs ) and tabulate_phis_tetrahedron( n , xs ) in
    expansions.py. It computes the basis values from all the previously tabulated variables."""

    code = []

    # Prefetch formats to speed up code generation
    format_multiply         = format["multiply"]
    format_secondary_index  = format["secondary index"]
    format_psitilde_a       = format["psitilde_a"]
    format_psitilde_bs      = format["psitilde_bs"]
    format_psitilde_cs      = format["psitilde_cs"]
    format_scalings         = format["scalings"]
    format_y                = format["y coordinate"]
    format_z                = format["z coordinate"]
    format_const_float      = format["const float declaration"]
    format_basisvalue       = format["basisvalue"]

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

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

    # Get the element shape
    element_shape = element.cell_shape()

    # Currently only 2D and 3D is supported
    # 2D
    if (element_shape == 2):
        count = 0
        for k in range(0,element.degree() + 1):
            for i in range(0,k + 1):
                ii = k-i
                jj = i
                factor = math.sqrt( (ii+0.5)*(ii+jj+1.0) )

                symbol = format_multiply([format_psitilde_a + format_secondary_index(ii),\
                         format_scalings(format_y, ii), format_psitilde_bs(ii) + format_secondary_index(jj)])

                # Declare variable
                name = format_const_float + format_basisvalue(count)

                # Let inner_product handle format of factor
                value = inner_product([factor],[symbol], format)

#                var += [format["multiply"](["psitilde_a_%d" % ii, "scalings_y_%d" % ii,\
#                        "psitilde_bs_%d_%d" %(ii, jj), format["floating point"](factor)])]

                code += [(Indent.indent(name), value)]
                count += 1
        if (count != poly_dim):
            raise RuntimeError, "The number of basis values must be the same as the polynomium dimension of the base"

    # 3D
    elif (element_shape == 3):
        count = 0
        for k in range(0, element.degree()+1):  # loop over degree
            for i in range(0, k+1):
                for j in range(0, k - i + 1):
                    ii = k-i-j
                    jj = j
                    kk = i
                    factor = math.sqrt( (ii+0.5) * (ii+jj+1.0) * (ii+jj+kk+1.5) )

#                    var += [format["multiply"](["psitilde_a_%d" % ii, "scalings_y_%d" % ii,\
#                        "psitilde_bs_%d_%d" % (ii, jj), "scalings_z_%d" % (ii+jj),\
#                        "psitilde_cs_%d%d_%d" % (ii, jj, kk), format["floating point"](factor)])]

                    symbol = format_multiply([format_psitilde_a + format_secondary_index(ii),\
                             format_scalings(format_y, ii), format_psitilde_bs(ii) + format_secondary_index(jj),\
                             format_scalings(format_z, (ii+jj)),\
                             format_psitilde_cs(ii, jj) + format_secondary_index(kk)])

                    # Declare variable
                    name = format_const_float + format_basisvalue(count)

                    # Let inner_product handle format of factor
                    value = inner_product([factor],[symbol], format)

                    code += [(Indent.indent(name), value)]
                    count += 1
        if (count != poly_dim):
            raise RuntimeError, "The number of basis values must be the same as the polynomium dimension of the base"
    else:
        raise RuntimeError, "Cannot compute basis values for shape: %d" % elemet_shape

    # Debug basis
#    code += ["std::cout" + "".join([" << basisvalue%d << " % i + '" "' for i in range(poly_dim)]) + " << std::endl;"]

    return code + [""]

def extract_elements(element):
    """This function extracts the basis elements recursively from vector elements and mixed elements.
    Example, the following mixed element:

    element1 = FiniteElement("Lagrange", "triangle", 1)
    element2 = VectorElement("Lagrange", "triangle", 2)

    element  = element2 + element1, has the structure:
    mixed-element[mixed-element[Lagrange order 2, Lagrange order 2], Lagrange order 1]

    This function returns the list of basis elements:
    elements = [Lagrange order 2, Lagrange order 2, Lagrange order 1]"""

    elements = []

    # If the element is not mixed (a basis element, add to list)
    if isinstance(element, FiniteElement):
        elements += [element]
    # Else call this function again for each subelement
    else:
        for i in range(element.num_sub_elements()):
            elements += extract_elements(element.sub_element(i))

    return elements

def eval_jacobi_batch_scalar(a, b, n, variables, format):
    """Implementation of FIAT function eval_jacobi_batch(a,b,n,xs) from jacobi.py"""

    # Prefetch formats to speed up code generation
    format_secondary_index  = format["secondary index"]
    format_float            = format["floating point"]
    format_epsilon          = format["epsilon"]
    
    # Format variables
    access = lambda i: variables[1] + format_secondary_index(i)
    coord = variables[0]

    if n == 0:
        return format["floating point"](1.0)
    if n == 1:
        # Results for entry 1, of type (a + b * coordinate) (coordinate = x, y or z)
        res0 = 0.5 * (a - b)
        res1 = 0.5 * ( a + b + 2.0 )

        val1 = inner_product([res1], [coord], format)

        if (abs(res0) > format_epsilon): # Only include if the value is not zero
            val0 = format_float(res0)
            if val1:
                if res0 > 0:
                    return format["add"]([val1, format_float(res0)])
                else:
                    return format["subtract"]([val1, format_float(-res0)])
            else:
                return val0
        else:
            return val1

    else:
        apb = a + b
        # Compute remaining entries, of type (a + b * coordinate) * psitilde[n-1] - c * psitilde[n-2])
        a1 = 2.0 * n * ( n + apb ) * ( 2.0 * n + apb - 2.0 )
        a2 = ( 2.0 * n + apb - 1.0 ) * ( a * a - b * b )
        a3 = ( 2.0 * n + apb - 2.0 ) * ( 2.0 * n + apb - 1.0 ) * ( 2.0 * n + apb )
        a4 = 2.0 * ( n + a - 1.0 ) * ( n + b - 1.0 ) * ( 2.0 * n + apb )
        a2 = a2 / a1
        a3 = a3 / a1
        # Note:  we subtract the value of a4!
        a4 = -a4 / a1

        float_numbers = [a2, a3, a4]
        symbols = [access(n-1), format["multiply"]([coord, access(n-1)]), access(n-2)]

        return inner_product(float_numbers, symbols, format)