Add typing generics

python/demo/demo_lagrange_variants.py

@@ -26,6 +26,7 @@
 from mpi4py import MPI
 
 import matplotlib.pylab as plt
+import numpy as np
 
 import basix
 import basix.ufl
@@ -142,11 +143,11 @@ def saw_tooth(x):
     uh.interpolate(lambda x: saw_tooth(x[0]))
     if MPI.COMM_WORLD.size == 1:  # Skip this plotting in parallel
         pts: list[list[float]] = []
-        cells: list[int] = []
+        cells = np.empty((0,), dtype=np.int32)
         for cell in range(N):
             for i in range(51):
                 pts.append([cell / N + i / 50 / N, 0, 0])
-                cells.append(cell)
+                cells = np.append(cells, [cell])
         values = uh.eval(pts, cells)
         plt.plot(pts, [saw_tooth(i[0]) for i in pts], "k--")
         plt.plot(pts, values, "r-")

python/dolfinx/__init__.py

@@ -29,7 +29,7 @@
 
 from dolfinx import common
 from dolfinx import cpp as _cpp
-from dolfinx import fem, geometry, graph, io, jit, la, log, mesh, nls, plot
+from dolfinx import fem, geometry, graph, io, jit, la, log, mesh, nls, plot, typing
 
 from dolfinx.common import (
     git_commit_hash,
@@ -79,6 +79,7 @@ def get_include(user=False):
     "mesh",
     "nls",
     "plot",
+    "typing",
     "git_commit_hash",
     "hardware_concurrency",
     "has_adios2",

python/dolfinx/fem/bcs.py

@@ -13,13 +13,15 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Iterable
+from typing import Generic
 
 import numpy as np
 import numpy.typing as npt
 
 import dolfinx
 from dolfinx import cpp as _cpp
 from dolfinx.fem.function import Constant, Function, FunctionSpace
+from dolfinx.typing import Scalar
 
 
 def locate_dofs_geometrical(
@@ -90,7 +92,7 @@
     return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)
 
 
-class DirichletBC:
+class DirichletBC(Generic[Scalar]):
     """Representation of Dirichlet boundary conditions.
 
     The conditions are imposed on a linear system.
@@ -119,8 +121,9 @@
         self._cpp_object = bc
 
     @property
-    def g(self) -> Function | Constant | np.ndarray:
+    def g(self) -> Function | Constant:
         """The boundary condition value(s)."""
+        # TODO: needs to be wrapped into Function or Constant
         return self._cpp_object.value
 
     @property
@@ -129,7 +132,7 @@
         return self._cpp_object.function_space
 
     def set(
-        self, x: npt.NDArray, x0: npt.NDArray[np.int32] | None = None, alpha: float = 1
+        self, x: npt.NDArray[Scalar], x0: npt.NDArray[Scalar] | None = None, alpha: float = 1
     ) -> None:
         """Set array entries that are constrained by a Dirichlet condition.
 
@@ -171,10 +174,10 @@
 
 
 def dirichletbc(
-    value: Function | Constant | np.ndarray | float | complex,
+    value: Function | Constant | npt.NDArray[Scalar] | float | complex,
     dofs: npt.NDArray[np.int32],
     V: dolfinx.fem.FunctionSpace | None = None,
-) -> DirichletBC:
+) -> DirichletBC[Scalar]:
     """Representation of Dirichlet boundary condition.
 
     Args:
@@ -231,8 +234,8 @@
 
 
 def bcs_by_block(
-    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC]
-) -> list[list[DirichletBC]]:
+    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC[Scalar]]
+) -> list[list[DirichletBC[Scalar]]]:
     """Arrange boundary conditions by the space that they constrain.
 
     Given a sequence of function spaces ``spaces`` and a sequence of

python/dolfinx/fem/element.py

@@ -6,16 +6,18 @@
 """Finite elements."""
 
 from functools import singledispatch
+from typing import Generic
 
 import numpy as np
 import numpy.typing as npt
 
 import basix
 import basix.ufl
 from dolfinx import cpp as _cpp
+from dolfinx.typing import Real
 
 
-class CoordinateElement:
+class CoordinateElement(Generic[Real]):
     """Coordinate element describing the geometry map for mesh cells."""
 
     _cpp_object: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64
@@ -61,10 +63,8 @@ def create_dof_layout(self) -> _cpp.fem.ElementDofLayout:
         return self._cpp_object.create_dof_layout()
 
     def push_forward(
-        self,
-        X: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
+        self, X: npt.NDArray[Real], cell_geometry: npt.NDArray[Real]
+    ) -> npt.NDArray[Real]:
         """Push points on the reference cell forward to the physical cell.
 
         Args:
@@ -82,11 +82,11 @@ def push_forward(
 
     def pull_back(
         self,
-        x: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
+        x: npt.NDArray[Real],
+        cell_geometry: npt.NDArray[Real],
         tol: float = 1.0e-6,
         maxit: int = 15,
-    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
+    ) -> npt.NDArray[Real]:
         """Pull points on the physical cell back to the reference cell.
 
         For non-affine cells, the pull-back is a nonlinear operation.
@@ -130,7 +130,7 @@ def coordinate_element(
     degree: int,
     variant=int(basix.LagrangeVariant.unset),
     dtype: npt.DTypeLike = np.float64,
-):
+) -> CoordinateElement:
     """Create a Lagrange CoordinateElement from element metadata.
 
     Coordinate elements are typically used to create meshes.
@@ -153,7 +153,7 @@ def coordinate_element(
 
 
 @coordinate_element.register(basix.finite_element.FiniteElement)
-def _(e: basix.finite_element.FiniteElement):
+def _(e: basix.finite_element.FiniteElement) -> CoordinateElement:
     """Create a Lagrange CoordinateElement from a Basix finite element.
 
     Coordinate elements are typically used when creating meshes.
@@ -170,7 +170,7 @@ def _(e: basix.finite_element.FiniteElement):
         return CoordinateElement(_cpp.fem.CoordinateElement_float64(e._e))
 
 
-class FiniteElement:
+class FiniteElement(Generic[Real]):
     """A finite element."""
 
     _cpp_object: _cpp.fem.FiniteElement_float32 | _cpp.fem.FiniteElement_float64
@@ -224,7 +224,7 @@ def value_shape(self) -> npt.NDArray[np.integer]:
         return self._cpp_object.value_shape
 
     @property
-    def interpolation_points(self) -> npt.NDArray[np.floating]:
+    def interpolation_points(self) -> npt.NDArray[Real]:
         """Points at which to evaluate the function to be interpolated.
 
         Interpolation point coordinates on the reference cell, returning
@@ -282,7 +282,7 @@ def signature(self) -> str:
         return self._cpp_object.signature
 
     def T_apply(
-        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
+        self, x: npt.NDArray[Real], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Transform basis from reference to physical ordering/orientation.
 
@@ -305,7 +305,7 @@ def T_apply(
         self._cpp_object.T_apply(x, cell_permutations, dim)
 
     def Tt_apply(
-        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
+        self, x: npt.NDArray[Real], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Apply the transpose of the operator applied by T_apply().
 
@@ -319,7 +319,7 @@ def Tt_apply(
         self._cpp_object.Tt_apply(x, cell_permutations, dim)
 
     def Tt_inv_apply(
-        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
+        self, x: npt.NDArray[Real], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Apply the inverse transpose of T_apply().
 
@@ -336,7 +336,7 @@ def Tt_inv_apply(
 def finiteelement(
     cell_type: _cpp.mesh.CellType,
     ufl_e: basix.ufl._ElementBase,
-    FiniteElement_dtype: np.dtype,
+    FiniteElement_dtype: npt.DTypeLike,
 ) -> FiniteElement:
     """Create a DOLFINx element from a basix.ufl element.
 

python/dolfinx/fem/forms.py

@@ -25,6 +25,7 @@
 from dolfinx import default_scalar_type, jit
 from dolfinx.fem import IntegralType
 from dolfinx.fem.function import Constant, Function, FunctionSpace
+from dolfinx.typing import Scalar
 
 if typing.TYPE_CHECKING:
     # import dolfinx.mesh just when doing type checking to avoid
@@ -33,7 +34,7 @@
     from dolfinx.mesh import Mesh, MeshTags
 
 
-class Form:
+class Form(typing.Generic[Scalar]):
     """A finite element form."""
 
     _cpp_object: (