Coverage for pySDC/implementations/datatype_classes/petsc_vec.py: 95%

40 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-29 09:02 +0000

1from petsc4py import PETSc 

2 

3from pySDC.core.Errors import DataError 

4 

5 

6class petsc_vec(PETSc.Vec): 

7 __array_priority__ = 1000 # otherwise rmul with float64 does not work (don't ask, won't tell) 

8 

9 def __new__(cls, init=None, val=0.0): 

10 if isinstance(init, petsc_vec) or isinstance(init, PETSc.Vec): 

11 obj = PETSc.Vec.__new__(cls) 

12 init.copy(obj) 

13 elif isinstance(init, PETSc.DMDA): 

14 tmp = init.createGlobalVector() 

15 obj = petsc_vec(tmp) 

16 objarr = init.getVecArray(obj) 

17 objarr[:] = val 

18 else: 

19 obj = PETSc.Vec.__new__(cls) 

20 return obj 

21 

22 def __abs__(self): 

23 """ 

24 Overloading the abs operator 

25 

26 Returns: 

27 float: absolute maximum of all vec values 

28 """ 

29 # take absolute values of the mesh values (INF = 3) 

30 return self.norm(3) 

31 

32 def isend(self, dest=None, tag=None, comm=None): 

33 """ 

34 Routine for sending data forward in time (non-blocking) 

35 

36 Args: 

37 dest (int): target rank 

38 tag (int): communication tag 

39 comm: communicator 

40 

41 Returns: 

42 request handle 

43 """ 

44 return comm.Issend(self.getArray(), dest=dest, tag=tag) 

45 

46 def irecv(self, source=None, tag=None, comm=None): 

47 """ 

48 Routine for receiving in time 

49 

50 Args: 

51 source (int): source rank 

52 tag (int): communication tag 

53 comm: communicator 

54 

55 Returns: 

56 None 

57 """ 

58 return comm.Irecv(self.getArray(), source=source, tag=tag) 

59 

60 def bcast(self, root=None, comm=None): 

61 """ 

62 Routine for broadcasting values 

63 

64 Args: 

65 root (int): process with value to broadcast 

66 comm: communicator 

67 

68 Returns: 

69 broadcasted values 

70 """ 

71 comm.Bcast(self.getArray(), root=root) 

72 return self 

73 

74 

75class petsc_vec_imex(object): 

76 """ 

77 RHS data type for Vec with implicit and explicit components 

78 

79 This data type can be used to have RHS with 2 components (here implicit and explicit) 

80 

81 Attributes: 

82 impl (petsc_vec): implicit part 

83 expl (petsc_vec): explicit part 

84 """ 

85 

86 def __init__(self, init, val=0.0): 

87 """ 

88 Initialization routine 

89 

90 Args: 

91 init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested) 

92 or another imex_mesh object 

93 val (float): an initial number (default: 0.0) 

94 Raises: 

95 DataError: if init is none of the types above 

96 """ 

97 

98 if isinstance(init, type(self)): 

99 self.impl = petsc_vec(init.impl) 

100 self.expl = petsc_vec(init.expl) 

101 elif isinstance(init, PETSc.DMDA): 

102 self.impl = petsc_vec(init, val=val) 

103 self.expl = petsc_vec(init, val=val) 

104 # something is wrong, if none of the ones above hit 

105 else: 

106 raise DataError('something went wrong during %s initialization' % type(self)) 

107 

108 

109class petsc_vec_comp2(object): 

110 """ 

111 RHS data type for Vec with implicit and explicit components 

112 

113 This data type can be used to have RHS with 2 components (here implicit and explicit) 

114 

115 Attributes: 

116 impl (petsc_vec): implicit part 

117 expl (petsc_vec): explicit part 

118 """ 

119 

120 def __init__(self, init, val=0.0): 

121 """ 

122 Initialization routine 

123 

124 Args: 

125 init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested) 

126 or another imex_mesh object 

127 val (float): an initial number (default: 0.0) 

128 Raises: 

129 DataError: if init is none of the types above 

130 """ 

131 

132 if isinstance(init, type(self)): 

133 self.comp1 = petsc_vec(init.comp1) 

134 self.comp2 = petsc_vec(init.comp2) 

135 elif isinstance(init, PETSc.DMDA): 

136 self.comp1 = petsc_vec(init, val=val) 

137 self.comp2 = petsc_vec(init, val=val) 

138 # something is wrong, if none of the ones above hit 

139 else: 

140 raise DataError('something went wrong during %s initialization' % type(self))