Ruff formatting

Author: jorgensd

python/basix/ufl.py

@@ -194,9 +194,7 @@ def sub_elements(self) -> list[_AbstractFiniteElement]:
 
     # Basix specific functions
     @_abstractmethod
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -208,9 +206,7 @@ def tabulate(
         """
 
     @_abstractmethod
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_typing.Any, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_typing.Any, int, int]:
         """Get element that represents a component, and the offset and stride of the component.
 
         For example, for a mixed element, this will return the
@@ -431,9 +427,7 @@ def basix_hash(self) -> _typing.Optional[int]:
         """Return the hash of the Basix element if this is a standard Basix element."""
         return self._element.hash()
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -448,9 +442,7 @@ def tabulate(
         # TODO: update FFCx to remove the need for transposing here
         return tab.transpose((0, 1, 3, 2)).reshape((tab.shape[0], tab.shape[1], -1))  # type: ignore
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -678,9 +670,7 @@ def __hash__(self) -> int:
         """Return a hash."""
         return super().__hash__()
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -700,24 +690,17 @@ def tabulate(
             elif len(self._element._reference_value_shape) == 1:
                 output.append(tbl[:, self._component, :])
             elif len(self._element._reference_value_shape) == 2:
-                if (
-                    isinstance(self._element, _BlockedElement)
-                    and self._element._has_symmetry
-                ):
+                if isinstance(self._element, _BlockedElement) and self._element._has_symmetry:
                     # FIXME: check that this behaves as expected
                     output.append(tbl[:, self._component, :])
                 else:
                     vs0 = self._element._reference_value_shape[0]
-                    output.append(
-                        tbl[:, self._component // vs0, self._component % vs0, :]
-                    )
+                    output.append(tbl[:, self._component // vs0, self._component % vs0, :])
             else:
                 raise NotImplementedError()
         return np.asarray(output, dtype=np.float64)
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -897,9 +880,7 @@ def __init__(self, sub_elements: list[_ElementBase]):
 
     def __eq__(self, other) -> bool:
         """Check if two elements are equal."""
-        if isinstance(other, _MixedElement) and len(self._sub_elements) == len(
-            other._sub_elements
-        ):
+        if isinstance(other, _MixedElement) and len(self._sub_elements) == len(other._sub_elements):
             for i, j in zip(self._sub_elements, other._sub_elements):
                 if i != j:
                     return False
@@ -925,9 +906,7 @@ def degree(self) -> int:
         """Degree of the element."""
         return max((e.degree for e in self._sub_elements), default=-1)
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -951,9 +930,7 @@ def tabulate(
             tables.append(new_table)
         return np.asarray(tables, dtype=np.float64)
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -1044,10 +1021,7 @@ def num_entity_dofs(self) -> list[list[int]]:
         """Number of DOFs associated with each entity."""
         data = [e.num_entity_dofs for e in self._sub_elements]
         return [
-            [
-                sum(d[tdim][entity_n] for d in data)
-                for entity_n, _ in enumerate(entities)
-            ]
+            [sum(d[tdim][entity_n] for d in data) for entity_n, _ in enumerate(entities)]
             for tdim, entities in enumerate(data[0])
         ]
 
@@ -1070,19 +1044,15 @@ def num_entity_closure_dofs(self) -> list[list[int]]:
         """Number of DOFs associated with the closure of each entity."""
         data = [e.num_entity_closure_dofs for e in self._sub_elements]
         return [
-            [
-                sum(d[tdim][entity_n] for d in data)
-                for entity_n, _ in enumerate(entities)
-            ]
+            [sum(d[tdim][entity_n] for d in data) for entity_n, _ in enumerate(entities)]
             for tdim, entities in enumerate(data[0])
         ]
 
     @property
     def entity_closure_dofs(self) -> list[list[list[int]]]:
         """DOF numbers associated with the closure of each entity."""
         dofs: list[list[list[int]]] = [
-            [[] for i in entities]
-            for entities in self._sub_elements[0].entity_closure_dofs
+            [[] for i in entities] for entities in self._sub_elements[0].entity_closure_dofs
         ]
         start_dof = 0
         for e in self._sub_elements:
@@ -1160,9 +1130,7 @@ def custom_quadrature(
                     custom_q = e.custom_quadrature()
                 else:
                     p, w = e.custom_quadrature()
-                    if not np.allclose(p, custom_q[0]) or not np.allclose(
-                        w, custom_q[1]
-                    ):
+                    if not np.allclose(p, custom_q[0]) or not np.allclose(w, custom_q[1]):
                         raise ValueError(
                             "Subelements of mixed element use different quadrature rules"
                         )
@@ -1207,13 +1175,9 @@ def __init__(
             )
         if symmetry is not None:
             if len(shape) != 2:
-                raise ValueError(
-                    "symmetry argument can only be passed to elements of rank 2."
-                )
+                raise ValueError("symmetry argument can only be passed to elements of rank 2.")
             if shape[0] != shape[1]:
-                raise ValueError(
-                    "symmetry argument can only be passed to square shaped elements."
-                )
+                raise ValueError("symmetry argument can only be passed to square shaped elements.")
 
         if symmetry:
             block_size = shape[0] * (shape[0] + 1) // 2
@@ -1281,9 +1245,7 @@ def is_quadrature(self) -> bool:
         """Is this a quadrature element?"""
         return self._sub_element.is_quadrature
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -1295,19 +1257,15 @@ def tabulate(
 
         """
         assert len(self._block_shape) == 1  # TODO: block shape
-        assert (
-            self.reference_value_size == self._block_size
-        )  # TODO: remove this assumption
+        assert self.reference_value_size == self._block_size  # TODO: remove this assumption
         output = []
         for table in self._sub_element.tabulate(nderivs, points):  # type: ignore
             # Repeat sub element horizontally
             assert len(table.shape) == 2  # type: ignore
             new_table = np.zeros(
                 (table.shape[0], *self._block_shape, self._block_size * table.shape[1])  # type: ignore
             )
-            for i, j in enumerate(
-                _itertools.product(*[range(s) for s in self._block_shape])
-            ):
+            for i, j in enumerate(_itertools.product(*[range(s) for s in self._block_shape])):
                 if len(j) == 1:
                     new_table[:, j[0], i :: self._block_size] = table
                 elif len(j) == 2:
@@ -1317,9 +1275,7 @@ def tabulate(
             output.append(new_table)
         return np.asarray(output, dtype=np.float64)
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -1368,29 +1324,23 @@ def dim(self) -> int:
     @property
     def num_entity_dofs(self) -> list[list[int]]:
         """Number of DOFs associated with each entity."""
-        return [
-            [j * self._block_size for j in i] for i in self._sub_element.num_entity_dofs
-        ]
+        return [[j * self._block_size for j in i] for i in self._sub_element.num_entity_dofs]
 
     @property
     def entity_dofs(self) -> list[list[list[int]]]:
         """DOF numbers associated with each entity."""
         # TODO: should this return this, or should it take blocks into
         # account?
         return [
-            [
-                [k * self._block_size + b for k in j for b in range(self._block_size)]
-                for j in i
-            ]
+            [[k * self._block_size + b for k in j for b in range(self._block_size)] for j in i]
             for i in self._sub_element.entity_dofs
         ]
 
     @property
     def num_entity_closure_dofs(self) -> list[list[int]]:
         """Number of DOFs associated with the closure of each entity."""
         return [
-            [j * self._block_size for j in i]
-            for i in self._sub_element.num_entity_closure_dofs
+            [j * self._block_size for j in i] for i in self._sub_element.num_entity_closure_dofs
         ]
 
     @property
@@ -1399,10 +1349,7 @@ def entity_closure_dofs(self) -> list[list[list[int]]]:
         # TODO: should this return this, or should it take blocks into
         # account?
         return [
-            [
-                [k * self._block_size + b for k in j for b in range(self._block_size)]
-                for j in i
-            ]
+            [[k * self._block_size + b for k in j for b in range(self._block_size)] for j in i]
             for i in self._sub_element.entity_closure_dofs
         ]
 
@@ -1620,9 +1567,7 @@ def __hash__(self) -> int:
         """Return a hash."""
         return super().__hash__()
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -1637,14 +1582,10 @@ def tabulate(
 
         if points.shape != self._points.shape:
             raise ValueError("Mismatch of tabulation points and element points.")
-        tables = np.asarray(
-            [np.eye(points.shape[0], points.shape[0])], dtype=points.dtype
-        )
+        tables = np.asarray([np.eye(points.shape[0], points.shape[0])], dtype=points.dtype)
         return tables
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -1810,9 +1751,7 @@ def __init__(self, cell: _basix.CellType, value_shape: tuple[int, ...]):
         self._cell_type = cell
         tdim = len(_basix.topology(cell)) - 1
 
-        super().__init__(
-            f"RealElement({cell.name}, {value_shape})", cell.name, value_shape, 0
-        )
+        super().__init__(f"RealElement({cell.name}, {value_shape})", cell.name, value_shape, 0)
 
         self._entity_counts = []
         if tdim >= 1:
@@ -1840,9 +1779,7 @@ def dtype(self) -> _npt.DTypeLike:
         """Element float type."""
         raise NotImplementedError()
 
-    def tabulate(
-        self, nderivs: int, points: _npt.NDArray[np.floating]
-    ) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -1858,9 +1795,7 @@ def tabulate(
             out[0, :, self.reference_value_size * v + v] = 1.0
         return out
 
-    def get_component_element(
-        self, flat_component: int
-    ) -> tuple[_ElementBase, int, int]:
+    def get_component_element(self, flat_component: int) -> tuple[_ElementBase, int, int]:
         """Get element that represents a component.
 
         Element that represents a component of the element, and the
@@ -2160,9 +2095,7 @@ def enriched_element(
         map_type = elements[0].map_type
         for e in elements:
             if e.map_type != map_type:
-                raise ValueError(
-                    "Enriched elements on different map types not supported."
-                )
+                raise ValueError("Enriched elements on different map types not supported.")
 
     dtype = e.dtype
     hcd = min(e.embedded_subdegree for e in elements)
@@ -2175,13 +2108,9 @@ def enriched_element(
         if e.cell_type != ct:
             raise ValueError("Enriched elements on different cell types not supported.")
         if e.polyset_type != ptype:
-            raise ValueError(
-                "Enriched elements on different polyset types not supported."
-            )
+            raise ValueError("Enriched elements on different polyset types not supported.")
         if e.reference_value_shape != vshape or e.reference_value_size != vsize:
-            raise ValueError(
-                "Enriched elements on different value shapes not supported."
-            )
+            raise ValueError("Enriched elements on different value shapes not supported.")
         if e.dtype != dtype:
             raise ValueError("Enriched elements with different dtypes no supported.")
     nderivs = max(e.interpolation_nderivs for e in elements)
@@ -2200,9 +2129,7 @@ def enriched_element(
             pt = 0
             dof = 0
             for mat in M_parts:
-                new_M[
-                    dof : dof + mat.shape[0], :, pt : pt + mat.shape[2], : mat.shape[3]
-                ] = mat
+                new_M[dof : dof + mat.shape[0], :, pt : pt + mat.shape[2], : mat.shape[3]] = mat
                 dof += mat.shape[0]
                 pt += mat.shape[2]
             M_row.append(new_M)
@@ -2410,9 +2337,7 @@ def blocked_element(
         A blocked finite element.
     """
     if len(sub_element.reference_value_shape) != 0:
-        raise ValueError(
-            "Cannot create a blocked element containing a non-scalar element."
-        )
+        raise ValueError("Cannot create a blocked element containing a non-scalar element.")
 
     return _BlockedElement(sub_element, shape=shape, symmetry=symmetry)