simplify Indexed(ComponentTensor(Indexed(ListTensor(...), ...), ...), ...) on construction

Co-authored-by: Pablo Brubeck <[email protected]>

Author: ksagiyam

test/test_simplify.py

@@ -14,6 +14,7 @@
     acos,
     as_tensor,
     as_ufl,
+    as_vector,
     asin,
     atan,
     cos,
@@ -234,7 +235,7 @@ def test_untangle_indexed_component_tensor(self):
     A = as_tensor(Indexed(C, MultiIndex(kk)), jj)
     assert A is not C
 
-    ii = kk
+    ii = indices(len(A.ufl_shape))
     expr = Indexed(A, MultiIndex(ii))
     assert isinstance(expr, Indexed)
     B, ll = expr.ufl_operands
@@ -266,3 +267,24 @@ def test_simplify_indexed(self):
     # ComponentTensor + ListTensor
     c = ComponentTensor(Indexed(ll, MultiIndex((i, j))), MultiIndex((j, i)))
     assert Indexed(c, MultiIndex((FixedIndex(1), FixedIndex(2)))) == l2[1]
+
+
+def test_simplify_indexed_componenttensor_indexed_listtensor():
+    list_item_0 = as_vector([10.0, 11.0, 12.0])
+    list_item_1 = as_vector([20.0, 21.0, 22.0])
+    i = Index()
+    j = Index()
+    value = Indexed(
+        ComponentTensor(
+            Indexed(
+                ListTensor(
+                    Indexed(list_item_0, MultiIndex((j,))),
+                    Indexed(list_item_1, MultiIndex((j,))),
+                ),
+                MultiIndex((i,)),
+            ),
+            MultiIndex((i, j)),
+        ),
+        MultiIndex((FixedIndex(1), FixedIndex(2))),
+    )
+    assert value == list_item_1[2]

ufl/tensors.py

@@ -235,14 +235,24 @@ def _simplify_indexed(self, multiindex):
         """Return a simplified Expr used in the constructor of Indexed(self, multiindex)."""
         # Untangle as_tensor(C[kk], jj)[ii] -> C[ll]
         B, jj = self.ufl_operands
+        if len(multiindex) != len(jj):
+            raise ValueError(f"len(multiindex) ({len(multiindex)}) != len(jj) ({len(jj)})")
+        rep = dict(zip(jj, multiindex))
+        # Avoid recursion and just attempt to simplify some common patterns
+        # as the result of this method is not cached.
+        if isinstance(B, Indexed):
+            C, kk = B.ufl_operands
+            if isinstance(C, ListTensor) and len(kk) == 1 and isinstance(rep[kk[0]], FixedIndex):
+                (k,) = kk
+                B = C.ufl_operands[int(rep[k])]
+                jj = MultiIndex(tuple(j for j in jj if j != k))
+                multiindex = MultiIndex(tuple(rep[j] for j in jj))
+                rep = dict(zip(jj, multiindex))
         if isinstance(B, Indexed):
             C, kk = B.ufl_operands
             if all(j in kk for j in jj):
-                ii = tuple(multiindex)
-                rep = dict(zip(jj, ii))
                 Cind = tuple(rep.get(k, k) for k in kk)
                 return Indexed(C, MultiIndex(Cind))
-
         return Operator._simplify_indexed(self, multiindex)
 
     def indices(self):