diff --git a/python/dolfinx/__init__.py b/python/dolfinx/__init__.py index 97755fe44f9..2fcd8b0b020 100644 --- a/python/dolfinx/__init__.py +++ b/python/dolfinx/__init__.py @@ -50,7 +50,7 @@ del _cpp, sys -def get_include(user=False): +def get_include(user: bool = False) -> str: import os d = os.path.dirname(__file__) diff --git a/python/dolfinx/common.py b/python/dolfinx/common.py index 6172ffad2b5..ba1ff085d86 100644 --- a/python/dolfinx/common.py +++ b/python/dolfinx/common.py @@ -7,8 +7,13 @@ import functools import typing +from types import TracebackType -from dolfinx import cpp as _cpp +from mpi4py import MPI as _MPI + +from dolfinx.cpp import Reduction as _Reduction +from dolfinx.cpp import Timer as _Timer +from dolfinx.cpp import TimingTyp as _TimingType from dolfinx.cpp.common import ( IndexMap, git_commit_hash, @@ -23,6 +28,8 @@ has_slepc, ufcx_signature, ) +from dolfinx.cpp.common import list_timings as _list_timings +from dolfinx.cpp.common import timing as _timing __all__ = [ "IndexMap", @@ -41,20 +48,22 @@ "ufcx_signature", ] -TimingType = _cpp.common.TimingType -Reduction = _cpp.common.Reduction +TimingType = _TimingType +Reduction = _Reduction -def timing(task: str): - return _cpp.common.timing(task) +def timing(task: str) -> tuple[int, float, float, float]: + return _timing(task) -def list_timings(comm, timing_types: list, reduction=Reduction.max): +def list_timings( + comm: _MPI.Comm, timing_types: list, reduction: _Reduction = Reduction.max +) -> None: """Print out a summary of all Timer measurements, with a choice of wall time, system time or user time. When used in parallel, a reduction is applied across all processes. By default, the maximum time is shown.""" - _cpp.common.list_timings(comm, timing_types, reduction) + _list_timings(comm, timing_types, reduction) class Timer: @@ -91,37 +100,42 @@ class Timer: list_timings(comm, [TimingType.wall, TimingType.user]) """ - _cpp_object: _cpp.common.Timer + _cpp_object: _Timer - def __init__(self, name: typing.Optional[str] = None): + def __init__(self, name: typing.Optional[str] = None) -> None: self._cpp_object = _cpp.common.Timer(name) - def __enter__(self): + def __enter__(self) -> typing.Self: self._cpp_object.start() return self - def __exit__(self, *args): + def __exit__( + self, + exception_type: typing.Optional[BaseException], + exception_value: typing.Optional[BaseException], + traceback: typing.Optional[TracebackType], + ) -> None: self._cpp_object.stop() - def start(self): + def start(self) -> None: self._cpp_object.start() - def stop(self): + def stop(self) -> float: return self._cpp_object.stop() - def resume(self): + def resume(self) -> None: self._cpp_object.resume() - def elapsed(self): + def elapsed(self) -> float: return self._cpp_object.elapsed() -def timed(task: str): +def timed(task: str) -> typing.Callable[..., typing.Any]: """Decorator for timing functions.""" - def decorator(func): + def decorator(func: typing.Callable[..., typing.Any]) -> typing.Callable[..., typing.Any]: @functools.wraps(func) - def wrapper(*args, **kwargs): + def wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: with Timer(task): return func(*args, **kwargs) diff --git a/python/dolfinx/fem/bcs.py b/python/dolfinx/fem/bcs.py index 737ca62d1e4..b6952917bbf 100644 --- a/python/dolfinx/fem/bcs.py +++ b/python/dolfinx/fem/bcs.py @@ -98,7 +98,15 @@ class DirichletBC: _cpp.fem.DirichletBC_float64, ] - def __init__(self, bc): + def __init__( + self, + bc: typing.Union[ + _cpp.fem.DirichletBC_complex64, + _cpp.fem.DirichletBC_complex128, + _cpp.fem.DirichletBC_float32, + _cpp.fem.DirichletBC_float64, + ], + ): """Representation of Dirichlet boundary condition which is imposed on a linear system. @@ -210,8 +218,8 @@ def dirichletbc( def bcs_by_block( - spaces: typing.Iterable[typing.Union[dolfinx.fem.FunctionSpace, None]], - bcs: typing.Iterable[DirichletBC], + spaces: list[typing.Union[dolfinx.fem.FunctionSpace, None]], + bcs: list[DirichletBC], ) -> list[list[DirichletBC]]: """Arrange Dirichlet boundary conditions by the function space that they constrain. @@ -222,8 +230,10 @@ def bcs_by_block( ``space[i]``. """ - def _bc_space(V, bcs): + def _bc_space( + V: list[typing.Union[dolfinx.fem.FunctionSpace, None]], bcs: list[DirichletBC] + ) -> list[DirichletBC]: """Return list of bcs that have the same space as V""" - return [bc for bc in bcs if V.contains(bc.function_space)] + return [bc for bc in bcs if bc.function_space in V] return [_bc_space(V, bcs) if V is not None else [] for V in spaces] diff --git a/python/dolfinx/fem/dofmap.py b/python/dolfinx/fem/dofmap.py index 0355a2cd860..059c694a742 100644 --- a/python/dolfinx/fem/dofmap.py +++ b/python/dolfinx/fem/dofmap.py @@ -4,7 +4,12 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import numpy.typing as npt + from dolfinx import cpp as _cpp +from dolfinx.common import IndexMap +from dolfinx.fem import ElementDofLayout class DofMap: @@ -19,7 +24,7 @@ class DofMap: def __init__(self, dofmap: _cpp.fem.DofMap): self._cpp_object = dofmap - def cell_dofs(self, cell_index: int): + def cell_dofs(self, cell_index: int) -> npt.NDArray[np.int32]: """Cell local-global dof map Args: @@ -31,26 +36,26 @@ def cell_dofs(self, cell_index: int): return self._cpp_object.cell_dofs(cell_index) @property - def bs(self): + def bs(self) -> int: """Returns the block size of the dofmap""" return self._cpp_object.bs @property - def dof_layout(self): + def dof_layout(self) -> ElementDofLayout: """Layout of dofs on an element.""" return self._cpp_object.dof_layout @property - def index_map(self): + def index_map(self) -> IndexMap: """Index map that described the parallel distribution of the dofmap.""" return self._cpp_object.index_map @property - def index_map_bs(self): + def index_map_bs(self) -> int: """Block size of the index map.""" return self._cpp_object.index_map_bs @property - def list(self): + def list(self) -> npt.NDArray[np.int32]: """Adjacency list with dof indices for each cell.""" return self._cpp_object.map() diff --git a/python/dolfinx/fem/element.py b/python/dolfinx/fem/element.py index 7e88427c12b..83df51f8d5c 100644 --- a/python/dolfinx/fem/element.py +++ b/python/dolfinx/fem/element.py @@ -51,9 +51,9 @@ def dtype(self) -> np.dtype: def coordinate_element( celltype: _cpp.mesh.CellType, degree: int, - variant=int(basix.LagrangeVariant.unset), + variant: int = 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. @@ -76,7 +76,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. diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 59821fd1817..b1a25ac2854 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -254,25 +254,22 @@ def _form(form): constants = [c._cpp_object for c in form.constants()] # Make map from integral_type to subdomain id - subdomain_ids = {type: [] for type in sd.get(domain).keys()} + subdomain_ids: dict[str, list[int]] = {type: [] for type in sd.get(domain).keys()} for integral in form.integrals(): if integral.subdomain_data() is not None: # Subdomain ids can be strings, its or tuples with strings and ints if integral.subdomain_id() != "everywhere": try: - ids = [sid for sid in integral.subdomain_id() if sid != "everywhere"] + subdomain_ids[integral.integral_type()] += [ + sid for sid in integral.subdomain_id() if sid != "everywhere" + ] except TypeError: # If not tuple, but single integer id - ids = [integral.subdomain_id()] - else: - ids = [] - subdomain_ids[integral.integral_type()].append(ids) - - # Chain and sort subdomain ids - for itg_type, marker_ids in subdomain_ids.items(): - flattened_ids = list(chain.from_iterable(marker_ids)) - flattened_ids.sort() - subdomain_ids[itg_type] = flattened_ids + subdomain_ids[integral.integral_type()].append(integral.subdomain_id()) + + # Sort subdomain ids + for val in subdomain_ids.values(): + val.sort() # Subdomain markers (possibly empty list for some integral types) subdomains = { diff --git a/python/dolfinx/fem/function.py b/python/dolfinx/fem/function.py index a04ce990c67..fc499a93d5e 100644 --- a/python/dolfinx/fem/function.py +++ b/python/dolfinx/fem/function.py @@ -7,9 +7,10 @@ from __future__ import annotations -import typing import warnings +from collections.abc import Sequence from functools import singledispatch +from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Union import numpy as np import numpy.typing as npt @@ -20,15 +21,16 @@ from dolfinx import default_scalar_type, jit, la from dolfinx.fem import dofmap from dolfinx.geometry import PointOwnershipData +from dolfinx.mesh import Mesh -if typing.TYPE_CHECKING: +if TYPE_CHECKING: from mpi4py import MPI as _MPI from dolfinx.mesh import Mesh class Constant(ufl.Constant): - _cpp_object: typing.Union[ + _cpp_object: Union[ _cpp.fem.Constant_complex64, _cpp.fem.Constant_complex128, _cpp.fem.Constant_float32, @@ -36,7 +38,9 @@ class Constant(ufl.Constant): ] def __init__( - self, domain, c: typing.Union[np.ndarray, typing.Sequence, np.floating, np.complexfloating] + self, + domain: Union[ufl.Mesh, Mesh], + c: Union[np.ndarray, Sequence, np.floating, np.complexfloating], ): """A constant with respect to a domain. @@ -61,25 +65,25 @@ def __init__( raise AttributeError("Constant value must have a dtype attribute.") @property - def value(self): + def value(self) -> float: """The value of the constant""" return self._cpp_object.value @value.setter - def value(self, v): + def value(self, v: float) -> None: np.copyto(self._cpp_object.value, np.asarray(v)) @property def dtype(self) -> np.dtype: return np.dtype(self._cpp_object.dtype) - def __float__(self): + def __float__(self) -> float: if self.ufl_shape or self.ufl_free_indices: raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.") else: return float(self.value) - def __complex__(self): + def __complex__(self) -> complex: if self.ufl_shape or self.ufl_free_indices: raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.") else: @@ -91,10 +95,10 @@ def __init__( self, e: ufl.core.expr.Expr, X: np.ndarray, - comm: typing.Optional[_MPI.Comm] = None, - form_compiler_options: typing.Optional[dict] = None, - jit_options: typing.Optional[dict] = None, - dtype: typing.Optional[npt.DTypeLike] = None, + comm: Optional[_MPI.Comm] = None, + form_compiler_options: Optional[dict] = None, + jit_options: Optional[dict] = None, + dtype: Optional[npt.DTypeLike] = None, ): """Create a DOLFINx Expression. @@ -173,7 +177,7 @@ def __init__( else: raise RuntimeError("Expressions with more that one Argument not allowed.") - def _create_expression(dtype): + def _create_expression(dtype: npt.DTypeLike) -> _cpp.fem.Expression: if np.issubdtype(dtype, np.float32): return _cpp.fem.create_expression_float32 elif np.issubdtype(dtype, np.float64): @@ -197,7 +201,7 @@ def eval( self, mesh: Mesh, entities: np.ndarray, - values: typing.Optional[np.ndarray] = None, + values: Optional[np.ndarray] = None, ) -> np.ndarray: """Evaluate Expression on entities. @@ -249,7 +253,7 @@ def X(self) -> np.ndarray: return self._cpp_object.X() @property - def ufl_expression(self): + def ufl_expression(self) -> Expression: """Original UFL Expression""" return self._ufl_expression @@ -259,12 +263,12 @@ def value_size(self) -> int: return self._cpp_object.value_size @property - def argument_function_space(self) -> typing.Optional[FunctionSpace]: + def argument_function_space(self) -> Optional[FunctionSpace]: """The argument function space if expression has argument""" return self._argument_function_space @property - def ufcx_expression(self): + def ufcx_expression(self) -> tuple[Any, Any | None, tuple[None, None] | tuple[str, str] | Any]: """The compiled ufcx_expression object""" return self._ufcx_expression @@ -283,7 +287,7 @@ class Function(ufl.Coefficient): (domain, element and dofmap) and a vector holding the degrees-of-freedom.""" - _cpp_object: typing.Union[ + _cpp_object: Union[ _cpp.fem.Function_complex64, _cpp.fem.Function_complex128, _cpp.fem.Function_float32, @@ -293,9 +297,9 @@ class Function(ufl.Coefficient): def __init__( self, V: FunctionSpace, - x: typing.Optional[la.Vector] = None, - name: typing.Optional[str] = None, - dtype: typing.Optional[npt.DTypeLike] = None, + x: Optional[la.Vector] = None, + name: Optional[str] = None, + dtype: Optional[npt.DTypeLike] = None, ): """Initialize a finite element Function. @@ -321,7 +325,14 @@ def __init__( ), "Incompatible FunctionSpace dtype and requested dtype." # Create cpp Function - def functiontype(dtype): + def functiontype( + dtype: npt.DTypeLike, + ) -> Union[ + _cpp.fem.Function_float32, + _cpp.fem.Function_float64, + _cpp.fem.Function_complex64, + _cpp.fem.Function_complex128, + ]: if np.issubdtype(dtype, np.float32): return _cpp.fem.Function_float32 elif np.issubdtype(dtype, np.float64): @@ -358,7 +369,9 @@ def function_space(self) -> FunctionSpace: """The FunctionSpace that the Function is defined on""" return self._V - def eval(self, x: npt.ArrayLike, cells: npt.ArrayLike, u=None) -> np.ndarray: + def eval( + self, x: npt.ArrayLike, cells: npt.ArrayLike, u: Optional[npt.NDArray[np.floating]] = None + ) -> npt.NDArray[np.floating]: """Evaluate Function at points x. Points where x has shape (num_points, 3), and cells has shape @@ -411,9 +424,9 @@ def interpolate_nonmatching( def interpolate( self, - u0: typing.Union[typing.Callable, Expression, Function], - cells0: typing.Optional[np.ndarray] = None, - cells1: typing.Optional[np.ndarray] = None, + u0: Union[Callable, Expression, Function], + cells0: Optional[np.ndarray] = None, + cells1: Optional[np.ndarray] = None, ) -> None: """Interpolate an expression. @@ -436,22 +449,22 @@ def interpolate( cells1 = np.arange(0, dtype=np.int32) @singledispatch - def _interpolate(u0): + def _interpolate(u0: Any) -> None: """Interpolate a cpp.fem.Function.""" self._cpp_object.interpolate(u0, cells0, cells1) # type: ignore @_interpolate.register(Function) - def _(u0: Function): + def _(u0: Function) -> None: """Interpolate a fem.Function.""" self._cpp_object.interpolate(u0._cpp_object, cells0, cells1) # type: ignore @_interpolate.register(int) - def _(u0_ptr: int): + def _(u0_ptr: int) -> None: """Interpolate using a pointer to a function f(x).""" self._cpp_object.interpolate_ptr(u0_ptr, cells0) # type: ignore @_interpolate.register(Expression) - def _(e0: Expression): + def _(e0: Expression) -> None: """Interpolate a fem.Expression.""" self._cpp_object.interpolate(e0._cpp_object, cells0, cells1) # type: ignore @@ -483,7 +496,7 @@ def x(self) -> la.Vector: return self._x @property - def vector(self): + def vector(self) -> Any: """PETSc vector holding the degrees-of-freedom. Upon first call, this function creates a PETSc ``Vec`` object @@ -511,10 +524,10 @@ def name(self) -> str: return self._cpp_object.name # type: ignore @name.setter - def name(self, name): + def name(self, name: str) -> None: self._cpp_object.name = name - def __str__(self): + def __str__(self) -> str: """Pretty print representation.""" return self.name @@ -562,7 +575,7 @@ def collapse(self) -> Function: return Function(V_collapsed, la.Vector(u_collapsed.x)) -class ElementMetaData(typing.NamedTuple): +class ElementMetaData(NamedTuple): """Data for representing a finite element :param family: Element type. @@ -574,16 +587,16 @@ class ElementMetaData(typing.NamedTuple): family: str degree: int - shape: typing.Optional[tuple[int, ...]] = None - symmetry: typing.Optional[bool] = None + shape: Optional[tuple[int, ...]] = None + symmetry: Optional[bool] = None def _create_dolfinx_element( - comm: _MPI.Intracomm, + comm: _MPI.Comm, cell_type: _cpp.mesh.CellType, ufl_e: ufl.FiniteElementBase, dtype: np.dtype, -) -> typing.Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]: +) -> Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]: """Create a DOLFINx element from a basix.ufl element.""" if np.issubdtype(dtype, np.float32): CppElement = _cpp.fem.FiniteElement_float32 @@ -606,9 +619,9 @@ def _create_dolfinx_element( def functionspace( mesh: Mesh, - element: typing.Union[ufl.FiniteElementBase, ElementMetaData, tuple[str, int, tuple, bool]], - form_compiler_options: typing.Optional[dict[str, typing.Any]] = None, - jit_options: typing.Optional[dict[str, typing.Any]] = None, + element: Union[ufl.FiniteElementBase, ElementMetaData, tuple[str, int, tuple, bool]], + form_compiler_options: Optional[dict[str, Any]] = None, + jit_options: Optional[dict[str, Any]] = None, ) -> FunctionSpace: """Create a finite element function space. @@ -667,14 +680,14 @@ def functionspace( class FunctionSpace(ufl.FunctionSpace): """A space on which Functions (fields) can be defined.""" - _cpp_object: typing.Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64] + _cpp_object: Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64] _mesh: Mesh def __init__( self, mesh: Mesh, element: ufl.FiniteElementBase, - cppV: typing.Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64], + cppV: Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64], ): """Create a finite element function space. @@ -756,11 +769,11 @@ def sub(self, i: int) -> FunctionSpace: cppV_sub = self._cpp_object.sub([i]) # type: ignore return FunctionSpace(self._mesh, sub_element, cppV_sub) - def component(self): + def component(self) -> npt.NDArray[np.integer]: """Return the component relative to the parent space.""" - return self._cpp_object.component() # type: ignore + return self._cpp_object.component() - def contains(self, V) -> bool: + def contains(self, V: FunctionSpace) -> bool: """Check if a space is contained in, or is the same as (identity), this space. @@ -773,11 +786,11 @@ def contains(self, V) -> bool: """ return self._cpp_object.contains(V._cpp_object) # type: ignore - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: """Comparison for equality.""" return super().__eq__(other) and self._cpp_object == other._cpp_object - def __ne__(self, other): + def __ne__(self, other: Any) -> bool: """Comparison for inequality.""" return super().__ne__(other) or self._cpp_object != other._cpp_object @@ -788,7 +801,7 @@ def ufl_function_space(self) -> ufl.FunctionSpace: @property def element( self, - ) -> typing.Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]: + ) -> Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]: """Function space finite element.""" return self._cpp_object.element # type: ignore diff --git a/python/dolfinx/geometry.py b/python/dolfinx/geometry.py index 03cd0870b24..84794520611 100644 --- a/python/dolfinx/geometry.py +++ b/python/dolfinx/geometry.py @@ -9,11 +9,13 @@ import typing +from mpi4py import MPI as _MPI + import numpy as np import numpy.typing as npt if typing.TYPE_CHECKING: - from dolfinx.cpp.graph import AdjacencyList_int32 + from dolfinx.cpp.graph import AdjacencyList_int32, AdjacencyList_int64 from dolfinx.mesh import Mesh @@ -40,7 +42,12 @@ class PointOwnershipData: _cpp.geometry.PointOwnershipData_float32, _cpp.geometry.PointOwnershipData_float64 ] - def __init__(self, ownership_data): + def __init__( + self, + ownership_data: typing.Union[ + _cpp.geometry.PointOwnershipData_float32, _cpp.geometry.PointOwnershipData_float64 + ], + ) -> None: """Wrap a C++ PointOwnershipData.""" self._cpp_object = ownership_data @@ -68,7 +75,12 @@ class BoundingBoxTree: _cpp.geometry.BoundingBoxTree_float32, _cpp.geometry.BoundingBoxTree_float64 ] - def __init__(self, tree): + def __init__( + self, + tree: typing.Union[ + _cpp.geometry.BoundingBoxTree_float32, _cpp.geometry.BoundingBoxTree_float64 + ], + ) -> None: """Wrap a C++ BoundingBoxTree. Note: @@ -83,7 +95,7 @@ def num_bboxes(self) -> int: """Number of bounding boxes.""" return self._cpp_object.num_bboxes - def get_bbox(self, i) -> npt.NDArray[np.floating]: + def get_bbox(self, i: int) -> npt.NDArray[np.floating]: """Get lower and upper corners of the ith bounding box. Args: @@ -96,7 +108,7 @@ def get_bbox(self, i) -> npt.NDArray[np.floating]: """ return self._cpp_object.get_bbox(i) - def create_global_tree(self, comm) -> BoundingBoxTree: + def create_global_tree(self, comm: _MPI.Comm) -> BoundingBoxTree: return BoundingBoxTree(self._cpp_object.create_global_tree(comm)) @@ -217,7 +229,7 @@ def create_midpoint_tree(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]) def compute_colliding_cells( mesh: Mesh, candidates: AdjacencyList_int32, x: npt.NDArray[np.floating] -): +) -> typing.Union[AdjacencyList_int32, AdjacencyList_int64]: """From a mesh, find which cells collide with a set of points. Args: @@ -234,7 +246,9 @@ def compute_colliding_cells( return _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates, x) -def squared_distance(mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]): +def squared_distance( + mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating] +) -> npt.NDArray[np.floating]: """Compute the squared distance between a point and a mesh entity. The distance is computed between the ith input points and the ith diff --git a/python/dolfinx/graph.py b/python/dolfinx/graph.py index dd16514a79a..18714b71540 100644 --- a/python/dolfinx/graph.py +++ b/python/dolfinx/graph.py @@ -7,7 +7,10 @@ from __future__ import annotations +from typing import Optional, Union + import numpy as np +import numpy.typing as npt from dolfinx import cpp as _cpp from dolfinx.cpp.graph import partitioner @@ -31,7 +34,9 @@ __all__ = ["adjacencylist", "partitioner"] -def adjacencylist(data: np.ndarray, offsets=None): +def adjacencylist( + data: np.ndarray, offsets: Optional[npt.NDArray[Union[np.int32, np.int64]]] = None +) -> Union[_cpp.graph.AdjacencyList_int32, _cpp.graph.AdjacencyList_int64]: """Create an AdjacencyList for int32 or int64 datasets. Args: diff --git a/python/dolfinx/io/utils.py b/python/dolfinx/io/utils.py index a479f898998..e1ae5dcb159 100644 --- a/python/dolfinx/io/utils.py +++ b/python/dolfinx/io/utils.py @@ -8,6 +8,7 @@ import typing from pathlib import Path +from types import TracebackType from mpi4py import MPI as _MPI @@ -26,11 +27,13 @@ __all__ = ["VTKFile", "XDMFFile", "cell_perm_gmsh", "cell_perm_vtk", "distribute_entity_data"] -def _extract_cpp_objects(functions: typing.Union[Mesh, Function, tuple[Function], list[Function]]): +def _extract_cpp_objects( + functions: typing.Union[Mesh, Function, tuple[Function], list[Function]], +) -> list[typing.Any | Function]: """Extract C++ objects""" if isinstance(functions, (list, tuple)): return [getattr(u, "_cpp_object", u) for u in functions] - else: + else: # TODO: remove case can not happen as indicated by typing return [getattr(functions, "_cpp_object", functions)] @@ -103,16 +106,21 @@ def __init__( comm, filename, _extract_cpp_objects(output), engine, mesh_policy ) # type: ignore[arg-type] - def __enter__(self): + def __enter__(self) -> typing.Self: return self - def __exit__(self, exception_type, exception_value, traceback): + def __exit__( + self, + exception_type: typing.Optional[BaseException], + exception_value: typing.Optional[BaseException], + traceback: typing.Optional[TracebackType], + ) -> None: self.close() - def write(self, t: float): + def write(self, t: float) -> None: self._cpp_object.write(t) - def close(self): + def close(self) -> None: self._cpp_object.close() class FidesWriter: @@ -174,16 +182,21 @@ def __init__( comm, filename, _extract_cpp_objects(output), engine, mesh_policy ) # type: ignore[arg-type] - def __enter__(self): + def __enter__(self) -> typing.Self: return self - def __exit__(self, exception_type, exception_value, traceback): + def __exit__( + self, + exception_type: typing.Optional[BaseException], + exception_value: typing.Optional[BaseException], + traceback: typing.Optional[TracebackType], + ) -> None: self.close() - def write(self, t: float): + def write(self, t: float) -> None: self._cpp_object.write(t) - def close(self): + def close(self) -> None: self._cpp_object.close() @@ -196,10 +209,15 @@ class VTKFile(_cpp.io.VTKFile): """ - def __enter__(self): + def __enter__(self) -> typing.Self: return self - def __exit__(self, exception_type, exception_value, traceback): + def __exit__( + self, + exception_type: typing.Optional[BaseException], + exception_value: typing.Optional[BaseException], + traceback: typing.Optional[TracebackType], + ) -> None: self.close() def write_mesh(self, mesh: Mesh, t: float = 0.0) -> None: @@ -212,10 +230,15 @@ def write_function(self, u: typing.Union[list[Function], Function], t: float = 0 class XDMFFile(_cpp.io.XDMFFile): - def __enter__(self): + def __enter__(self) -> typing.Self: return self - def __exit__(self, exception_type, exception_value, traceback): + def __exit__( + self, + exception_type: typing.Optional[BaseException], + exception_value: typing.Optional[BaseException], + traceback: typing.Optional[TracebackType], + ) -> None: self.close() def write_mesh(self, mesh: Mesh, xpath: str = "/Xdmf/Domain") -> None: @@ -233,8 +256,11 @@ def write_meshtags( super().write_meshtags(tags._cpp_object, x, geometry_xpath, xpath) def write_function( - self, u: Function, t: float = 0.0, mesh_xpath="/Xdmf/Domain/Grid[@GridType='Uniform'][1]" - ): + self, + u: Function, + t: float = 0.0, + mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]", + ) -> None: """Write function to file for a given time. Note: @@ -252,7 +278,10 @@ def write_function( super().write_function(getattr(u, "_cpp_object", u), t, mesh_xpath) def read_mesh( - self, ghost_mode=GhostMode.shared_facet, name="mesh", xpath="/Xdmf/Domain" + self, + ghost_mode: GhostMode = GhostMode.shared_facet, + name: str = "mesh", + xpath: str = "/Xdmf/Domain", ) -> Mesh: """Read mesh data from file.""" cell_shape, cell_degree = super().read_cell_type(name, xpath) @@ -276,7 +305,7 @@ def read_mesh( ) return Mesh(msh, domain) - def read_meshtags(self, mesh, name, xpath="/Xdmf/Domain"): + def read_meshtags(self, mesh: Mesh, name: str, xpath: str = "/Xdmf/Domain") -> MeshTags: mt = super().read_meshtags(mesh._cpp_object, name, xpath) return MeshTags(mt) diff --git a/python/dolfinx/jit.py b/python/dolfinx/jit.py index d27093599bc..9bb37c32fda 100644 --- a/python/dolfinx/jit.py +++ b/python/dolfinx/jit.py @@ -10,9 +10,9 @@ import os import sys from pathlib import Path -from typing import Optional +from typing import Any, Callable, Optional -from mpi4py import MPI +from mpi4py import MPI as _MPI import ffcx import ffcx.codegeneration.jit @@ -44,7 +44,9 @@ ) -def mpi_jit_decorator(local_jit, *args, **kwargs): +def mpi_jit_decorator( + local_jit: Callable[..., Any], *args: Any, **kwargs: Any +) -> Callable[..., Any]: """A decorator for jit compilation. Use this function as a decorator to any jit compiler function. In a @@ -56,7 +58,7 @@ def mpi_jit_decorator(local_jit, *args, **kwargs): """ @functools.wraps(local_jit) - def mpi_jit(comm, *args, **kwargs): + def mpi_jit(comm: _MPI.Comm, *args: Any, **kwargs: Any) -> Callable[..., Any]: # Just call JIT compiler when running in serial if comm.size == 1: return local_jit(*args, **kwargs) @@ -83,7 +85,7 @@ def mpi_jit(comm, *args, **kwargs): # Wait for the compiling process to finish and get status # TODO: Would be better to broadcast the status from root but # this works. - global_status = comm.allreduce(status, op=MPI.MAX) + global_status = comm.allreduce(status, op=_MPI.MAX) if global_status == 0: # Success, call jit on all other processes (this should just # read the cache) @@ -103,7 +105,7 @@ def mpi_jit(comm, *args, **kwargs): @functools.cache -def _load_options(): +def _load_options() -> tuple[dict, dict]: """Loads options from JSON files.""" user_config_file = os.getenv("XDG_CONFIG_HOME", default=Path.home().joinpath(".config")) / Path( "dolfinx", "dolfinx_jit_options.json" @@ -157,8 +159,10 @@ def get_options(priority_options: Optional[dict] = None) -> dict: @mpi_jit_decorator def ffcx_jit( - ufl_object, form_compiler_options: Optional[dict] = None, jit_options: Optional[dict] = None -): + ufl_object: ufl.Form, + form_compiler_options: Optional[dict] = None, + jit_options: Optional[dict] = None, +) -> tuple[Any, Any, tuple[str, str]]: """Compile UFL object with FFCx and CFFI. Args: diff --git a/python/dolfinx/la.py b/python/dolfinx/la.py index 9998be3fac1..9779133b9e8 100644 --- a/python/dolfinx/la.py +++ b/python/dolfinx/la.py @@ -5,7 +5,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Linear algebra functionality""" -import typing +from typing import Any, Union import numpy as np import numpy.typing as npt @@ -28,7 +28,7 @@ class MatrixCSR: - _cpp_object: typing.Union[ + _cpp_object: Union[ _cpp.la.MatrixCSR_float32, _cpp.la.MatrixCSR_float64, _cpp.la.MatrixCSR_complex64, @@ -37,7 +37,7 @@ class MatrixCSR: def __init__( self, - A: typing.Union[ + A: Union[ _cpp.la.MatrixCSR_float32, _cpp.la.MatrixCSR_float64, _cpp.la.MatrixCSR_complex64, @@ -64,7 +64,7 @@ def index_map(self, i: int) -> IndexMap: return self._cpp_object.index_map(i) @property - def block_size(self): + def block_size(self) -> tuple[int, int]: """Block sizes for the matrix.""" return self._cpp_object.bs @@ -132,7 +132,7 @@ def to_dense(self) -> npt.NDArray[np.floating]: """ return self._cpp_object.to_dense() - def to_scipy(self, ghosted=False): + def to_scipy(self, ghosted: bool = False) -> Any: """Convert to a SciPy CSR/BSR matrix. Data is shared. Note: @@ -174,7 +174,9 @@ def to_scipy(self, ghosted=False): def matrix_csr( - sp: _cpp.la.SparsityPattern, block_mode=BlockMode.compact, dtype: npt.DTypeLike = np.float64 + sp: _cpp.la.SparsityPattern, + block_mode: BlockMode = BlockMode.compact, + dtype: npt.DTypeLike = np.float64, ) -> MatrixCSR: """Create a distributed sparse matrix. @@ -203,7 +205,7 @@ def matrix_csr( class Vector: - _cpp_object: typing.Union[ + _cpp_object: Union[ _cpp.la.Vector_float32, _cpp.la.Vector_float64, _cpp.la.Vector_complex64, @@ -215,7 +217,7 @@ class Vector: def __init__( self, - x: typing.Union[ + x: Union[ _cpp.la.Vector_float32, _cpp.la.Vector_float64, _cpp.la.Vector_complex64, @@ -237,7 +239,7 @@ def __init__( self._cpp_object = x self._petsc_x = None - def __del__(self): + def __del__(self) -> None: if self._petsc_x is not None: self._petsc_x.destroy() @@ -257,7 +259,7 @@ def array(self) -> np.ndarray: return self._cpp_object.array @property - def petsc_vec(self): + def petsc_vec(self) -> Any: """PETSc vector holding the entries of the vector. Upon first call, this function creates a PETSc ``Vec`` object @@ -286,7 +288,7 @@ def scatter_reverse(self, mode: InsertMode) -> None: self._cpp_object.scatter_reverse(mode) -def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector: +def vector(map: IndexMap, bs: int = 1, dtype: npt.DTypeLike = np.float64) -> Vector: """Create a distributed vector. Args: @@ -318,7 +320,7 @@ def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector: return Vector(vtype(map, bs)) -def create_petsc_vector_wrap(x: Vector): +def create_petsc_vector_wrap(x: Vector) -> Any: """Wrap a distributed DOLFINx vector as a PETSc vector. Note: @@ -343,7 +345,7 @@ def create_petsc_vector_wrap(x: Vector): return PETSc.Vec().createGhostWithArray(ghosts, x.array, size=size, bsize=bs, comm=map.comm) # type: ignore -def create_petsc_vector(map, bs: int): +def create_petsc_vector(map: IndexMap, bs: int) -> Any: """Create a distributed PETSc vector. Note: @@ -367,7 +369,7 @@ def create_petsc_vector(map, bs: int): return PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=map.comm) # type: ignore -def orthonormalize(basis: list[Vector]): +def orthonormalize(basis: list[Vector]) -> None: """Orthogonalise set of PETSc vectors in-place.""" _cpp.la.orthonormalize([x._cpp_object for x in basis]) diff --git a/python/dolfinx/mesh.py b/python/dolfinx/mesh.py index cf04e746fe6..d2b2b76bb8c 100644 --- a/python/dolfinx/mesh.py +++ b/python/dolfinx/mesh.py @@ -19,10 +19,18 @@ import ufl from dolfinx import cpp as _cpp from dolfinx import default_real_type +from dolfinx.cpp.graph import AdjacencyList_int32, AdjacencyList_int64 from dolfinx.cpp.mesh import ( CellType, DiagonalType, + Geometry_float32, + Geometry_float64, GhostMode, + MeshTags_float64, + MeshTags_int8, + MeshTags_int32, + MeshTags_int64, + Topology, build_dual_graph, cell_dim, create_cell_partitioner, @@ -66,6 +74,11 @@ "entities_to_geometry", ] +CellPartionerType = typing.Callable[ + [_MPI.Comm, int, list[CellType], list[npt.NDArray[np.int64]]], + AdjacencyList_int32 | AdjacencyList_int64, +] + class Mesh: """A mesh.""" @@ -73,7 +86,11 @@ class Mesh: _mesh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64] _ufl_domain: typing.Optional[ufl.Mesh] - def __init__(self, mesh, domain: typing.Optional[ufl.Mesh]): + def __init__( + self, + mesh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64], + domain: typing.Optional[ufl.Mesh], + ): """Initialize mesh from a C++ mesh. Args: @@ -90,15 +107,15 @@ def __init__(self, mesh, domain: typing.Optional[ufl.Mesh]): self._ufl_domain._ufl_cargo = self._cpp_object # type: ignore @property - def comm(self): + def comm(self) -> _MPI.Comm: return self._cpp_object.comm @property - def name(self): + def name(self) -> str: return self._cpp_object.name @name.setter - def name(self, value): + def name(self, value: str) -> None: self._cpp_object.name = value def ufl_cell(self) -> ufl.Cell: @@ -138,12 +155,12 @@ def h(self, dim: int, entities: npt.NDArray[np.int32]) -> npt.NDArray[np.float64 return _cpp.mesh.h(self._cpp_object, dim, entities) @property - def topology(self): + def topology(self) -> Topology: "Mesh topology." return self._cpp_object.topology @property - def geometry(self): + def geometry(self) -> typing.Union[Geometry_float32, Geometry_float64]: "Mesh geometry." return self._cpp_object.geometry @@ -151,7 +168,10 @@ def geometry(self): class MeshTags: """Mesh tags associate data (markers) with a subset of mesh entities of a given dimension.""" - def __init__(self, meshtags): + def __init__( + self, + meshtags: typing.Union[MeshTags_int8, MeshTags_int32, MeshTags_int64, MeshTags_float64], + ): """Initialize tags from a C++ MeshTags object. Args: @@ -189,7 +209,7 @@ def indices(self) -> npt.NDArray[np.int32]: return self._cpp_object.indices @property - def values(self): + def values(self) -> npt.NDArray[np.int32]: """Values associated with tagged mesh entities.""" return self._cpp_object.values @@ -199,10 +219,10 @@ def name(self) -> str: return self._cpp_object.name @name.setter - def name(self, value): + def name(self, value: np.int32) -> None: self._cpp_object.name = value - def find(self, value) -> npt.NDArray[np.int32]: + def find(self, value: np.int32) -> npt.NDArray[np.int32]: """Get a list of all entity indices with a given value. Args: @@ -214,11 +234,15 @@ def find(self, value) -> npt.NDArray[np.int32]: return self._cpp_object.find(value) -def compute_incident_entities(topology, entities: npt.NDArray[np.int32], d0: int, d1: int): +def compute_incident_entities( + topology: Topology, entities: npt.NDArray[np.int32], d0: int, d1: int +) -> npt.NDArray[np.int32]: return _cpp.mesh.compute_incident_entities(topology, entities, d0, d1) -def compute_midpoints(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]): +def compute_midpoints( + mesh: Mesh, dim: int, entities: npt.NDArray[np.int32] +) -> npt.NDArray[np.floating]: return _cpp.mesh.compute_midpoints(mesh._cpp_object, dim, entities) @@ -551,7 +575,7 @@ def meshtags( def meshtags_from_entities( mesh: Mesh, dim: int, entities: _cpp.graph.AdjacencyList_int32, values: npt.NDArray[typing.Any] -): +) -> MeshTags: """Create a :class:dolfinx.mesh.MeshTags` object that associates data with a subset of mesh entities, where the entities are defined by their vertices. @@ -585,8 +609,8 @@ def create_interval( nx: int, points: npt.ArrayLike, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, ) -> Mesh: """Create an interval mesh. @@ -621,8 +645,8 @@ def create_unit_interval( comm: _MPI.Comm, nx: int, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, ) -> Mesh: """Create a mesh on the unit interval. @@ -647,10 +671,10 @@ def create_rectangle( comm: _MPI.Comm, points: npt.ArrayLike, n: npt.ArrayLike, - cell_type=CellType.triangle, + cell_type: CellType = CellType.triangle, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, diagonal: DiagonalType = DiagonalType.right, ) -> Mesh: """Create a rectangle mesh. @@ -689,10 +713,10 @@ def create_unit_square( comm: _MPI.Comm, nx: int, ny: int, - cell_type=CellType.triangle, + cell_type: CellType = CellType.triangle, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, diagonal: DiagonalType = DiagonalType.right, ) -> Mesh: """Create a mesh of a unit square. @@ -729,10 +753,10 @@ def create_box( comm: _MPI.Comm, points: list[npt.ArrayLike], n: list, - cell_type=CellType.tetrahedron, + cell_type: CellType = CellType.tetrahedron, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, ) -> Mesh: """Create a box mesh. @@ -768,10 +792,10 @@ def create_unit_cube( nx: int, ny: int, nz: int, - cell_type=CellType.tetrahedron, + cell_type: CellType = CellType.tetrahedron, dtype: npt.DTypeLike = default_real_type, - ghost_mode=GhostMode.shared_facet, - partitioner=None, + ghost_mode: GhostMode = GhostMode.shared_facet, + partitioner: typing.Optional[CellPartionerType] = None, ) -> Mesh: """Create a mesh of a unit cube. @@ -803,7 +827,7 @@ def create_unit_cube( def entities_to_geometry( - mesh: Mesh, dim: int, entities: npt.NDArray[np.int32], permute=False + mesh: Mesh, dim: int, entities: npt.NDArray[np.int32], permute: bool = False ) -> npt.NDArray[np.int32]: """Compute the geometric DOFs associated with the closure of the given mesh entities. diff --git a/python/dolfinx/nls/petsc.py b/python/dolfinx/nls/petsc.py index 33f1f22a7fa..0173419581f 100644 --- a/python/dolfinx/nls/petsc.py +++ b/python/dolfinx/nls/petsc.py @@ -41,11 +41,11 @@ def __init__(self, comm: MPI.Intracomm, problem: NonlinearProblem): self.setF(problem.F, self._b) self.set_form(problem.form) - def __del__(self): + def __del__(self) -> None: self._A.destroy() self._b.destroy() - def solve(self, u: fem.Function): + def solve(self, u: fem.Function) -> tuple[int, bool]: """Solve non-linear problem into function u. Returns the number of iterations and if the solver converged.""" n, converged = super().solve(u.x.petsc_vec) diff --git a/python/dolfinx/pkgconfig.py b/python/dolfinx/pkgconfig.py index 41150a9fbef..075f3941a57 100644 --- a/python/dolfinx/pkgconfig.py +++ b/python/dolfinx/pkgconfig.py @@ -14,7 +14,7 @@ import subprocess -def _pkgconfig_query(s): +def _pkgconfig_query(s: str) -> tuple[int, str]: pkg_config_exe = os.environ.get("PKG_CONFIG", None) or "pkg-config" cmd = [pkg_config_exe, *s.split()] proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) @@ -23,12 +23,12 @@ def _pkgconfig_query(s): return (rc, out.rstrip().decode("utf-8")) -def exists(package) -> bool: +def exists(package: str) -> bool: """Test for the existence of a pkg-config file for a named package.""" return _pkgconfig_query("--exists " + package)[0] == 0 -def parse(package): +def parse(package: str) -> dict[str, list[str]]: """Return a dict containing compile-time definitions.""" parse_map = { "-D": "define_macros", @@ -37,7 +37,7 @@ def parse(package): "-l": "libraries", } - result = {x: [] for x in parse_map.values()} + result: dict[str, list[str]] = {x: [] for x in parse_map.values()} # Execute the query to pkg-config and clean the result out = _pkgconfig_query(package + " --cflags --libs")[1] diff --git a/python/dolfinx/plot.py b/python/dolfinx/plot.py index ef90cb1c2ca..c57395e4d14 100644 --- a/python/dolfinx/plot.py +++ b/python/dolfinx/plot.py @@ -6,9 +6,10 @@ """Support functions for plotting""" import functools -import typing +from typing import Optional import numpy as np +import numpy.typing as npt from dolfinx import cpp as _cpp from dolfinx import fem, mesh @@ -31,7 +32,9 @@ @functools.singledispatch -def vtk_mesh(msh: mesh.Mesh, dim: typing.Optional[int] = None, entities=None): +def vtk_mesh( + msh: mesh.Mesh, dim: Optional[int] = None, entities: Optional[npt.NDArray[np.int32]] = None +) -> tuple[mesh.Topology, npt.NDArray[np.int8], npt.NDArray[np.floating]]: """Create vtk mesh topology data for mesh entities of a given dimension. The vertex indices in the returned topology array are the indices for the associated entry in the mesh geometry. @@ -75,7 +78,9 @@ def vtk_mesh(msh: mesh.Mesh, dim: typing.Optional[int] = None, entities=None): @vtk_mesh.register(fem.FunctionSpace) -def _(V: fem.FunctionSpace, entities=None): +def _( + V: fem.FunctionSpace, entities: Optional[npt.NDArray[np.int8] | range] = None +) -> tuple[mesh.Topology, npt.NDArray[np.int8], npt.NDArray[np.floating]]: """Creates a VTK mesh topology (topology array and array of cell types) that is based on the degree-of-freedom coordinates. diff --git a/python/dolfinx/wrappers/__init__.py b/python/dolfinx/wrappers/__init__.py index 9757a1d6da5..54aadac0876 100644 --- a/python/dolfinx/wrappers/__init__.py +++ b/python/dolfinx/wrappers/__init__.py @@ -4,9 +4,9 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later +from pathlib import Path -def get_include_path(): - """Return path to nanobind wrapper header files""" - import pathlib - return pathlib.Path(__file__).parent +def get_include_path() -> Path: + """Return path to nanobind wrapper header files""" + return Path(__file__).parent diff --git a/python/pyproject.toml b/python/pyproject.toml index baed60b2948..ad58679e5dc 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -71,9 +71,9 @@ markers = [ [tool.mypy] # Suggested at https://blog.wolt.com/engineering/2021/09/30/professional-grade-mypy-configuration/ # Goal would be to make all of the below True long-term -disallow_untyped_defs = false +disallow_untyped_defs = true disallow_any_unimported = false -no_implicit_optional = false +no_implicit_optional = true check_untyped_defs = false warn_return_any = false warn_unused_ignores = false