Fix regression in FiniteElement Python wrapper (#3552)

* Fix documentation in finite element wrapper and inverse permutation

* Further typing fixes

* Add dof inverse permutation check

Author: jorgensd

python/dolfinx/fem/element.py

@@ -263,13 +263,15 @@ def signature(self) -> str:
         """String identifying the finite element."""
         return self._cpp_object.signature
 
-    def T_apply(self, x: npt.NDArray[np.floating], cell_permutations: np.int32, dim: int) -> None:
+    def T_apply(
+        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
+    ) -> None:
         """Transform basis functions from the reference element ordering and orientation to the
         globally consistent physical element ordering and orientation.
 
         Args:
-            x: Data to transform (in place). The shape is ``(m, n)``, where `m` is the number of
-            dgerees-of-freedom and the storage is row-major.
+            x: Data to transform (in place). The shape is ``(num_cells, n, dim)``, where ``n`` is
+                the number degrees-of-freedom and the data is flattened (row-major).
             cell_permutations: Permutation data for the cell.
             dim: Number of columns in ``data``.
 
@@ -279,37 +281,31 @@ def T_apply(self, x: npt.NDArray[np.floating], cell_permutations: np.int32, dim:
         """
         self._cpp_object.T_apply(x, cell_permutations, dim)
 
-    def Tt_apply(self, x: npt.NDArray[np.floating], cell_permutations: np.int32, dim: int) -> None:
+    def Tt_apply(
+        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
+    ) -> None:
         """Apply the transpose of the operator applied by T_apply().
 
         Args:
-            x: Data to transform (in place). The shape is ``(m, n)``, where `m` is the number of
-            dgerees-of-freedom and the storage is row-major.
-            cell_permutations: Permutation data for the cell.
-            dim: Number of columns in `data`.
-
-        Note:
-            Exposed for testing. Function is not vectorised across multiple cells. Please see
-            `basix.numba_helpers` for performant versions.
+            x: Data to transform (in place). The shape is ``(num_cells, n, dim)``, where ``n`` is
+                the number degrees-of-freedom and the data is flattened (row-major).
+            cell_permutations: Permutation data for the cells
+            dim: Number of columns in ``data``.
         """
         self._cpp_object.Tt_apply(x, cell_permutations, dim)
 
     def Tt_inv_apply(
-        self, x: npt.NDArray[np.floating], cell_permutations: np.int32, dim: int
+        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Apply the inverse transpose of the operator applied by T_apply().
 
         Args:
-            x: Data to transform (in place). The shape is ``(m, n)``, where ``m`` is the number of
-            dgerees-of-freedom and the storage is row-major.
-            cell_permutations: Permutation data for the cell.
-            dim: Number of columns in `data`.
-
-        Note:
-            Exposed for testing. Function is not vectorised across multiple cells. Please see
-            ``basix.numba_helpers`` for performant versions.
+            x: Data to transform (in place). The shape is ``(num_cells, n, dim)``, where ``n`` is
+                the number degrees-of-freedom and the data is flattened (row-major).
+            cell_permutations: Permutation data for the cells
+            dim: Number of columns in ``data``.
         """
-        self._cpp_object.Tt_apply(x, cell_permutations, dim)
+        self._cpp_object.Tt_inv_apply(x, cell_permutations, dim)
 
 
 def finiteelement(

python/test/unit/fem/test_dof_permuting.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2009-2020 Garth N. Wells, Matthew W. Scroggs and Jorgen S. Dokken
+# Copyright (C) 2009-2024 Garth N. Wells, Matthew W. Scroggs and Jorgen S. Dokken
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -16,7 +16,7 @@
 from basix.ufl import element
 from dolfinx import default_real_type
 from dolfinx.fem import Function, assemble_scalar, form, functionspace
-from dolfinx.mesh import create_mesh
+from dolfinx.mesh import create_mesh, create_unit_cube
 
 
 def randomly_ordered_mesh(cell_type):
@@ -392,3 +392,31 @@ def tangent2(x):
 
             value = assemble_scalar(form(_form))
             assert np.isclose(value, 0.0, rtol=1.0e-6, atol=1.0e-6)
+
+
+@pytest.mark.parametrize(
+    "data_types",
+    [
+        (np.float32, np.float32),
+        (np.float64, np.float64),
+        (np.float32, np.complex64),
+        (np.float64, np.complex128),
+    ],
+)
+@pytest.mark.parametrize("space_order", range(3, 5))
+def test_permutation_wrappers(space_order, data_types):
+    s_type, d_type = data_types
+    domain = create_unit_cube(MPI.COMM_WORLD, 5, 3, 4, dtype=s_type)
+    V = functionspace(domain, ("N1curl", space_order))
+    u = Function(V, name="u", dtype=d_type)
+    u.interpolate(lambda x: np.array([x[0], x[1], x[2]]))
+
+    arr = u.x.array[V.dofmap.list]
+
+    domain.topology.create_entity_permutations()
+    cell_perm = domain.topology.get_cell_permutation_info()
+    org_data = arr.copy()
+    V.element.Tt_apply(arr.reshape(-1), cell_perm, 1)
+    V.element.Tt_inv_apply(arr.reshape(-1), cell_perm, 1)
+    eps = 100 * np.finfo(s_type).eps
+    np.testing.assert_allclose(org_data.reshape(-1), arr.reshape(-1), atol=eps)