diff --git a/python/dune/perftool/ufl/dimensionindex.py b/python/dune/perftool/ufl/dimensionindex.py
index e5fb3e9ffe92a02c0aab9563f3190b7f0467bf4f..62651b210884a167441152bd3c29376160499c78 100644
--- a/python/dune/perftool/ufl/dimensionindex.py
+++ b/python/dune/perftool/ufl/dimensionindex.py
@@ -9,6 +9,7 @@ class _DimensionIndexMapping(MultiFunction):
     def __call__(self, o):
         self.shape = 0
         self.dimension_index_dict = {}
+        self.component = None
         self.call(o)
 
         return self.dimension_index_dict
@@ -18,12 +19,15 @@ class _DimensionIndexMapping(MultiFunction):
             self.call(op)
 
     def function_view(self, o):
-        from ufl.functionview import select_subelement
-        subelement = select_subelement(o.ufl_operands[0].ufl_element(), o.ufl_operands[1])
-        self.shape = len(subelement.value_shape())
+        self.component = o.ufl_operands[1]
+        self.call(o.ufl_operands[0])
+        self.component = None
 
-    def coefficient(self, o):
-        self.shape = len(o.ufl_element().value_shape())
+    def argument(self, o):
+        if self.component:
+            from ufl.functionview import select_subelement
+            subelement = select_subelement(o.ufl_element(), self.component)
+            self.shape = len(subelement.value_shape())
 
     def indexed(self, o):
         self.call(o.ufl_operands[0])