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

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) -> _PETSc.Mat: # type: ignore 

28 """ 

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

30 and Dirichlet boundary conditions. 

31 

32 Args: 

33 form: The compiled bilinear variational form 

34 constraint: The multi point constraint 

35 bcs: Sequence of Dirichlet boundary conditions 

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

37 A: PETSc matrix to assemble into 

38 

39 Returns: 

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

41 """ 

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

43 if not isinstance(constraint, Sequence): 

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

45 constraint = (constraint, constraint) 

46 

47 # Generate matrix with MPC sparsity pattern 

48 if A is None: 

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

50 A.zeroEntries() 

51 

52 # Assemble matrix in C++ 

53 cpp.mpc.assemble_matrix(A, form._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object, bcs, diagval) 

54 

55 # Add one on diagonal for Dirichlet boundary conditions 

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

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

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

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

60 

61 A.assemble() 

62 return A 

63 

64 

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

66 """ 

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

68 

69 Args: 

70 form: The form 

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

72 axis 0 & 1, respectively 

73 """ 

74 if isinstance(mpc, Sequence): 

75 assert len(mpc) == 2 

76 for mpc_ in mpc: 

77 mpc_._not_finalized() # type: ignore 

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

79 else: 

80 mpc._not_finalized() # type: ignore 

81 return cpp.mpc.create_sparsity_pattern( 

82 form._cpp_object, 

83 mpc._cpp_object, # type: ignore 

84 mpc._cpp_object, # type: ignore 

85 ) # type: ignore 

86 

87 

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

89 """ 

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

91 given the provided multi points constraints 

92 

93 Args: 

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

95 constraints: An ordered list of multi point constraints 

96 """ 

97 assert len(constraints) == len(a) 

98 

99 A_ = [[None for _ in range(len(a[0]))] for _ in range(len(a))] 

100 

101 for i, a_row in enumerate(a): 

102 for j, a_block in enumerate(a_row): 

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

104 continue 

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

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

107 ) 

108 

109 A = _PETSc.Mat().createNest( # type: ignore 

110 A_, comm=constraints[0].function_space.mesh.comm 

111 ) 

112 return A 

113 

114 

115def assemble_matrix_nest( 

116 A: _PETSc.Mat, # type: ignore 

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

118 constraints: Sequence[MultiPointConstraint], 

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

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

121): 

122 """ 

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

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

125 conditions. 

126 

127 Args: 

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

129 constraints: An ordered list of multi point constraints 

130 bcs: Sequence of Dirichlet boundary conditions 

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

132 A: PETSc matrix to assemble into 

133 """ 

134 for i, a_row in enumerate(a): 

135 for j, a_block in enumerate(a_row): 

136 if a_block is not None: 

137 Asub = A.getNestSubMatrix(i, j) 

138 assemble_matrix(a_block, (constraints[i], constraints[j]), bcs=bcs, diagval=diagval, A=Asub)