Coverage for src/airball/environments.py: 70%

179 statements  

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

1import rebound as _rebound 

2import numpy as _np 

3from scipy.stats import uniform as _uniform 

4from scipy.stats import maxwell as _maxwell 

5from scipy.stats import expon as _exponential 

6from scipy.optimize import fminbound as _fminbound 

7 

8from . import units as _u 

9from . import imf as _imf 

10from .imf import IMF as _IMF 

11from .stars import Star as _Star 

12from .stars import Stars as _Stars 

13from . import analytic as _analytic 

14from . import tools as _tools 

15from .tools import UnitSet as _UnitSet 

16 

17class StellarEnvironment: 

18 ''' 

19 This is the AIRBALL StellarEnvironment class. 

20 It encapsulates the relevant data for a static stellar environment. 

21 

22 # Example 

23 my_env = airball.StellarEnvironment(stellar_density=10, velocity_dispersion=20, lower_mass_limit=0.08, upper_mass_limit=8, name='My Environment') 

24 my_star = my_env.random_star() 

25 

26 If a `maximum_impact_parameter` is not given, AIRBALL attempts to estimate a reasonable one.  

27 The Maximum Impact Parameter is radius defining the outer limit of the sphere of influence around a stellar system. 

28 There are predefined subclasses for the LocalNeighborhood, a generic OpenCluster, a generic GlobularCluster, and the Milky Way center GalacticBulge and GalacticCore. 

29 ''' 

30 def __init__(self, stellar_density, velocity_dispersion, lower_mass_limit, upper_mass_limit, mass_function=None, maximum_impact_parameter=None, name=None, UNIT_SYSTEM=[], object_name=None, seed=None): 

31 

32 # Check to see if an stars object unit is defined in the given UNIT_SYSTEM and if the user defined a different name for the objects. 

33 self.units = _UnitSet(UNIT_SYSTEM) 

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

35 if objectUnit == [] and object_name is not None: self.units.object = _u.def_unit(object_name, _u.stars) 

36 elif objectUnit == [] and object_name is None: self.units.object = _u.stars 

37 else: self.units.object = objectUnit[0] 

38 

39 self.density = stellar_density 

40 self.velocity_dispersion = velocity_dispersion 

41 

42 self._upper_mass_limit = upper_mass_limit.to(self.units['mass']) if _tools.isQuantity(upper_mass_limit) else upper_mass_limit * self.units['mass'] 

43 self._lower_mass_limit = lower_mass_limit.to(self.units['mass']) if _tools.isQuantity(lower_mass_limit) else lower_mass_limit * self.units['mass'] 

44 self._IMF = _IMF(min_mass=self._lower_mass_limit, max_mass=self._upper_mass_limit, mass_function=mass_function, unit=self.units['mass']) 

45 self._median_mass = self.IMF.median_mass 

46 self.maximum_impact_parameter = maximum_impact_parameter 

47 

48 self.name = name if name is not None else 'Stellar Environment' 

49 self.seed = seed if seed is not None else None #_np.random.randint(0, int(2**32 - 1)) 

50 

51 def random_star(self, size=1, include_orientation=True, maximum_impact_parameter=None, **kwargs): 

52 ''' Alias for `random_stars`.''' 

53 return self.random_stars(size=size, include_orientation=include_orientation, maximum_impact_parameter=maximum_impact_parameter, **kwargs) 

54 

55 def random_stars(self, size=1, include_orientation=True, maximum_impact_parameter=None, **kwargs): 

56 ''' 

57 Computes a random star from a stellar environment. 

58 Returns: airball.Star() or airball.Stars() if size > 1. 

59 ''' 

60 if isinstance(size, tuple): size = tuple([int(i) for i in size]) 

61 else: size = int(size) 

62 

63 self.seed = kwargs.get('seed') 

64 if self.seed != None: _np.random.seed(self.seed) 

65 

66 v = _maxwell.rvs(scale=_tools.maxwell_boltzmann_scale_from_dispersion(self.velocity_dispersion), size=size) # Relative velocity of the star at infinity. 

67 

68 max_impact = maximum_impact_parameter if maximum_impact_parameter is not None else self.maximum_impact_parameter 

69 b = max_impact * _np.sqrt(_uniform.rvs(size=size)) # Impact parameter of the star. 

70 

71 m = self.IMF.random_mass(size=size) # Mass of the star. 

72 

73 zeros = _np.zeros(size) 

74 inc = 2.0*_np.pi * _uniform.rvs(size=size) - _np.pi if include_orientation else zeros 

75 ω = 2.0*_np.pi * _uniform.rvs(size=size) - _np.pi if include_orientation else zeros 

76 Ω = 2.0*_np.pi * _uniform.rvs(size=size) - _np.pi if include_orientation else zeros 

77 

78 if isinstance(size, tuple): return _Stars(m=m, b=b, v=v, inc=inc, omega=ω, Omega=Ω, UNIT_SYSTEM=self.UNIT_SYSTEM) 

79 elif size > 1: return _Stars(m=m, b=b, v=v, inc=inc, omega=ω, Omega=Ω, UNIT_SYSTEM=self.UNIT_SYSTEM)#, environment=self) 

80 else: return _Star(m, b[0], v[0], inc[0], ω[0], Ω[0], UNIT_SYSTEM=self.UNIT_SYSTEM) 

81 

82 def stats(self): 

83 ''' 

84 Prints a summary of the current stats of the Stellar Environment. 

85 ''' 

86 s = self.name 

87 s += "\n------------------------------------------\n" 

88 s += "{1} Density: {0:12.4g} \n".format(self.density, "Stellar" if self.object_unit.to_string() == _u.stars.to_string() else "Object") 

89 s += "Velocity Scale: {0:12.4g} \n".format(self.velocity_dispersion) 

90 s += "Mass Range: {0:6.4g} - {1:1.4g}\n".format(self.lower_mass_limit.value, self.upper_mass_limit) 

91 s += "Median Mass: {0:12.4g} \n".format(self.median_mass) 

92 s += "Max Impact Param: {0:12.4g} \n".format(self.maximum_impact_parameter) 

93 s += "Encounter Rate: {0:12.4g} \n".format(self.encounter_rate) 

94 s += "------------------------------------------" 

95 print(s) 

96 

97 @property 

98 def object_unit(self): 

99 return self.units['object'] 

100 

101 @property 

102 def object_name(self): 

103 return self.units['object'].to_string() 

104 

105 @object_name.setter 

106 def object_name(self, value): 

107 self.units.object = _u.def_unit(value, _u.stars) 

108 

109 @property 

110 def UNIT_SYSTEM(self): 

111 return self.units.UNIT_SYSTEM 

112 

113 @UNIT_SYSTEM.setter 

114 def UNIT_SYSTEM(self, UNIT_SYSTEM): 

115 self.units.UNIT_SYSTEM = UNIT_SYSTEM 

116 

117 @property 

118 def median_mass(self): 

119 ''' 

120 The median mass of the environment's IMF 

121 ''' 

122 return self.IMF.median_mass.to(self.units['mass']) 

123 

124 @property 

125 def maximum_impact_parameter(self): 

126 ''' 

127 The largest impact parameter to affect a stellar system in the environment. 

128 ''' 

129 return self._maximum_impact_parameter.to(self.units['length']) 

130 

131 @maximum_impact_parameter.setter 

132 def maximum_impact_parameter(self, value): 

133 if value is not None: 

134 self._maximum_impact_parameter = value.to(self.units['length']) if _tools.isQuantity(value) else value * self.units['length'] 

135 else: 

136 # TODO: Convert from fminbound to interpolation 

137 sim = _rebound.Simulation() 

138 sim.add(m=1.0) 

139 sim.add(m=5.2e-05, a=30.2, e=0.013) 

140 _f = lambda b: _np.log10(_np.abs(1e-16 - _np.abs(_analytic.relative_energy_change(sim, _Star(self.upper_mass_limit, b * self.units['length'], _np.sqrt(2.0)*_tools.maxwell_boltzmann_mean_from_dispersion(self.velocity_dispersion)))))) 

141 bs = _np.logspace(1, 6, 1000) 

142 b0 = bs[_np.argmin(_f(bs))] 

143 self._maximum_impact_parameter = _fminbound(_f, b0/5, 5*b0) * self.units['length'] 

144 

145 @property 

146 def density(self): 

147 ''' 

148 The number density of the environment. 

149 Default units: pc^{-3}. 

150 ''' 

151 return self._density.to(self.units['density']) 

152 

153 @density.setter 

154 def density(self, value): 

155 ''' 

156 The number density of the environment. 

157 Default units: pc^{-3}. 

158 ''' 

159 if _tools.isQuantity(value): 

160 if value.unit.is_equivalent(_u.stars/_u.m**3): self._density = value.to(self.units['density']) 

161 elif value.unit.is_equivalent(1/_u.m**3): self._density = (value * self.units['object']).to(self.units['density']) 

162 else: raise AssertionError('The given density units are not compatible.') 

163 else: self._density = value * self.units['density'] 

164 

165 @property 

166 def velocity_dispersion(self): 

167 ''' 

168 Return the velocity dispersion of the environment. 

169 Default units: km/s. 

170 ''' 

171 return self._velocity.to(self.units['velocity']) 

172 

173 @velocity_dispersion.setter 

174 def velocity_dispersion(self, value): 

175 ''' 

176 The velocity dispersion of the environment. 

177 Default units: km/s. 

178 ''' 

179 self._velocity = value.to(self.units['velocity']) if _tools.isQuantity(value) else value * self.units['velocity'] 

180 

181 @property 

182 def velocity_mean(self): 

183 ''' 

184 Return the velocity dispersion of the environment. 

185 Default units: km/s. 

186 ''' 

187 return _tools.maxwell_boltzmann_mean_from_dispersion(self.velocity_dispersion).to(self.units['velocity']) 

188 

189 @property 

190 def velocity_mode(self): 

191 ''' 

192 Return the velocity dispersion of the environment. 

193 Default units: km/s. 

194 ''' 

195 return _tools.maxwell_boltzmann_mode_from_dispersion(self.velocity_dispersion).to(self.units['velocity']) 

196 

197 @property 

198 def velocity_rms(self): 

199 ''' 

200 Return the velocity dispersion of the environment. 

201 Default units: km/s. 

202 ''' 

203 

204 v = _maxwell.rvs(scale=_tools.maxwell_boltzmann_scale_from_dispersion(self.velocity_dispersion), size=int(1e6)) 

205 return _tools.verify_unit(_np.sqrt(_np.mean(v**2)), self.units['velocity']) 

206 

207 @property 

208 def lower_mass_limit(self): 

209 ''' 

210 Return the lower mass limit of the IMF of the environment. 

211 Default units: solMass 

212 ''' 

213 return self.IMF.min_mass.to(self.units['mass']) 

214 

215 @lower_mass_limit.setter 

216 def lower_mass_limit(self, value): 

217 ''' 

218 The lower mass limit of the IMF of the environment. 

219 Default units: solMass 

220 ''' 

221 self.IMF.min_mass = value 

222 

223 @property 

224 def upper_mass_limit(self): 

225 ''' 

226 Return the lower mass limit of the IMF of the environment. 

227 Default units: solMass 

228 ''' 

229 return self.IMF.max_mass.to(self.units['mass']) 

230 

231 @upper_mass_limit.setter 

232 def upper_mass_limit(self, value): 

233 ''' 

234 The lower mass limit of the IMF of the environment. 

235 Default units: solMass 

236 ''' 

237 self.IMF.max_mass = value 

238 

239 @property 

240 def IMF(self): 

241 ''' 

242 Return the IMF of the environment. 

243 ''' 

244 return self._IMF 

245 

246 @IMF.setter 

247 def IMF(self, value): 

248 ''' 

249 The IMF of the environment. 

250 ''' 

251 if isinstance(value, _IMF): self._IMF = _IMF(value.min_mass, value.max_mass, value.imf, value.unit, value.number_samples, value.seed) 

252 else: raise AssertionError('Initial Mass Function (IMF) must be an airball.IMF object.') 

253 

254 @property 

255 def encounter_rate(self): 

256 ''' 

257 Compute the expected flyby encounter rate Γ = ⟨nσv⟩ for the stellar environment in units of flybys per year. 

258 The inverse of the encouter rate will give the average number of years until a flyby. 

259 

260 n : stellar number density. Default units: pc^{-3} 

261 σ : interaction cross section. Default units: AU^2 

262 v : velocity dispersion. Default units: km/s 

263 

264 The interaction cross section σ = πb^2 considers gravitational focussing b = q√[1 + (2GM)/(q v∞^2)] and considers 

265 - the median mass of the environment 

266 - the maximum impact parameter 

267 - the relative velocity at infinity derived from the velocity dispersion 

268 ''' 

269 return _tools.encounter_rate(self._density, _tools.maxwell_boltzmann_mean_from_dispersion(self.velocity_dispersion), self._maximum_impact_parameter, self.median_mass, unit_set=self.units).to(self.units['object']/self.units['time']) 

270 

271 def cumulative_encounter_times(self, size=None): 

272 ''' 

273 Returns the cumulative time from t=0 for when to the expect the next flyby encounters. 

274 Assumes a Poisson Process and uses an Exponential distribution with the encounter rate. 

275 ''' 

276 if isinstance(size, tuple): 

277 size = tuple([int(i) for i in size]) 

278 result = _np.cumsum(_exponential.rvs(scale=1/self.encounter_rate, size=size), axis=1) << self.units['time'] 

279 result -= result[:, 0][:, None] 

280 return result 

281 else: 

282 size = int(size) 

283 result = _np.cumsum(_exponential.rvs(scale=1/self.encounter_rate, size=size)) << self.units['time'] 

284 result -= result[0] 

285 return result 

286 

287 def encounter_times(self, size=None): 

288 ''' 

289 Returns the time between encounters for when to the expect the next flyby encounters. 

290 Assumes a Poisson Process and uses an Exponential distribution with the encounter rate. 

291 ''' 

292 if isinstance(size, tuple): 

293 size = tuple([int(i) for i in size]) 

294 return _exponential.rvs(scale=1/self.encounter_rate, size=size) << self.units['time'] 

295 else: 

296 size = int(size) 

297 return _exponential.rvs(scale=1/self.encounter_rate, size=size) << self.units['time'] 

298 

299 def time_to_next_encounter(self): 

300 ''' 

301 Draw a time to the next expected flyby encounter. 

302 Assumes a Poisson Process and uses an Exponential distribution with the encounter rate. 

303 ''' 

304 return _exponential.rvs(scale=1/self.encounter_rate) * self.units['time'] 

305 

306 

307 

308class LocalNeighborhood(StellarEnvironment): 

309 ''' 

310 This is a AIRBALL StellarEnvironment subclass for the Local Neighborhood. 

311 It encapsulates the relevant data for a static stellar environment representing the local neighborhood of the solar system. 

312 

313 The stellar density is 0.14 pc^-3 defined by Bovy (2017). 

314 The velocity dispersion is 20 km/s, defined by Binnery & Tremaine (2008) v_rms ~50 km/s and Bailer-Jones et al. (2018) so that 90% of stars have v < 100 km/s with an encounter rate of ~20 stars/Myr within 1 pc. A more accurate representation of the velocity distribution in the solar neighborhood is a triaxial Gaussian distribution. 

315 The mass limits is defined to between 0.08-8 solar masses using Chabrier (2003) for single stars when m < 1 and a power-law model from Bovy (2017) for stars m ≥ 1 to account for depleted stars due to stellar evolution. 

316 

317 # Example 

318 my_local = airball.LocalNeighborhood() 

319 my_10stars = my_local.random_star(size=10) 

320 # returns a Stars object with the masses, impact parameters, velocities, and orientation of the 10 Star objects in a heliocentric model. 

321 ''' 

322 short_name = 'Local' 

323 def local_mass_function(x): 

324 ''' 

325 This defined using Chabrier (2003) for single stars when m < 1 and a power-law model from Bovy (2017) for stars m ≥ 1 to account for depleted stars due to stellar evolution. 

326 ''' 

327 return _imf.chabrier_2003_single(1, 0.0567) * (x)**-4.7 if x > 1 else _imf.chabrier_2003_single(x, 0.0567) 

328 

329 def __init__(self, stellar_density = 0.14 * _u.stars/_u.pc**3, velocity_dispersion = 20.8 * _u.km/_u.s, lower_mass_limit=0.08 * _u.solMass, upper_mass_limit = 8 * _u.solMass, maximum_impact_parameter=10000 * _u.au, UNIT_SYSTEM=[], mass_function=local_mass_function, object_name=None): 

330 super().__init__(stellar_density=stellar_density, velocity_dispersion=velocity_dispersion, lower_mass_limit=lower_mass_limit, upper_mass_limit=upper_mass_limit, mass_function=mass_function, maximum_impact_parameter=maximum_impact_parameter, UNIT_SYSTEM=UNIT_SYSTEM, name = 'Local Neighborhood', object_name=object_name) 

331 

332class OpenCluster(StellarEnvironment): 

333 ''' 

334 This is a AIRBALL StellarEnvironment subclass for a generic Open Cluster. 

335 It encapsulates the relevant data for a static stellar environment representing a generic open cluster. 

336 

337 The stellar density is 100 pc^-3 informed by Adams (2010). 

338 The velocity scale is 1 km/s informed by Adams (2010) and Malmberg et al. (2011). 

339 The mass limit is defined to between 0.08-100 solar masses using Chabrier (2003) for single stars when m < 1 and Salpeter (1955) for stars m ≥ 1. 

340 

341 # Example 

342 my_open = airball.OpenCluster() 

343 my_10stars = my_open.random_star(size=10) 

344 # returns a Stars object with the masses, impact parameters, velocities, and orientation of the 10 Star objects in a heliocentric model. 

345 ''' 

346 short_name = 'Open' 

347 

348 def __init__(self, stellar_density = 100 * _u.stars * _u.pc**-3, velocity_dispersion = 1 * _u.km/_u.s, lower_mass_limit=0.08 * _u.solMass, upper_mass_limit = 100 * _u.solMass, maximum_impact_parameter=1000 * _u.au, UNIT_SYSTEM=[], object_name=None): 

349 super().__init__(stellar_density=stellar_density, velocity_dispersion=velocity_dispersion, lower_mass_limit=lower_mass_limit, upper_mass_limit=upper_mass_limit, mass_function=None, maximum_impact_parameter=maximum_impact_parameter, UNIT_SYSTEM=UNIT_SYSTEM, name = 'Open Cluster', object_name=object_name) 

350 

351class GlobularCluster(StellarEnvironment): 

352 short_name = 'Globular' 

353 

354 def __init__(self, stellar_density = 1000 * _u.stars * _u.pc**-3, velocity_dispersion = 10 * _u.km/_u.s, lower_mass_limit=0.08 * _u.solMass, upper_mass_limit = 1 * _u.solMass, maximum_impact_parameter=5000 * _u.au, UNIT_SYSTEM=[], object_name=None): 

355 super().__init__(stellar_density=stellar_density, velocity_dispersion=velocity_dispersion, lower_mass_limit=lower_mass_limit, upper_mass_limit=upper_mass_limit, mass_function=None, maximum_impact_parameter=maximum_impact_parameter, UNIT_SYSTEM=UNIT_SYSTEM, name = 'Globular Cluster', object_name=object_name) 

356 

357class GalacticBulge(StellarEnvironment): 

358 short_name = 'Bulge' 

359 

360 def __init__(self, stellar_density = 50 * _u.stars * _u.pc**-3, velocity_dispersion = 120 * _u.km/_u.s, lower_mass_limit=0.08 * _u.solMass, upper_mass_limit = 10 * _u.solMass, maximum_impact_parameter=50000 * _u.au, UNIT_SYSTEM=[], object_name=None): 

361 super().__init__(stellar_density=stellar_density, velocity_dispersion=velocity_dispersion, lower_mass_limit=lower_mass_limit, upper_mass_limit=upper_mass_limit, mass_function=None, maximum_impact_parameter=maximum_impact_parameter, UNIT_SYSTEM=UNIT_SYSTEM, name = 'Milky Way Bulge', object_name=object_name) 

362 

363class GalacticCore(StellarEnvironment): 

364 short_name = 'Core' 

365 

366 def __init__(self, stellar_density = 10000 * _u.stars * _u.pc**-3, velocity_dispersion = 170 * _u.km/_u.s, lower_mass_limit=0.08 * _u.solMass, upper_mass_limit = 10 * _u.solMass, maximum_impact_parameter=50000 * _u.au, UNIT_SYSTEM=[_u.yr], object_name=None): 

367 super().__init__(stellar_density=stellar_density, velocity_dispersion=velocity_dispersion, lower_mass_limit=lower_mass_limit, upper_mass_limit=upper_mass_limit, mass_function=None, maximum_impact_parameter=maximum_impact_parameter, UNIT_SYSTEM=UNIT_SYSTEM, name = 'Milky Way Core', object_name=object_name)