Coverage for pySDC/implementations/datatype_classes/firedrake_mesh.py: 0%

1import firedrake as fd 

2 

3from pySDC.core.errors import DataError 

4from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator 

5 

6 

7class firedrake_mesh(object): 

8 """ 

9 Wrapper for firedrake function data. 

10 

11 Attributes: 

12 functionspace (firedrake.Function): firedrake data 

13 """ 

14 

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

16 if fd.functionspaceimpl.WithGeometry in type(init).__mro__: 

17 self.functionspace = fd.Function(init) 

18 self.functionspace.assign(val) 

19 elif fd.Function in type(init).__mro__: 

20 self.functionspace = fd.Function(init) 

21 elif type(init) == firedrake_mesh: 

22 self.functionspace = init.functionspace.copy(deepcopy=True) 

23 else: 

24 raise DataError('something went wrong during %s initialization' % type(init)) 

25 

26 def __getattr__(self, key): 

27 return getattr(self.functionspace, key) 

28 

29 @property 

30 def asnumpy(self): 

31 """ 

32 Get a numpy array of the values associated with this data 

33 """ 

34 return self.functionspace.dat._numpy_data 

35 

36 def __add__(self, other): 

37 if isinstance(other, type(self)): 

38 me = firedrake_mesh(other) 

39 me.functionspace.assign(self.functionspace + other.functionspace) 

40 return me 

41 else: 

42 raise DataError("Type error: cannot add %s to %s" % (type(other), type(self))) 

43 

44 def __sub__(self, other): 

45 if isinstance(other, type(self)): 

46 me = firedrake_mesh(other) 

47 me.functionspace.assign(self.functionspace - other.functionspace) 

48 return me 

49 else: 

50 raise DataError("Type error: cannot add %s to %s" % (type(other), type(self))) 

51 

52 def __rmul__(self, other): 

53 """ 

54 Overloading the right multiply by scalar factor 

55 

56 Args: 

57 other (float): factor 

58 Raises: 

59 DataError: if other is not a float 

60 Returns: 

61 fenics_mesh: copy of original values scaled by factor 

62 """ 

63 

64 try: 

65 me = firedrake_mesh(self) 

66 me.functionspace.assign(other * self.functionspace) 

67 return me 

68 except TypeError as e: 

69 raise DataError("Type error: cannot multiply %s to %s" % (type(other), type(self))) from e 

70 

71 def __abs__(self): 

72 """ 

73 Overloading the abs operator for mesh types 

74 

75 Returns: 

76 float: L2 norm 

77 """ 

78 

79 return fd.norm(self.functionspace, 'L2') 

80 

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

82 """ 

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

84 

85 Args: 

86 dest (int): target rank 

87 tag (int): communication tag 

88 comm: communicator 

89 

90 Returns: 

91 request handle 

92 """ 

93 assert ( 

94 type(comm) == FiredrakeEnsembleCommunicator 

95 ), f'Need to give a FiredrakeEnsembleCommunicator here, not {type(comm)}' 

96 return comm.Isend(self.functionspace, dest=dest, tag=tag) 

97 

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

99 """ 

100 Routine for receiving in time 

101 

102 Args: 

103 source (int): source rank 

104 tag (int): communication tag 

105 comm: communicator 

106 

107 Returns: 

108 None 

109 """ 

110 assert ( 

111 type(comm) == FiredrakeEnsembleCommunicator 

112 ), f'Need to give a FiredrakeEnsembleCommunicator here, not {type(comm)}' 

113 return comm.Irecv(self.functionspace, source=source, tag=tag) 

114 

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

116 """ 

117 Routine for broadcasting values 

118 

119 Args: 

120 root (int): process with value to broadcast 

121 comm: communicator 

122 

123 Returns: 

124 broadcasted values 

125 """ 

126 assert ( 

127 type(comm) == FiredrakeEnsembleCommunicator 

128 ), f'Need to give a FiredrakeEnsembleCommunicator here, not {type(comm)}' 

129 comm.Bcast(self.functionspace, root=root) 

130 return self 

131 

132 

133class IMEX_firedrake_mesh(object): 

134 """ 

135 Datatype for IMEX integration with firedrake data. 

136 

137 Attributes: 

138 impl (firedrake_mesh): implicit part 

139 expl (firedrake_mesh): explicit part 

140 """ 

141 

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

143 if type(init) == type(self): 

144 self.impl = firedrake_mesh(init.impl) 

145 self.expl = firedrake_mesh(init.expl) 

146 else: 

147 self.impl = firedrake_mesh(init, val=val) 

148 self.expl = firedrake_mesh(init, val=val) 

149 

150 def __add__(self, other): 

151 me = IMEX_firedrake_mesh(self) 

152 me.impl = self.impl + other.impl 

153 me.expl = self.expl + other.expl 

154 return me 

155 

156 def __sub__(self, other): 

157 me = IMEX_firedrake_mesh(self) 

158 me.impl = self.impl - other.impl 

159 me.expl = self.expl - other.expl 

160 return me 

161 

162 def __rmul__(self, other): 

163 me = IMEX_firedrake_mesh(self) 

164 me.impl = other * self.impl 

165 me.expl = other * self.expl 

166 return me