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

1 import numpy as _np

2 import joblib as _joblib

3 import warnings as _warnings

4 import astropy.constants as const

5 from . import units as _u

6 from scipy.stats import uniform

7 from scipy.stats import maxwell

8 from scipy.stats import expon

9

10 twopi = 2.*_np.pi

11

12 class 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

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

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

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

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

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

260 def numberOfElementsReturnedBySlice(start, stop, step):

261 return (stop - start) // step

262

263 def _integrate(sim, tmax):

264 sim.integrate(tmax)

265 return sim

266

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

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

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

305 def unit_vector(vector):

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

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

308

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

323 def reb_mod2pi(f):

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

325

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

598 def verify_unit(value, unit):

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

600

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

606 def isQuantity(var):

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

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

609

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

619 class InvalidKeyException(Exception):

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

621

622 class InvalidUnitException(Exception):

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

624