Coverage for pySDC/implementations/datatype_classes/particles.py: 79%

112 statements  

« prev     ^ index     » next       coverage.py v7.6.7, created at 2024-11-16 14:51 +0000

1import numpy as np 

2 

3from pySDC.implementations.datatype_classes.mesh import mesh 

4from pySDC.core.errors import DataError 

5 

6try: 

7 from mpi4py import MPI 

8except ImportError: 

9 MPI = None 

10 

11 

12class particles(object): 

13 """ 

14 Particle data type for particles in 3 dimensions 

15 

16 This data type can be used for particles in 3 dimensions with 3 position and 3 velocity values per particle 

17 

18 Attributes: 

19 pos: contains the positions of all particles 

20 vel: contains the velocities of all particles 

21 """ 

22 

23 class position(mesh): 

24 pass 

25 

26 class velocity(mesh): 

27 pass 

28 

29 def __init__(self, init=None, val=None): 

30 """ 

31 Initialization routine 

32 

33 Args: 

34 init: can either be a number or another particle object 

35 val: initial tuple of values for position and velocity (default: (None,None)) 

36 Raises: 

37 DataError: if init is none of the types above 

38 """ 

39 

40 # if init is another particles object, do a copy (init by copy) 

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

42 self.pos = particles.position(init.pos) 

43 self.vel = particles.velocity(init.vel) 

44 self.q = init.q.copy() 

45 self.m = init.m.copy() 

46 # if init is a number, create particles object and pick the corresponding initial values 

47 elif ( 

48 isinstance(init, tuple) 

49 and (init[1] is None or isinstance(init[1], MPI.Intracomm)) 

50 and isinstance(init[2], np.dtype) 

51 ): 

52 if isinstance(val, int) or isinstance(val, float) or val is None: 

53 self.pos = particles.position(init, val=val) 

54 self.vel = particles.velocity(init, val=val) 

55 if isinstance(init[0], tuple): 

56 self.q = np.zeros(init[0][-1]) 

57 self.m = np.zeros(init[0][-1]) 

58 elif isinstance(init[0], int): 

59 self.q = np.zeros(init[0]) 

60 self.m = np.zeros(init[0]) 

61 self.q[:] = 1.0 

62 self.m[:] = 1.0 

63 elif isinstance(val, tuple) and len(val) == 4: 

64 self.pos = particles.position(init, val=val[0]) 

65 self.vel = particles.velocity(init, val=val[1]) 

66 if isinstance(init[0], tuple): 

67 self.q = np.zeros(init[0][-1]) 

68 self.m = np.zeros(init[0][-1]) 

69 elif isinstance(init[0], int): 

70 self.q = np.zeros(init[0]) 

71 self.m = np.zeros(init[0]) 

72 self.q[:] = val[2] 

73 self.m[:] = val[3] 

74 else: 

75 raise DataError('type of val is wrong, got %s', val) 

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

77 else: 

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

79 

80 def __add__(self, other): 

81 """ 

82 Overloading the addition operator for particles types 

83 

84 Args: 

85 other (particles): particles object to be added 

86 Raises: 

87 DataError: if other is not a particles object 

88 Returns: 

89 particles: sum of caller and other values (self+other) 

90 """ 

91 

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

93 # always create new particles, since otherwise c = a + b changes a as well! 

94 p = particles(self) 

95 p.pos[:] = self.pos + other.pos 

96 p.vel[:] = self.vel + other.vel 

97 p.m = self.m 

98 p.q = self.q 

99 return p 

100 else: 

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

102 

103 def __sub__(self, other): 

104 """ 

105 Overloading the subtraction operator for particles types 

106 

107 Args: 

108 other (particles): particles object to be subtracted 

109 Raises: 

110 DataError: if other is not a particles object 

111 Returns: 

112 particles: differences between caller and other values (self-other) 

113 """ 

114 

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

116 # always create new particles, since otherwise c = a - b changes a as well! 

117 p = particles(self) 

118 p.pos[:] = self.pos - other.pos 

119 p.vel[:] = self.vel - other.vel 

120 p.m = self.m 

121 p.q = self.q 

122 return p 

123 else: 

124 raise DataError("Type error: cannot subtract %s from %s" % (type(other), type(self))) 

125 

126 def __rmul__(self, other): 

127 """ 

128 Overloading the right multiply by factor operator for particles types 

129 

130 Args: 

131 other (float): factor 

132 Raises: 

133 DataError: if other is not a particles object 

134 Returns: 

135 particles: scaled particle's velocity and position as new particle 

136 """ 

137 

138 if isinstance(other, float): 

139 # always create new particles 

140 p = particles(self) 

141 p.pos[:] = other * self.pos 

142 p.vel[:] = other * self.vel 

143 p.m = self.m 

144 p.q = self.q 

145 return p 

146 else: 

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

148 

149 def __abs__(self): 

150 """ 

151 Overloading the abs operator for particles types 

152 

153 Returns: 

154 float: absolute maximum of abs(pos) and abs(vel) for all particles 

155 """ 

156 abspos = abs(self.pos) 

157 absvel = abs(self.vel) 

158 return np.amax((abspos, absvel)) 

159 

160 def send(self, dest=None, tag=None, comm=None): 

161 """ 

162 Routine for sending data forward in time (blocking) 

163 

164 Args: 

165 dest (int): target rank 

166 tag (int): communication tag 

167 comm: communicator 

168 

169 Returns: 

170 None 

171 """ 

172 

173 comm.send(self, dest=dest, tag=tag) 

174 return None 

175 

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

177 """ 

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

179 

180 Args: 

181 dest (int): target rank 

182 tag (int): communication tag 

183 comm: communicator 

184 

185 Returns: 

186 request handle 

187 """ 

188 return comm.isend(self, dest=dest, tag=tag) 

189 

190 def recv(self, source=None, tag=None, comm=None): 

191 """ 

192 Routine for receiving in time 

193 

194 Args: 

195 source (int): source rank 

196 tag (int): communication tag 

197 comm: communicator 

198 

199 Returns: 

200 None 

201 """ 

202 part = comm.recv(source=source, tag=tag) 

203 self.pos[:] = part.pos.copy() 

204 self.vel[:] = part.vel.copy() 

205 self.m = part.m.copy() 

206 self.q = part.q.copy() 

207 return None 

208 

209 

210class acceleration(mesh): 

211 pass 

212 

213 

214class fields(object): 

215 """ 

216 Field data type for 3 dimensions 

217 

218 This data type can be used for electric and magnetic fields in 3 dimensions 

219 

220 Attributes: 

221 elec: contains the electric field 

222 magn: contains the magnetic field 

223 """ 

224 

225 class electric(mesh): 

226 pass 

227 

228 class magnetic(mesh): 

229 pass 

230 

231 def __init__(self, init=None, val=None): 

232 """ 

233 Initialization routine 

234 

235 Args: 

236 init: can either be a number or another fields object 

237 val: initial tuple of values for electric and magnetic (default: (None,None)) 

238 Raises: 

239 DataError: if init is none of the types above 

240 """ 

241 

242 # if init is another fields object, do a copy (init by copy) 

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

244 self.elec = fields.electric(init.elec) 

245 self.magn = fields.magnetic(init.magn) 

246 # if init is a number, create fields object and pick the corresponding initial values 

247 elif ( 

248 isinstance(init, tuple) 

249 and (init[1] is None or isinstance(init[1], MPI.Intracomm)) 

250 and isinstance(init[2], np.dtype) 

251 ): 

252 if isinstance(val, int) or isinstance(val, float) or val is None: 

253 self.elec = fields.electric(init, val=val) 

254 self.magn = fields.magnetic(init, val=val) 

255 elif isinstance(val, tuple) and len(val) == 2: 

256 self.elec = fields.electric(init, val=val[0]) 

257 self.magn = fields.magnetic(init, val=val[1]) 

258 else: 

259 raise DataError('wrong type of val, got %s' % val) 

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

261 else: 

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

263 

264 def __add__(self, other): 

265 """ 

266 Overloading the addition operator for fields types 

267 

268 Args: 

269 other (fields): fields object to be added 

270 Raises: 

271 DataError: if other is not a fields object 

272 Returns: 

273 fields: sum of caller and other values (self+other) 

274 """ 

275 

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

277 # always create new fields, since otherwise c = a - b changes a as well! 

278 p = fields(self) 

279 p.elec[:] = self.elec + other.elec 

280 p.magn[:] = self.magn + other.magn 

281 return p 

282 else: 

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

284 

285 def __sub__(self, other): 

286 """ 

287 Overloading the subtraction operator for fields types 

288 

289 Args: 

290 other (fields): fields object to be subtracted 

291 Raises: 

292 DataError: if other is not a fields object 

293 Returns: 

294 fields: differences between caller and other values (self-other) 

295 """ 

296 

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

298 # always create new fields, since otherwise c = a - b changes a as well! 

299 p = fields(self) 

300 p.elec[:] = self.elec - other.elec 

301 p.magn[:] = self.magn - other.magn 

302 return p 

303 else: 

304 raise DataError("Type error: cannot subtract %s from %s" % (type(other), type(self))) 

305 

306 def __rmul__(self, other): 

307 """ 

308 Overloading the multiply with factor from right operator for fields types 

309 

310 Args: 

311 other (float): factor 

312 Raises: 

313 DataError: if other is not a fields object 

314 Returns: 

315 fields: scaled fields 

316 """ 

317 

318 if isinstance(other, float): 

319 # always create new fields, since otherwise c = a - b changes a as well! 

320 p = fields(self) 

321 p.elec[:] = other * self.elec 

322 p.magn[:] = other * self.magn 

323 return p 

324 else: 

325 raise DataError("Type error: cannot multiply %s with %s" % (type(other), type(self)))