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

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

2# 

3# This file is part of DOLFINX_MPC 

4# 

5# SPDX-License-Identifier: MIT 

6from __future__ import annotations 

7 

8from collections.abc import Sequence 

9from typing import Optional, Union 

10 

11from petsc4py import PETSc as _PETSc 

12 

13import dolfinx.cpp as _cpp 

14import dolfinx.fem as _fem 

15 

16from dolfinx_mpc import cpp 

17 

18from .multipointconstraint import MultiPointConstraint 

19 

20 

21def assemble_matrix( 

22 form: _fem.Form, 

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

24 bcs: Optional[Sequence[_fem.DirichletBC]] = None, 

25 diagval: _PETSc.ScalarType = 1, # type: ignore 

26 A: Optional[_PETSc.Mat] = None, # type: ignore 

27 num_threads: Optional[int] = 1, 

28) -> _PETSc.Mat: # type: ignore 

29 """ 

30 Assemble a compiled DOLFINx bilinear form into a PETSc matrix with corresponding multi point constraints 

31 and Dirichlet boundary conditions. 

32 

33 Args: 

34 form: The compiled bilinear variational form 

35 constraint: The multi point constraint 

36 bcs: Sequence of Dirichlet boundary conditions 

37 diagval: Value to set on the diagonal of the matrix 

38 A: PETSc matrix to assemble into 

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

40 Returns: 

41 _PETSc.Mat: The matrix with the assembled bi-linear form #type: ignore 

42 """ 

43 bcs = [] if bcs is None else [bc._cpp_object for bc in bcs] 

44 if not isinstance(constraint, Sequence): 

45 assert form.function_spaces[0] == form.function_spaces[1] 

46 constraint = (constraint, constraint) 

47 

48 # Generate matrix with MPC sparsity pattern 

49 if A is None: 

50 A = cpp.mpc.create_matrix(form._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object) 

51 A.zeroEntries() 

52 

53 # Assemble matrix in C++ 

54 cpp.mpc.assemble_matrix( 

55 A, form._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object, bcs, diagval, num_threads 

56 ) 

57 

58 # Add one on diagonal for Dirichlet boundary conditions 

59 if form.function_spaces[0] is form.function_spaces[1]: 

60 A.assemblyBegin(_PETSc.Mat.AssemblyType.FLUSH) # type: ignore 

61 A.assemblyEnd(_PETSc.Mat.AssemblyType.FLUSH) # type: ignore 

62 _cpp.fem.petsc.insert_diagonal(A, form.function_spaces[0], bcs, diagval) 

63 

64 A.assemble() 

65 return A 

66 

67 

68def create_sparsity_pattern(form: _fem.Form, mpc: Union[MultiPointConstraint, Sequence[MultiPointConstraint]]): 

69 """ 

70 Create sparsity-pattern for MPC given a compiled DOLFINx form 

71 

72 Args: 

73 form: The form 

74 mpc: For square forms, the MPC. For rectangular forms a list of 2 MPCs on 

75 axis 0 & 1, respectively 

76 """ 

77 if isinstance(mpc, Sequence): 

78 assert len(mpc) == 2 

79 for mpc_ in mpc: 

80 mpc_._not_finalized() # type: ignore 

81 return cpp.mpc.create_sparsity_pattern(form._cpp_object, mpc[0]._cpp_object, mpc[1]._cpp_object) 

82 else: 

83 mpc._not_finalized() # type: ignore 

84 return cpp.mpc.create_sparsity_pattern( 

85 form._cpp_object, 

86 mpc._cpp_object, # type: ignore 

87 mpc._cpp_object, # type: ignore 

88 ) # type: ignore 

89 

90 

91def create_matrix_nest(a: Sequence[Sequence[_fem.Form]], constraints: Sequence[MultiPointConstraint]): 

92 """ 

93 Create a PETSc matrix of type "nest" with appropriate sparsity pattern 

94 given the provided multi points constraints 

95 

96 Args: 

97 a: The compiled bilinear variational form provided in a rank 2 list 

98 constraints: An ordered list of multi point constraints 

99 """ 

100 assert len(constraints) == len(a) 

101 

102 A_: list[list[_PETSc.Mat | None]] = [[None for _ in range(len(a[0]))] for _ in range(len(a))] 

103 

104 for i, a_row in enumerate(a): 

105 for j, a_block in enumerate(a_row): 

106 if a[i][j] is None: 

107 continue 

108 A_[i][j] = cpp.mpc.create_matrix( 

109 a[i][j]._cpp_object, constraints[i]._cpp_object, constraints[j]._cpp_object 

110 ) 

111 

112 A = _PETSc.Mat().createNest( 

113 A_, # type: ignore 

114 comm=constraints[0].function_space.mesh.comm, 

115 ) 

116 return A 

117 

118 

119def assemble_matrix_nest( 

120 A: _PETSc.Mat, # type: ignore 

121 a: Sequence[Sequence[_fem.Form]], 

122 constraints: Sequence[MultiPointConstraint], 

123 bcs: Sequence[_fem.DirichletBC] = [], 

124 diagval: _PETSc.ScalarType = 1, # type: ignore 

125 num_threads: Optional[int] = 1, 

126): 

127 """ 

128 Assemble a compiled DOLFINx bilinear form into a PETSc matrix of type 

129 "nest" with corresponding multi point constraints and Dirichlet boundary 

130 conditions. 

131 

132 Args: 

133 a: The compiled bilinear variational form provided in a rank 2 list 

134 constraints: An ordered list of multi point constraints 

135 bcs: Sequence of Dirichlet boundary conditions 

136 diagval: Value to set on the diagonal of the matrix (Default 1) 

137 A: PETSc matrix to assemble into 

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

139 """ 

140 for i, a_row in enumerate(a): 

141 for j, a_block in enumerate(a_row): 

142 if a_block is not None: 

143 Asub = A.getNestSubMatrix(i, j) 

144 assemble_matrix( 

145 a_block, (constraints[i], constraints[j]), bcs=bcs, diagval=diagval, A=Asub, num_threads=num_threads 

146 )