diff --git a/.github/workflows/ccpp.yml b/.github/workflows/ccpp.yml index fcfad14945..48add837a9 100644 --- a/.github/workflows/ccpp.yml +++ b/.github/workflows/ccpp.yml @@ -138,7 +138,7 @@ jobs: mypy --config-file python/pyproject.toml -p dolfinx # package mode, type checks installed dolfinx (working around name collision in python/ directory) cd python/ mypy test - #mypy demo # TODO: Enable this in future + mypy demo - name: Install gmsh and pyvista (and dependencies) run: | diff --git a/.github/workflows/fenicsx-refs.env b/.github/workflows/fenicsx-refs.env index 22c7e44a7e..7db9f61af8 100644 --- a/.github/workflows/fenicsx-refs.env +++ b/.github/workflows/fenicsx-refs.env @@ -1,5 +1,5 @@ basix_repository=FEniCS/basix -basix_ref=main +basix_ref=schnellerhase/mypy-dolfinx ufl_repository=FEniCS/ufl ufl_ref=main ffcx_repository=FEniCS/ffcx diff --git a/python/demo/demo_hdg.py b/python/demo/demo_hdg.py index e8b5a77e7a..2634ffcc0a 100644 --- a/python/demo/demo_hdg.py +++ b/python/demo/demo_hdg.py @@ -98,8 +98,8 @@ def u_e(x): # Create the mesh -n = 8 # Number of elements in each direction -msh = mesh.create_unit_cube(comm, n, n, n, ghost_mode=mesh.GhostMode.none) +n_el = 8 # Number of elements in each direction +msh = mesh.create_unit_cube(comm, n_el, n_el, n_el, ghost_mode=mesh.GhostMode.none) # We need to create a broken Lagrange space defined over the facets of # the mesh. To do so, we require a sub-mesh of the all facets. We begin diff --git a/python/demo/demo_interpolation-io.py b/python/demo/demo_interpolation-io.py index 0145691d96..5b30352855 100644 --- a/python/demo/demo_interpolation-io.py +++ b/python/demo/demo_interpolation-io.py @@ -119,7 +119,7 @@ pl.add_text("y-component", font_size=12, color="black", position="upper_edge") pl.add_mesh(grid.copy(), component=1, show_edges=True) - pl.view_xy() + pl.view_xy() # type: ignore pl.link_views() # If pyvista environment variable is set to off-screen (static) diff --git a/python/demo/demo_lagrange-variants.py b/python/demo/demo_lagrange-variants.py index 2f7687fa32..b615ef6fc7 100644 --- a/python/demo/demo_lagrange-variants.py +++ b/python/demo/demo_lagrange-variants.py @@ -65,8 +65,8 @@ if MPI.COMM_WORLD.size == 1: for i in range(values.shape[1]): plt.plot(lattice, values[:, i]) - plt.plot(element._element.points, [0] * 11, "ko") - plt.ylim([-1, 6]) + plt.plot(element.basix_element.points, [0] * 11, "ko") + plt.ylim((-1, 6)) plt.savefig("demo_lagrange_variants_equispaced_10.png") plt.clf() # - @@ -97,8 +97,8 @@ if MPI.COMM_WORLD.size == 1: # Skip this plotting in parallel for i in range(values.shape[1]): plt.plot(lattice, values[:, i]) - plt.plot(element._element.points, [0] * 11, "ko") - plt.ylim([-1, 6]) + plt.plot(element.basix_element.points, [0] * 11, "ko") + plt.ylim((-1, 6)) plt.savefig("demo_lagrange_variants_gll_10.png") plt.clf() # - @@ -152,7 +152,7 @@ def saw_tooth(x): plt.plot(pts, [saw_tooth(i[0]) for i in pts], "k--") plt.plot(pts, values, "r-") plt.legend(["function", "approximation"]) - plt.ylim([-0.1, 0.4]) + plt.ylim((-0.1, 0.4)) plt.title(variant.name) plt.savefig(f"demo_lagrange_variants_interpolation_{variant.name}.png") plt.clf() diff --git a/python/demo/demo_mixed-topology.py b/python/demo/demo_mixed-topology.py index f725b90465..f6a49cb71e 100644 --- a/python/demo/demo_mixed-topology.py +++ b/python/demo/demo_mixed-topology.py @@ -193,9 +193,9 @@ def marker(x): A_scipy = A.to_scipy() b_scipy = b.array -x = spsolve(A_scipy, b_scipy) +x_scipy = spsolve(A_scipy, b_scipy) -print(f"Solution vector norm {np.linalg.norm(x)}") +print(f"Solution vector norm {np.linalg.norm(x_scipy)}") # Mixed-topology I/O # We manually build a ASCII XDMF file to store the mesh @@ -235,8 +235,8 @@ def marker(x): - - {" ".join(str(val) for val in x)} + + {" ".join(str(val) for val in x_scipy)} """ diff --git a/python/demo/demo_navier-stokes.py b/python/demo/demo_navier-stokes.py index 07e2a956ed..d2613bc0b3 100644 --- a/python/demo/demo_navier-stokes.py +++ b/python/demo/demo_navier-stokes.py @@ -242,7 +242,7 @@ def f_expr(x): # We define some simulation parameters -n = 16 +n_el = 16 num_time_steps = 25 t_end = 10 Re = 25 # Reynolds Number @@ -254,7 +254,7 @@ def f_expr(x): # interpolate into for artifact free visualisation. # + -msh = mesh.create_unit_square(MPI.COMM_WORLD, n, n) +msh = mesh.create_unit_square(MPI.COMM_WORLD, n_el, n_el) # Function spaces for the velocity and for the pressure V = fem.functionspace(msh, ("Raviart-Thomas", k + 1)) @@ -413,7 +413,7 @@ def jump(phi, n): # We perform the time-stepping as a for-loop # + -for n in range(num_time_steps): +for _ in range(num_time_steps): t += delta_t.value navier_stokes_problem.solve() diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index b9990405e3..73a41375a3 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -476,10 +476,23 @@ def _create_form(form): return _create_form(form) +@typing.overload +def extract_function_spaces(forms: Form, index: int = 0) -> FunctionSpace | None: ... + + +@typing.overload def extract_function_spaces( - forms: Form | Sequence[Form] | Sequence[Sequence[Form]], - index: int = 0, -) -> FunctionSpace | list[None | FunctionSpace]: + forms: Sequence[Form], index: int = 0 +) -> list[FunctionSpace | None]: ... + + +@typing.overload +def extract_function_spaces( + forms: Sequence[Sequence[Form]], index: int = 0 +) -> list[list[FunctionSpace | None]]: ... + + +def extract_function_spaces(forms, index: int = 0): """Extract common function spaces from an array of forms. If ``forms`` is a list of linear forms, this function returns of list diff --git a/python/dolfinx/fem/petsc.py b/python/dolfinx/fem/petsc.py index bb1eae1404..d1f9e36a34 100644 --- a/python/dolfinx/fem/petsc.py +++ b/python/dolfinx/fem/petsc.py @@ -468,10 +468,10 @@ def _( "Cannot have a entire {'row' if index == 0 else 'column'} of a full of None" ) is0 = _cpp.la.petsc.create_index_sets( - [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[0]] # type: ignore[union-attr] + [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[0]] # type: ignore ) is1 = _cpp.la.petsc.create_index_sets( - [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[1]] # type: ignore[union-attr] + [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[1]] # type: ignore ) _bcs = [bc._cpp_object for bc in bcs] if bcs is not None else [] diff --git a/python/dolfinx/io/utils.py b/python/dolfinx/io/utils.py index a6c4276d29..68838a11cd 100644 --- a/python/dolfinx/io/utils.py +++ b/python/dolfinx/io/utils.py @@ -213,6 +213,7 @@ def read_mesh( # Get coordinate element, special handling for second order # serendipity. + basix_el: basix.ufl._BasixElement | basix.ufl._BlockedElement num_nodes_per_cell = cells.shape[1] if (cell_shape == CellType.quadrilateral and num_nodes_per_cell == 8) or ( cell_shape == CellType.hexahedron and num_nodes_per_cell == 20 diff --git a/python/dolfinx/plot.py b/python/dolfinx/plot.py index 2f39658223..34b80a25af 100644 --- a/python/dolfinx/plot.py +++ b/python/dolfinx/plot.py @@ -6,8 +6,10 @@ """Support functions for plotting.""" import functools +from typing import overload import numpy as np +import numpy.typing as npt from dolfinx import cpp as _cpp from dolfinx import fem, mesh @@ -29,8 +31,18 @@ } +@overload +def vtk_mesh( + msh: mesh.Mesh, dim: int | None = None, entities: npt.NDArray[np.int32] | None = None +): ... + + +@overload +def vtk_mesh(V: fem.FunctionSpace, entities: npt.NDArray[np.int32] | None = None): ... + + @functools.singledispatch -def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities=None): +def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities: npt.NDArray[np.int32] | None = None): """Create VTK mesh topology data for mesh entities. The vertex indices in the returned topology array are the indices @@ -74,8 +86,8 @@ def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities=None): return topology.reshape(-1), cell_types, msh.geometry.x -@vtk_mesh.register -def _(V: fem.FunctionSpace, entities=None): # type: ignore +@vtk_mesh.register # type: ignore[attr-defined] +def _(V: fem.FunctionSpace, entities: npt.NDArray[np.int32] | None = None): # type: ignore """Create VTK mesh topology based on the degree-of-freedom coordinates. This function supports visualisation when the degree of the finite @@ -113,7 +125,7 @@ def _(V: fem.FunctionSpace, entities=None): # type: ignore msh = V.mesh tdim = msh.topology.dim if entities is None: - entities = range(msh.topology.index_map(tdim).size_local) + entities = np.arange(msh.topology.index_map(tdim).size_local, dtype=np.int32) dofmap = V.dofmap num_dofs_per_cell = V.dofmap.dof_layout.num_dofs @@ -123,9 +135,9 @@ def _(V: fem.FunctionSpace, entities=None): # type: ignore vtk_type = ( _first_order_vtk[cell_type] if degree == 1 else _cpp.io.get_vtk_cell_type(cell_type, tdim) ) - cell_types = np.full(len(entities), vtk_type) + cell_types = np.full(entities.size, vtk_type) - topology = np.zeros((len(entities), num_dofs_per_cell + 1), dtype=np.int32) + topology = np.zeros((entities.size, num_dofs_per_cell + 1), dtype=np.int32) topology[:, 0] = num_dofs_per_cell dofmap_ = dofmap.list topology[:, 1:] = dofmap_[entities][:, perm]