better reindexing

firedrakeproject · Nov 6, 2023 · 890adb7 · 890adb7
1 parent f9f88a7
commit 890adb7
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/FIAT/expansions.py b/FIAT/expansions.py
@@ -270,10 +270,7 @@ def tabulate_jet(self, n, pts, order=1):
         for r in range(1, order+1):
             v = numpy.zeros((D,)*r + v0.shape, dtype=v0.dtype)
             for index in zip(*[range(D) for k in range(r)]):
-                alpha = [0] * D
-                for i in index:
-                    alpha[i] += 1
-                v[index] = vals[tuple(alpha)]
+                v[index] = vals[tuple(map(index.count, range(D)))]
             data.append(v.transpose((r, r+1) + tuple(range(r))))
         return data