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
« 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
7from . import tools as _tools
8from . import units as _u
10############################################################
11################# Flyby Helper Functions ###################
12############################################################
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
28def add_star_to_sim(sim, star, hash, **kwargs):
29 '''Adds a Star to a REBOUND Simulation.
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.
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.
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)
46 plane = kwargs.get('plane')
47 if plane is not None: rotation = _rotate_into_plane(sim, plane)
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.
54 if plane is not None: sim.rotate(rotation.inverse())
55 sim.move_to_com()
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']}
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.
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.
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'])
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
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):
90def _integrate_with_ias15(sim, tmax):
91 sim.integrator = 'ias15'
92 sim.gravity = 'basic'
93 sim.integrate(tmax)
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()
106############################################################
107#################### Flyby Functions #######################
108############################################################
110def hybrid_flyby(sim, star, **kwargs):
111 '''
112 Simulate a stellar flyby to a REBOUND simulation.
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
119 Parameters
120 ----------
121 sim : the REBOUND Simulation (star and planets) that will experience the flyby star
122 star: a AIRBALL Star object
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 '''
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)
137 star_vars = add_star_to_sim(sim, star, rmax=kwargs.get('rmax', 1e5*_u.au), plane=kwargs.get('plane'), hash=hash)
139 tperi = sim.particles[hash].T - sim.t # Compute the time to periapsis for the flyby star from the current time.
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
148 dt = sim.dt
149 dt_frac = sim.dt/sim.particles[1].P
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)
155 else: _integrate_with_whckl(sim, tmax=(sim.t + 2*tperi), dt=sim.dt, dt_frac=sim.dt/sim.particles[1].P)
157 # Remove the flyby star.
158 remove_star_from_sim(sim, hash=hash)
160 return sim
162def hybrid_flybys(sims, stars, **kwargs):
163 '''
164 Run serial flybys in parallel.
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.
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'.
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.
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)]
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]
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]
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']
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)
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)
236 return sim_results
238def flyby(sim, star, **kwargs):
239 '''
240 Simulate a stellar flyby to a REBOUND simulation.
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.
246 Parameters
247 ----------
248 sim : the REBOUND Simulation (star and planets) that will experience the flyby star
249 star: a airball.Star object
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
268def flybys(sims, stars, **kwargs):
269 '''
270 Run serial flybys in parallel.
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.
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.
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.
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)]
306 try: inds = kwargs['inds']
307 except KeyError: inds = _np.arange(Nruns)
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.')
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)
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
332def successive_flybys(sim, stars, **kwargs):
333 '''
334 Simulate a stellar flyby to a REBOUND simulation.
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
341 Parameters
342 ----------
343 sim : the REBOUND Simulation (star and planets) that will experience the flyby star
344 star: a AIRBALL Star object
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 '''
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
367def concurrent_flybys(sim, stars, start_times, **kwargs):
368 '''
369 Simulate a stellar flyby to a REBOUND simulation.
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.
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.
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)
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)
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)])
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]
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)
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
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.
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
435 Parameters
436 ----------
437 sim : the REBOUND Simulation (star and planets) that will experience the flyby star
438 star: a AIRBALL Star object
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 '''
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)
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)
461 tperi = sim.particles[hash].T - sim.t # Compute the time to periapsis for the flyby star from the current time.
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
470 dt = sim.dt
471 dt_frac = sim.dt/sim.particles[1].P
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)
479 else: _integrate_with_whckl(sim, tmax=(sim.t + 2*tperi), dt=sim.dt, dt_frac=sim.dt/sim.particles[1].P)
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
487def _hybrid_successive_flybys_parallel(sims, stars, **kwargs):
488 '''
489 Run serial flybys in parallel.
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.
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.
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.
516 Returns
517 -------------
518 hybrid_flybys : list
519 List of REBOUND simulations that experienced a flyby.
520 '''
522 _warnings.warn("This function has not been thoroughly tested. Use with caution.", RuntimeWarning)
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
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]
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]
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)
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)
557 return sim_results
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.
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
568 Parameters
569 ----------
570 sim : the REBOUND Simulation (star and planets) that will experience the flyby star
571 star: a AIRBALL Star object
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 '''
581 _warnings.warn("This function has not been thoroughly tested. Use with caution.", RuntimeWarning)
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)
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)
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)
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]])
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
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
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)
646 return output