Add typing generics

.github/workflows/ccpp.yml

@@ -7,8 +7,8 @@ on:
     tags:
       - "v*"
   pull_request:
-    branches:
-      - main
+    # branches:
+    #   - main
   merge_group:
     branches:
       - main
@@ -132,15 +132,9 @@ jobs:
           pip install --check-build-dependencies --no-build-isolation --config-settings=cmake.build-type="Developer" 'python/[test]'
           python -c "from mpi4py import MPI; import dolfinx; assert not dolfinx.has_petsc; assert not dolfinx.has_petsc4py; assert dolfinx.has_superlu_dist"
 
-      - name: Run mypy  # Use a venv to avoid NumPy upgrades that are incompatible with numba
+      - name: Run mypy
         working-directory: python
-        run: |
-          python -m venv mypy-env
-          . ./mypy-env/bin/activate
-          pip install mypy types-cffi scipy-stubs
-          mypy -p dolfinx
-        # mypy test
-        # mypy demo
+        run: make lint-mypy
 
       - name: Install gmsh and pyvista (and dependencies)
         run: |

python/Makefile

@@ -0,0 +1,14 @@
+# Note: mypy is currenlty on CI tested without petsc4py!
+lint-mypy:
+	python -m venv dolfinx-venv-mypy
+	. dolfinx-venv-mypy/bin/activate && pip install mypy types-cffi scipy-stubs
+	. dolfinx-venv-mypy/bin/activate && mypy -p dolfinx
+	. dolfinx-venv-mypy/bin/activate && mypy test
+	. dolfinx-venv-mypy/bin/activate && mypy demo
+
+lint-mypy-clean:
+	rm -rf dolfinx-venv-mypy
+
+lint: lint-mypy
+
+clean: lint-mypy-clean

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/demo/demo_pml.py

@@ -214,8 +214,8 @@ def generate_mesh_wire(
 
 def compute_a(nu: int, m: complex, alpha: float) -> float:
     """Compute the Mie coefficient a_nu for a cylinder."""
-    J_nu_alpha = jv(nu, alpha)
-    J_nu_malpha = jv(nu, m * alpha)
+    J_nu_alpha = jv(nu, alpha)  # type: ignore
+    J_nu_malpha = jv(nu, m * alpha)  # type: ignore
     J_nu_alpha_p = jvp(nu, alpha, 1)
     J_nu_malpha_p = jvp(nu, m * alpha, 1)
 

python/demo/demo_pyamg.py

@@ -63,7 +63,7 @@ def poisson_problem(dtype: npt.DTypeLike, solver_type: str) -> None:
         solver_type: pyamg solver type, either "ruge_stuben" or
             "smoothed_aggregation"
     """
-    real_type = np.real(dtype(0)).dtype
+    real_type = np.real(np.zeros(0, dtype=dtype)).dtype
     mesh = create_box(
         comm=MPI.COMM_WORLD,
         points=[(0.0, 0.0, 0.0), (3.0, 2.0, 1.0)],
@@ -84,7 +84,7 @@ def poisson_problem(dtype: npt.DTypeLike, solver_type: str) -> None:
 
     dofs = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets)
 
-    bc = dirichletbc(value=dtype(0), dofs=dofs, V=V)
+    bc = dirichletbc(value=dtype(0.0), dofs=dofs, V=V)  # type: ignore
 
     u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
     x = ufl.SpatialCoordinate(mesh)
@@ -110,14 +110,14 @@ def poisson_problem(dtype: npt.DTypeLike, solver_type: str) -> None:
     print(ml)
 
     # Solve linear systems
-    print(f"\nSolve Poisson equation: {dtype.__name__}")
+    print(f"\nSolve Poisson equation: {dtype!s}")
     res: list[float] = []
     tol = 1e-10 if real_type == np.float64 else 1e-6
     uh.x.array[:] = ml.solve(b.array, tol=tol, residuals=res, accel="cg")
     for i, q in enumerate(res):
         print(f"Convergence history: iteration {i}, residual= {q}")
 
-    with io.XDMFFile(mesh.comm, f"out_pyamg/poisson_{dtype.__name__}.xdmf", "w") as file:
+    with io.XDMFFile(mesh.comm, f"out_pyamg/poisson_{dtype!s}.xdmf", "w") as file:
         file.write_mesh(mesh)
         file.write_function(uh)
 
@@ -126,7 +126,7 @@ def poisson_problem(dtype: npt.DTypeLike, solver_type: str) -> None:
 
 
 # +
-def nullspace_elasticty(Q: fem.FunctionSpace) -> list[np.ndarray]:
+def nullspace_elasticty(Q: fem.FunctionSpace) -> npt.NDArray:
     """Create the elasticity (near)nulspace.
 
     Args:

python/demo/demo_scattering_boundary_conditions.py

@@ -220,8 +220,8 @@ def generate_mesh_wire(
 # +
 def compute_a(nu: int, m: complex, alpha: float) -> float:
     """Compute the a_nu coefficient."""
-    J_nu_alpha = jv(nu, alpha)
-    J_nu_malpha = jv(nu, m * alpha)
+    J_nu_alpha = jv(nu, alpha)  # type: ignore
+    J_nu_malpha = jv(nu, m * alpha)  # type: ignore
     J_nu_alpha_p = jvp(nu, alpha, 1)
     J_nu_malpha_p = jvp(nu, m * alpha, 1)
 

python/dolfinx/fem/bcs.py

@@ -12,8 +12,8 @@
 
 from __future__ import annotations
 
-import numbers
 from collections.abc import Callable, Iterable
+from typing import Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -91,7 +91,10 @@ def locate_dofs_topological(
     return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)
 
 
-class DirichletBC:
+_T = TypeVar("_T", np.float32, np.float64, np.complex64, np.complex128)
+
+
+class DirichletBC(Generic[_T]):
     """Representation of Dirichlet boundary conditions.
 
     The conditions are imposed on a linear system.
@@ -120,18 +123,17 @@ class initialiser. This class is combined with different
         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
     def function_space(self) -> dolfinx.fem.FunctionSpace:
         """Function space on which the boundary condition is defined."""
         return self._cpp_object.function_space
 
-    def set(
-        self, x: npt.NDArray, x0: npt.NDArray[np.int32] | None = None, alpha: float = 1
-    ) -> None:
+    def set(self, x: npt.NDArray[_T], x0: npt.NDArray[_T] | None = None, alpha: float = 1) -> None:
         """Set array entries that are constrained by a Dirichlet condition.
 
         Entries in ``x`` that are constrained by a Dirichlet boundary
@@ -172,10 +174,10 @@ def dof_indices(self) -> tuple[npt.NDArray[np.int32], int]:
 
 
 def dirichletbc(
-    value: Function | Constant | np.ndarray,
+    value: Function | Constant | npt.NDArray[_T] | float | complex,
     dofs: npt.NDArray[np.int32],
     V: dolfinx.fem.FunctionSpace | None = None,
-) -> DirichletBC:
+) -> DirichletBC[_T]:
     """Representation of Dirichlet boundary condition.
 
     Args:
@@ -193,7 +195,7 @@ def dirichletbc(
         A representation of the boundary condition for modifying linear
         systems.
     """
-    if isinstance(value, numbers.Number):
+    if isinstance(value, float | complex):
         value = np.asarray(value)
 
     try:
@@ -232,8 +234,8 @@ def dirichletbc(
 
 
 def bcs_by_block(
-    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC]
-) -> list[list[DirichletBC]]:
+    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC[_T]]
+) -> list[list[DirichletBC[_T]]]:
     """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,6 +6,7 @@
 """Finite elements."""
 
 from functools import singledispatch
+from typing import Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -14,8 +15,10 @@
 import basix.ufl
 from dolfinx import cpp as _cpp
 
+_T = TypeVar("_T", np.float32, np.float64)
 
-class CoordinateElement:
+
+class CoordinateElement(Generic[_T]):
     """Coordinate element describing the geometry map for mesh cells."""
 
     _cpp_object: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64
@@ -60,11 +63,7 @@ def create_dof_layout(self) -> _cpp.fem.ElementDofLayout:
         """Compute and return the dof layout."""
         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]:
+    def push_forward(self, X: npt.NDArray[_T], cell_geometry: npt.NDArray[_T]) -> npt.NDArray[_T]:
         """Push points on the reference cell forward to the physical cell.
 
         Args:
@@ -82,11 +81,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[_T],
+        cell_geometry: npt.NDArray[_T],
         tol: float = 1.0e-6,
         maxit: int = 15,
-    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
+    ) -> npt.NDArray[_T]:
         """Pull points on the physical cell back to the reference cell.
 
         For non-affine cells, the pull-back is a nonlinear operation.
@@ -130,7 +129,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 +152,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 +169,7 @@ def _(e: basix.finite_element.FiniteElement):
         return CoordinateElement(_cpp.fem.CoordinateElement_float64(e._e))
 
 
-class FiniteElement:
+class FiniteElement(Generic[_T]):
     """A finite element."""
 
     _cpp_object: _cpp.fem.FiniteElement_float32 | _cpp.fem.FiniteElement_float64
@@ -224,7 +223,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[_T]:
         """Points at which to evaluate the function to be interpolated.
 
         Interpolation point coordinates on the reference cell, returning
@@ -282,7 +281,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[_T], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Transform basis from reference to physical ordering/orientation.
 
@@ -305,7 +304,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[_T], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Apply the transpose of the operator applied by T_apply().
 
@@ -319,7 +318,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[_T], cell_permutations: npt.NDArray[np.uint32], dim: int
     ) -> None:
         """Apply the inverse transpose of T_apply().
 
@@ -336,7 +335,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

@@ -32,8 +32,10 @@
     from dolfinx.mesh import EntityMap as _EntityMap
     from dolfinx.mesh import Mesh, MeshTags
 
+_S = typing.TypeVar("_S", np.float32, np.float64, np.complex64, np.complex128)  # scalar
 
-class Form:
+
+class Form(typing.Generic[_S]):
     """A finite element form."""
 
     _cpp_object: (