address review comments

FIAT/mardal_tai_winther.py

@@ -66,9 +66,9 @@ def __init__(self, cell, degree):
         dofs.extend(_dofs)
 
         for entity_id in range(3):
-            dof_ids[1][entity_id] = dof_ids[1][entity_id] + _edge_dof_ids[entity_id]
+            dof_ids[1][entity_id].extend(_edge_dof_ids[entity_id])
 
-        dof_ids[2][0] = dof_ids[2][0] + _cell_dof_ids
+        dof_ids[2][0].extend(_cell_dof_ids)
 
         super(MardalTaiWintherDual, self).__init__(dofs, cell, dof_ids)
 
@@ -118,18 +118,21 @@ def _generate_constraint_dofs(cell, degree, offset):
         as described in the FIAT paper.
         """
         dofs = []
-
         edge_dof_ids = {}
+
+        start_order = 2
+        stop_order = 3
+        qdegree = degree + stop_order
         for entity_id in range(3):
-            dofs.append(IntegralLegendreNormalMoment(cell, entity_id, 2, degree+3))
-            dofs.append(IntegralLegendreNormalMoment(cell, entity_id, 3, degree+3))
+            cur = len(dofs)
+            dofs.extend(IntegralLegendreNormalMoment(cell, entity_id, order, qdegree)
+                        for order in range(start_order, stop_order+1))
 
-            edge_dof_ids[entity_id] = [offset, offset+1]
-            offset += 2
+            edge_dof_ids[entity_id] = list(range(offset+cur, offset+len(dofs)))
 
-        cell_dofs = DivergenceDubinerMoments(cell, 1, 2, 6)
-        dofs.extend(cell_dofs)
-        cell_dof_ids = list(range(offset, offset+len(cell_dofs)))
+        cur = len(dofs)
+        dofs.extend(DivergenceDubinerMoments(cell, start_order-1, stop_order-1, degree))
+        cell_dof_ids = list(range(offset+cur, offset+len(dofs)))
 
         return (dofs, edge_dof_ids, cell_dof_ids)