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 List, Optional, Sequence 

10 

11from petsc4py import PETSc as _PETSc 

12 

13import dolfinx.cpp as _cpp 

14import dolfinx.fem as _fem 

15import numpy 

16import ufl 

17from dolfinx import default_scalar_type 

18from dolfinx.common import Timer 

19from dolfinx.la.petsc import create_vector 

20 

21import dolfinx_mpc.cpp 

22 

23from .multipointconstraint import MultiPointConstraint, _float_classes 

24 

25 

26def apply_lifting( 

27 b: _PETSc.Vec, # type: ignore 

28 form: List[_fem.Form], 

29 bcs: List[List[_fem.DirichletBC]], 

30 constraint: MultiPointConstraint, 

31 x0: List[_PETSc.Vec] = [], # type: ignore 

32 scale: _float_classes = default_scalar_type(1.0), 

33): # type: ignore 

34 """ 

35 Apply lifting to vector b, i.e. 

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

37 

38 Args: 

39 b: PETSc vector to assemble into 

40 form: The linear form 

41 bcs: List of Dirichlet boundary conditions 

42 constraint: The multi point constraint 

43 x0: List of vectors 

44 scale: Scaling for lifting 

45 """ 

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

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

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

49 with contextlib.ExitStack() as stack: 

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

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

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

53 _forms = [f._cpp_object for f in form] 

54 _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs] 

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

56 t.stop() 

57 

58 

59def assemble_vector( 

60 form: ufl.form.Form, 

61 constraint: MultiPointConstraint, 

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

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

64 """ 

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

66 

67 Args: 

68 form: The linear form 

69 constraint: The multi point constraint 

70 b: PETSc vector to assemble 

71 

72 Returns: 

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

74 """ 

75 

76 if b is None: 

77 b = create_vector( 

78 constraint.function_space.dofmap.index_map, 

79 constraint.function_space.dofmap.index_map_bs, 

80 ) 

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

82 with b.localForm() as b_local: 

83 b_local.set(0.0) 

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

85 t.stop() 

86 return b 

87 

88 

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

90 """ 

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

92 point constraints 

93 

94 Args: 

95 L: A sequence of linear forms 

96 constraints: An ordered list of multi point constraints 

97 

98 Returns: 

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

100 """ 

101 assert len(constraints) == len(L) 

102 

103 maps = [ 

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

105 for constraint in constraints 

106 ] 

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

108 

109 

110def assemble_vector_nest( 

111 b: _PETSc.Vec, # type: ignore 

112 L: Sequence[_fem.Form], 

113 constraints: Sequence[MultiPointConstraint], 

114): 

115 """ 

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

117 

118 Args: 

119 b: A PETSc vector of type "nest" 

120 L: A sequence of linear forms 

121 constraints: An ordered list of multi point constraints 

122 """ 

123 assert len(constraints) == len(L) 

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

125 

126 b_sub_vecs = b.getNestSubVecs() 

127 for i, L_row in enumerate(L): 

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