Coverage for src/airball/core.py: 15%

304 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-12-08 10:31 +0900

1import rebound as _rebound 

2import numpy as _np 

3import joblib as _joblib 

4import warnings as _warnings 

5import tempfile as _tempfile 

6 

7from . import tools as _tools 

8from . import units as _u 

9 

10############################################################ 

11################# Flyby Helper Functions ################### 

12############################################################ 

13 

14def _rotate_into_plane(sim, plane): 

15 '''Rotates the simulation into the specified plane.''' 

16 int_types = int, _np.integer 

17 rotation = _rebound.Rotation.to_new_axes(newz=[0,0,1]) 

18 if plane is not None: 

19 # Move the system into the chosen plane of reference. TODO: Make sure the angular momentum calculations don't include other flyby stars. 

20 if plane == 'invariable': rotation = _rebound.Rotation.to_new_axes(newz=sim.angular_momentum()) 

21 elif plane == 'ecliptic': rotation = _rebound.Rotation.to_new_axes(newz=_tools.calculate_angular_momentum(sim)[3]) # Assumes Earth is particle 3. 0-Sun, 1-Mecury, 2-Venus, 3-Earth, ... 

22 elif isinstance(plane, int_types): 

23 p = sim.particles[int(plane)] 

24 rotation = (_rebound.Rotation.orbit(Omega=p.Omega, inc=p.inc, omega=p.omega)).inverse() 

25 sim.rotate(rotation) 

26 return rotation 

27 

28def add_star_to_sim(sim, star, hash, **kwargs): 

29 '''Adds a Star to a REBOUND Simulation. 

30 

31 Parameters 

32 ---------- 

33 sim : the REBOUND Simulation (star and planets). 

34 star: an AIRBALL Star object. 

35 hash: a string to refer to the Star object. 

36 

37 rmax : the starting distance of the flyby star in units of AU; if rmax=0, then the star will be placed at perihelion. 

38 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

39 ''' 

40 # Because REBOUND Simulations are C structs underneath Python, this function passes the simulation by reference. 

41 

42 units = _tools.rebound_units(sim) 

43 rmax = _tools.verify_unit(kwargs.get('rmax', 1e5*_u.au), units['length']) 

44 stellar_elements, semilatus_rectum = _tools.initial_conditions_from_stellar_params(sim, star, rmax) 

45 

46 plane = kwargs.get('plane') 

47 if plane is not None: rotation = _rotate_into_plane(sim, plane) 

48 

49 sim.add(**stellar_elements, hash=hash, primary=sim.particles[0]) 

50 # Because a new particle was added, we need to tell REBOUND to recalculate the coordinates if WHFast is being used. 

51 if sim.integrator == 'whfast': sim.ri_whfast.recalculate_coordinates_this_timestep = 1 

52 sim.integrator_synchronize() # For good measure. 

53 

54 if plane is not None: sim.rotate(rotation.inverse()) 

55 sim.move_to_com() 

56 

57 # Because REBOUND Simulations are C structs underneath Python, this function passes the simulation by reference. 

58 return {'m':stellar_elements['m'] * units['mass'], 'a':stellar_elements['a'] * units['length'], 'e':stellar_elements['e'], 'l':semilatus_rectum * units['length']} 

59 

60def remove_star_from_sim(sim, hash): 

61 # Because REBOUND Simulations are C structs underneath Python, this function passes the simulation by reference. 

62 sim.remove(hash=hash) 

63 # Because a particle was removed, we need to tell REBOUND to recalculate the coordinates if WHFast is being used and to synchronize. 

64 if sim.integrator == 'whfast': sim.ri_whfast.recalculate_coordinates_this_timestep = 1 

65 sim.integrator_synchronize() 

66 sim.move_to_com() # Readjust the system back into the centre of mass/momentum frame for integrating. 

67 # Because REBOUND Simulations are C structs underneath Python, this function passes the simulation by reference. 

68 

69def _time_to_periapsis_from_crossover_point(sim, sim_units, crossoverFactor, index, star_elements): 

70 ''' 

71 Compute the time to periapsis from crossover point. 

72 ''' 

73 rCrossOver = crossoverFactor * sim.particles[index].a * sim_units['length'] # This is the distance to switch integrators 

74 q = star_elements['a'] * (1 - star_elements['e']) 

75 if q < rCrossOver: 

76 f = _np.arccos((star_elements['l']/rCrossOver-1.)/star_elements['e']) # Compute the true anomaly for the cross-over point. 

77 

78 G = (sim.G * sim_units['length']**3 / sim_units['mass'] / sim_units['time']**2) 

79 mu = G * (_tools.system_mass(sim) * sim_units['mass'] + star_elements['m']) 

80 

81 # Compute the time to periapsis from the switching point (-a because the semi-major axis is negative). 

82 with _u.set_enabled_equivalencies(_u.dimensionless_angles()): 

83 E = _np.arccosh((_np.cos(f)+star_elements['e'])/(1.+star_elements['e']*_np.cos(f))) # Compute the eccentric anomaly 

84 M = star_elements['e'] * _np.sinh(E)-E # Compute the mean anomaly 

85 return True, M/_np.sqrt(mu/(-star_elements['a']*star_elements['a']*star_elements['a'])) 

86 else: return False, None 

87 

88# old signature flyby(sim, star=None, m=0.3, b=1000, v=40, e=None, inc='uniform', omega='uniform', Omega='uniform', rmax=2.5e5, hybrid=True, crossoverFactor=30, overwrite=False): 

89 

90def _integrate_with_ias15(sim, tmax): 

91 sim.integrator = 'ias15' 

92 sim.gravity = 'basic' 

93 sim.integrate(tmax) 

94 

95def _integrate_with_whckl(sim, tmax, dt, dt_frac): 

96 sim.integrator = 'whckl' 

97 sim.ri_whfast.safe_mode = 0 

98 sim.ri_whfast.recalculate_coordinates_this_timestep = 1 

99 sim.integrator_synchronize() 

100 if sim.particles[1].P > 0: sim.dt = dt_frac*sim.particles[1].P 

101 else: sim.dt = dt 

102 sim.integrate(tmax) 

103 sim.ri_whfast.recalculate_coordinates_this_timestep = 1 

104 sim.integrator_synchronize() 

105 

106############################################################ 

107#################### Flyby Functions ####################### 

108############################################################ 

109 

110def hybrid_flyby(sim, star, **kwargs): 

111 ''' 

112 Simulate a stellar flyby to a REBOUND simulation. 

113 

114 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

115 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

116 This function assumes that you are using the WHCKL integrator with REBOUND. 

117 Uses IAS15 (instead of WHCKL) for the closest approach if b < planet_a * crossoverFactor 

118 

119 Parameters 

120 ---------- 

121 sim : the REBOUND Simulation (star and planets) that will experience the flyby star 

122 star: a AIRBALL Star object 

123 

124 rmax : the starting distance of the flyby star in units of AU 

125 crossoverFactor: the value for when to switch integrators, i.e. 30 times the semi-major axis of particle 1. Default is 30x. 

126 particle_index: the particle index to consider for the crossoverFactor. Default is 1. 

127 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

128 integrator: sets the integrator for before and after the hybrid switch (for example, if you want to use WHCKL instead of WHFast) 

129 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

130 ''' 

131 

132 overwrite = kwargs.get('overwrite', True) 

133 if not overwrite: sim = sim.copy() 

134 hash = kwargs.get('hash', 'flybystar') 

135 sim_units = _tools.rebound_units(sim) 

136 

137 star_vars = add_star_to_sim(sim, star, rmax=kwargs.get('rmax', 1e5*_u.au), plane=kwargs.get('plane'), hash=hash) 

138 

139 tperi = sim.particles[hash].T - sim.t # Compute the time to periapsis for the flyby star from the current time. 

140 

141 # Integrate the flyby. Start at the current time and go to twice the time to periapsis. 

142 switch, tIAS15 = _time_to_periapsis_from_crossover_point(sim, sim_units, crossoverFactor=kwargs.get('crossoverFactor', 30), index=kwargs.get('particle_index', 1), star_elements=star_vars) 

143 if switch: 

144 t_switch = sim.t + tperi - tIAS15.value 

145 t_switch_back = sim.t + tperi + tIAS15.value 

146 t_end = sim.t + 2*tperi 

147 

148 dt = sim.dt 

149 dt_frac = sim.dt/sim.particles[1].P 

150 

151 _integrate_with_whckl(sim, t_switch, dt, dt_frac) 

152 _integrate_with_ias15(sim, t_switch_back) 

153 _integrate_with_whckl(sim, t_end, dt, dt_frac) 

154 

155 else: _integrate_with_whckl(sim, tmax=(sim.t + 2*tperi), dt=sim.dt, dt_frac=sim.dt/sim.particles[1].P) 

156 

157 # Remove the flyby star. 

158 remove_star_from_sim(sim, hash=hash) 

159 

160 return sim 

161 

162def hybrid_flybys(sims, stars, **kwargs): 

163 ''' 

164 Run serial flybys in parallel. 

165 

166 Parameters 

167 --------------- 

168 sims : A list of REBOUND Simulations. 

169 REBOUND simulations to integrate flybys with. If only one simulation is given, then AIRBALL will duplicate it to match the number of Stars given. Required. 

170 stars : AIRBALL Stars. 

171 The objects that will flyby the given REBOUND simulations. Required. 

172 

173 crossoverFactor : Float. 

174 The value for when to switch to IAS15 as a multiple of sim.particles[1].a Default is 30. 

175 overwrite : True/False. 

176 Sets whether or not to return new simulation objects or overwrite the given ones. Default is True, meaning the same simulation object will be returned. This keeps all original pointers attached to it. 

177 rmax : Float. 

178 The starting distance of the flyby object (in units of the REBOUND Simulation). Default is 1e5. 

179 plane : String/Int. 

180 The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

181 particle_index : Int. 

182 The simulation particle index to define the crossoverFactor with respect to. Default is 1. 

183 hashes : array_like. 

184 A list of hash values for adding and removing stars from simulations. Default is 'flybystar'. 

185 

186 inds : array_like 

187 An array of indices to determine which sims and stars to integrate. Default is all of them. 

188 n_jobs : Integer. 

189 The number of jobs per CPU to run in parallel. Default is -1. 

190 verbose : Integer. 

191 The amount of details to display for the parallel jobs. Default is 0. 

192 

193 Returns 

194 ------------- 

195 hybrid_flybys : list 

196 List of REBOUND simulations that experienced a flyby. 

197 ''' 

198 Nruns = 0 

199 try: 

200 Nruns = len(sims) 

201 if Nruns != len(stars): raise Exception('Sims and stars are unequal lengths') 

202 except: 

203 Nruns = len(stars) 

204 sims = [sims.copy() for _ in range(Nruns)] 

205 

206 try: 

207 rmax = kwargs['rmax'] 

208 if not _tools.isList(rmax): rmax = Nruns * [rmax] 

209 elif len(rmax) != Nruns: raise Exception('List arguments must be same length.') 

210 except KeyError: rmax = Nruns * [1e5] 

211 

212 try: 

213 crossoverFactor = kwargs['crossoverFactor'] 

214 if not _tools.isList(crossoverFactor): crossoverFactor = Nruns * [crossoverFactor] 

215 elif len(crossoverFactor) != Nruns: raise Exception('List arguments must be same length.') 

216 except KeyError: crossoverFactor = Nruns * [30] 

217 

218 try: 

219 hashes = kwargs['hashes'] 

220 if not _tools.isList(hashes): hashes = Nruns * [hashes] 

221 elif len(hashes) != Nruns: raise Exception('List arguments must be same length.') 

222 except KeyError: hashes = Nruns * ['flybystar'] 

223 

224 inds = kwargs.get('inds', _np.arange(Nruns)) 

225 overwrite = kwargs.get('overwrite', True) 

226 n_jobs = kwargs.get('n_jobs', -1) 

227 verbose = kwargs.get('verbose', 0) 

228 particle_index = kwargs.get('particle_index', 1) 

229 plane = kwargs.get('plane', None) 

230 

231 sim_results = _joblib.Parallel(n_jobs=n_jobs, verbose=verbose)( 

232 _joblib.delayed(hybrid_flyby)( 

233 sim=sims[int(i)], star=stars[i], rmax=rmax[i], crossoverFactor=crossoverFactor[i], overwrite=overwrite, particle_index=particle_index, plane=plane, hash=hashes[i]) 

234 for i in inds) 

235 

236 return sim_results 

237 

238def flyby(sim, star, **kwargs): 

239 ''' 

240 Simulate a stellar flyby to a REBOUND simulation. 

241 

242 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

243 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

244 This function assumes that you are using a WHFAST integrator with REBOUND. 

245 

246 Parameters 

247 ---------- 

248 sim : the REBOUND Simulation (star and planets) that will experience the flyby star 

249 star: a airball.Star object 

250 

251 rmax : the starting distance of the flyby star in units of AU 

252 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

253 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

254 hash: String. The name for the flyby star. Default is `flybystar`. 

255 ''' 

256 if kwargs.get('hybrid', False): return hybrid_flyby(sim, star, **kwargs) 

257 else: 

258 if sim.integrator == 'whfast': _warnings.warn("Did you intend to use the hybrid method with WHFast?", RuntimeWarning) 

259 overwrite = kwargs.get('overwrite', True) 

260 if not overwrite: sim = sim.copy() 

261 hash = kwargs.get('hash', 'flybystar') 

262 add_star_to_sim(sim, star, hash, rmax=kwargs.get('rmax', 1e5*_u.au), plane=kwargs.get('plane')) 

263 tperi = sim.particles[hash].T - sim.t # Compute the time to periapsis for the flyby star from the current time. 

264 sim.integrate(sim.t + 2*tperi) 

265 remove_star_from_sim(sim, hash) 

266 return sim 

267 

268def flybys(sims, stars, **kwargs): 

269 ''' 

270 Run serial flybys in parallel. 

271 

272 Parameters 

273 --------------- 

274 sims : A list of REBOUND Simulations. 

275 REBOUND simulations to integrate flybys with. If only one simulation is given, then AIRBALL will duplicate it to match the number of Stars given. Required. 

276 stars : AIRBALL Stars. 

277 The objects that will flyby the given REBOUND simulations. Required. 

278 

279 overwrite : True/False. 

280 Sets whether or not to return new simulation objects or overwrite the given ones. Default is True, meaning the same simulation object will be returned. This keeps all original pointers attached to it. 

281 rmax : Float. 

282 The starting distance of the flyby object (in units of the REBOUND Simulation). Default is 1e5. 

283 plane : String/Int. 

284 The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

285 

286 inds : array_like 

287 An array of indices to determine which sims and stars to integrate. Default is all of them. 

288 n_jobs : Integer. 

289 The number of jobs per CPU to run in parallel. Default is -1. 

290 verbose : Integer. 

291 The amount of details to display for the parallel jobs. Default is 0. 

292 

293 Returns 

294 ------------- 

295 flybys : list 

296 List of REBOUND simulations that experienced a flyby. 

297 ''' 

298 Nruns = 0 

299 try: 

300 Nruns = len(sims) 

301 if Nruns != len(stars): raise Exception('Sims and stars are unequal lengths') 

302 except: 

303 Nruns = len(stars) 

304 sims = [sims.copy() for _ in range(Nruns)] 

305 

306 try: inds = kwargs['inds'] 

307 except KeyError: inds = _np.arange(Nruns) 

308 

309 try: 

310 rmax = kwargs['rmax'] 

311 rmax = _tools.verify_unit(rmax, _u.au) 

312 if len(rmax.shape) == 0: rmax = _np.array(stars.N * [rmax.value]) << rmax.unit 

313 elif len(rmax) != stars.N: raise Exception('List arguments must be same length.') 

314 except KeyError: rmax = _np.array(Nruns * [1e5]) << _u.au 

315 if _np.any(stars.b > rmax): raise Exception('Some stellar impact parameters are greater than the stellar starting distance, rmax.') 

316 

317 heartbeat = kwargs.get('heartbeat', None) 

318 inds = kwargs.get('inds', _np.arange(Nruns)) 

319 overwrite = kwargs.get('overwrite', True) 

320 n_jobs = kwargs.get('n_jobs', -1) 

321 verbose = kwargs.get('verbose', 0) 

322 require = kwargs.get('require') 

323 plane = kwargs.get('plane', None) 

324 hybrid = kwargs.get('hybrid', False) 

325 

326 sim_results = _joblib.Parallel(n_jobs=n_jobs, verbose=verbose, require=require)( 

327 _joblib.delayed(flyby)( 

328 sim=sims[int(i)], star=stars[i], rmax=rmax[i], overwrite=overwrite, heartbeat=heartbeat, plane=plane, hybrid=hybrid) 

329 for i in inds) 

330 return sim_results 

331 

332def successive_flybys(sim, stars, **kwargs): 

333 ''' 

334 Simulate a stellar flyby to a REBOUND simulation. 

335 

336 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

337 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

338 This function assumes that you are using a WHFAST integrator with REBOUND. 

339 Uses IAS15 (instead of WHFast) for the closest approach if b < planet_a * crossoverFactor 

340 

341 Parameters 

342 ---------- 

343 sim : the REBOUND Simulation (star and planets) that will experience the flyby star 

344 star: a AIRBALL Star object 

345 

346 rmax : the starting distance of the flyby star in units of AU 

347 crossoverFactor: the value for when to switch integrators if hybrid=True 

348 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

349 integrator: sets the integrator for before and after the hybrid switch (for example, if you want to use WHCKL instead of WHFast) 

350 heartbeat: sets a heartbeat function 

351 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

352 ''' 

353 

354 # Do not overwrite given sim. 

355 overwrite = kwargs.get('overwrite', True) 

356 if overwrite == False: sim = sim.copy() 

357 hashes = kwargs.get('hashes', [f'flybystar{i}' for i in range(stars.N)]) 

358 saveSnapshots = kwargs.get('snapshots', False) 

359 if saveSnapshots: snapshots = [sim.copy()] 

360 for i,star in enumerate(stars): 

361 if overwrite == True: flyby(sim, star, hash=hashes[i], **kwargs) 

362 else: sim = flyby(sim, star, hash=hashes[i], **kwargs) 

363 if saveSnapshots: snapshots.append(sim.copy()) 

364 if saveSnapshots: return snapshots 

365 else: return sim 

366 

367def concurrent_flybys(sim, stars, start_times, **kwargs): 

368 ''' 

369 Simulate a stellar flyby to a REBOUND simulation. 

370 

371 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

372 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

373 

374 Parameters 

375 ---------- 

376 sim : the REBOUND Simulation (star and planets) that will experience the flyby star. 

377 stars: an AIRBALL Stars object (containing multiple stars). 

378 times: an array of times for the stars to be added to the sim. 

379 

380 rmax : the starting distance of the flyby star in units of AU 

381 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

382 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

383 ''' 

384 message = "Integrating flybys concurrently may give unintuitive results. Use with caution." 

385 _warnings.warn(message, RuntimeWarning) 

386 

387 # Do not overwrite given sim. 

388 overwrite = kwargs.get('overwrite', True) 

389 if not overwrite: sim = sim.copy() 

390 sim_units = _tools.rebound_units(sim) 

391 

392 rmax = kwargs.get('rmax', 1e5*_u.au) 

393 plane = kwargs.get('plane') 

394 start_times = _tools.verify_unit(start_times, sim_units['time']).value 

395 hashes = kwargs.get('hashes', [f'flybystar{i}' for i in range(stars.N)]) 

396 

397 # Using the sim and the start times, compute the end times for the flyby stars. 

398 all_times = _np.zeros((stars.N, 2)) 

399 for star_number, star in enumerate(stars): 

400 tmp_sim = sim.copy() 

401 hash = hashes[star_number] 

402 add_star_to_sim(tmp_sim, star, hash=hash, rmax=rmax, plane=plane) 

403 # Compute the time to periapsis for the flyby star from the current simulation time. 

404 tperi = tmp_sim.particles[hash].T - tmp_sim.t 

405 end_time = start_times[star_number] + tmp_sim.t + 2*tperi 

406 all_times[star_number] = [start_times[star_number], end_time] 

407 

408 # Sort the event times sequentially. 

409 all_times = all_times.flatten() 

410 event_order = _np.argsort(all_times) 

411 max_event_number = len(all_times) 

412 

413 # Integrate the flybys, adding and removing them at the appropriate times. 

414 event_number = 0 

415 while event_number < max_event_number: 

416 event_index = event_order[event_number] 

417 star_number = event_index//2 

418 sim.integrate(all_times[event_index]) 

419 if event_index%2 == 0: add_star_to_sim(sim, stars[star_number], hash=hashes[star_number], rmax=rmax, plane=plane) 

420 else: remove_star_from_sim(sim, hash=hashes[star_number]) 

421 event_number += 1 

422 return sim 

423 

424 

425 

426def _hybrid_successive_flybys(sim, stars, rmax=1e5, crossoverFactor=30, overwrite=False, heartbeat=None, particle_index=1, plane=None): 

427 ''' 

428 Simulate a stellar flyby to a REBOUND simulation. 

429 

430 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

431 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

432 This function assumes that you are using a WHFAST integrator with REBOUND. 

433 Uses IAS15 (instead of WHFast) for the closest approach if b < planet_a * crossoverFactor 

434 

435 Parameters 

436 ---------- 

437 sim : the REBOUND Simulation (star and planets) that will experience the flyby star 

438 star: a AIRBALL Star object 

439 

440 rmax : the starting distance of the flyby star in units of AU 

441 crossoverFactor: the value for when to switch integrators if hybrid=True 

442 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

443 integrator: sets the integrator for before and after the hybrid switch (for example, if you want to use WHCKL instead of WHFast) 

444 heartbeat: sets a heartbeat function 

445 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

446 ''' 

447 

448 # Do not overwrite given sim. 

449 if not overwrite: sim = sim.copy() 

450 if heartbeat is not None: sim.heartbeat = heartbeat 

451 sim_units = _tools.rebound_units(sim) 

452 

453 output = None 

454 with _tempfile.NamedTemporaryFile() as tmp: 

455 sim.simulationarchive_snapshot(tmp.name, deletefile=True) 

456 for star_number, star in enumerate(stars): 

457 index = int(particle_index) 

458 hash = f'flybystar{star_number}' 

459 star_vars = add_star_to_sim(sim, star, rmax=rmax, plane=plane, hash=hash) 

460 

461 tperi = sim.particles[hash].T - sim.t # Compute the time to periapsis for the flyby star from the current time. 

462 

463 # Integrate the flyby. Start at the current time and go to twice the time to periapsis. 

464 switch, tIAS15 = _time_to_periapsis_from_crossover_point(sim, sim_units, crossoverFactor, index, star_vars) 

465 if switch: 

466 t_switch = sim.t + tperi - tIAS15.value 

467 t_switch_back = sim.t + tperi + tIAS15.value 

468 t_end = sim.t + 2*tperi 

469 

470 dt = sim.dt 

471 dt_frac = sim.dt/sim.particles[1].P 

472 

473 _integrate_with_whckl(sim, t_switch, dt, dt_frac) 

474 sim.simulationarchive_snapshot(tmp.name, deletefile=False) 

475 _integrate_with_ias15(sim, t_switch_back) 

476 sim.simulationarchive_snapshot(tmp.name, deletefile=False) 

477 _integrate_with_whckl(sim, t_end, dt, dt_frac) 

478 

479 else: _integrate_with_whckl(sim, tmax=(sim.t + 2*tperi), dt=sim.dt, dt_frac=sim.dt/sim.particles[1].P) 

480 

481 # Remove the flyby star. 

482 remove_star_from_sim(sim, hash=hash) 

483 sim.simulationarchive_snapshot(tmp.name, deletefile=False) 

484 output = _rebound.SimulationArchive(tmp.name) 

485 return output 

486 

487def _hybrid_successive_flybys_parallel(sims, stars, **kwargs): 

488 ''' 

489 Run serial flybys in parallel. 

490 

491 Parameters 

492 --------------- 

493 sims : A list of REBOUND Simulations. 

494 REBOUND simulations to integrate flybys with. If only one simulation is given, then AIRBALL will duplicate it to match the number of Stars given. Required. 

495 stars : A list of AIRBALL Stars. 

496 The objects that will flyby the given REBOUND simulations. Required. 

497 

498 crossoverFactor : Float. 

499 The value for when to switch to IAS15 as a multiple of sim.particles[1].a Default is 30. 

500 overwrite : True/False. 

501 Sets whether or not to return new simulation objects or overwrite the given ones. Default is True, meaning the same simulation object will be returned. This keeps all original pointers attached to it. 

502 rmax : Float. 

503 The starting distance of the flyby object (in units of the REBOUND Simulation). Default is 1e5. 

504 plane : String/Int. 

505 The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

506 particle_index : Int. 

507 The simulation particle index to define the crossoverFactor with respect to. Default is 1. 

508 

509 inds : array_like 

510 An array of indices to determine which sims and stars to integrate. Default is all of them. 

511 n_jobs : Integer. 

512 The number of jobs per CPU to run in parallel. Default is -1. 

513 verbose : Integer. 

514 The amount of details to display for the parallel jobs. Default is 0. 

515 

516 Returns 

517 ------------- 

518 hybrid_flybys : list 

519 List of REBOUND simulations that experienced a flyby. 

520 ''' 

521 

522 _warnings.warn("This function has not been thoroughly tested. Use with caution.", RuntimeWarning) 

523 

524 Nruns = 0 

525 try: 

526 Nruns = len(sims) 

527 if Nruns != len(stars): raise Exception('Sims and stars are unequal lengths') 

528 except Exception as err: 

529 # TypeError: object of type 'Simulation' has no len() 

530 raise err 

531 

532 try: 

533 rmax = kwargs['rmax'] 

534 if not _tools.isList(rmax): rmax = Nruns * [rmax] 

535 elif len(rmax) != Nruns: raise Exception('List arguments must be same length.') 

536 except KeyError: rmax = Nruns * [1e5] 

537 

538 try: 

539 crossoverFactor = kwargs['crossoverFactor'] 

540 if not _tools.isList(crossoverFactor): crossoverFactor = Nruns * [crossoverFactor] 

541 elif len(crossoverFactor) != Nruns: raise Exception('List arguments must be same length.') 

542 except KeyError: crossoverFactor = Nruns * [30] 

543 

544 heartbeat = kwargs.get('heartbeat', None) 

545 inds = kwargs.get('inds', _np.arange(Nruns)) 

546 overwrite = kwargs.get('overwrite', True) 

547 n_jobs = kwargs.get('n_jobs', -1) 

548 verbose = kwargs.get('verbose', 0) 

549 particle_index = kwargs.get('particle_index', 1) 

550 plane = kwargs.get('plane', None) 

551 

552 sim_results = _joblib.Parallel(n_jobs=n_jobs, verbose=verbose)( 

553 _joblib.delayed(_hybrid_successive_flybys)( 

554 sim=sims[int(i)], stars=stars[int(i)], rmax=rmax[i], crossoverFactor=crossoverFactor[i], overwrite=overwrite, particle_index=particle_index, plane=plane, heartbeat=heartbeat) 

555 for i in inds) 

556 

557 return sim_results 

558 

559def _hybrid_concurrent_flybys(sim, stars, times, rmax=1e5, crossoverFactor=30, overwrite=False, heartbeat=None, particle_index=1, plane=None, verbose=False): 

560 ''' 

561 Simulate a stellar flyby to a REBOUND simulation. 

562 

563 Because REBOUND Simulations are C structs underneath the Python, this function can pass the simulation by reference. 

564 Meaning, any changes made inside this function to the REBOUND simulation are permanent. This can be avoided by specifying overwrite=False. 

565 This function assumes that you are using a WHFAST integrator with REBOUND. 

566 Uses IAS15 (instead of WHFast) for the closest approach if b < planet_a * crossoverFactor 

567 

568 Parameters 

569 ---------- 

570 sim : the REBOUND Simulation (star and planets) that will experience the flyby star 

571 star: a AIRBALL Star object 

572 

573 rmax : the starting distance of the flyby star in units of AU 

574 crossoverFactor: the value for when to switch integrators if hybrid=True 

575 overwrite: determines whether or not to return a copy of sim (overwrite=False) or integrate using the original sim (overwrite=True) 

576 integrator: sets the integrator for before and after the hybrid switch (for example, if you want to use WHCKL instead of WHFast) 

577 heartbeat: sets a heartbeat function 

578 plane: String/Int. The plane defining the orientation of the star, None, 'invariable', 'ecliptic', or Int. Default is None. 

579 ''' 

580 

581 _warnings.warn("This function has not been thoroughly tested. Use with caution.", RuntimeWarning) 

582 

583 # Do not overwrite given sim. 

584 if not overwrite: sim = sim.copy() 

585 if heartbeat is not None: sim.heartbeat = heartbeat 

586 sim_units = _tools.rebound_units(sim) 

587 index = int(particle_index) 

588 

589 times = _tools.verify_unit(times, sim_units['time']).value 

590 all_times = [] 

591 for star_number, star in enumerate(stars): 

592 these_times = [] 

593 tmp_sim = sim.copy() 

594 hash = f'tmp{star_number}' 

595 star_vars = add_star_to_sim(tmp_sim, star, rmax=rmax, plane=plane, hash=hash) 

596 

597 tperi = times[star_number] + tmp_sim.particles[hash].T - tmp_sim.t # Compute the time to periapsis for the flyby star from the current time. 

598 these_times.append(times[star_number]) 

599 # Integrate the flyby. Start at the current time and go to twice the time to periapsis. 

600 switch, tIAS15 = _time_to_periapsis_from_crossover_point(tmp_sim, sim_units, crossoverFactor, index, star_vars) 

601 if switch: 

602 these_times.append(times[star_number] + tmp_sim.t + tperi - tIAS15.value) 

603 these_times.append(times[star_number] + tmp_sim.t + tperi + tIAS15.value) 

604 else: 

605 these_times.append(_np.nan) 

606 these_times.append(_np.nan) 

607 

608 these_times.append(times[star_number] + tmp_sim.t + 2*tperi) 

609 all_times.append(these_times) 

610 all_times = _np.array(all_times).flatten() 

611 max_event_number = len(all_times) - _np.sum(_np.isnan(all_times)) 

612 event_order = _np.argsort(all_times) 

613 if verbose: 

614 tmpdic = {0 : f'ADD', 1 : f'start IAS15', 2 : f'end IAS15', 3 : f'REMOVE'} 

615 print([f'{tmpdic[i%4]} {i//4}' for i in event_order[:max_event_number]]) 

616 

617 useIAS15 = _np.array([False] * stars.N) 

618 def startUsingIAS15(i, IAS15_array): IAS15_array[i//4] = True 

619 def stopUsingIAS15(i, IAS15_array): IAS15_array[i//4] = False 

620 

621 def function_map( i, v, sim, star, IAS15_array, plane, hash): 

622 if not _np.isnan(v): 

623 map_i = i%4 

624 if map_i == 0: add_star_to_sim(sim, star, plane=plane, hash=hash) 

625 elif map_i == 1: startUsingIAS15(i, IAS15_array) 

626 elif map_i == 2: stopUsingIAS15(i, IAS15_array) 

627 elif map_i == 3: remove_star_from_sim(sim, hash=hash) 

628 else: pass 

629 

630 output = None 

631 event_number = 0 

632 dt = sim.dt 

633 dt_frac = sim.dt/sim.particles[1].P 

634 with _tempfile.NamedTemporaryFile() as tmp: 

635 sim.simulationarchive_snapshot(tmp.name, deletefile=True) 

636 while event_number < max_event_number: 

637 event_index = event_order[event_number] 

638 star_number = event_index//4 

639 if _np.any(useIAS15): _integrate_with_ias15(sim, all_times[event_index]) 

640 else: _integrate_with_whckl(sim, all_times[event_index], dt, dt_frac) 

641 function_map(event_index, all_times[event_index], sim, stars[star_number], useIAS15, plane, hash=f'flybystar{star_number}') 

642 sim.simulationarchive_snapshot(tmp.name, deletefile=False) 

643 event_number += 1 

644 output = _rebound.SimulationArchive(tmp.name) 

645 

646 return output