Tabulate Dubiner basis using Duffy transform

FIAT/expansions.py

@@ -432,3 +432,209 @@ def polynomial_dimension(ref_el, degree):
         return max(0, (degree + 1) * (degree + 2) * (degree + 3) // 6)
     else:
         raise ValueError("Unknown reference element type.")
+
+
+def eta_square(xi):
+    """Maps from the (-1,1) reference triangle to [-1,1]^2."""
+    xi1, xi2 = xi
+    with numpy.errstate(divide='ignore', invalid='ignore'):
+        eta1 = 2. * (1. + xi1) / (1. - xi2) - 1.
+    eta2 = xi2
+    if eta1.dtype != object:
+        eta1[numpy.logical_not(numpy.isfinite(eta1))] = 1.
+    return eta1, eta2
+
+
+def eta_cube(xi):
+    """Maps from the (-1,1) reference tetrahedron to [-1,1]^3."""
+    xi1, xi2, xi3 = xi
+    with numpy.errstate(divide='ignore', invalid='ignore'):
+        eta1 = 2. * (1. + xi1) / (-xi2 - xi3) - 1.
+        eta2 = 2. * (1. + xi2) / (1. - xi3) - 1.
+    eta3 = xi3
+    if eta1.dtype != object:
+        eta1[numpy.logical_not(numpy.isfinite(eta1))] = 1.
+    if eta2.dtype != object:
+        eta2[numpy.logical_not(numpy.isfinite(eta2))] = 1.
+    return eta1, eta2, eta3
+
+
+from operator import mul
+from functools import reduce
+from math import prod
+
+
+def chain_rule(eta, dphi_dxi):
+    dim = len(eta)
+    Jii = 1.
+    dphi_deta = dphi_dxi
+    for i in reversed(range(dim)):
+        iupper = range(i + 1, dim)
+        offdiag = (prod(((1. - eta[k])*0.5 for k in iupper if k != j), (1. + eta[i])*0.5) for j in iupper)
+        dphi_deta[i] = sum(reduce(mul, dphi_dxi[i+1:], offdiag), dphi_deta[i]) * (1./Jii)
+        Jii *= (1. - eta[i])*0.5
+    return dphi_deta
+
+
+def flat_index(i, j):
+    return (i + j) * (i + j + 1) // 2 + j
+
+
+def dubiner_1d(order, dim, x):
+    if dim == 0:
+        return jacobi.eval_jacobi_batch(0, 0, degree, x[:, None])
+    sd = (order + 1) * (order + 2) // 2
+    phi = numpy.zeros((sd, x.size), dtype=x.dtype)
+    xhat = (1. - x) * 0.5
+    for j in range(order+1):
+        n = order - j
+        alpha = 2 * j + dim
+        results = jacobi.eval_jacobi_batch(alpha, 0, n, x[:, None])
+        if j > 0:
+            results *= xhat ** j
+        s = [flat_index(i, j) for i in range(n + 1)]
+        phi[s, :] = results
+    return phi
+
+
+def dubiner_deriv_1d(order, dim, x):
+    if dim == 0:
+        return jacobi.eval_jacobi_deriv_batch(0, 0, degree, x[:, None])
+    sd = (order + 1) * (order + 2) // 2
+    dphi = numpy.zeros((sd, x.size), dtype=x.dtype)
+    xhat = (1. - x) * 0.5
+    for j in range(order):
+        n = order - j
+        alpha = 2 * j + dim
+        derivs = jacobi.eval_jacobi_deriv_batch(alpha, 0, n, x[:, None])
+        if j > 0:
+            results = jacobi.eval_jacobi_batch(alpha, 0, n, x[:, None])
+            derivs *= xhat
+            derivs -= results * (0.5*j)
+            if j > 1:
+                derivs *= xhat ** (j - 1)
+        s = [flat_index(i, j) for i in range(n + 1)]
+        dphi[s, :] = derivs
+    return dphi
+
+
+def dubiner_2d(order, xi, alphas=None):
+    if alphas is None:
+        alphas = [(0,) * 2]
+    sd = (order + 1) * (order + 2) // 2
+    eta = eta_square(numpy.transpose(xi))
+    B = [dubiner_1d(order, k, x) for k, x in enumerate(eta)]
+    D = [None] * len(B)
+    if any(sum(alpha) > 0 for alpha in alphas):
+        D = [dubiner_deriv_1d(order, k, x) for k, x in enumerate(eta)]
+    def idx(p, q):
+        return (p + q) * (p + q + 1) // 2 + q
+
+    tabulations = {}
+    for alpha in alphas:
+        T = [Dj if aj else Bj for aj, Bj, Dj in zip(alpha, B, D)]
+        phi = numpy.zeros((sd, T[0].shape[1]), dtype=T[0].dtype)
+        for i in range(order + 1):
+            Ti = T[0][i]
+            for j in range(order + 1 - i):
+                scale = ((i + 0.5) * (i + j + 1.0)) ** 0.5
+                phi[idx(i, j)] = T[1][flat_index(j, i)] * Ti * scale
+        tabulations[alpha] = phi
+    return tabulations
+
+
+def dubiner_3d(order, xi, alphas=None):
+    if alphas is None:
+        alphas = [(0,) * 3]
+    sd = (order + 1) * (order + 2) * (order + 3) // 6
+    eta = eta_cube(numpy.transpose(xi))
+    B = [dubiner_1d(order, k, x) for k, x in enumerate(eta)]
+    D = [None] * len(B)
+    if any(sum(alpha) > 0 for alpha in alphas):
+        D = [dubiner_deriv_1d(order, k, x) for k, x in enumerate(eta)]
+    def idx(p, q, r):
+        return (p + q + r)*(p + q + r + 1)*(p + q + r + 2)//6 + (q + r)*(q + r + 1)//2 + r
+
+    tabulations = {}
+    for alpha in alphas:
+        T = [Dj if aj else Bj for aj, Bj, Dj in zip(alpha, B, D)]
+        phi = numpy.zeros((sd, T[0].shape[1]), dtype=T[0].dtype)
+        for i in range(order + 1):
+            Ti = T[0][i]
+            for j in range(order + 1 - i):
+                Tij = T[1][flat_index(j, i)] * Ti
+                for k in range(order + 1 - i - j):
+                    scale = ((i + 0.5) * (i + j + 1.0) * (i + j + k + 1.5)) ** 0.5
+                    phi[idx(i, j, k)] = T[2][flat_index(k, i + j)] * Tij * scale
+        tabulations[alpha] = phi
+    return tabulations
+
+
+if __name__ == "__main__":
+    def symmetric_simplex(dim):
+        s = reference_element.ufc_simplex(dim)
+        r = lambda x: x ** 0.5
+        if dim == 2:
+            s.vertices = [(0.0, 0.0), (-1.0, -r(3.0)), (1.0, -r(3.0))]
+        elif dim == 3:
+            s.vertices = [(r(3.0)/3, 0.0, 0.0), (-r(3.0)/6, 0.5, 0.0),
+                          (-r(3.0)/6, -0.5, 0.0), (0.0, 0.0, r(6.0)/3)]
+        return s
+
+    dim = 2
+    degree = 2
+    tabulate = [lambda n, x: jacobi.eval_jacobi_batch(0, 0, n, x), dubiner_2d, dubiner_3d][dim-1]
+
+    # ref_el = symmetric_simplex(dim)
+    ref_el = reference_element.ufc_simplex(dim)
+    expansion_set = get_expansion_set(ref_el)
+
+    if dim == 1:
+        base_ref_el = reference_element.DefaultInterval()
+    elif dim == 2:
+        base_ref_el = reference_element.DefaultTriangle()
+    elif dim == 3:
+        base_ref_el = reference_element.DefaultTetrahedron()
+
+    v1 = ref_el.get_vertices()
+    v2 = base_ref_el.get_vertices()
+    A, b = reference_element.make_affine_mapping(v1, v2)
+    mapping = lambda x: numpy.dot(x, A.T) + b
+
+    if 1:
+        alphas = [(0,) * dim]
+        alphas.extend(tuple(row) for row in numpy.eye(dim))
+        simplify = lambda x: numpy.array(sympy.simplify(x))
+        X = [tuple(map(sympy.Symbol, ("x", "y", "z")[:dim]))]
+        Tnew = tabulate(degree, mapping(X), alphas=alphas)
+        Told = expansion_set.tabulate(degree, X)
+        print("New")
+        print(simplify(Tnew[(0,) * dim]))
+        print("Old")
+        print(simplify(Told))
+
+        dz = tabulate(degree, mapping(X), alphas=alphas)
+        dy = expansion_set.tabulate_derivatives(degree, X)
+        for alpha, Xi in zip(alphas, X):
+            print("New")
+            print(simplify(Tnew[alpha]))
+            print("Old")
+            print(simplify(sympy.diff(Told, Xi)))
+
+    else:
+        import FIAT
+        from matplotlib import pyplot as plt
+
+        line = reference_element.ufc_simplex(1)
+        lr = FIAT.quadrature.GaussLobattoLegendreQuadratureLineRule
+        point_set = FIAT.recursive_points.RecursivePointSet(lambda n: lr(line, n+1).get_points() if n else None)
+        points = point_set.recursive_points(ref_el.get_vertices(), degree*5)
+        phi = tabulate(degree, mapping(points))
+        z = phi[(0,) * dim]
+        y = expansion_set.tabulate(degree, points)
+
+        x = numpy.array(points)
+        fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
+        ax.plot_trisurf(x[:, 0], x[:, 1], z[-1], linewidth=0.2, antialiased=True)
+        ax.plot_trisurf(x[:, 0], x[:, 1], y[-1], linewidth=0.2, antialiased=True)
+        plt.show()

FIAT/jacobi.py

@@ -90,7 +90,7 @@ def eval_jacobi_deriv_batch(a, b, n, xs):
     Returns a two-dimensional array of points, where the
     rows correspond to the Jacobi polynomials and the
     columns correspond to the points."""
-    results = numpy.zeros((n + 1, len(xs)), "d")
+    results = numpy.zeros((n + 1, len(xs)), xs.dtype)
     if n == 0:
         return results
     else: