gistfile1.py

from ufl import HDiv
from firedrake.fiat_utils import *

class CellIndirection(object):

    def __init__(self, fs):
        ufl_ele = fs.ufl_element()
        # Unwrap non HDiv'd element if necessary
        if isinstance(ufl_ele, HDiv):
            ele = ufl_ele._element
        else:
            ele = ufl_ele

        element = fiat_from_ufl_element(ele)
        self._element = element
        self._n_h = element.A.space_dimension()
        self._n_v = element.B.space_dimension()
        self._ndof = self._n_h * self._n_v
        # Where are the nodes?
        pe = [p.get_point_dict().keys()[0] for p in element.dual_basis()]

        # Permutation from reference-element ordering into "layers in cell" ordering
        # i.e. all the dofs with z=0, then z=1, z=2, ...
        self._permutation = zip(*sorted(zip(range(self._n_h*self._n_v), pe),
                                        key=lambda x: x[1][::-1]))[0]


    def ki_to_idx(self, k, i):
        if i > self._n_h:
            raise RuntimeError("i out of range [0, %d)" % self._n_h)
        return (k / self._n_v) * self._ndof + self._permutation[k % self._n_v + i]

    def c_permutation(self, name):
        return "int %s[%d] = {%s};" % (name, len(self._permutation),
                                       ', '.join("%d" % p for p in self._permutation))

    def c_ki_to_idx(self, k, i, name):
        return "((%(k)s / %(nv)d) * %(nd)d + %(name)s[(%(k)s %% %(nv)d) + %(i)s])" % \
            {'k': k,
             'i': i,
             'name': name,
             'nv': self._n_v,
             'nd': self._ndof}