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

1 import rebound as _rebound

2 import numpy as _np

3 import joblib as _joblib

4 import warnings as _warnings

5 import tempfile as _tempfile

6

7 from . import tools as _tools

8 from . import units as _u

9

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

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

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

13

14 def _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

28 def 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

60 def 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

69 def _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

90 def _integrate_with_ias15(sim, tmax):

91 sim.integrator = 'ias15'

92 sim.gravity = 'basic'

93 sim.integrate(tmax)

94

95 def _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

110 def 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

162 def 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

238 def 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

268 def 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

332 def 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

367 def 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

426 def _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

487 def _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

559 def _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