Coverage for src/airball/tools.py: 49%

328 statements  

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

1import numpy as _np 

2import joblib as _joblib 

3import warnings as _warnings 

4import astropy.constants as const 

5from . import units as _u 

6from scipy.stats import uniform 

7from scipy.stats import maxwell 

8from scipy.stats import expon 

9 

10twopi = 2.*_np.pi 

11 

12class UnitSet(): 

13 

14 def __init__(self, UNIT_SYSTEM=[]) -> None: 

15 self._units = {'length': _u.au, 'time': _u.Myr, 'mass': _u.solMass, 'angle': _u.rad, 'velocity': _u.km/_u.s, 'object': _u.stars, 'density': _u.stars/_u.pc**3} 

16 self.UNIT_SYSTEM = UNIT_SYSTEM 

17 pass 

18 

19 @property 

20 def units(self): 

21 return self._units 

22 

23 @property 

24 def UNIT_SYSTEM(self): 

25 return self._UNIT_SYSTEM 

26 

27 def __getitem__(self, key): 

28 if isinstance(key, str): return self.units[key] 

29 else: raise InvalidKeyException() 

30 

31 def __setitem__(self, key, value): 

32 if isinstance(key, str): 

33 if isUnit(value): self.units[key] = value 

34 else: raise InvalidUnitException() 

35 else: raise InvalidKeyException() 

36 

37 def __str__(self): 

38 s = '{' 

39 for key in self.units: 

40 s += f'{key}: {self.units[key].to_string()}, ' 

41 s = s[:-2] + '}' 

42 return s 

43 

44 def __repr__(self): 

45 s = '{' 

46 for key in self.units: 

47 s += f'{key}: {self.units[key].to_string()},\n' 

48 s = s[:-2] + '}' 

49 return s 

50 

51 def __iter__(self): 

52 for k in self.units: 

53 yield self.units[k] 

54 

55 def __eq__(self, other): 

56 '''Determines if two UnitSets are equivalent, not necessarily identical.''' 

57 if isinstance(other, UnitSet): 

58 return _np.all([u1.is_equivalent(u2) for u1,u2 in zip(self, other)]) 

59 return NotImplemented 

60 

61 def __ne__(self, other): 

62 """Overrides the default implementation (unnecessary in Python 3)""" 

63 x = self.__eq__(other) 

64 if x is not NotImplemented: 

65 return not x 

66 return NotImplemented 

67 

68 def __hash__(self): 

69 """Overrides the default implementation""" 

70 return hash(tuple(sorted(self.__dict__.items()))) 

71 

72 

73 @property 

74 def length(self): 

75 """Units of LENGTH.""" 

76 return self._units['length'] 

77 

78 @length.setter 

79 def length(self, value): 

80 """ 

81 Setter for the Units of LENGTH. 

82 

83 Parameters: 

84 - value: An Astropy Quantity describing LENGTH. 

85 """ 

86 self.UNIT_SYSTEM = [value] 

87 

88 @property 

89 def time(self): 

90 """Units of TIME.""" 

91 return self._units['time'] 

92 

93 @time.setter 

94 def time(self, value): 

95 """ 

96 Setter for the Units of TIME. 

97 

98 Parameters: 

99 - value: An Astropy Quantity describing TIME. 

100 """ 

101 self.UNIT_SYSTEM = [value] 

102 

103 @property 

104 def mass(self): 

105 """Units of MASS.""" 

106 return self._units['mass'] 

107 

108 @mass.setter 

109 def mass(self, value): 

110 """ 

111 Setter for the Units of MASS. 

112 

113 Parameters: 

114 - value: An Astropy Quantity describing MASS. 

115 """ 

116 self.UNIT_SYSTEM = [value] 

117 

118 @property 

119 def angle(self): 

120 """Units of ANGLE.""" 

121 return self._units['angle'] 

122 

123 @angle.setter 

124 def angle(self, value): 

125 """ 

126 Setter for the Units of ANGLE. 

127 

128 Parameters: 

129 - value: An Astropy Quantity describing ANGLE. 

130 """ 

131 self.UNIT_SYSTEM = [value] 

132 

133 @property 

134 def velocity(self): 

135 """Units of VELOCITY.""" 

136 return self._units['velocity'] 

137 

138 @velocity.setter 

139 def velocity(self, value): 

140 """ 

141 Setter for the Units of VELOCITY. 

142 

143 Parameters: 

144 - value: An Astropy Quantity describing VELOCITY. 

145 """ 

146 self.UNIT_SYSTEM = [value] 

147 

148 @property 

149 def density(self): 

150 """Units of DENSITY.""" 

151 return self._units['density'] 

152 

153 @density.setter 

154 def density(self, value): 

155 """ 

156 Setter for the Units of DENSITY . 

157 

158 Parameters: 

159 - value: An Astropy Quantity describing DENSITY. 

160 """ 

161 self.UNIT_SYSTEM = [value] 

162 

163 @property 

164 def object(self): 

165 """Units of an OBJECT (such as a star).""" 

166 return self._units['object'] 

167 

168 @object.setter 

169 def object(self, value): 

170 """ 

171 Setter for the Units of an OBJECT (such as a star). 

172 

173 Parameters: 

174 - value: An Astropy Quantity describing an OBJECT (such as a star). 

175 """ 

176 self.UNIT_SYSTEM = [value] 

177 

178 @UNIT_SYSTEM.setter 

179 def UNIT_SYSTEM(self, UNIT_SYSTEM): 

180 if UNIT_SYSTEM != []: 

181 lengthUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.m)] 

182 self._units['length'] = lengthUnit[0] if lengthUnit != [] else self._units['length'] 

183 

184 timeUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.s)] 

185 self._units['time'] = timeUnit[0] if timeUnit != [] else self._units['time'] 

186 

187 velocityUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.km/_u.s)] 

188 if velocityUnit == [] and timeUnit != [] and lengthUnit != []: velocityUnit = [lengthUnit[0]/timeUnit[0]] 

189 self._units['velocity'] = velocityUnit[0] if velocityUnit != [] else self._units['velocity'] 

190 

191 massUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.kg)] 

192 self._units['mass'] = massUnit[0] if massUnit != [] else self._units['mass'] 

193 

194 angleUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.rad)] 

195 self._units['angle'] = angleUnit[0] if angleUnit != [] else self._units['angle'] 

196 

197 objectUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.stars)] 

198 self._units['object'] = objectUnit[0] if objectUnit != [] else _u.stars 

199 

200 densityUnit = [this for this in UNIT_SYSTEM if this.is_equivalent(_u.stars/_u.m**3)] 

201 densityUnit2 = [this for this in UNIT_SYSTEM if this.is_equivalent(1/_u.m**3)] 

202 if densityUnit == [] and densityUnit2 != []: 

203 densityUnit = [self._units['object'] * densityUnit2[0]] 

204 elif densityUnit == [] and objectUnit != [] and lengthUnit != []: 

205 densityUnit = [self._units['object']/self._units['length']**3] 

206 elif densityUnit == [] and densityUnit2 == [] and objectUnit != []: 

207 densityLength = [this for this in self._units['density'].bases if this.is_equivalent(_u.m)][0] 

208 densityUnit = [self._units['object']/densityLength**3] 

209 self._units['density'] = densityUnit[0] if densityUnit != [] else self._units['density'] 

210 

211 self._UNIT_SYSTEM = list(self._units.values()) 

212 

213 

214 

215############################################################ 

216################### Helper Functions ####################### 

217############################################################ 

218 

219# Implemented from StackOverflow: https://stackoverflow.com/a/14314054 

220def moving_average(a, n=3, method=None) : 

221 '''Compute the moving average of an array of numbers using the nearest n elements.''' 

222 if method == 'nan': ret = _np.nancumsum(a) 

223 elif method == 'nn': 

224 bool = _np.isnan(a) 

225 inds = _np.arange(len(a))[bool] 

226 ret = a.copy() 

227 for i in inds: 

228 ret[i] = (ret[i-1 if i-1 > 0 else i+1] + ret[i+1 if i+1 < len(a) else i-1])/2.0 

229 ret = _np.cumsum(ret) 

230 else: ret = _np.cumsum(a) 

231 ret[n:] = ret[n:] - ret[:-n] 

232 return ret[n - 1:] / n 

233 

234# Implemented from StackOverflow: https://stackoverflow.com/a/33585850 

235def moving_median(arr, n=3): 

236 '''Compute the moving median of an array of numbers using the nearest n elements.''' 

237 idx = _np.arange(n) + _np.arange(len(arr)-n+1)[:,None] 

238 return _np.median(arr[idx], axis=1) 

239 

240def save_as_simulationarchive(filename, sims, deletefile=True): 

241 ''' 

242 Saves a list of REBOUND Simulations as a SimulationArchive. 

243 ''' 

244 for i,s in enumerate(sims): 

245 s.simulationarchive_snapshot(filename, deletefile=(deletefile if i == 0 else False)) 

246 

247# From REBOUND particle.py 

248def notNone(a): 

249 """ 

250 Returns True if array a contains at least one element that is not None. Returns False otherwise. 

251 """ 

252 return a.count(None) != len(a) 

253 

254def hasTrue(a): 

255 """ 

256 Returns True if array a contains at least one element that is True. Returns False otherwise. 

257 """ 

258 return a.count(True) > 0 

259 

260def numberOfElementsReturnedBySlice(start, stop, step): 

261 return (stop - start) // step 

262 

263def _integrate(sim, tmax): 

264 sim.integrate(tmax) 

265 return sim 

266 

267def integrate(sims, tmaxs, n_jobs=-1, verbose=0): 

268 sim_results = _joblib.Parallel(n_jobs=n_jobs, verbose=verbose, require='sharedmem')( 

269 _joblib.delayed(_integrate)(sim=sims[int(i)], tmax=tmaxs[int(i)]) for i in range(len(sims))) 

270 return sim_results 

271 

272def hist(arr, bins=10, normalize=False, density=False, wfac=1): 

273 # https://stackoverflow.com/questions/30551694/logarithmic-multi-sequenz-plot-with-equal-bar-widths/30555229#30555229 

274 # """Return pairwise geometric means of adjacent elements.""" 

275 geometric_means = lambda a: _np.sqrt(a[1:] * a[:-1]) 

276 

277 astart = _np.min(arr) 

278 aend = _np.max(arr) 

279 arange = _np.linspace(astart, aend, bins+1, endpoint=True) 

280 

281 y,b = _np.histogram(arr, bins=arange, density=density) 

282 x = geometric_means(b) 

283 w = wfac * _np.mean(x[1:] - x[:-1]) 

284 

285 if normalize: return x, y/_np.trapz(y,x), w 

286 else: return x,y,w 

287 

288def hist10(arr, bins=10, normalize=False, density=False, wfac=1): 

289 # https://stackoverflow.com/questions/30551694/logarithmic-multi-sequenz-plot-with-equal-bar-widths/30555229#30555229 

290 # """Return pairwise geometric means of adjacent elements.""" 

291 geometric_means = lambda a: _np.sqrt(a[1:] * a[:-1]) 

292 

293 astart = _np.log10(_np.min(arr)/2) 

294 aend = _np.log10(_np.max(arr)*2) 

295 arange = _np.logspace(astart, aend, bins+1, endpoint=True) 

296 

297 y,b = _np.histogram(arr, bins=arange, density=density) 

298 x = geometric_means(b) 

299 w = wfac * x*_np.mean((x[1:] - x[:-1])/x[:-1]) 

300 

301 if normalize: return x, y/_np.trapz(y,x), w 

302 else: return x,y,w 

303 

304# https://stackoverflow.com/a/13849249/71522 

305def unit_vector(vector): 

306 """ Returns the unit vector of the vector. """ 

307 return vector / _np.linalg.norm(vector) 

308 

309def angle_between(v1, v2): 

310 """ Returns the angle in radians between vectors 'v1' and 'v2':: 

311 

312 >>> angle_between((1, 0, 0), (0, 1, 0)) 

313 1.5707963267948966 

314 >>> angle_between((1, 0, 0), (1, 0, 0)) 

315 0.0 

316 >>> angle_between((1, 0, 0), (-1, 0, 0)) 

317 3.141592653589793 

318 """ 

319 v1_u = unit_vector(v1) 

320 v2_u = unit_vector(v2) 

321 return _np.arccos(_np.clip(_np.dot(v1_u, v2_u), -1.0, 1.0)) 

322 

323def reb_mod2pi(f): 

324 return _np.mod(twopi + _np.mod(f, twopi), twopi) 

325 

326def reb_M_to_E(e, M): 

327 E = 0 

328 if e < 1.0 : 

329 M = reb_mod2pi(M); # avoid numerical artefacts for negative numbers 

330 E = M if e < 0.8 else _np.pi 

331 F = E - e*_np.sin(E) - M 

332 for i in range(100): 

333 E = E - F/(1.-e*_np.cos(E)) 

334 F = E - e*_np.sin(E) - M 

335 if _np.all(_np.abs(F) < 1.0e-16) : break 

336 E = reb_mod2pi(E) 

337 return E 

338 else: 

339 E = M/_np.abs(M)*_np.log(2.*_np.abs(M)/e + 1.8) 

340 F = E - e*_np.sinh(E) + M 

341 for i in range(100): 

342 E = E - F/(1.0 - e*_np.cosh(E)) 

343 F = E - e*_np.sinh(E) + M 

344 if _np.all(_np.abs(F) < 1.0e-16): break 

345 return E 

346 

347def reb_E_to_f(e, E): 

348 if e > 1. :return reb_mod2pi(2.*_np.arctan(_np.sqrt((1.+e)/(e-1.))*_np.tanh(0.5*E))); 

349 else: return reb_mod2pi(2.*_np.arctan(_np.sqrt((1.+e)/(1.-e))*_np.tan(0.5*E))); 

350 

351 

352def reb_M_to_f(e, M): 

353 E = reb_M_to_E(e, M) 

354 return reb_E_to_f(e, E) 

355 

356 

357############################################################ 

358############### Properties and Elements #################### 

359############################################################ 

360 

361def calculate_angular_momentum(sim): 

362 L = _np.zeros((sim.N, 3)) 

363 L[0] = sim.angular_momentum() 

364 for i,p in enumerate(sim.particles[1:]): 

365 L[i+1,0] = p.m*(p.y*p.vz - p.z*p.vy) 

366 L[i+1,1] = p.m*(p.z*p.vx - p.x*p.vz) 

367 L[i+1,2] = p.m*(p.x*p.vy - p.y*p.vx) 

368 return L 

369 

370def vinf_and_b_to_e(mu, star_b, star_v): 

371 ''' 

372 Using the impact parameter to convert from the relative velocity at infinity between the two stars to the eccentricity of the flyby star. 

373 Equation (2) from Spurzem et al. (2009) https://ui.adsabs.harvard.edu/abs/2009ApJ...697..458S/abstract 

374 

375 Parameters 

376 ---------- 

377 mu : the total mass of the system (Sun, planets, and flyby star) times the gravitational constant G 

378 star_b : impact parameter of the flyby star 

379 star_v : the relative velocity at infinity between the central star and the flyby star (hyperbolic excess velocity) 

380 ''' 

381 

382 star_b = verify_unit(star_b, _u.au) 

383 star_v = verify_unit(star_v, _u.km/_u.s) 

384 

385 numerator = star_b * star_v**2. 

386 return _np.sqrt(1 + (numerator/mu)**2.) * _u.dimensionless_unscaled 

387 

388def vinf_and_q_to_e(mu, star_q, star_v): 

389 ''' 

390 Using the perihelion to convert from the relative velocity at infinity between the two stars to the eccentricity of the flyby star. 

391 

392 Parameters 

393 ---------- 

394 mu : the total mass of the system (Sun, planets, and flyby star) times the gravitational constant G 

395 star_q : perihelion of the flyby star 

396 star_v : the relative velocity at infinity between the central star and the flyby star (hyperbolic excess velocity) 

397 ''' 

398 

399 star_q = verify_unit(star_q, _u.au) 

400 star_vinf = verify_unit(star_v, _u.km/_u.s) 

401 return (1 + star_q * star_vinf * star_vinf / mu) * _u.dimensionless_unscaled 

402 

403def vinf_and_q_to_b(mu, star_q, star_v): 

404 ''' 

405 Using the perihelion to convert from the relative velocity at infinity between the two stars to the eccentricity of the flyby star. 

406 

407 Parameters 

408 ---------- 

409 mu : the total mass of the system (Sun, planets, and flyby star) times the gravitational constant G 

410 star_q : perihelion of the flyby star 

411 star_v : the relative velocity at infinity between the central star and the flyby star (hyperbolic excess velocity) 

412 ''' 

413 

414 mu = verify_unit(mu, (_u.au**3)/(_u.yr2pi**2)) 

415 star_q = verify_unit(star_q, _u.au) 

416 star_vinf = verify_unit(star_v, _u.km/_u.s) 

417 star_e = 1 + star_q * star_vinf * star_vinf / mu 

418 return verify_unit(star_q * _np.sqrt((star_e + 1.0)/(star_e - 1.0)), _u.au) 

419 

420def gravitational_mu(sim, star=None, star_mass=None): 

421 # Convert the units of the REBOUND Simulation into Astropy Units. 

422 units = rebound_units(sim) 

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

424 if star is not None and star_mass is not None: raise Exception('Cannot define both star and star_mass.') 

425 elif star is not None and star_mass is None: star_mass = verify_unit(star.mass, units['mass']) 

426 elif star is None and star_mass is not None: star_mass = verify_unit(star_mass, units['mass']) 

427 else: raise Exception('Either star or star_mass must be defined.') 

428 return G * (system_mass(sim) * units['mass'] + star_mass) 

429 

430def star_q(sim, star): 

431 ''' 

432 Using the impact parameter and the relative velocity at infinity between the two stars convert to the perhelion of the flyby star. 

433 

434 Parameters 

435 ---------- 

436 mu : the total mass of the system (Sun, planets, and flyby star) times the gravitational constant G 

437 star_q : perihelion of the flyby star 

438 star_v : the relative velocity at infinity between the central star and the flyby star (hyperbolic excess velocity) 

439 ''' 

440 

441 units = rebound_units(sim) 

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

443 mu = G * (system_mass(sim) * units['mass'] + star.m) 

444 

445 star_e = vinf_and_b_to_e(mu, star.b, star.v) 

446 return star.b * _np.sqrt((star_e - 1.0)/(star_e + 1.0)) 

447 

448def system_mass(sim): 

449 ''' 

450 The total bound mass of the system. 

451 ''' 

452 total_mass = 0 

453 for i,p in enumerate(sim.particles): 

454 if i == 0: total_mass += p.m 

455 elif p.a > 0: total_mass += p.m 

456 else: pass 

457 return total_mass 

458 

459def determine_eccentricity(sim, star_mass, star_b, star_v): 

460 '''Calculate the eccentricity of the flyby star. ''' 

461 # Convert the units of the REBOUND Simulation into Astropy Units. 

462 units = rebound_units(sim) 

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

464 star_mass = verify_unit(star_mass, units['mass']) 

465 star_b = verify_unit(star_b, units['length']) 

466 star_v = verify_unit(star_v, units['length']/units['time']) 

467 

468 mu = G * (system_mass(sim) * units['mass'] + star_mass) 

469 return vinf_and_b_to_e(mu=mu, star_b=star_b, star_v=star_v) 

470 

471def initial_conditions_from_stellar_params(sim, star, rmax): 

472 ''' 

473 Calculate the flyby star's initial conditions based on the provided Simulation and starting distance (rmax). 

474 ''' 

475 e = determine_eccentricity(sim, star.m, star.b, star.v) 

476 a = -star.b/_np.sqrt(e**2. - 1.) # Compute the semi-major axis of the flyby star 

477 l = -a*(e*e-1.) # Compute the semi-latus rectum of the hyperbolic orbit to get the true anomaly (-a because the semi-major axis is negative) 

478 if rmax == 0 * _u.au: f = 0 * _u.dimensionless_unscaled 

479 else: f = _np.arccos((l/rmax-1.)/e) # Compute the true anomaly 

480 

481 return {'m':star.m.value, 'a':a.value, 'e':e.value, 'inc':star.inc.value, 'omega':star.omega.value, 'Omega':star.Omega.value, 'f':-f.value}, l.value 

482 

483def hyperbolic_elements_from_stellar_params(sim, star, rmax): 

484 ''' 

485 Calculate the flyby star's hyperbolic orbital elements based on the provided Simulation and starting distance (rmax). 

486 ''' 

487 sim_units = rebound_units(sim) 

488 e = determine_eccentricity(sim, star.m, star.b, star.v) 

489 a = -star.b/_np.sqrt(e**2. - 1.) # Compute the semi-major axis of the flyby star 

490 l = a*(1.0 - e*e) # Compute the semi-latus rectum of the hyperbolic orbit to get the true anomaly 

491 if rmax == 0 * _u.au: f = 0 * _u.dimensionless_unscaled 

492 else: 

493 with _warnings.catch_warnings(): 

494 _warnings.simplefilter("error") # Make sure that the value for rmax is sufficient. 

495 try: f = _np.arccos((l/rmax-1.)/e) # Compute the true anomaly 

496 except RuntimeWarning as err: raise RuntimeError(f'{err}, rmax={rmax:1.6g} likely not larger than impact parameter, b={star.b:1.6g}.') from err 

497 

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

499 mu = G * (system_mass(sim) * sim_units['mass'] + star.m) 

500 

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

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

503 E = _np.arccosh((_np.cos(f)+e)/(1.+e*_np.cos(f))) # Compute the eccentric anomaly 

504 M = e * _np.sinh(E)-E # Compute the mean anomaly 

505 Tperi = M/_np.sqrt(mu/(-a*a*a)) 

506 

507 return {'m':star.m, 'a':a, 'e':e, 'inc':star.inc, 'omega':star.omega, 'Omega':star.Omega, 'f':-f, 'T':Tperi, 'l':l} 

508 

509def impulse_gradient(star): 

510 '''Calculate the impulse gradient for a flyby star.''' 

511 G = (1 * _u.au**3 / _u.solMass / _u.yr2pi**2) 

512 return ((2.0 * G * star.m) / (star.v * star.b**2.0)).to(_u.km/_u.s/_u.au) 

513 

514############################################################ 

515############# Stellar Environment Functions ################ 

516############################################################ 

517 

518 

519def maxwell_boltzmann_scale_from_dispersion(sigma): 

520 ''' 

521 Converts velocity dispersion (sigma) to scale factor for Maxwell-Boltzmann distributions. 

522 ''' 

523 return _np.sqrt((_np.pi*_np.square(sigma))/(3.0*_np.pi - 8.0)) 

524 

525def maxwell_boltzmann_scale_from_mean(mu): 

526 ''' 

527 Converts mean (mu) to scale factor for Maxwell-Boltzmann distributions. 

528 ''' 

529 return _np.sqrt(_np.pi/2.0) * (mu / 2.0) 

530 

531def maxwell_boltzmann_mean_from_dispersion(sigma): 

532 ''' 

533 Converts velocity dispersion (sigma) to mean (mu) for Maxwell-Boltzmann distributions. 

534 ''' 

535 scale = maxwell_boltzmann_scale_from_dispersion(sigma) 

536 return (2.0 * scale) * _np.sqrt(2.0/_np.pi) 

537 

538def maxwell_boltzmann_mode_from_dispersion(sigma): 

539 ''' 

540 Converts velocity dispersion (sigma) to mode (most common or typical value) for Maxwell-Boltzmann distributions. 

541 ''' 

542 scale = maxwell_boltzmann_scale_from_dispersion(sigma) 

543 return scale * _np.sqrt(2.0) 

544 

545def cross_section(M, R, v, unit_set=UnitSet()): 

546 ''' 

547 The cross-section with gravitational focusing. 

548  

549 Parameters 

550 ---------- 

551 M : the mass of flyby star (default units: solMass) 

552 R : the maximum interaction radius (default units: AU) 

553 v : the typical velocity from the distribution (default units: km/s) 

554 ''' 

555 

556 # Newton's gravitational constant in units of Msun, AU, and Years/2pi (G ~ 1). 

557 G = const.G.decompose(unit_set.UNIT_SYSTEM) 

558 sun_mass = 1 * _u.solMass # mass of the Sun in units of Msun 

559 

560 v = verify_unit(v, unit_set.units['velocity']) 

561 R = verify_unit(R, unit_set.units['length']) 

562 M = verify_unit(M, unit_set.units['mass']) 

563 

564 return (_np.pi * R**2) * (1 + 2*G*(sun_mass + M)/(R * v**2)) 

565 

566def encounter_rate(n, v, R, M=(1 * _u.solMass), unit_set=UnitSet()): 

567 ''' 

568 The expected flyby encounter rate within an stellar environment 

569  

570 Parameters 

571 ---------- 

572 n : stellar number density (default units: pc^{-3}) 

573 v : average velocity (default units: km/s) 

574 R : interaction radius (default units: AU) 

575 M : mass of a typical flyby star (default units: solMass) 

576 ''' 

577 n = verify_unit(n, unit_set.units['density']) 

578 v = verify_unit(v, unit_set.units['velocity']) 

579 R = verify_unit(R, unit_set.units['length']) 

580 M = verify_unit(M, unit_set.units['mass']) 

581 

582 return n * v * cross_section(M, R, v, unit_set) 

583 

584 

585############################################################ 

586################### Units Functions ######################## 

587############################################################ 

588 

589def rebound_units(sim): 

590 defrebunits = {'length': _u.au, 'mass': _u.solMass, 'time': _u.yr2pi} 

591 simunits = sim.units 

592 

593 for unit in simunits: 

594 if simunits[unit] == None: simunits[unit] = defrebunits[unit] 

595 else: simunits[unit] = _u.Unit(simunits[unit]) 

596 return simunits 

597 

598def verify_unit(value, unit): 

599 return value.to(unit) if isQuantity(value) else value * unit 

600 

601def isList(l): 

602 '''Determines if an object is a list or numpy array. Used for flyby parallelization.''' 

603 if isinstance(l,(list,_np.ndarray)): return True 

604 else: return False 

605 

606def isQuantity(var): 

607 '''Determines if an object is an Astropy Quantity. Used for Stellar Environment initializations.''' 

608 return isinstance(var, _u.quantity.Quantity) 

609 

610def isUnit(var): 

611 '''Determines if an object is an Astropy Quantity. Used for Stellar Environment initializations.''' 

612 return isinstance(var, (_u.core.IrreducibleUnit, _u.core.CompositeUnit, _u.Unit)) 

613 

614 

615############################################################ 

616###################### Exceptions ########################## 

617############################################################ 

618 

619class InvalidKeyException(Exception): 

620 def __init__(self): super().__init__('Invalid key type.') 

621 

622class InvalidUnitException(Exception): 

623 def __init__(self): super().__init__('Value is not a valid unit type.') 

624