Remove functional.index_iterator

FIAT/barycentric_interpolation.py

@@ -8,7 +8,6 @@
 
 import numpy
 from FIAT import reference_element, expansions, polynomial_set
-from FIAT.functional import index_iterator
 
 
 def get_lagrange_points(nodes):
@@ -94,33 +93,28 @@ def _tabulate_on_cell(self, n, pts, order=0, cell=0, direction=None):
 
 class LagrangePolynomialSet(polynomial_set.PolynomialSet):
 
-    def __init__(self, ref_el, pts, shape=tuple()):
+    def __init__(self, ref_el, pts, shape=()):
         if ref_el.get_shape() != reference_element.LINE:
             raise ValueError("Invalid reference element type.")
 
         expansion_set = LagrangeLineExpansionSet(ref_el, pts)
         degree = expansion_set.degree
-        if shape == tuple():
-            num_components = 1
-        else:
-            flat_shape = numpy.ravel(shape)
-            num_components = numpy.prod(flat_shape)
+        num_components = numpy.prod(shape, dtype=int)
         num_exp_functions = expansion_set.get_num_members(degree)
         num_members = num_components * num_exp_functions
         embedded_degree = degree
 
         # set up coefficients
-        if shape == tuple():
+        if shape == ():
             coeffs = numpy.eye(num_members, dtype="d")
         else:
             coeffs_shape = (num_members, *shape, num_exp_functions)
             coeffs = numpy.zeros(coeffs_shape, "d")
-            # use functional's index_iterator function
-            cur_bf = 0
-            for idx in index_iterator(shape):
-                for exp_bf in range(num_exp_functions):
-                    cur_idx = (cur_bf, *idx, exp_bf)
-                    coeffs[cur_idx] = 1.0
-                    cur_bf += 1
+            cur = 0
+            exp_bf = range(num_exp_functions)
+            for idx in numpy.ndindex(shape):
+                cur_bf = range(cur, cur+num_exp_functions)
+                coeffs[(cur_bf, *idx, exp_bf)] = 1.0
+                cur += num_exp_functions
 
         super().__init__(ref_el, degree, embedded_degree, expansion_set, coeffs)

FIAT/functional.py

@@ -18,13 +18,6 @@
 from FIAT import polynomial_set, jacobi, quadrature_schemes
 
 
-def index_iterator(shp):
-    """Constructs a generator iterating over all indices in
-    shp in generalized column-major order  So if shp = (2,2), then we
-    construct the sequence (0,0),(0,1),(1,0),(1,1)"""
-    return numpy.ndindex(shp)
-
-
 class Functional(object):
     r"""Abstract class representing a linear functional.
     All FIAT functionals are discrete in the sense that
@@ -384,7 +377,7 @@ def __init__(self, ref_el, Q, f_at_qpts, nm=None):
         self.f_at_qpts = f_at_qpts
         qpts, qwts = Q.get_points(), Q.get_weights()
         weights = numpy.transpose(numpy.multiply(f_at_qpts, qwts), (-1,) + tuple(range(len(shp))))
-        alphas = list(index_iterator(shp))
+        alphas = list(numpy.ndindex(shp))
 
         pt_dict = {tuple(pt): [(wt[alpha], alpha) for alpha in alphas] for pt, wt in zip(qpts, weights)}
         Functional.__init__(self, ref_el, shp, pt_dict, {}, nm or "FrobeniusIntegralMoment")
@@ -494,7 +487,7 @@ def __init__(self, ref_el, Q, f_at_qpts):
         weights = numpy.multiply(f_at_qpts, Q.get_weights()).T
 
         alphas = tuple(map(tuple, numpy.eye(sd, dtype=int)))
-        dpt_dict = {tuple(pt): [(wt[i], alphas[j], (i, j)) for i, j in index_iterator(shp)]
+        dpt_dict = {tuple(pt): [(wt[i], alphas[j], (i, j)) for i, j in numpy.ndindex(shp)]
                     for pt, wt in zip(points, weights)}
 
         super().__init__(ref_el, tuple(), {}, dpt_dict, "IntegralMomentOfDivergence")
@@ -655,7 +648,7 @@ def __init__(self, ref_el, v, w, pt):
         wvT = numpy.outer(w, v)
         shp = wvT.shape
 
-        pt_dict = {tuple(pt): [(wvT[idx], idx) for idx in index_iterator(shp)]}
+        pt_dict = {tuple(pt): [(wvT[idx], idx) for idx in numpy.ndindex(shp)]}
 
         super().__init__(ref_el, shp, pt_dict, {}, "PointwiseInnerProductEval")
 

FIAT/polynomial_set.py

@@ -18,7 +18,6 @@
 import numpy
 from itertools import chain
 from FIAT import expansions
-from FIAT.functional import index_iterator
 
 
 def mis(m, n):
@@ -113,25 +112,21 @@ class ONPolynomialSet(PolynomialSet):
     identity matrix of coefficients.  Can be used to specify ON bases
     for vector- and tensor-valued sets as well.
     """
-    def __init__(self, ref_el, degree, shape=tuple(), **kwargs):
+    def __init__(self, ref_el, degree, shape=(), **kwargs):
         expansion_set = expansions.ExpansionSet(ref_el, **kwargs)
-        if shape == tuple():
-            num_components = 1
-        else:
-            flat_shape = numpy.ravel(shape)
-            num_components = numpy.prod(flat_shape)
+        num_components = numpy.prod(shape, dtype=int)
         num_exp_functions = expansion_set.get_num_members(degree)
         num_members = num_components * num_exp_functions
         embedded_degree = degree
 
         # set up coefficients
-        if shape == tuple():
+        if shape == ():
             coeffs = numpy.eye(num_members)
         else:
             coeffs = numpy.zeros((num_members, *shape, num_exp_functions))
             cur = 0
             exp_bf = range(num_exp_functions)
-            for idx in index_iterator(shape):
+            for idx in numpy.ndindex(shape):
                 cur_bf = range(cur, cur+num_exp_functions)
                 coeffs[(cur_bf, *idx, exp_bf)] = 1.0
                 cur += num_exp_functions
@@ -243,7 +238,7 @@ def __init__(self, ref_el, degree, size=None, **kwargs):
         coeffs = numpy.zeros((num_members, *shape, num_exp_functions))
         cur = 0
         exp_bf = range(num_exp_functions)
-        for i, j in index_iterator(shape):
+        for i, j in numpy.ndindex(shape):
             if i > j:
                 continue
             cur_bf = range(cur, cur+num_exp_functions)
@@ -275,7 +270,7 @@ def __init__(self, ref_el, degree, size=None, **kwargs):
         coeffs = numpy.zeros((num_members, *shape, num_exp_functions))
         cur = 0
         exp_bf = range(num_exp_functions)
-        for i, j in index_iterator(shape):
+        for i, j in numpy.ndindex(shape):
             if i == size-1 and j == size-1:
                 continue
             cur_bf = range(cur, cur+num_exp_functions)