DO NOT MERGE

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)

finat/physically_mapped.py

@@ -1,4 +1,5 @@
 from abc import ABCMeta, abstractmethod
+from collections.abc import Mapping
 
 import gem
 import numpy
@@ -260,6 +261,46 @@ class NeedsCoordinateMappingElement(metaclass=ABCMeta):
     pass
 
 
+class MappedTabulation(Mapping):
+    """A lazy tabulation dict that applies the basis transformation only
+    on the requested derivatives."""
+
+    def __init__(self, M, ref_tabulation):
+        self.M = M
+        self.ref_tabulation = ref_tabulation
+        # we expect M to be sparse with O(1) nonzeros per row
+        # for each row, get the column index of each nonzero entry
+        csr = [[j for j in range(M.shape[1]) if not isinstance(M.array[i, j], gem.Zero)]
+               for i in range(M.shape[0])]
+        self.csr = csr
+        self._tabulation_cache = {}
+
+    def matvec(self, table):
+        # basis recombination using hand-rolled sparse-dense matrix multiplication
+        ii = gem.indices(len(table.shape)-1)
+        phi = [gem.Indexed(table, (j, *ii)) for j in range(self.M.shape[1])]
+        # the sum approach is faster than calling numpy.dot or gem.IndexSum
+        exprs = [gem.ComponentTensor(gem.Sum(*(self.M.array[i, j] * phi[j] for j in js)), ii)
+                 for i, js in enumerate(self.csr)]
+
+        val = gem.ListTensor(exprs)
+        # val = self.M @ table
+        return gem.optimise.aggressive_unroll(val)
+
+    def __getitem__(self, alpha):
+        try:
+            return self._tabulation_cache[alpha]
+        except KeyError:
+            result = self.matvec(self.ref_tabulation[alpha])
+            return self._tabulation_cache.setdefault(alpha, result)
+
+    def __iter__(self):
+        return iter(self.ref_tabulation)
+
+    def __len__(self):
+        return len(self.ref_tabulation)
+
+
 class PhysicallyMappedElement(NeedsCoordinateMappingElement):
     """A mixin that applies a "physical" transformation to tabulated
     basis functions."""
@@ -277,24 +318,10 @@ def basis_transformation(self, coordinate_mapping):
         :arg coordinate_mapping: Object providing physical geometry."""
         pass
 
-    def map_tabulation(self, tabulation, coordinate_mapping):
+    def map_tabulation(self, ref_tabulation, coordinate_mapping):
         assert coordinate_mapping is not None
-
         M = self.basis_transformation(coordinate_mapping)
-        # we expect M to be sparse with O(1) nonzeros per row
-        # for each row, get the column index of each nonzero entry
-        csr = [[j for j in range(M.shape[1]) if not isinstance(M.array[i, j], gem.Zero)]
-               for i in range(M.shape[0])]
-
-        def matvec(table):
-            # basis recombination using hand-rolled sparse-dense matrix multiplication
-            table = [gem.partial_indexed(table, (j,)) for j in range(M.shape[1])]
-            # the sum approach is faster than calling numpy.dot or gem.IndexSum
-            expressions = [sum(M.array[i, j] * table[j] for j in js) for i, js in enumerate(csr)]
-            val = gem.ListTensor(expressions)
-            return gem.optimise.aggressive_unroll(val)
-
-        return {alpha: matvec(tabulation[alpha]) for alpha in tabulation}
+        return MappedTabulation(M, ref_tabulation)
 
     def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         result = super().basis_evaluation(order, ps, entity=entity)

gem/coffee.py

@@ -4,8 +4,7 @@
 This file is NOT for code generation as a COFFEE AST.
 """
 
-from collections import OrderedDict
-import itertools
+from itertools import chain, repeat
 import logging
 
 import numpy
@@ -58,10 +57,10 @@ def find_optimal_atomics(monomials, linear_indices):
 
     :returns: list of atomic GEM expressions
     """
-    atomics = tuple(OrderedDict.fromkeys(itertools.chain(*(monomial.atomics for monomial in monomials))))
+    atomics = tuple(dict.fromkeys(chain.from_iterable(monomial.atomics for monomial in monomials)))
 
     def cost(solution):
-        extent = sum(map(lambda atomic: index_extent(atomic, linear_indices), solution))
+        extent = sum(map(index_extent, solution, repeat(linear_indices)))
         # Prefer shorter solutions, but larger extents
         return (len(solution), -extent)
 

gem/gem.py

@@ -15,7 +15,7 @@
 """
 
 from abc import ABCMeta
-from itertools import chain
+from itertools import chain, repeat
 from functools import reduce
 from operator import attrgetter
 from numbers import Integral, Number
@@ -55,8 +55,8 @@ def __call__(self, *args, **kwargs):
 
         # Set free_indices if not set already
         if not hasattr(obj, 'free_indices'):
-            obj.free_indices = unique(chain(*[c.free_indices
-                                              for c in obj.children]))
+            obj.free_indices = unique(chain.from_iterable(c.free_indices
+                                                          for c in obj.children))
         # Set dtype if not set already.
         if not hasattr(obj, 'dtype'):
             obj.dtype = obj.inherit_dtype_from_children(obj.children)
@@ -118,9 +118,8 @@ def __matmul__(self, other):
             raise ValueError(f"Mismatching shapes {self.shape} and {other.shape} in matmul")
         *i, k = indices(len(self.shape))
         _, *j = indices(len(other.shape))
-        expr = Product(Indexed(self, tuple(i) + (k, )),
-                       Indexed(other, (k, ) + tuple(j)))
-        return ComponentTensor(IndexSum(expr, (k, )), tuple(i) + tuple(j))
+        expr = Product(Indexed(self, (*i, k)), Indexed(other, (k, *j)))
+        return ComponentTensor(IndexSum(expr, (k, )), (*i, *j))
 
     def __rmatmul__(self, other):
         return as_gem(other).__matmul__(self)
@@ -341,7 +340,7 @@ def __new__(cls, *args):
             return a
 
         if isinstance(a, Constant) and isinstance(b, Constant):
-            return Literal(a.value + b.value, dtype=Node.inherit_dtype_from_children([a, b]))
+            return Literal(a.value + b.value, dtype=Node.inherit_dtype_from_children((a, b)))
 
         self = super(Sum, cls).__new__(cls)
         self.children = a, b
@@ -370,7 +369,7 @@ def __new__(cls, *args):
             return a
 
         if isinstance(a, Constant) and isinstance(b, Constant):
-            return Literal(a.value * b.value, dtype=Node.inherit_dtype_from_children([a, b]))
+            return Literal(a.value * b.value, dtype=Node.inherit_dtype_from_children((a, b)))
 
         self = super(Product, cls).__new__(cls)
         self.children = a, b
@@ -394,7 +393,7 @@ def __new__(cls, a, b):
             return a
 
         if isinstance(a, Constant) and isinstance(b, Constant):
-            return Literal(a.value / b.value, dtype=Node.inherit_dtype_from_children([a, b]))
+            return Literal(a.value / b.value, dtype=Node.inherit_dtype_from_children((a, b)))
 
         self = super(Division, cls).__new__(cls)
         self.children = a, b
@@ -407,7 +406,7 @@ class FloorDiv(Scalar):
     def __new__(cls, a, b):
         assert not a.shape
         assert not b.shape
-        dtype = Node.inherit_dtype_from_children([a, b])
+        dtype = Node.inherit_dtype_from_children((a, b))
         if dtype != uint_type:
             raise ValueError(f"dtype ({dtype}) != unit_type ({uint_type})")
         # Constant folding
@@ -430,7 +429,7 @@ class Remainder(Scalar):
     def __new__(cls, a, b):
         assert not a.shape
         assert not b.shape
-        dtype = Node.inherit_dtype_from_children([a, b])
+        dtype = Node.inherit_dtype_from_children((a, b))
         if dtype != uint_type:
             raise ValueError(f"dtype ({dtype}) != uint_type ({uint_type})")
         # Constant folding
@@ -453,7 +452,7 @@ class Power(Scalar):
     def __new__(cls, base, exponent):
         assert not base.shape
         assert not exponent.shape
-        dtype = Node.inherit_dtype_from_children([base, exponent])
+        dtype = Node.inherit_dtype_from_children((base, exponent))
 
         # Constant folding
         if isinstance(base, Zero):
@@ -569,7 +568,7 @@ def __new__(cls, condition, then, else_):
         self = super(Conditional, cls).__new__(cls)
         self.children = condition, then, else_
         self.shape = then.shape
-        self.dtype = Node.inherit_dtype_from_children([then, else_])
+        self.dtype = Node.inherit_dtype_from_children((then, else_))
         return self
 
 
@@ -932,7 +931,7 @@ def is_equal(self, other):
         """Common subexpression eliminating equality predicate."""
         if type(self) is not type(other):
             return False
-        if (self.array == other.array).all():
+        if numpy.array_equal(self.array, other.array):
             self.array = other.array
             return True
         return False
@@ -973,7 +972,7 @@ class Delta(Scalar, Terminal):
     def __new__(cls, i, j, dtype=None):
         if isinstance(i, tuple) and isinstance(j, tuple):
             # Handle multiindices
-            return Product(*map(Delta, i, j))
+            return Product(*map(Delta, i, j, repeat(dtype)))
         assert isinstance(i, IndexBase)
         assert isinstance(j, IndexBase)
 
@@ -985,26 +984,18 @@ def __new__(cls, i, j, dtype=None):
         if isinstance(i, Integral) and isinstance(j, Integral):
             return one if i == j else Zero()
 
-        if isinstance(i, Integral):
-            return Indexed(Literal(numpy.eye(j.extent)[i]), (j,))
-
-        if isinstance(j, Integral):
-            return Indexed(Literal(numpy.eye(i.extent)[j]), (i,))
-
         self = super(Delta, cls).__new__(cls)
         self.i = i
         self.j = j
         # Set up free indices
-        free_indices = []
-        for index in (i, j):
-            if isinstance(index, Index):
-                free_indices.append(index)
-            elif isinstance(index, VariableIndex):
-                raise NotImplementedError("Can not make Delta with VariableIndex")
+        free_indices = [index for index in (i, j) if isinstance(index, Index)]
         self.free_indices = tuple(unique(free_indices))
         self._dtype = dtype
         return self
 
+    def reconstruct(self, *args):
+        return Delta(*args, dtype=self.dtype)
+
 
 class Inverse(Node):
     """The inverse of a square matrix."""