diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index 0cdc292313..e9a1d094f9 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -23,7 +23,8 @@ jobs: env: PETSC_ARCH: arch-darwin-c-opt PETSC_DIR: ${{ github.workspace }}/petsc - PETSC_VERSION: 3.25.3 + # PETSC_VERSION: "3.25.3" + PETSC_VERSION: "petsc4py-stub-fixes" PYTEST_ADDOPTS: "-W error" steps: @@ -64,7 +65,7 @@ jobs: run: | export PATH="$(brew --prefix gfortran)/bin:$(brew --prefix bison)/bin:$PATH" export PATH="$(brew --prefix make)/libexec/gnubin:$PATH" - git clone --depth 1 --branch v${PETSC_VERSION} https://gitlab.com/petsc/petsc.git petsc + git clone --depth 1 --branch ${PETSC_VERSION} https://gitlab.com/paul.kuehner/petsc.git petsc cd petsc ./configure \ --with-64-bit-indices=no \ @@ -131,6 +132,14 @@ jobs: python -m scikit_build_core.build requires | python -c "import sys, json; print(' '.join(json.load(sys.stdin)))" | xargs pip install pip install --check-build-dependencies --no-build-isolation --config-settings=cmake.build-type="Developer" './[test]' + - name: Run mypy + working-directory: python + run: | + pip install mypy types-cffi scipy-stubs + mypy -p dolfinx + mypy test + mypy demo + - name: Basic test run: | mpiexec -np 1 python -c "from mpi4py import MPI; import dolfinx; assert dolfinx.has_parmetis and dolfinx.has_petsc and dolfinx.has_superlu_dist and dolfinx.has_ptscotch" diff --git a/python/dolfinx/fem/petsc.py b/python/dolfinx/fem/petsc.py index d1f9e36a34..e080f4b459 100644 --- a/python/dolfinx/fem/petsc.py +++ b/python/dolfinx/fem/petsc.py @@ -265,7 +265,7 @@ def assemble_vector( @assemble_vector.register def _( - b: PETSc.Vec, # type: ignore[name-defined] + b: PETSc.Vec, L: Form | Sequence[Form], constants: npt.NDArray | Sequence[npt.NDArray] | None = None, coeffs: ( @@ -273,7 +273,7 @@ def _( | Sequence[dict[tuple[IntegralType, int], npt.NDArray]] | None ) = None, -) -> PETSc.Vec: # type: ignore[name-defined] +) -> PETSc.Vec: """Assemble linear form(s) into a PETSc vector. The vector ``b`` must have been initialized with a size/layout that @@ -304,7 +304,7 @@ def _( Returns: Assembled vector. """ - if b.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + if b.getType() == PETSc.Vec.Type.NEST: if not isinstance(L, Sequence): raise ValueError("Must provide a sequence of forms when assembling a nest vector") constants = [None] * len(L) if constants is None else constants # type: ignore[list-item] @@ -329,7 +329,7 @@ def _( offset1[1:], strict=True, ): - bx_ = np.zeros((off1 - off0) + (offg1 - offg0), dtype=PETSc.ScalarType) # type: ignore[attr-defined] + bx_ = np.zeros((off1 - off0) + (offg1 - offg0), dtype=PETSc.ScalarType) _assemble_vector_array(bx_, L_, const, coeff) # type: ignore[arg-type] size = off1 - off0 b_l.array_w[off0:off1] += bx_[:size] @@ -500,7 +500,7 @@ def _( ) # Flush to enable switch from add to set in the matrix - A.assemble(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined] + A.assemble(PETSc.Mat.AssemblyType.FLUSH) # Set diagonal for i, a_row in enumerate(a): @@ -516,8 +516,8 @@ def _( _bcs = [bc._cpp_object for bc in bcs] if bcs is not None else [] _cpp.fem.petsc.assemble_matrix(A, a._cpp_object, constants, coeffs, _bcs) if a.function_spaces[0] is a.function_spaces[1]: - A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined] - A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined] + A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH) + A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH) _cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diag) return A @@ -527,10 +527,10 @@ def _( def apply_lifting( - b: PETSc.Vec, # type: ignore[name-defined] + b: PETSc.Vec, a: Sequence[Form] | Sequence[Sequence[Form]], bcs: Sequence[DirichletBC] | Sequence[Sequence[DirichletBC]] | None, - x0: Sequence[PETSc.Vec] | None = None, # type: ignore[name-defined] + x0: PETSc.Vec | Sequence[PETSc.Vec] | None = None, alpha: float = 1, constants: Sequence[npt.NDArray] | Sequence[Sequence[npt.NDArray]] | None = None, coeffs: ( @@ -592,8 +592,10 @@ def apply_lifting( function. Use :func:`dolfinx.fem.DirichletBC.set` to set values in ``b``. """ - if b.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] - x0 = [] if x0 is None else x0.getNestSubVecs() # type: ignore[attr-defined] + if b.getType() == PETSc.Vec.Type.NEST: + assert isinstance(x0, PETSc.Vec | None) + x0 = [] if x0 is None else x0.getNestSubVecs() + constants = [pack_constants(forms) for forms in a] if constants is None else constants # type: ignore[assignment] coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs # type: ignore[misc] for b_sub, a_sub, const, coeff in zip( @@ -604,10 +606,7 @@ def apply_lifting( strict=True, ): const_ = list( - map( - lambda x: np.array([], dtype=PETSc.ScalarType) if x is None else x, - const, - ) # type: ignore[attr-defined, call-overload] + map(lambda x: np.array([], dtype=PETSc.ScalarType) if x is None else x, const) # type: ignore[call-overload] ) apply_lifting(b_sub, a_sub, bcs, x0, alpha, const_, coeff) # type: ignore[arg-type] else: @@ -652,9 +651,9 @@ def apply_lifting( def set_bc( - b: PETSc.Vec, # type: ignore[name-defined] + b: PETSc.Vec, bcs: Sequence[DirichletBC] | Sequence[Sequence[DirichletBC]], - x0: PETSc.Vec | None = None, # type: ignore[name-defined] + x0: PETSc.Vec | None = None, alpha: float = 1, ) -> None: """Set constraint (Dirchlet boundary condition) values in an vector. @@ -686,7 +685,7 @@ def set_bc( x0 = x0.array_r if x0 is not None else None for bc in bcs: bc.set(b.array_w, x0, alpha) # type: ignore[union-attr] - elif b.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + elif b.getType() == PETSc.Vec.Type.NEST: _b = b.getNestSubVecs() x0 = len(_b) * [None] if x0 is None else x0.getNestSubVecs() for b_sub, bc, x_sub in zip(_b, bcs, x0, strict=True): # type: ignore[assignment, arg-type] @@ -810,14 +809,14 @@ def __init__( """ self._a = _create_form( a, - dtype=PETSc.ScalarType, # type: ignore[attr-defined] + dtype=PETSc.ScalarType, form_compiler_options=form_compiler_options, jit_options=jit_options, entity_maps=entity_maps, ) self._L = _create_form( L, - dtype=PETSc.ScalarType, # type: ignore[attr-defined] + dtype=PETSc.ScalarType, form_compiler_options=form_compiler_options, jit_options=jit_options, entity_maps=entity_maps, @@ -825,7 +824,7 @@ def __init__( self._A = create_matrix(self._a, kind=kind) self._preconditioner = _create_form( P, # type: ignore[arg-type] - dtype=PETSc.ScalarType, # type: ignore[attr-defined] + dtype=PETSc.ScalarType, form_compiler_options=form_compiler_options, jit_options=jit_options, entity_maps=entity_maps, @@ -837,7 +836,7 @@ def __init__( ) # For nest matrices kind can be a nested list. - kind = "nest" if self.A.getType() == PETSc.Mat.Type.NEST else kind # type: ignore[attr-defined] + kind = "nest" if self.A.getType() == PETSc.Mat.Type.NEST else kind assert kind is None or isinstance(kind, str) self._b = create_vector(_extract_function_spaces(self.L), kind=kind) # type: ignore self._x = create_vector(_extract_function_spaces(self.L), kind=kind) # type: ignore @@ -855,7 +854,7 @@ def __init__( self.bcs = [] if bcs is None else bcs - self._solver = PETSc.KSP().create(self.A.comm) # type: ignore[attr-defined] + self._solver = PETSc.KSP().create(self.A.comm) self.solver.setOperators(self.A, self.P_mat) if petsc_options_prefix == "": @@ -871,7 +870,7 @@ def __init__( # Set options on KSP only if petsc_options is not None: - opts = PETSc.Options() # type: ignore[attr-defined] + opts = PETSc.Options() opts.prefixPush(self.solver.getOptionsPrefix()) for k, v in petsc_options.items(): @@ -942,8 +941,8 @@ def solve(self) -> _Function | Sequence[_Function]: apply_lifting(self.b, self.a, bcs=bcs1) # type: ignore[arg-type] dolfinx.la.petsc._ghost_update( self.b, - PETSc.InsertMode.ADD, # type: ignore[attr-defined] - PETSc.ScatterMode.REVERSE, # type: ignore[attr-defined] + PETSc.InsertMode.ADD, + PETSc.ScatterMode.REVERSE, ) bcs0 = _bcs_by_block(_extract_function_spaces(self.L), self.bcs) # type: ignore[arg-type] dolfinx.fem.petsc.set_bc(self.b, bcs0) @@ -951,17 +950,17 @@ def solve(self) -> _Function | Sequence[_Function]: apply_lifting(self.b, [self.a], bcs=[self.bcs]) # type: ignore[arg-type] dolfinx.la.petsc._ghost_update( self.b, - PETSc.InsertMode.ADD, # type: ignore[attr-defined] - PETSc.ScatterMode.REVERSE, # type: ignore[attr-defined] + PETSc.InsertMode.ADD, + PETSc.ScatterMode.REVERSE, ) for bc in self.bcs: bc.set(self.b.array_w) else: - dolfinx.la.petsc._ghost_update(self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # Solve linear system and update ghost values in the solution self.solver.solve(self.b, self.x) - dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) dolfinx.fem.petsc.assign(self.x, self.u) return self.u @@ -1064,16 +1063,16 @@ def assemble_residual( format of this argument. """ # Update input vector before assigning - dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # Assign the input vector to the unknowns assign(x, u) # Assign block data if block assembly is requested - if isinstance(residual, Sequence) and b.getType() != PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + if isinstance(residual, Sequence) and b.getType() != PETSc.Vec.Type.NEST: assert _blocks is not None, "Block data must be provided for block assembly." - b.setAttr("_blocks", _blocks) # type: ignore[attr-defined] - x.setAttr("_blocks", _blocks) # type: ignore[attr-defined] + b.setAttr("_blocks", _blocks) + x.setAttr("_blocks", _blocks) # Assemble the residual dolfinx.la.petsc._zero_vector(b) @@ -1084,15 +1083,15 @@ def assemble_residual( # Nest and blocked lifting bcs1 = _bcs_by_block(_extract_function_spaces(jacobian, 1), bcs) # type: ignore[arg-type] apply_lifting(b, jacobian, bcs=bcs1, x0=x, alpha=-1.0) - dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) bcs0 = _bcs_by_block(_extract_function_spaces(residual), bcs) # type: ignore[arg-type] set_bc(b, bcs0, x0=x, alpha=-1.0) else: # Single form lifting apply_lifting(b, [jacobian], bcs=[bcs], x0=[x], alpha=-1.0) - dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) set_bc(b, bcs, x0=x, alpha=-1.0) - dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) def assemble_jacobian( @@ -1135,7 +1134,7 @@ def assemble_jacobian( """ # Copy existing solution into the function used in the residual and # Jacobian - dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) assign(x, u) # Assemble Jacobian @@ -1286,14 +1285,14 @@ def __init__( self._P_mat = None # Determine the vector kind based on the matrix type - kind = "nest" if self._A.getType() == PETSc.Mat.Type.NEST else kind # type: ignore[attr-defined] + kind = "nest" if self._A.getType() == PETSc.Mat.Type.NEST else kind assert kind is None or isinstance(kind, str) self._b = create_vector(_extract_function_spaces(self.F), kind=kind) # type: ignore self._x = create_vector(_extract_function_spaces(self.F), kind=kind) # type: ignore # Create the SNES solver and attach the corresponding Jacobian and # residual computation functions - self._snes = PETSc.SNES().create(self.A.comm) # type: ignore[attr-defined] + self._snes = PETSc.SNES().create(self.A.comm) jacobian_ctx = { "u": self.u, "jacobian": self.J, @@ -1320,7 +1319,7 @@ def __init__( # Set options for SNES only if petsc_options is not None: - opts = PETSc.Options() # type: ignore[attr-defined] + opts = PETSc.Options() opts.prefixPush(self.solver.getOptionsPrefix()) for k, v in petsc_options.items(): @@ -1367,7 +1366,7 @@ def solve(self) -> _Function | Sequence[_Function]: # Solve problem self.solver.solve(None, self.x) - dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # Copy solution back to function assign(self.x, self.u) @@ -1398,22 +1397,22 @@ def preconditioner(self) -> Form | Sequence[Sequence[Form]] | None: return self._preconditioner @property - def A(self) -> PETSc.Mat: # type: ignore[name-defined] + def A(self) -> PETSc.Mat: """Jacobian matrix.""" return self._A @property - def P_mat(self) -> PETSc.Mat | None: # type: ignore[name-defined] + def P_mat(self) -> PETSc.Mat | None: """Preconditioner matrix.""" return self._P_mat @property - def b(self) -> PETSc.Vec: # type: ignore[name-defined] + def b(self) -> PETSc.Vec: """Residual vector.""" return self._b @property - def x(self) -> PETSc.Vec: # type: ignore[name-defined] + def x(self) -> PETSc.Vec: """Solution vector. Note: @@ -1423,7 +1422,7 @@ def x(self) -> PETSc.Vec: # type: ignore[name-defined] return self._x @property - def solver(self) -> PETSc.SNES: # type: ignore[name-defined] + def solver(self) -> PETSc.SNES: """The SNES solver.""" return self._snes @@ -1521,7 +1520,7 @@ def a(self) -> Form: """The compiled bilinear form (the Jacobian form).""" return self._a - def form(self, x: PETSc.Vec) -> None: # type: ignore[name-defined] + def form(self, x: PETSc.Vec) -> None: """Function called before the residual or Jacobian is computed. This is usually used to update ghost values. @@ -1529,9 +1528,9 @@ def form(self, x: PETSc.Vec) -> None: # type: ignore[name-defined] Args: x: The vector containing the latest solution. """ - x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined] + x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) - def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None: # type: ignore[name-defined] + def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None: """Assemble the residual F into the vector b. Args: @@ -1545,12 +1544,12 @@ def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None: # type: ignore[name-defined] # Apply boundary condition if self.bcs is not None: apply_lifting(b, [self._a], bcs=[self.bcs], x0=[x], alpha=-1.0) - b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined] + b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(b, self.bcs, x, -1.0) else: - b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined] + b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) - def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None: # type: ignore[name-defined] + def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None: """Assemble the Jacobian matrix. Args: @@ -1565,7 +1564,7 @@ def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None: # type: ignore[name-defined] # -- Additional free helper functions (interpolations, assignments etc.) -- -def discrete_curl(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat: # type: ignore[name-defined] +def discrete_curl(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat: """Assemble a discrete curl operator. Args: @@ -1578,7 +1577,7 @@ def discrete_curl(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat: return _discrete_curl(space0._cpp_object, space1._cpp_object) -def discrete_gradient(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat: # type: ignore[name-defined] +def discrete_gradient(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat: """Assemble a discrete gradient operator. The discrete gradient operator interpolates the gradient of a H1 @@ -1596,7 +1595,7 @@ def discrete_gradient(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.M return _discrete_gradient(space0._cpp_object, space1._cpp_object) -def interpolation_matrix(V0: _FunctionSpace, V1: _FunctionSpace) -> PETSc.Mat: # type: ignore[name-defined] +def interpolation_matrix(V0: _FunctionSpace, V1: _FunctionSpace) -> PETSc.Mat: """Create an interpolation operator between finite element spaces. Consider is the vector of degrees-of-freedom :math:`u_{i}` @@ -1635,7 +1634,7 @@ def assign(u: _Function | Sequence[_Function], x: PETSc.Vec): # type: ignore[na u: ``Function`` (s) to assign degree-of-freedom value from. x: Vector to assign degree-of-freedom values in ``u`` to. """ - if x.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined] + if x.getType() == PETSc.Vec.Type().NEST: dolfinx.la.petsc.assign([v.x.array for v in u], x) else: if isinstance(u, Sequence): @@ -1651,7 +1650,7 @@ def assign(u: _Function | Sequence[_Function], x: PETSc.Vec): # type: ignore[na @assign.register -def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-defined] +def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): """Assign vector entries to :class:`Function` degrees-of-freedom. Assigns values in ``x`` to the degrees-of-freedom of ``u``, which is @@ -1664,7 +1663,7 @@ def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-de x: Vector with values to assign values from. u: ``Function`` (s) to assign degree-of-freedom values to. """ - if x.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined] + if x.getType() == PETSc.Vec.Type().NEST: dolfinx.la.petsc.assign(x, [v.x.array for v in u]) else: if isinstance(u, Sequence): @@ -1692,7 +1691,7 @@ def get_petsc_lib() -> pathlib.Path: import petsc4py as _petsc4py petsc_dir = _petsc4py.get_config()["PETSC_DIR"] - petsc_arch = _petsc4py.lib.getPathArchPETSc()[1] + petsc_arch = _petsc4py.lib.getPathArchPETSc()[1] # type: ignore candidate_paths = [ os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.so"), os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.dylib"), @@ -1743,9 +1742,9 @@ def set_vals(A: int, _llvmlite.binding.load_library_permanently(str(get_petsc_lib())) - _int = _numba.from_dtype(_PETSc.IntType) # type: ignore - _scalar = _numba.from_dtype(_PETSc.ScalarType) # type: ignore - _real = _numba.from_dtype(_PETSc.RealType) # type: ignore + _int = _numba.from_dtype(_PETSc.IntType) + _scalar = _numba.from_dtype(_PETSc.ScalarType) + _real = _numba.from_dtype(_PETSc.RealType) _int_ptr = _numba.core.types.CPointer(_int) _scalar_ptr = _numba.core.types.CPointer(_scalar) _MatSetValues_sig = _numba.core.typing.signature( @@ -1899,8 +1898,8 @@ def set_vals(A: int, np.longlong: "long long", } - _c_int_t = _CTYPES[_PETSc.IntType] # type: ignore - _c_scalar_t = _CTYPES[_PETSc.ScalarType] # type: ignore + _c_int_t = _CTYPES[_PETSc.IntType] + _c_scalar_t = _CTYPES[_PETSc.ScalarType] _ffi.cdef( f""" int MatSetValuesLocal(void* mat, {_c_int_t} nrow, const {_c_int_t}* irow, diff --git a/python/dolfinx/la/petsc.py b/python/dolfinx/la/petsc.py index cbdfc2937b..1196847160 100644 --- a/python/dolfinx/la/petsc.py +++ b/python/dolfinx/la/petsc.py @@ -32,9 +32,9 @@ __all__ = ["assign", "create_vector", "create_vector_wrap"] -def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): # type: ignore[name-defined] +def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): """Helper function for ghost updating PETSc vectors.""" - if x.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + if x.getType() == PETSc.Vec.Type.NEST: for x_sub in x.getNestSubVecs(): x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode) x_sub.destroy() @@ -42,9 +42,9 @@ def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PET x.ghostUpdate(addv=insert_mode, mode=scatter_mode) -def _zero_vector(x: PETSc.Vec): # type: ignore[name-defined] +def _zero_vector(x: PETSc.Vec): """Helper function for zeroing out PETSc vectors.""" - if x.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + if x.getType() == PETSc.Vec.Type.NEST: for x_sub in x.getNestSubVecs(): with x_sub.localForm() as x_sub_local: x_sub_local.set(0.0) @@ -54,7 +54,7 @@ def _zero_vector(x: PETSc.Vec): # type: ignore[name-defined] x_local.set(0.0) -def create_vector_wrap(x: Vector) -> PETSc.Vec: # type: ignore[name-defined] +def create_vector_wrap(x: Vector) -> PETSc.Vec: """Wrap a distributed DOLFINx vector as a PETSc vector. Args: @@ -67,14 +67,15 @@ def create_vector_wrap(x: Vector) -> PETSc.Vec: # type: ignore[name-defined] ghosts = index_map.ghosts.astype(PETSc.IntType) # type: ignore[attr-defined] bs = x.block_size size = (index_map.size_local * bs, index_map.size_global * bs) - return PETSc.Vec().createGhostWithArray( # type: ignore[attr-defined] + + return PETSc.Vec().createGhostWithArray( ghosts, x.array, size=size, bsize=bs, comm=index_map.comm ) def create_vector( maps: typing.Sequence[tuple[IndexMap, int]], kind: str | None = None -) -> PETSc.Vec: # type: ignore[name-defined] +) -> PETSc.Vec: """Create a PETSc vector from a sequence of maps and blocksizes. Three cases are supported: @@ -123,16 +124,16 @@ def create_vector( index_map, bs = maps[0] ghosts = index_map.ghosts.astype(PETSc.IntType) # type: ignore[attr-defined] size = (index_map.size_local * bs, index_map.size_global * bs) - b = PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm) # type: ignore - if kind == PETSc.Vec.Type.MPI: # type: ignore[attr-defined] + b = PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm) + if kind == PETSc.Vec.Type.MPI: _assign_block_data(maps, b) return b - if kind is None or kind == PETSc.Vec.Type.MPI: # type: ignore[attr-defined] + if kind is None or kind == PETSc.Vec.Type.MPI: b = dolfinx.cpp.fem.petsc.create_vector_block(maps) _assign_block_data(maps, b) return b - elif kind == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + elif kind == PETSc.Vec.Type.NEST: return dolfinx.cpp.fem.petsc.create_vector_nest(maps) else: raise NotImplementedError( @@ -144,8 +145,8 @@ def create_vector( @functools.singledispatch def assign( - x0: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]], - x1: PETSc.Vec, # type: ignore[name-defined] + x0: PETSc.Vec | npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]], + x1: PETSc.Vec, ): """Assign ``x0`` values to a PETSc vector ``x1``. @@ -172,8 +173,9 @@ def assign( x0: An array or list of arrays that will be assigned to ``x1``. x1: Vector to assign values to. """ - if x1.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined] + if x1.getType() == PETSc.Vec.Type().NEST: x1_nest = x1.getNestSubVecs() + assert isinstance(x0, PETSc.Vec) for _x0, _x1 in zip(x0, x1_nest, strict=True): with _x1.localForm() as x: x.array_w[:] = _x0 @@ -191,7 +193,7 @@ def assign( @assign.register def _( - x0: PETSc.Vec, # type: ignore[name-defined] + x0: PETSc.Vec, x1: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]], ): """Assign PETSc vector ``x0`` values to (blocked) array(s) ``x1``. @@ -204,11 +206,11 @@ def _( x0: Vector that will have its values assigned to ``x1``. x1: An array or list of arrays to assign to. """ - if x0.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined] + if x0.getType() == PETSc.Vec.Type().NEST: x0_nest = x0.getNestSubVecs() for _x0, _x1 in zip(x0_nest, x1, strict=True): with _x0.localForm() as x: - _x1[:] = x.array_r[:] # type: ignore[index] + _x1[:] = x.array_r[:] else: with x0.localForm() as _x0: if isinstance(x1, Sequence): @@ -221,7 +223,7 @@ def _( x1[:] = _x0.array_r[:] -def _assign_block_data(maps: Iterable[tuple[IndexMap, int]], vec: PETSc.Vec): # type: ignore[name-defined] +def _assign_block_data(maps: Iterable[tuple[IndexMap, int]], vec: PETSc.Vec): """Assign block data to a PETSc vector. Args: