Simplify nested Indexed objects

ufl/algorithms/formsplitter.py

@@ -74,7 +74,7 @@ def indexed(self, o, child, multiindex):
             if len(indices) == 1:
                 return child[indices[0]]
             elif len(indices) == len(child.ufl_operands) and all(
-                k == i for k, i in enumerate(indices)
+                k == int(i) for k, i in enumerate(indices)
             ):
                 return child
             else:

ufl/core/operator.py

@@ -39,9 +39,9 @@ def _ufl_signature_data_(self):
 
     def _ufl_compute_hash_(self):
         """Compute a hash code for this expression. Used by sets and dicts."""
-        return hash((self._ufl_typecode_,) + tuple(hash(o) for o in self.ufl_operands))
+        return hash((self._ufl_typecode_,) + tuple(map(hash, self.ufl_operands)))
 
     def __repr__(self):
         """Default repr string construction for operators."""
         # This should work for most cases
-        return f"{self._ufl_class_.__name__}({', '.join(repr(op) for op in self.ufl_operands)})"
+        return f"{self._ufl_class_.__name__}({', '.join(map(repr, self.ufl_operands))})"

ufl/corealg/multifunction.py

@@ -7,15 +7,17 @@
 #
 # Modified by Massimiliano Leoni, 2016
 
-import inspect
+from functools import cache
+from inspect import signature
 
 from ufl.core.expr import Expr
 from ufl.core.ufl_type import UFLType
 
 
+@cache
 def get_num_args(function):
     """Return the number of arguments accepted by *function*."""
-    sig = inspect.signature(function)
+    sig = signature(function)
     return len(sig.parameters) + 1
 
 

ufl/indexed.py

@@ -26,15 +26,30 @@ class Indexed(Operator):
 
     def __new__(cls, expression, multiindex):
         """Create a new Indexed."""
+        # cyclic import
         from ufl.tensors import ListTensor
 
+        simpler = False
         indices = multiindex.indices()
-        while isinstance(expression, ListTensor) and isinstance(indices[0], FixedIndex):
+        while (
+            len(indices) > 0
+            and isinstance(expression, ListTensor)
+            and isinstance(indices[0], FixedIndex)
+        ):
             # Simplify indexed ListTensor objects
             expression = expression[indices[0]]
             indices = indices[1:]
+            simpler = True
+
+        if isinstance(expression, Indexed):
+            # Simplify nested Indexed objects
+            indices = expression.ufl_operands[1].indices() + indices
+            expression = expression.ufl_operands[0]
+            simpler = True
 
-        if isinstance(expression, Zero):
+        if len(indices) == 0:
+            return expression
+        elif isinstance(expression, Zero):
             # Zero-simplify indexed Zero objects
             shape = expression.ufl_shape
             efi = expression.ufl_free_indices
@@ -49,8 +64,8 @@ def __new__(cls, expression, multiindex):
             else:
                 fi, fid = (), ()
             return Zero(shape=(), free_indices=fi, index_dimensions=fid)
-        elif expression.ufl_shape == () and multiindex == ():
-            return expression
+        elif simpler:
+            return Indexed(expression, MultiIndex(indices))
         else:
             return Operator.__new__(cls)
 

ufl/tensors.py

@@ -55,7 +55,7 @@ def __new__(cls, *expressions):
 
         # Simplify to Zero if possible
         if all(isinstance(e, Zero) for e in expressions):
-            shape = (len(expressions),) + sh
+            shape = (len(expressions), *sh)
             return Zero(shape, fi, fid)
 
         # Simplify [v[0], v[1], ..., v[k]] -> v
@@ -132,7 +132,7 @@ def substring(expressions, indent):
                 s = (",\n" + ind).join(substrings)
                 return "%s[\n%s%s\n%s]" % (ind, ind, s, ind)
             else:
-                s = ", ".join(str(e) for e in expressions)
+                s = ", ".join(map(str, expressions))
                 return "%s[%s]" % (ind, s)
 
         return substring(self.ufl_operands, 0)