Skip to content

Commit 8c8a410

Browse files
schnellerhaseschnellerhase
committed
fix: petsc...
1 parent 15e3fe9 commit 8c8a410

File tree

1 file changed

+52
-46
lines changed

1 file changed

+52
-46
lines changed

‎python/dolfinx/fem/petsc.py‎

Lines changed: 52 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ def create_matrix(
195195
else: # Array of 'kind' types
196196
return _cpp.fem.petsc.create_matrix_nest(_a, kind)
197197
else: # Single form
198-
return _cpp.fem.petsc.create_matrix(a._cpp_object, kind)
198+
return _cpp.fem.petsc.create_matrix(a._cpp_object, kind) # type: ignore
199199

200200

201201
# -- Vector assembly ------------------------------------------------------
@@ -265,7 +265,7 @@ def assemble_vector(
265265

266266

267267
@assemble_vector.register
268-
def _(
268+
def _( # type: ignore
269269
b: PETSc.Vec, # type: ignore[name-defined]
270270
L: Form | Sequence[Form],
271271
constants: npt.NDArray | Sequence[npt.NDArray] | None = None,
@@ -316,10 +316,16 @@ def _(
316316
elif isinstance(L, Sequence):
317317
constants = pack_constants(L) if constants is None else constants
318318
coeffs = pack_coefficients(L) if coeffs is None else coeffs
319-
offset0, offset1 = b.getAttr("_blocks")
319+
offset0, offset1 = b.getAttr("_blocks") # type: ignore
320320
with b.localForm() as b_l:
321321
for L_, const, coeff, off0, off1, offg0, offg1 in zip(
322-
L, constants, coeffs, offset0, offset0[1:], offset1, offset1[1:]
322+
L,
323+
constants,
324+
coeffs,
325+
offset0, # type: ignore
326+
offset0[1:], # type: ignore
327+
offset1, # type: ignore
328+
offset1[1:], # type: ignore
323329
):
324330
bx_ = np.zeros((off1 - off0) + (offg1 - offg0), dtype=PETSc.ScalarType) # type: ignore[attr-defined]
325331
_assemble_vector_array(bx_, L_, const, coeff) # type: ignore[arg-type]
@@ -403,8 +409,8 @@ def assemble_matrix(
403409

404410

405411
@assemble_matrix.register
406-
def _(
407-
A: PETSc.Mat, # type: ignore[name-defined]
412+
def _( # type: ignore
413+
A: PETSc.Mat,
408414
a: Form | Sequence[Sequence[Form]],
409415
bcs: Sequence[DirichletBC] | None = None,
410416
diag: float = 1,
@@ -468,7 +474,7 @@ def _(
468474
for i, a_row in enumerate(a):
469475
for j, a_sub in enumerate(a_row):
470476
if a_sub is not None:
471-
Asub = A.getLocalSubMatrix(is0[i], is1[j])
477+
Asub = A.getLocalSubMatrix(is0[i], is1[j]) # type: ignore
472478
_cpp.fem.petsc.assemble_matrix(
473479
Asub,
474480
a_sub._cpp_object,
@@ -477,7 +483,7 @@ def _(
477483
_bcs,
478484
True,
479485
)
480-
A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
486+
A.restoreLocalSubMatrix(is0[i], is1[j], Asub) # type: ignore
481487
elif i == j:
482488
for bc in _bcs:
483489
row_forms = [row_form for row_form in a_row if row_form is not None]
@@ -496,15 +502,15 @@ def _(
496502
for i, a_row in enumerate(a):
497503
for j, a_sub in enumerate(a_row):
498504
if a_sub is not None:
499-
Asub = A.getLocalSubMatrix(is0[i], is1[j])
505+
Asub = A.getLocalSubMatrix(is0[i], is1[j]) # type: ignore
500506
if a_sub.function_spaces[0] is a_sub.function_spaces[1]:
501507
_cpp.fem.petsc.insert_diagonal(Asub, a_sub.function_spaces[0], _bcs, diag)
502-
A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
508+
A.restoreLocalSubMatrix(is0[i], is1[j], Asub) # type: ignore
503509
else: # Non-blocked
504510
constants = pack_constants(a) if constants is None else constants # type: ignore[assignment]
505511
coeffs = pack_coefficients(a) if coeffs is None else coeffs # type: ignore[assignment]
506512
_bcs = [bc._cpp_object for bc in bcs] if bcs is not None else []
507-
_cpp.fem.petsc.assemble_matrix(A, a._cpp_object, constants, coeffs, _bcs)
513+
_cpp.fem.petsc.assemble_matrix(A, a._cpp_object, constants, coeffs, _bcs) # type: ignore
508514
if a.function_spaces[0] is a.function_spaces[1]:
509515
A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined]
510516
A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined]
@@ -606,7 +612,7 @@ def apply_lifting(
606612
else:
607613
xlocal = None
608614

609-
offset0, offset1 = b.getAttr("_blocks")
615+
offset0, offset1 = b.getAttr("_blocks") # type: ignore
610616
with b.localForm() as b_l:
611617
for i, (a_, off0, off1, offg0, offg1) in enumerate(
612618
zip(a, offset0, offset0[1:], offset1, offset1[1:])
@@ -629,8 +635,6 @@ def apply_lifting(
629635
b_local = stack.enter_context(b.localForm())
630636
_apply_lifting(b_local.array_w, a, bcs, x0_r, alpha, constants, coeffs) # type: ignore[arg-type]
631637

632-
return b
633-
634638

635639
def set_bc(
636640
b: PETSc.Vec, # type: ignore[name-defined]
@@ -664,19 +668,19 @@ def set_bc(
664668
return
665669

666670
if not isinstance(bcs[0], Sequence):
667-
x0 = x0.array_r if x0 is not None else None
671+
x0 = x0.array_r if x0 is not None else None # type: ignore
668672
for bc in bcs:
669-
bc.set(b.array_w, x0, alpha) # type: ignore[union-attr]
673+
bc.set(b.array_w, x0, alpha) # type: ignore
670674
elif b.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined]
671675
_b = b.getNestSubVecs()
672-
x0 = len(_b) * [None] if x0 is None else x0.getNestSubVecs()
673-
for b_sub, bc, x_sub in zip(_b, bcs, x0): # type: ignore[assignment, arg-type]
676+
x0 = len(_b) * [None] if x0 is None else x0.getNestSubVecs() # type: ignore
677+
for b_sub, bc, x_sub in zip(_b, bcs, x0): # type: ignore
674678
set_bc(b_sub, bc, x_sub, alpha) # type: ignore[arg-type]
675679
else: # block vector
676-
offset0, _ = b.getAttr("_blocks")
680+
offset0, _ = b.getAttr("_blocks") # type: ignore
677681
b_array = b.getArray(readonly=False)
678682
x_array = x0.getArray(readonly=True) if x0 is not None else None
679-
for bcs, off0, off1 in zip(bcs, offset0, offset0[1:]): # type: ignore[assignment]
683+
for bcs, off0, off1 in zip(bcs, offset0, offset0[1:]): # type: ignore
680684
x0_sub = x_array[off0:off1] if x0 is not None else None # type: ignore[index]
681685
for bc in bcs:
682686
bc.set(b_array[off0:off1], x0_sub, alpha) # type: ignore[arg-type, union-attr]
@@ -855,13 +859,13 @@ def __init__(
855859
opts.prefixPush(self.solver.getOptionsPrefix())
856860

857861
for k, v in petsc_options.items():
858-
opts[k] = v
862+
opts[k] = v # type: ignore
859863

860864
self.solver.setFromOptions()
861865

862866
# Tidy up global options
863867
for k in petsc_options.keys():
864-
del opts[k]
868+
del opts[k] # type: ignore
865869

866870
opts.prefixPop()
867871

@@ -921,17 +925,17 @@ def solve(self) -> _Function | Sequence[_Function]:
921925
apply_lifting(self.b, self.a, bcs=bcs1) # type: ignore[arg-type]
922926
dolfinx.la.petsc._ghost_update(
923927
self.b,
924-
PETSc.InsertMode.ADD, # type: ignore[attr-defined]
925-
PETSc.ScatterMode.REVERSE, # type: ignore[attr-defined]
928+
PETSc.InsertMode.ADD, # type: ignore
929+
PETSc.ScatterMode.REVERSE, # type: ignore
926930
)
927931
bcs0 = _bcs_by_block(_extract_function_spaces(self.L), self.bcs) # type: ignore[arg-type]
928932
dolfinx.fem.petsc.set_bc(self.b, bcs0)
929933
else: # single form
930934
apply_lifting(self.b, [self.a], bcs=[self.bcs]) # type: ignore[arg-type]
931935
dolfinx.la.petsc._ghost_update(
932936
self.b,
933-
PETSc.InsertMode.ADD, # type: ignore[attr-defined]
934-
PETSc.ScatterMode.REVERSE, # type: ignore[attr-defined]
937+
PETSc.InsertMode.ADD, # type: ignore
938+
PETSc.ScatterMode.REVERSE, # type: ignore
935939
)
936940
for bc in self.bcs:
937941
bc.set(self.b.array_w)
@@ -940,8 +944,8 @@ def solve(self) -> _Function | Sequence[_Function]:
940944

941945
# Solve linear system and update ghost values in the solution
942946
self.solver.solve(self.b, self.x)
943-
dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined]
944-
dolfinx.fem.petsc.assign(self.x, self.u)
947+
dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
948+
dolfinx.fem.petsc.assign(self.x, self.u) # type: ignore
945949
return self.u
946950

947951
@property
@@ -965,7 +969,7 @@ def A(self) -> PETSc.Mat: # type: ignore[name-defined]
965969
return self._A
966970

967971
@property
968-
def P_mat(self) -> PETSc.Mat: # type: ignore[name-defined]
972+
def P_mat(self) -> PETSc.Mat | None: # type: ignore[name-defined]
969973
"""Preconditioner matrix."""
970974
return self._P_mat
971975

@@ -1055,10 +1059,10 @@ def assemble_residual(
10551059
b: Vector to assemble the residual into.
10561060
"""
10571061
# Update input vector before assigning
1058-
dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined]
1062+
dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
10591063

10601064
# Assign the input vector to the unknowns
1061-
assign(x, u)
1065+
assign(x, u) # type: ignore
10621066

10631067
# Assign block data if block assembly is requested
10641068
if isinstance(residual, Sequence) and b.getType() != PETSc.Vec.Type.NEST: # type: ignore[attr-defined]
@@ -1073,16 +1077,16 @@ def assemble_residual(
10731077
if isinstance(jacobian, Sequence):
10741078
# Nest and blocked lifting
10751079
bcs1 = _bcs_by_block(_extract_function_spaces(jacobian, 1), bcs) # type: ignore[arg-type]
1076-
apply_lifting(b, jacobian, bcs=bcs1, x0=x, alpha=-1.0)
1077-
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined]
1080+
apply_lifting(b, jacobian, bcs=bcs1, x0=x, alpha=-1.0) # type: ignore
1081+
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
10781082
bcs0 = _bcs_by_block(_extract_function_spaces(residual), bcs) # type: ignore[arg-type]
10791083
set_bc(b, bcs0, x0=x, alpha=-1.0)
10801084
else:
10811085
# Single form lifting
10821086
apply_lifting(b, [jacobian], bcs=[bcs], x0=[x], alpha=-1.0)
1083-
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[attr-defined]
1087+
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
10841088
set_bc(b, bcs, x0=x, alpha=-1.0)
1085-
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined]
1089+
dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
10861090

10871091

10881092
def assemble_jacobian(
@@ -1123,8 +1127,8 @@ def assemble_jacobian(
11231127
"""
11241128
# Copy existing soultion into the function used in the residual and
11251129
# Jacobian
1126-
dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined]
1127-
assign(x, u)
1130+
dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
1131+
assign(x, u) # type: ignore
11281132

11291133
# Assemble Jacobian
11301134
J.zeroEntries()
@@ -1155,6 +1159,8 @@ class NonlinearProblem:
11551159
``.destroy()`` on returned PETSc objects.
11561160
""" # noqa: D301
11571161

1162+
_P_mat: PETSc.Mat | None
1163+
11581164
def __init__(
11591165
self,
11601166
F: ufl.form.Form | Sequence[ufl.form.Form],
@@ -1302,13 +1308,13 @@ def __init__(
13021308
opts.prefixPush(self.solver.getOptionsPrefix())
13031309

13041310
for k, v in petsc_options.items():
1305-
opts[k] = v
1311+
opts[k] = v # type: ignore
13061312

13071313
self.solver.setFromOptions()
13081314

13091315
# Tidy up global options
13101316
for k in petsc_options.keys():
1311-
del opts[k]
1317+
del opts[k] # type: ignore
13121318

13131319
opts.prefixPop()
13141320

@@ -1345,10 +1351,10 @@ def solve(self) -> _Function | Sequence[_Function]:
13451351

13461352
# Solve problem
13471353
self.solver.solve(None, self.x)
1348-
dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore[attr-defined]
1354+
dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
13491355

13501356
# Copy solution back to function
1351-
assign(self.x, self.u)
1357+
assign(self.x, self.u) # type: ignore
13521358

13531359
return self.u
13541360

@@ -1517,7 +1523,7 @@ def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None: # type: ignore[name-defined]
15171523
"""
15181524
# Reset the residual vector
15191525
dolfinx.la.petsc._zero_vector(b)
1520-
assemble_vector(b, self._L)
1526+
assemble_vector(b, self._L) # type: ignore
15211527

15221528
# Apply boundary condition
15231529
if self.bcs is not None:
@@ -1628,7 +1634,7 @@ def assign(u: _Function | Sequence[_Function], x: PETSc.Vec): # type: ignore[na
16281634

16291635

16301636
@assign.register
1631-
def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-defined]
1637+
def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-defined, misc]
16321638
"""Assign vector entries to :class:`Function` degrees-of-freedom.
16331639
16341640
Assigns values in ``x`` to the degrees-of-freedom of ``u``, which is
@@ -1642,7 +1648,7 @@ def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-de
16421648
u: ``Function`` (s) to assign degree-of-freedom values to.
16431649
"""
16441650
if x.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined]
1645-
dolfinx.la.petsc.assign(x, [v.x.array for v in u])
1651+
dolfinx.la.petsc.assign(x, [v.x.array for v in u]) # type: ignore
16461652
else:
16471653
if isinstance(u, Sequence):
16481654
data0, data1 = [], []
@@ -1651,9 +1657,9 @@ def _(x: PETSc.Vec, u: _Function | Sequence[_Function]): # type: ignore[name-de
16511657
n = v.function_space.dofmap.index_map.size_local
16521658
data0.append(v.x.array[: bs * n])
16531659
data1.append(v.x.array[bs * n :])
1654-
dolfinx.la.petsc.assign(x, data0 + data1)
1660+
dolfinx.la.petsc.assign(x, data0 + data1) # type: ignore
16551661
else:
1656-
dolfinx.la.petsc.assign(x, u.x.array)
1662+
dolfinx.la.petsc.assign(x, u.x.array) # type: ignore
16571663

16581664

16591665
def get_petsc_lib() -> pathlib.Path:

0 commit comments

Comments
 (0)