From c5d3ada5c20b0f1fd8b7001b3890bad99f7ef8d3 Mon Sep 17 00:00:00 2001
From: David Ham <David.Ham@imperial.ac.uk>
Date: Tue, 26 Apr 2016 22:20:13 +0100
Subject: [PATCH 1/3] AffineIndex and IndexIterator

Implement affine index groups and an iterator for multi-indices over sets of Index and AffineIndex.
---
 gem/gem.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 68 insertions(+), 3 deletions(-)

diff --git a/gem/gem.py b/gem/gem.py
index afef65e0..33a25664 100644
--- a/gem/gem.py
+++ b/gem/gem.py
@@ -28,9 +28,9 @@
 __all__ = ['Node', 'Identity', 'Literal', 'Zero', 'Variable', 'Sum',
            'Product', 'Division', 'Power', 'MathFunction', 'MinValue',
            'MaxValue', 'Comparison', 'LogicalNot', 'LogicalAnd',
-           'LogicalOr', 'Conditional', 'Index', 'VariableIndex',
-           'Indexed', 'ComponentTensor', 'IndexSum', 'ListTensor',
-           'Delta', 'partial_indexed']
+           'LogicalOr', 'Conditional', 'Index', 'AffineIndex',
+           'VariableIndex', 'Indexed', 'ComponentTensor', 'IndexSum',
+           'ListTensor', 'Delta', 'IndexIterator', 'affine_index_group', 'partial_indexed']
 
 
 class NodeMeta(type):
@@ -379,6 +379,22 @@ def __repr__(self):
         return "Index(%r)" % self.name
 
 
+class AffineIndex(Index):
+    """An index in an affine_index_group. Do not instantiate directly but
+    instead call :func:`affine_index_group`."""
+    __slots__ = ('name', 'extent', 'count', 'group')
+
+    def __str__(self):
+        if self.name is None:
+            return "i_%d" % self.count
+        return self.name
+
+    def __repr__(self):
+        if self.name is None:
+            return "AffineIndex(%r)" % self.count
+        return "AffineIndex(%r)" % self.name
+
+
 class VariableIndex(IndexBase):
     """An index that is constant during a single execution of the
     kernel, but whose value is not known at compile time."""
@@ -566,6 +582,55 @@ def __new__(cls, i, j):
         self.free_indices = tuple(unique(free_indices))
         return self
 
+    
+class IndexIterator(object):
+    """An iterator whose value is a multi-index (tuple) iterating over the
+    extent of the supplied :class:`.Index` objects in a last index varies
+    fastest (ie 'c') ordering.
+
+    :arg *indices: the indices over whose extent to iterate."""
+    def __init__(self, *indices):
+
+        self.affine_groups = set()
+        for i in indices:
+            if isinstance(i, AffineIndex):
+                try:
+                    pos = tuple(indices.index(g) for g in i.group)
+                except ValueError:
+                    raise ValueError("Only able to iterate over all indices in an affine group at once")
+                self.affine_groups.add((i.group, pos))
+        
+        self.ndindex = numpy.ndindex(tuple(i.extent for i in indices))
+
+    def _affine_groups_legal(self, multiindex):
+        for group, pos in self.affine_groups:
+            if sum(multiindex[p] for p in pos) >= group[0].extent:
+                return False
+        return True
+        
+    def __iter__(self):
+        # Fix this for affine index groups.
+        while True:
+            multiindex = self.ndindex.next()
+            if self._affine_groups_legal(multiindex):
+                yield multiindex
+
+
+def affine_index_group(n, extent):
+    """A set of indices whose values are constrained to lie in a simplex
+    subset of the iteration space.
+
+    :arg n: the number of indices in the group.
+    :arg extent: sum(indices) < extent
+    """
+
+    group = tuple(AffineIndex(extent=extent) for i in range(n))
+
+    for i in range(n):
+        group[i].group = group
+
+    return group
+
 
 def partial_indexed(tensor, indices):
     """Generalised indexing into a tensor.  The number of indices may

From 800ddff6edb878fb0f465ae4e18b0b12a476fcf6 Mon Sep 17 00:00:00 2001
From: David Ham <david.ham@imperial.ac.uk>
Date: Wed, 27 Apr 2016 11:29:49 +0100
Subject: [PATCH 2/3] FInAT dependency

---
 requirements-git.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements-git.txt b/requirements-git.txt
index 89251708..53541e68 100644
--- a/requirements-git.txt
+++ b/requirements-git.txt
@@ -1,3 +1,4 @@
 git+https://github.com/coneoproject/COFFEE#egg=COFFEE
 git+https://github.com/firedrakeproject/ufl.git#egg=ufl
 git+https://github.com/firedrakeproject/fiat.git#egg=fiat
+git+https://github.com/firedrakeproject/FInAT.git#egg=finat

From a85e6ba9dbb3071e44ce8432653dbfef63df8c84 Mon Sep 17 00:00:00 2001
From: David Ham <david.ham@imperial.ac.uk>
Date: Wed, 27 Apr 2016 11:39:39 +0100
Subject: [PATCH 3/3] Slightly more ideomatic python

---
 gem/gem.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gem/gem.py b/gem/gem.py
index 33a25664..2345300d 100644
--- a/gem/gem.py
+++ b/gem/gem.py
@@ -626,8 +626,8 @@ def affine_index_group(n, extent):
 
     group = tuple(AffineIndex(extent=extent) for i in range(n))
 
-    for i in range(n):
-        group[i].group = group
+    for g in group:
+        g.group = group
 
     return group