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

1import firedrake as fd 

2 

3from pySDC.core.errors import DataError 

4from pySDC.implementations.datatype_classes.container import MultiComponentContainer 

5from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator 

6 

7 

8class firedrake_mesh(object): 

9 """ 

10 Wrapper for firedrake function data. 

11 

12 Attributes: 

13 functionspace (firedrake.Function): firedrake data 

14 """ 

15 

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

17 """ 

18 Args: 

19 init (firedrake.FunctionSpace, firedrake.Function or firedrake_mesh): Function space to create a function 

20 on, or 

21 data to copy 

22 val (float): Value to assign if `init` is a function space, ignored otherwise 

23 

24 Raises: 

25 DataError: If `init` has none of these types 

26 """ 

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

28 self.functionspace = fd.Function(init) 

29 self.functionspace.assign(val) 

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

31 self.functionspace = fd.Function(init) 

32 elif type(init) == firedrake_mesh: 

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

34 else: 

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

36 

37 def __getattr__(self, key): 

38 return getattr(self.functionspace, key) 

39 

40 @property 

41 def asnumpy(self): 

42 """ 

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

44 """ 

45 return self.functionspace.dat._numpy_data 

46 

47 def __add__(self, other): 

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

49 me = firedrake_mesh(other) 

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

51 return me 

52 else: 

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

54 

55 def __sub__(self, other): 

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

57 me = firedrake_mesh(other) 

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

59 return me 

60 else: 

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

62 

63 def __rmul__(self, other): 

64 """ 

65 Overloading the right multiply by scalar factor 

66 

67 Args: 

68 other (float): factor 

69 Raises: 

70 DataError: if other is not a float 

71 Returns: 

72 fenics_mesh: copy of original values scaled by factor 

73 """ 

74 

75 try: 

76 me = firedrake_mesh(self) 

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

78 return me 

79 except TypeError as e: 

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

81 

82 def __abs__(self): 

83 """ 

84 Overloading the abs operator for mesh types 

85 

86 Returns: 

87 float: L2 norm 

88 """ 

89 

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

91 

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

93 """ 

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

95 

96 Args: 

97 dest (int): target rank 

98 tag (int): communication tag 

99 comm: communicator 

100 

101 Returns: 

102 request handle 

103 """ 

104 assert ( 

105 type(comm) == FiredrakeEnsembleCommunicator 

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

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

108 

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

110 """ 

111 Routine for receiving in time 

112 

113 Args: 

114 source (int): source rank 

115 tag (int): communication tag 

116 comm: communicator 

117 

118 Returns: 

119 None 

120 """ 

121 assert ( 

122 type(comm) == FiredrakeEnsembleCommunicator 

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

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

125 

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

127 """ 

128 Routine for broadcasting values 

129 

130 Args: 

131 root (int): process with value to broadcast 

132 comm: communicator 

133 

134 Returns: 

135 broadcasted values 

136 """ 

137 assert ( 

138 type(comm) == FiredrakeEnsembleCommunicator 

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

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

141 return self 

142 

143 

144class IMEX_firedrake_mesh(MultiComponentContainer): 

145 """ 

146 Datatype for IMEX integration with firedrake data. 

147 

148 Attributes: 

149 impl (firedrake_mesh): implicit part 

150 expl (firedrake_mesh): explicit part 

151 """ 

152 

153 components = ['impl', 'expl'] 

154 component_type = firedrake_mesh