Move FiniteElement.components to FunctionSpace.components (#307)

* missing reference_

* remove IPython embed

* remove components property

* ruff

* move components to function space

* import numpy as np

---------

Co-authored-by: Jørgen Schartum Dokken <jsd55@cam.ac.uk>

ufl/algorithms/expand_indices.py

@@ -49,7 +49,7 @@ def form_argument(self, x):
         if sh == ():
             return x
         else:
-            e = x.ufl_element()
+            space = x.ufl_function_space()
             r = len(sh)
 
             # Get component
@@ -58,8 +58,8 @@ def form_argument(self, x):
                 raise ValueError("Component size mismatch.")
 
             # Map it through an eventual symmetry mapping
-            if len(e.components) > 1:
-                c = min(i for i, j in e.components.items() if j == e.components[c])
+            if len(space.components) > 1:
+                c = min(i for i, j in space.components.items() if j == space.components[c])
                 if r != len(c):
                     raise ValueError("Component size mismatch after symmetry mapping.")
 

ufl/finiteelement.py

@@ -15,8 +15,6 @@
 import abc as _abc
 import typing as _typing
 
-import numpy as np
-
 from ufl.cell import Cell as _Cell
 from ufl.pullback import AbstractPullback as _AbstractPullback
 from ufl.pullback import IdentityPullback as _IdentityPullback
@@ -130,30 +128,6 @@ def _ufl_signature_data_(self) -> str:
         """Return UFL signature data."""
         return repr(self)
 
-    @property
-    def components(self) -> _typing.Dict[_typing.Tuple[int, ...], int]:
-        """Get the numbering of the components of the element.
-
-        Returns:
-            A map from the components of the values on a physical cell (eg (0, 1))
-            to flat component numbers on the reference cell (eg 1)
-        """
-        if isinstance(self.pullback, _SymmetricPullback):
-            return self.pullback._symmetry
-
-        if len(self.sub_elements) == 0:
-            return {(): 0}
-
-        components = {}
-        offset = 0
-        c_offset = 0
-        for e in self.sub_elements:
-            for i, j in enumerate(np.ndindex(e.reference_value_shape)):
-                components[(offset + i,)] = c_offset + e.components[j]
-            c_offset += max(e.components.values()) + 1
-            offset += e.value_size
-        return components
-
     @property
     def reference_value_size(self) -> int:
         """Return the integer product of the reference value shape."""

ufl/functionspace.py

@@ -11,6 +11,8 @@
 
 import typing
 
+import numpy as np
+
 from ufl.core.ufl_type import UFLObject
 from ufl.domain import join_domains
 from ufl.duals import is_dual, is_primal
@@ -60,6 +62,32 @@ def __init__(self, domain, element, label=""):
         self._ufl_domain = domain
         self._ufl_element = element
 
+    @property
+    def components(self) -> typing.Dict[typing.Tuple[int, ...], int]:
+        """Get the numbering of the components of the element of this space.
+
+        Returns:
+            A map from the components of the values on a physical cell (eg (0, 1))
+            to flat component numbers on the reference cell (eg 1)
+        """
+        from ufl.pullback import SymmetricPullback
+
+        if isinstance(self._ufl_element.pullback, SymmetricPullback):
+            return self._ufl_element.pullback._symmetry
+
+        if len(self._ufl_element.sub_elements) == 0:
+            return {(): 0}
+
+        components = {}
+        offset = 0
+        c_offset = 0
+        for s in self.ufl_sub_spaces():
+            for i, j in enumerate(np.ndindex(s.value_shape)):
+                components[(offset + i,)] = c_offset + s.components[j]
+            c_offset += max(s.components.values()) + 1
+            offset += s.value_size
+        return components
+
     def label(self):
         """Return label of boundary domains to differentiate restricted and unrestricted."""
         return self._label