Coverage for python/src/dolfinx_mpc/assemble

3# This file is part of DOLFINX_MPC

5# SPDX-License-Identifier: MIT

6from __future__ import annotations

8import contextlib

9from typing import Optional, Sequence, Union, cast

11from petsc4py import PETSc as _PETSc

13import dolfinx.cpp as _cpp

14import dolfinx.fem as _fem

15import numpy

16from dolfinx import default_scalar_type

17from dolfinx.common import Timer

18from dolfinx.la.petsc import _zero_vector

19from dolfinx.la.petsc import create_vector as _create_petsc_vector

21import dolfinx_mpc.cpp

23from ._kind import Kind, blocked_layout, deprecated

24from .dirichletbc import BCData

25from .multipointconstraint import MultiPointConstraint, _float_classes

28def _block_maps(constraints: Sequence[MultiPointConstraint]):

29 """Extended index map and block size of each block, as the vector layout is built from."""

30 return [

31 (mpc.function_space.dofmap.index_map._cpp_object, mpc.function_space.dofmap.index_map_bs) for mpc in constraints

32 ]

35def _is_block_vector(b: _PETSc.Vec) -> bool: # type: ignore

36 """Whether `b` is a monolithic vector of several blocks, as made by :func:`create_vector_block`."""

37 return b.getType() != "nest" and b.getAttr("_blocks") is not None

40def _block_offsets(b: _PETSc.Vec) -> tuple[Sequence[int], Sequence[int]]: # type: ignore

41 """Offsets of the owned and of the ghost entries of each block of the monolithic vector `b`."""

42 return cast("tuple[Sequence[int], Sequence[int]]", b.getAttr("_blocks"))

45def _block_scratch(b: _PETSc.Vec, k: int) -> numpy.ndarray: # type: ignore

46 """A zeroed array with the `[owned, ghosts]` layout of block `k` of the monolithic vector `b`."""

47 off_owned, off_ghost = _block_offsets(b)

48 size = (off_owned[k + 1] - off_owned[k]) + (off_ghost[k + 1] - off_ghost[k])

49 return numpy.zeros(size, dtype=_PETSc.ScalarType) # type: ignore

52def _add_to_block(b_local: numpy.ndarray, b: _PETSc.Vec, k: int, values: numpy.ndarray): # type: ignore

53 """Add `values`, laid out `[owned, ghosts]`, to block `k` of the local array of a monolithic vector."""

54 off_owned, off_ghost = _block_offsets(b)

55 size = off_owned[k + 1] - off_owned[k]

56 b_local[off_owned[k] : off_owned[k + 1]] += values[:size]

57 b_local[off_ghost[k] : off_ghost[k + 1]] += values[size:]

60@contextlib.contextmanager

61def _block_arrays(b: _PETSc.Vec): # type: ignore

62 """

63 The writable local array, owned and ghost entries, of every block of a nest or monolithic

64 vector.

66 For a monolithic vector, whose blocks are not contiguous, these are zeroed scratch arrays that

67 are added to `b` on exit. All assembly into them only adds, so this equals working in place.

68 """

69 if b.getType() == "nest":

70 with contextlib.ExitStack() as stack:

71 yield [stack.enter_context(b_k.localForm()).array_w for b_k in b.getNestSubVecs()]

72 else:

73 num_blocks = len(_block_offsets(b)[0]) - 1

74 scratch = [_block_scratch(b, k) for k in range(num_blocks)]

75 yield scratch

76 with b.localForm() as b_local:

77 for k, values in enumerate(scratch):

78 _add_to_block(b_local.array_w, b, k, values)

81@contextlib.contextmanager

82def _block_arrays_read(x: Optional[_PETSc.Vec], constraints: Sequence[MultiPointConstraint]): # type: ignore

83 """The local array, owned and ghost entries, of every block of a nest or monolithic vector."""

84 if x is None:

85 yield []

86 elif x.getType() == "nest":

87 with contextlib.ExitStack() as stack:

88 yield [stack.enter_context(x_k.localForm()).array_r for x_k in x.getNestSubVecs()]

89 else:

90 yield _cpp.la.petsc.get_local_vectors(x, _block_maps(constraints))

93def apply_lifting(

94 b: _PETSc.Vec,

95 form: Union[Sequence[Sequence[Optional[_fem.Form]]], Sequence[Optional[_fem.Form]]],

96 bcs: Union[Sequence[_fem.DirichletBC], Sequence[Sequence[_fem.DirichletBC]]],

97 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],

98 x0: Optional[Sequence[_PETSc.Vec]] = None,

99 scale: _float_classes = default_scalar_type(1.0), # type: ignore

100 num_threads: Optional[int] = 1,

101 bc_data: Optional[BCData] = None,

102): # type: ignore

103 """

104 Apply lifting of Dirichlet conditions to the vector `b`.

105

106 For a single vector,

107

108 .. math::

109

110 b \\leftarrow b - \\mathrm{scale}\\, K^T \\sum_j A_j (g_j - x0_j),

111

112 where :math:`A_j` is assembled from `form[j]`, :math:`g_j` holds the

113 Dirichlet values on its trial space and :math:`K` is the reduction matrix

114 of `constraint`. For a nest vector, block row :math:`i` is

115

116 .. math::

117

118 b_i \\leftarrow b_i - \\mathrm{scale}\\, K_i^T \\sum_j A_{ij} (g_j - x0_j),

119

120 where :math:`A_{ij}` is assembled from `form[i][j]` and :math:`K_i` is the

121 reduction matrix of `constraint[i]`. A `None` form is skipped.

122

123 Args:

124 b: PETSc vector to assemble into

125 form: The bilinear forms, `form[j]` (single vector) or `form[i][j]`

126 (nest vector) as above

127 bcs: List of Dirichlet boundary conditions. A condition contributes to

128 :math:`g_j` if it is defined on the trial space of column `j` or a

129 subspace of it, so the conditions may be given per block or as one

130 flat list.

131 constraint: The multi point constraint

132 x0: List of vectors

133 scale: Scaling for lifting

134 num_threads: The number of threads to use for certain operations

135 bc_data: A :class:`BCData` cache. Built from `bcs` when not supplied;

136 pass one to reuse the dof markers across calls.

137 """

138 t = Timer("~MPC: Apply lifting (C++)")

139 if isinstance(scale, numpy.generic): # nanobind conversion of numpy dtypes to general Python types

140 scale = scale.item() # type: ignore

141 if bc_data is None:

142 bc_data = BCData([bc for bcs_j in bcs for bc in (bcs_j if isinstance(bcs_j, Sequence) else [bcs_j])])

143

144 # The values depend only on the trial space, so compute them once per space

145 # rather than once per block row.

146 lifting_data: dict[int, tuple[numpy.ndarray, numpy.ndarray]] = {}

147

148 def _lifting_data(forms):

149 markers, values = [], []

150 for a in forms:

151 if a is None:

152 markers.append(numpy.empty(0, dtype=numpy.int8))

153 values.append(numpy.empty(0, dtype=_PETSc.ScalarType)) # type: ignore

154 continue

155 V1 = a.function_spaces[1]

156 key = id(V1._cpp_object)

157 if key not in lifting_data:

158 lifting_data[key] = bc_data.lifting(V1, _PETSc.ScalarType) # type: ignore

159 markers.append(lifting_data[key][0])

160 values.append(lifting_data[key][1])

161 return markers, values

162

163 if b.getType() == "nest" or _is_block_vector(b):

164 # Each block row lifts into its own vector, and its constraint puts the masters into the

165 # vector of their block, so every block's array is passed

166 assert isinstance(form, Sequence) and isinstance(constraint, Sequence)

167 with _block_arrays(b) as arrays, _block_arrays_read(x0, constraint) as x0_arrays: # type: ignore[arg-type]

168 for i, (a_sub, mpc_i) in enumerate(zip(form, constraint)):

169 markers, values = _lifting_data(a_sub)

170 _a = [None if f is None else f._cpp_object for f in a_sub] # type:ignore

171 dolfinx_mpc.cpp.mpc.apply_lifting_blocks(

172 arrays, i, _a, markers, values, x0_arrays, scale, mpc_i._cpp_object, num_threads

173 )

174 else:

175 with contextlib.ExitStack() as stack:

176 if x0 is None:

177 x0 = []

178 else:

179 x0 = [stack.enter_context(x.localForm()) for x in x0]

180 x0_r = [x.array_r for x in x0]

181 b_local = stack.enter_context(b.localForm())

182 markers, values = _lifting_data(form)

183 _forms = [None if f is None else f._cpp_object for f in form] # type: ignore

184 assert isinstance(constraint, MultiPointConstraint)

185 dolfinx_mpc.cpp.mpc.apply_lifting(

186 b_local.array_w, _forms, markers, values, x0_r, scale, constraint._cpp_object, num_threads

187 )

188 t.stop()

189

190

191def apply_mpc_lifting(

192 b: _PETSc.Vec, # type: ignore

193 form: Sequence[_fem.Form],

194 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],

195 constraint1: Optional[Sequence[MultiPointConstraint]] = None,

196 scale: _float_classes = default_scalar_type(1.0), # type: ignore

197 num_threads: Optional[int] = 1,

198):

199 """

200 Lift the inhomogeneity of a multi point constraint into the vector `b`, i.e.

201

202 :math:`b = b - scale \\cdot K^T (A_j g_j)`

203

204 where :math:`g` is the constraint offset of the constraint on the trial space. This is

205 the term arising in :math:`K^T A K x_{red} = K^T (b - A g)` for the affine constraint

206 :math:`x = K x_{red} + g`, and is a no-op for a homogeneous constraint.

207

208 Note:

209 Only required when solving directly for :math:`x_{red}`. A residual assembled at an

210 iterate that already satisfies the constraint contains :math:`K^T A g` already, so

211 the Newton/SNES path must not call this.

212

213 Args:

214 b: PETSc vector to assemble into

215 form: The bilinear forms, one per block column

216 constraint: The multi point constraint for the rows of `b`

217 constraint1: The multi point constraints for the columns, one per block. Defaults

218 to `constraint`, which is correct for a square problem.

219 scale: Scaling for lifting

220 num_threads: The number of threads to use for certain operations

221 """

222 t = Timer("~MPC: Apply MPC lifting (C++)")

223 if isinstance(scale, numpy.generic): # nanobind conversion of numpy dtypes to general Python types

224 scale = scale.item() # type: ignore

225

226 if b.getType() == "nest" or _is_block_vector(b):

227 assert isinstance(form, Sequence) and isinstance(constraint, Sequence)

228 cols = constraint if constraint1 is None else constraint1

229 _mpc1 = [c._cpp_object for c in cols] # type: ignore

230 with _block_arrays(b) as arrays:

231 for i, (a_sub, mpc_i) in enumerate(zip(form, constraint)):

232 _a = [None if f is None else f._cpp_object for f in a_sub] # type: ignore

233 dolfinx_mpc.cpp.mpc.apply_mpc_lifting_blocks(

234 arrays, i, _a, scale, mpc_i._cpp_object, _mpc1, num_threads

235 )

236 else:

237 assert isinstance(constraint, MultiPointConstraint)

238 cols = [constraint] if constraint1 is None else constraint1

239 with b.localForm() as b_local:

240 _forms = [f._cpp_object for f in form] # type: ignore

241 _mpc1 = [c._cpp_object for c in cols] # type: ignore

242 dolfinx_mpc.cpp.mpc.apply_mpc_lifting(

243 b_local.array_w, _forms, scale, constraint._cpp_object, _mpc1, num_threads

244 )

245 t.stop()

246

247

248def create_vector(

249 L: Union[_fem.Form, Sequence[_fem.Form]],

250 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],

251 kind: Kind = None,

252) -> _PETSc.Vec: # type: ignore

253 """

254 Create a PETSc vector appropriate for a linear form, or a sequence of them, under multi point

255 constraints.

256

257 As in :func:`dolfinx.fem.petsc.create_vector`, a single form gives a ghosted vector, while a

258 sequence of forms with `kind` ``"nest"``, or a nested sequence of matrix types, gives a vector

259 of type ``nest`` and any other `kind` a single, monolithic vector. On each process the

260 monolithic vector is ``[b_0, b_1, ..., b_n, b_0g, b_1g, ..., b_ng]``, where ``b_i`` holds the

261 owned entries of block ``i`` and ``b_ig`` its ghosts, which include the masters added by the

262 constraint. The offsets of the blocks are in the attribute ``_blocks``, see

263 :func:`dolfinx.la.petsc.create_vector`.

264

265 Args:

266 L: A linear form, or a sequence of them, one per block

267 constraint: The multi point constraint, or one per block

268 kind: The kind of vector, as above

269

270 Returns:

271 A PETSc vector, not initialised to zero.

272 """

273 if not isinstance(L, Sequence):

274 assert isinstance(constraint, MultiPointConstraint)

275 return _create_petsc_vector(

276 [(constraint.function_space.dofmap.index_map, constraint.function_space.dofmap.index_map_bs)]

277 )

278 assert isinstance(constraint, Sequence)

279 if blocked_layout(kind)[0] == "nest":

280 return _create_vector_nest(L, constraint)

281 return _create_vector_block(L, constraint)

282

283

284def assemble_vector(

285 form: Union[_fem.Form, Sequence[_fem.Form]],

286 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],

287 b: Optional[_PETSc.Vec] = None, # type: ignore

288 num_threads: Optional[int] = 1,

289 kind: Kind = None,

290) -> _PETSc.Vec: # type: ignore

291 """

292 Assemble a linear form, or a sequence of them, into vector `b` with the corresponding multi

293 point constraints. The kind of vector is selected by `kind`, or by the type of `b` if it is

294 supplied.

295

296 Args:

297 form: The linear form, or a sequence of them, one per block

298 constraint: The multi point constraint, or one per block

299 b: PETSc vector to assemble into. Assembly is additive, so `b` is not

300 zeroed; use `dolfinx.la.petsc._zero_vector` first to discard its

301 contents. If not supplied a new, zeroed vector is created.

302 num_threads: The number of threads to use for certain operations

303 kind: The kind of vector to create when `b` is not supplied, see :func:`create_vector`.

304

305 Returns:

306 The vector with the assembled linear form (`b` if supplied)

307 """

308 if b is None:

309 b = create_vector(form, constraint, kind)

310 _zero_vector(b)

311 t = Timer("~MPC: Assemble vector (C++)")

312 if isinstance(form, Sequence):

313 assert isinstance(constraint, Sequence)

314 if b.getType() == "nest":

315 _assemble_vector_nest(b, form, constraint, num_threads)

316 else:

317 _assemble_vector_block(b, form, constraint, num_threads)

318 else:

319 assert isinstance(constraint, MultiPointConstraint)

320 _assemble_form(b, form, constraint, num_threads)

321 t.stop()

322 return b

323

324

325def _assemble_form(

326 b: _PETSc.Vec, # type: ignore

327 form: _fem.Form,

328 constraint: MultiPointConstraint,

329 num_threads: Optional[int] = 1,

330):

331 """

332 Assemble one compiled linear form into a vector.

333

334 Additive: `b` is not zeroed, following the convention of the DOLFINx

335 assemblers.

336 """

337 with b.localForm() as b_local:

338 dolfinx_mpc.cpp.mpc.assemble_vector(b_local.array_w, form._cpp_object, constraint._cpp_object, num_threads)

339

340

341def _create_vector_nest(L: Sequence[_fem.Form], constraints: Sequence[MultiPointConstraint]) -> _PETSc.Vec: # type: ignore

342 """Create a PETSc vector of type "nest" appropriate for the provided multi point constraints."""

343 assert len(constraints) == len(L)

344

345 maps = [

346 (constraint.function_space.dofmap.index_map._cpp_object, constraint.function_space.dofmap.index_map_bs)

347 for constraint in constraints

348 ]

349 return _cpp.fem.petsc.create_vector_nest(maps)

350

351

352def create_vector_nest(L: Sequence[_fem.Form], constraints: Sequence[MultiPointConstraint]) -> _PETSc.Vec: # type: ignore

353 """

354 Create a PETSc vector of type "nest" appropriate for the provided multi

355 point constraints

356

357 .. deprecated::

358 Use :func:`create_vector` with ``kind="nest"``.

359

360 Args:

361 L: A sequence of linear forms

362 constraints: An ordered list of multi point constraints

363

364 Returns:

365 PETSc.Vec: A PETSc vector of type "nest" #type: ignore

366 """

367 deprecated("create_vector_nest", "create_vector(L, constraints, kind='nest')")

368 return _create_vector_nest(L, constraints)

369

370

371def _assemble_vector_nest(

372 b: _PETSc.Vec, # type: ignore

373 L: Sequence[_fem.Form],

374 constraints: Sequence[MultiPointConstraint],

375 num_threads: Optional[int] = 1,

376):

377 """Assemble linear forms into a PETSc vector of type "nest". Additive, `b` is not zeroed."""

378 assert len(constraints) == len(L)

379 assert b.getType() == "nest"

380

381 _assemble_vector_blocks(b, L, constraints, num_threads)

382

383

384def _assemble_vector_blocks(

385 b: _PETSc.Vec, # type: ignore

386 L: Sequence[_fem.Form],

387 constraints: Sequence[MultiPointConstraint],

388 num_threads: Optional[int] = 1,

389):

390 """

391 Assemble linear forms into a nest or monolithic vector. Each form goes to the vector of its

392 block, and the masters of its constraint to the vector of their block. A `None` form is

393 skipped.

394 """

395 with _block_arrays(b) as arrays:

396 for i, (L_i, mpc_i) in enumerate(zip(L, constraints)):

397 # A block without a linear form, such as a block holding only masters, adds nothing

398 if L_i is None:

399 continue

400 dolfinx_mpc.cpp.mpc.assemble_vector_blocks(arrays, i, L_i._cpp_object, mpc_i._cpp_object, num_threads)

401

402

403def assemble_vector_nest(

404 b: _PETSc.Vec, # type: ignore

405 L: Sequence[_fem.Form],

406 constraints: Sequence[MultiPointConstraint],

407 num_threads: Optional[int] = 1,

408):

409 """

410 Assemble a linear form into a PETSc vector of type "nest"

411

412 .. deprecated::

413 Use :func:`assemble_vector`, which selects the layout from `b`.

414

415 Args:

416 b: A PETSc vector of type "nest" to assemble into. Assembly is additive,

417 so `b` is not zeroed; use `dolfinx.la.petsc._zero_vector` first to

418 discard its contents.

419 L: A sequence of linear forms

420 constraints: An ordered list of multi point constraints

421 """

422 deprecated("assemble_vector_nest", "assemble_vector(L, constraints, b)")

423 _assemble_vector_nest(b, L, constraints, num_threads)

424

425

426def _create_vector_block(L: Sequence[_fem.Form], constraints: Sequence[MultiPointConstraint]) -> _PETSc.Vec: # type: ignore

427 """

428 Create a monolithic PETSc vector appropriate for the provided multi point constraints.

429

430 On each process the vector is ``[b_0, b_1, ..., b_n, b_0g, b_1g, ..., b_ng]``, where ``b_i``

431 holds the owned entries of block ``i`` and ``b_ig`` its ghosts, which include the masters

432 added by the constraint. The offsets of the blocks are in the attribute ``_blocks``, see

433 :func:`dolfinx.la.petsc.create_vector`.

434

435 Args:

436 L: A sequence of linear forms, one per block

437 constraints: An ordered list of multi point constraints, one per block

438

439 Returns:

440 A PETSc vector with the layout above, not initialised to zero.

441 """

442 assert len(constraints) == len(L)

443 maps = [(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs) for mpc in constraints]

444 # A vector of one block gets the block layout too, rather than none

445 return _create_petsc_vector(maps, kind=_PETSc.Vec.Type.MPI) # type: ignore

446

447

448def _assemble_vector_block(

449 b: _PETSc.Vec, # type: ignore

450 L: Sequence[_fem.Form],

451 constraints: Sequence[MultiPointConstraint],

452 num_threads: Optional[int] = 1,

453):

454 """

455 Assemble linear forms into a monolithic PETSc vector.

456

457 Args:

458 b: A vector made by :func:`create_vector_block` to assemble into. Assembly is additive,

459 so `b` is not zeroed; use `dolfinx.la.petsc._zero_vector` first to discard its

460 contents.

461 L: A sequence of linear forms, one per block

462 constraints: An ordered list of multi point constraints, one per block

463 """

464 assert len(constraints) == len(L)

465 if not _is_block_vector(b):

466 raise ValueError("The vector must be created by create_vector_block")

467 _assemble_vector_blocks(b, L, constraints, num_threads)

Coverage for python/src/dolfinx_mpc/assemble_vector.py: 99%

161 statements