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

1# Copyright (C) 2020 Jørgen S. Dokken 

2# 

3# This file is part of DOLFINX_MPC 

4# 

5# SPDX-License-Identifier: MIT 

6from __future__ import annotations 

7 

8import contextlib 

9from typing import Iterable, Optional, Sequence, Union 

10 

11from petsc4py import PETSc as _PETSc 

12 

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 create_vector 

19 

20import dolfinx_mpc.cpp 

21 

22from .multipointconstraint import MultiPointConstraint, _float_classes 

23 

24 

25def apply_lifting( 

26 b: _PETSc.Vec, 

27 form: Union[Iterable[Sequence[_fem.Form]], Iterable[_fem.Form]], # type: ignore 

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

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

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

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

32): # type: ignore 

33 """ 

34 Apply lifting to vector b, i.e. 

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

36 

37 Args: 

38 b: PETSc vector to assemble into 

39 form: The linear form 

40 bcs: List of Dirichlet boundary conditions 

41 constraint: The multi point constraint 

42 x0: List of vectors 

43 scale: Scaling for lifting 

44 """ 

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

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

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

48 

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

50 try: 

51 bcs = _fem.bcs_by_block(_fem.extract_function_spaces(form, 1), bcs) # type: ignore 

52 except AttributeError: 

53 pass 

54 x0 = [] if x0 is None else x0.getNestSubVecs() # type: ignore 

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

56 for b_sub, a_sub, mpc_i in zip(b.getNestSubVecs(), form, constraint): 

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

58 _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs] # type: ignore 

59 dolfinx_mpc.cpp.mpc.apply_lifting(b_sub.array_w, _a, _bcs, x0, scale, mpc_i._cpp_object) 

60 else: 

61 with contextlib.ExitStack() as stack: 

62 if x0 is None: 

63 x0 = [] 

64 else: 

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

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

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

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

69 _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs] # type: ignore 

70 assert isinstance(constraint, MultiPointConstraint) 

71 dolfinx_mpc.cpp.mpc.apply_lifting(b_local.array_w, _forms, _bcs, x0_r, scale, constraint._cpp_object) 

72 t.stop() 

73 

74 

75def assemble_vector( 

76 form: _fem.Form, 

77 constraint: MultiPointConstraint, 

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

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

80 """ 

81 Assemble a linear form into vector `b` with corresponding multi point constraint 

82 

83 Args: 

84 form: The linear form 

85 constraint: The multi point constraint 

86 b: PETSc vector to assemble 

87 

88 Returns: 

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

90 """ 

91 

92 if b is None: 

93 b = create_vector([(constraint.function_space.dofmap.index_map, constraint.function_space.dofmap.index_map_bs)]) 

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

95 with b.localForm() as b_local: 

96 b_local.set(0.0) 

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

98 t.stop() 

99 return b 

100 

101 

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

103 """ 

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

105 point constraints 

106 

107 Args: 

108 L: A sequence of linear forms 

109 constraints: An ordered list of multi point constraints 

110 

111 Returns: 

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

113 """ 

114 assert len(constraints) == len(L) 

115 

116 maps = [ 

117 (constraint.function_space.dofmap.index_map, constraint.function_space.dofmap.index_map_bs) 

118 for constraint in constraints 

119 ] 

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

121 

122 

123def assemble_vector_nest( 

124 b: _PETSc.Vec, # type: ignore 

125 L: Sequence[_fem.Form], 

126 constraints: Sequence[MultiPointConstraint], 

127): 

128 """ 

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

130 

131 Args: 

132 b: A PETSc vector of type "nest" 

133 L: A sequence of linear forms 

134 constraints: An ordered list of multi point constraints 

135 """ 

136 assert len(constraints) == len(L) 

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

138 

139 b_sub_vecs = b.getNestSubVecs() 

140 for i, L_row in enumerate(L): 

141 assemble_vector(L_row, constraints[i], b=b_sub_vecs[i])