Switch to FInAT

gem/__init__.py

@@ -1,3 +1,4 @@
 from __future__ import absolute_import, print_function, division
 
 from gem.gem import *  # noqa
+from gem.optimise import select_expression  # noqa

gem/gem.py

@@ -16,12 +16,10 @@
 
 from __future__ import absolute_import, print_function, division
 from six import with_metaclass
-from six.moves import map
+from six.moves import range, zip
 
 from abc import ABCMeta
-import collections
 from itertools import chain
-import numbers
 from operator import attrgetter
 
 import numpy
@@ -34,9 +32,9 @@
            'Variable', 'Sum', 'Product', 'Division', 'Power',
            'MathFunction', 'MinValue', 'MaxValue', 'Comparison',
            'LogicalNot', 'LogicalAnd', 'LogicalOr', 'Conditional',
-           'Index', 'VariableIndex', 'Indexed', 'FlexiblyIndexed',
-           'ComponentTensor', 'IndexSum', 'ListTensor', 'as_gem',
-           'partial_indexed', 'reshape']
+           'Index', 'VariableIndex', 'Indexed', 'ComponentTensor',
+           'IndexSum', 'ListTensor', 'Delta', 'index_sum',
+           'partial_indexed', 'reshape', 'view']
 
 
 class NodeMeta(type):
@@ -129,10 +127,6 @@ def value(self):
         assert not self.shape
         return 0.0
 
-    @property
-    def array(self):
-        return numpy.zeros(self.shape)
-
 
 class Identity(Constant):
     """Identity matrix"""
@@ -206,12 +200,15 @@ def __new__(cls, a, b):
         assert not a.shape
         assert not b.shape
 
-        # Zero folding
+        # Constant folding
         if isinstance(a, Zero):
             return b
         elif isinstance(b, Zero):
             return a
 
+        if isinstance(a, Constant) and isinstance(b, Constant):
+            return Literal(a.value + b.value)
+
         self = super(Sum, cls).__new__(cls)
         self.children = a, b
         return self
@@ -224,10 +221,18 @@ def __new__(cls, a, b):
         assert not a.shape
         assert not b.shape
 
-        # Zero folding
+        # Constant folding
         if isinstance(a, Zero) or isinstance(b, Zero):
             return Zero()
 
+        if a == one:
+            return b
+        if b == one:
+            return a
+
+        if isinstance(a, Constant) and isinstance(b, Constant):
+            return Literal(a.value * b.value)
+
         self = super(Product, cls).__new__(cls)
         self.children = a, b
         return self
@@ -240,12 +245,18 @@ def __new__(cls, a, b):
         assert not a.shape
         assert not b.shape
 
-        # Zero folding
+        # Constant folding
         if isinstance(b, Zero):
             raise ValueError("division by zero")
         if isinstance(a, Zero):
             return Zero()
 
+        if b == one:
+            return a
+
+        if isinstance(a, Constant) and isinstance(b, Constant):
+            return Literal(a.value / b.value)
+
         self = super(Division, cls).__new__(cls)
         self.children = a, b
         return self
@@ -264,7 +275,7 @@ def __new__(cls, base, exponent):
                 raise ValueError("cannot solve 0^0")
             return Zero()
         elif isinstance(exponent, Zero):
-            return Literal(1)
+            return one
 
         self = super(Power, cls).__new__(cls)
         self.children = base, exponent
@@ -438,6 +449,10 @@ def __new__(cls, aggregate, multiindex):
             if isinstance(index, Index):
                 index.set_extent(extent)
 
+        # Empty multiindex
+        if not multiindex:
+            return aggregate
+
         # Zero folding
         if isinstance(aggregate, Zero):
             return Zero()
@@ -474,7 +489,7 @@ def __init__(self, variable, dim2idxs):
 
         For example, if ``variable`` is rank two, and ``dim2idxs`` is
 
-            ((1, ((i, 2), (j, 3), (k, 4))), (0, ()))
+            ((1, ((i, 12), (j, 4), (k, 1))), (0, ()))
 
         then this corresponds to the indexing:
 
@@ -484,31 +499,32 @@ def __init__(self, variable, dim2idxs):
         assert isinstance(variable, Variable)
         assert len(variable.shape) == len(dim2idxs)
 
-        indices = []
+        dim2idxs_ = []
+        free_indices = []
         for dim, (offset, idxs) in zip(variable.shape, dim2idxs):
-            strides = []
+            offset_ = offset
+            idxs_ = []
+            last = 0
             for idx in idxs:
                 index, stride = idx
-                strides.append(stride)
-
                 if isinstance(index, Index):
-                    if index.extent is None:
-                        index.set_extent(stride)
-                    elif not (index.extent <= stride):
-                        raise ValueError("Index extent cannot exceed stride")
-                    indices.append(index)
+                    assert index.extent is not None
+                    free_indices.append(index)
+                    idxs_.append((index, stride))
+                    last += (index.extent - 1) * stride
                 elif isinstance(index, int):
-                    if not (index <= stride):
-                        raise ValueError("Index cannot exceed stride")
+                    offset_ += index * stride
                 else:
                     raise ValueError("Unexpected index type for flexible indexing")
 
-            if dim is not None and offset + numpy.prod(strides) > dim:
+            if dim is not None and offset_ + last >= dim:
                 raise ValueError("Offset {0} and indices {1} exceed dimension {2}".format(offset, idxs, dim))
 
+            dim2idxs_.append((offset_, tuple(idxs_)))
+
         self.children = (variable,)
-        self.dim2idxs = dim2idxs
-        self.free_indices = unique(indices)
+        self.dim2idxs = tuple(dim2idxs_)
+        self.free_indices = unique(free_indices)
 
 
 class ComponentTensor(Node):
@@ -543,26 +559,36 @@ def __new__(cls, expression, multiindex):
 
 
 class IndexSum(Scalar):
-    __slots__ = ('children', 'index')
-    __back__ = ('index',)
+    __slots__ = ('children', 'multiindex')
+    __back__ = ('multiindex',)
 
-    def __new__(cls, summand, index):
+    def __new__(cls, summand, multiindex):
         # Sum zeros
         assert not summand.shape
         if isinstance(summand, Zero):
             return summand
 
-        # Sum a single expression
-        if index.extent == 1:
-            return Indexed(ComponentTensor(summand, (index,)), (0,))
+        # Unroll singleton sums
+        unroll = tuple(index for index in multiindex if index.extent <= 1)
+        if unroll:
+            assert numpy.prod([index.extent for index in unroll]) == 1
+            summand = Indexed(ComponentTensor(summand, unroll),
+                              (0,) * len(unroll))
+            multiindex = tuple(index for index in multiindex
+                               if index not in unroll)
+
+        # No indices case
+        multiindex = tuple(multiindex)
+        if not multiindex:
+            return summand
 
         self = super(IndexSum, cls).__new__(cls)
         self.children = (summand,)
-        self.index = index
+        self.multiindex = multiindex
 
         # Collect shape and free indices
-        assert index in summand.free_indices
-        self.free_indices = unique(set(summand.free_indices) - {index})
+        assert set(multiindex) <= set(summand.free_indices)
+        self.free_indices = unique(set(summand.free_indices) - set(multiindex))
 
         return self
 
@@ -621,6 +647,31 @@ def get_hash(self):
         return hash((type(self), self.shape, self.children))
 
 
+class Delta(Scalar, Terminal):
+    __slots__ = ('i', 'j')
+    __front__ = ('i', 'j')
+
+    def __new__(cls, i, j):
+        assert isinstance(i, IndexBase)
+        assert isinstance(j, IndexBase)
+
+        # \delta_{i,i} = 1
+        if i == j:
+            return one
+
+        # Fixed indices
+        if isinstance(i, int) and isinstance(j, int):
+            return Literal(int(i == j))
+
+        self = super(Delta, cls).__new__(cls)
+        self.i = i
+        self.j = j
+        # Set up free indices
+        free_indices = tuple(index for index in (i, j) if isinstance(index, Index))
+        self.free_indices = tuple(unique(free_indices))
+        return self
+
+
 def unique(indices):
     """Sorts free indices and eliminates duplicates.
 
@@ -630,18 +681,13 @@ def unique(indices):
     return tuple(sorted(set(indices), key=id))
 
 
-def as_gem(expr):
-    """Cast an expression to GEM."""
-    if isinstance(expr, Node):
-        return expr
-    elif isinstance(expr, numbers.Number):
-        return Literal(expr)
-    elif isinstance(expr, numpy.ndarray) and expr.dtype != object:
-        return Literal(expr)
-    elif isinstance(expr, collections.Iterable):
-        return ListTensor(list(map(as_gem, expr)))
-    else:
-        raise ValueError("Cannot cast %s to GEM." % (expr,))
+def index_sum(expression, indices):
+    """Eliminates indices from the free indices of an expression by
+    summing over them.  Skips any index that is not a free index of
+    the expression."""
+    multiindex = tuple(index for index in indices
+                       if index in expression.free_indices)
+    return IndexSum(expression, multiindex)
 
 
 def partial_indexed(tensor, indices):
@@ -667,20 +713,103 @@ def partial_indexed(tensor, indices):
         raise ValueError("More indices than rank!")
 
 
-def reshape(variable, *shapes):
+def strides_of(shape):
+    """Calculate cumulative strides from per-dimension capacities.
+
+    For example:
+
+        [2, 3, 4] ==> [12, 4, 1]
+
+    """
+    temp = numpy.flipud(numpy.cumprod(numpy.flipud(list(shape)[1:])))
+    return list(temp) + [1]
+
+
+def decompose_variable_view(expression):
+    """Extract ComponentTensor + FlexiblyIndexed view onto a variable."""
+    if isinstance(expression, Variable):
+        variable = expression
+        indexes = tuple(Index(extent=extent) for extent in expression.shape)
+        dim2idxs = tuple((0, ((index, 1),)) for index in indexes)
+    elif isinstance(expression, ComponentTensor) and isinstance(expression.children[0], FlexiblyIndexed):
+        variable = expression.children[0].children[0]
+        indexes = expression.multiindex
+        dim2idxs = expression.children[0].dim2idxs
+    else:
+        raise ValueError("Cannot handle {} objects.".format(type(expression).__name__))
+
+    return variable, dim2idxs, indexes
+
+
+def reshape(expression, *shapes):
     """Reshape a variable (splitting indices only).
 
-    :arg variable: a :py:class:`Variable`
+    :arg expression: view of a :py:class:`Variable`
     :arg shapes: one shape tuple for each dimension of the variable.
     """
-    dim2idxs = []
+    variable, dim2idxs, indexes = decompose_variable_view(expression)
+    assert len(indexes) == len(shapes)
+    shape_of = dict(zip(indexes, shapes))
+
+    dim2idxs_ = []
     indices = []
-    for shape in shapes:
-        idxs = []
-        for e in shape:
-            i = Index(extent=e)
-            idxs.append((i, e))
-            indices.append(i)
-        dim2idxs.append((0, tuple(idxs)))
-    expr = FlexiblyIndexed(variable, tuple(dim2idxs))
+    for offset, idxs in dim2idxs:
+        idxs_ = []
+        for idx in idxs:
+            index, stride = idx
+            assert isinstance(index, Index)
+            dim = index.extent
+            shape = shape_of[index]
+            if dim is not None and numpy.prod(shape) != dim:
+                raise ValueError("Shape {} does not match extent {}.".format(shape, dim))
+            strides = strides_of(shape)
+            for extent, stride_ in zip(shape, strides):
+                index = Index(extent=extent)
+                idxs_.append((index, stride_ * stride))
+                indices.append(index)
+        dim2idxs_.append((offset, tuple(idxs_)))
+
+    expr = FlexiblyIndexed(variable, tuple(dim2idxs_))
     return ComponentTensor(expr, tuple(indices))
+
+
+def view(expression, *slices):
+    """View a part of a variable.
+
+    :arg expression: view of a :py:class:`Variable`
+    :arg slices: one slice object for each dimension of the variable.
+    """
+    variable, dim2idxs, indexes = decompose_variable_view(expression)
+    assert len(indexes) == len(slices)
+    slice_of = dict(zip(indexes, slices))
+
+    dim2idxs_ = []
+    indices = []
+    for offset, idxs in dim2idxs:
+        offset_ = offset
+        idxs_ = []
+        for idx in idxs:
+            index, stride = idx
+            assert isinstance(index, Index)
+            dim = index.extent
+            s = slice_of[index]
+            start = s.start or 0
+            stop = s.stop or dim
+            if stop is None:
+                raise ValueError("Unknown extent!")
+            if dim is not None and stop > dim:
+                raise ValueError("Slice exceeds dimension extent!")
+            step = s.step or 1
+            offset_ += start * stride
+            extent = 1 + (stop - start - 1) // step
+            index_ = Index(extent=extent)
+            indices.append(index_)
+            idxs_.append((index_, step * stride))
+        dim2idxs_.append((offset_, tuple(idxs_)))
+
+    expr = FlexiblyIndexed(variable, tuple(dim2idxs_))
+    return ComponentTensor(expr, tuple(indices))
+
+
+# Static one object for quicker constant folding
+one = Literal(1)