Coverage for src/airball/imf.py: 81%

88 statements  

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

1import numpy as _np 

2from scipy.integrate import quad as _quad 

3from . import units as u 

4from . import tools 

5 

6 

7def chabrier_2003_single(x, A=0.158): 

8 """ 

9 Chabrier 2003 IMF for single stars. 

10 This function calculates the probability density for a given mass value (x) based on the Chabrier 2003 IMF equation. 

11 

12 Parameters: 

13 - x: Mass value (float). 

14 

15 Returns: 

16 - Probability density at the given mass value. 

17 """ 

18 return (A / x) * _np.exp(-((_np.log10(x) - _np.log10(0.079)) ** 2) / (2 * 0.69 ** 2)) 

19 

20def salpeter_1955(x, A): 

21 """ 

22 Salpeter 1955 IMF for single stars. 

23 This function calculates the probability density for a given mass value (x) based on the Salpeter 1955 IMF equation. 

24 

25 Parameters: 

26 - x: Mass value (float). 

27 

28 Returns: 

29 - Probability density at the given mass value. 

30 """ 

31 return A * x ** -2.3 

32 

33 

34class IMF(): 

35 """ 

36 Class representing an Initial Mass Function (IMF). 

37 It generates random masses based on a given mass function and provides various properties and methods for manipulating and analyzing the IMF. 

38 

39 Parameters: 

40 - min_mass: Minimum mass value of the IMF range (float or astropy.units.Quantity). 

41 - max_mass: Maximum mass value of the IMF range (float or astropy.units.Quantity). 

42 - mass_function: Mass function to use for the IMF (optional). 

43 - unit: Unit of mass (optional, defaults to solar masses). 

44 - number_samples: Number of samples to use for interpolating the CDF (default: 100). 

45 - seed: Value to seed the random number generator with (optional, int, or None to turn off). 

46 """ 

47 

48 def __init__(self, min_mass, max_mass, mass_function=None, unit=u.solMass, number_samples=100, seed=None): 

49 self._number_samples = int(number_samples) 

50 self._seed = seed 

51 self.unit = unit if tools.isUnit(unit) else u.solMass 

52 

53 # Convert min_mass and max_mass to specified unit if they are Quantity objects, otherwise assume they are already in the correct unit 

54 self._min_mass = min_mass.to(self.unit) if tools.isQuantity(min_mass) else min_mass * self.unit 

55 if self._min_mass.value <= 0: raise Exception('Cannot have minimum mass value be less than or equal to 0.') 

56 self._max_mass = max_mass.to(self.unit) if tools.isQuantity(max_mass) else max_mass * self.unit 

57 

58 # Determine the probability distribution function (PDF) based on the given mass function or default to a piecewise Chabrier 2003 and Salpeter 1955 

59 if mass_function is None: mass_function = lambda x: _np.where(x < 1, chabrier_2003_single(x), salpeter_1955(x, chabrier_2003_single(1))) 

60 self._imf = mass_function 

61 

62 # Recalculate the IMF properties based on the updated parameters 

63 self._recalculate() 

64 

65 def _recalculate(self): 

66 """ 

67 Recalculates the IMF properties based on the current parameters. 

68 This function updates the normalization factor, normalized PDF, cumulative distribution function (CDF), 

69 mass values, and IMF values based on the current min_mass, max_mass, and PDF. 

70 """ 

71 # Calculate the normalization factor for the PDF 

72 pdf = _np.vectorize(lambda x: _np.where(x < self.min_mass.value, 0, _np.where(x > self.max_mass.value, 0, self._imf(x)))) 

73 normalization_factor = _quad(pdf, self._min_mass.value, self._max_mass.value)[0] 

74 

75 # Create a normalized PDF 

76 

77 npdf = lambda x: pdf(x) / normalization_factor 

78 self._npdf = _np.vectorize(npdf) 

79 

80 # Create a cumulative distribution function (CDF) 

81 cdf = lambda x: _np.where(x < self.min_mass.value, 0, _np.where(x > self.max_mass.value, 1, _quad(self._npdf, self._min_mass.value, x)[0])) 

82 self._cdf = _np.vectorize(cdf) 

83 

84 # Generate logarithmically spaced mass values between min_mass and max_mass 

85 self._masses = _np.logspace(_np.log10(self._min_mass.value), _np.log10(self._max_mass.value), self.number_samples) 

86 

87 # Calculate the CDF and IMF values for the mass values 

88 self._CDF = self._cdf(self._masses) 

89 self._norm_imf = self._npdf(self._masses) 

90 

91 def random_mass(self, size=1, **kwargs): 

92 """ 

93 Generates random mass values from the IMF. 

94 

95 Parameters: 

96 - size: Number of random mass values to generate (default: 1). 

97 

98 Returns: 

99 - Randomly generated mass value(s) from the IMF. 

100 """ 

101 if 'seed' in kwargs: self.seed = kwargs['seed'] 

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

103 rand_masses = _np.interp(_np.random.uniform(size=size), self._CDF, self._masses) * self.unit 

104 if isinstance(size, tuple): return rand_masses 

105 elif size > 1: return rand_masses 

106 else: return rand_masses[0] 

107 

108 def masses(self, number_samples, endpoint=True): 

109 """ 

110 Generates an array of mass values logarithmically spanning the IMF range. 

111 

112 Parameters: 

113 - number_samples: Number of mass values to generate. 

114 

115 Returns: 

116 - numpy array of mass values logarithmically spanning the IMF range. 

117 """ 

118 return _np.logspace(_np.log10(self._min_mass.value), _np.log10(self._max_mass.value), number_samples, endpoint=endpoint) 

119 

120 @property 

121 def min_mass(self): 

122 """ 

123 Minimum mass value of the IMF range. 

124 """ 

125 return self._min_mass 

126 

127 @min_mass.setter 

128 def min_mass(self, value): 

129 """ 

130 Setter for the minimum mass value of the IMF range. 

131 Recalculates the IMF properties when the min_mass value is updated. 

132 

133 Parameters: 

134 - value: New minimum mass value (float or astropy.units.Quantity). 

135 """ 

136 value = value.to(self.unit) if tools.isQuantity(value) else value * self.unit 

137 if value.value <= 0: raise Exception('Cannot have minimum mass value be less than or equal to 0.') 

138 self._min_mass = value 

139 self._recalculate() 

140 

141 @property 

142 def max_mass(self): 

143 """ 

144 Maximum mass value of the IMF range. 

145 """ 

146 return self._max_mass 

147 

148 @max_mass.setter 

149 def max_mass(self, value): 

150 """ 

151 Setter for the maximum mass value of the IMF range. 

152 Recalculates the IMF properties when the max_mass value is updated. 

153 

154 Parameters: 

155 - value: New maximum mass value (float or astropy.units.Quantity). 

156 """ 

157 self._max_mass = value.to(self.unit) if tools.isQuantity(value) else value * self.unit 

158 self._recalculate() 

159 

160 @property 

161 def median_mass(self): 

162 """ 

163 Median mass value of the IMF. 

164 """ 

165 return _np.interp(0.5, self._CDF, self._masses) * self.unit 

166 

167 @property 

168 def seed(self): 

169 """ 

170 Seed for the random number generator. 

171 """ 

172 return self._seed 

173 

174 @seed.setter 

175 def seed(self, value): 

176 """ 

177 Setter for the seed for the random number generator  

178 

179 Parameters: 

180 - value: New seed for the random number generator (int, or None to turn off) 

181 """ 

182 self._seed = value 

183 

184 @property 

185 def number_samples(self): 

186 """ 

187 Seed for the random number generator. 

188 """ 

189 return self._number_samples 

190 

191 @number_samples.setter 

192 def number_samples(self, value): 

193 """ 

194 Setter for the seed for the random number generator  

195 

196 Parameters: 

197 - value: New seed for the random number generator (int, or None to turn off) 

198 """ 

199 self._number_samples = int(value) 

200 self._recalculate() 

201 

202 @property 

203 def cdf(self): 

204 """ 

205 Cumulative distribution function (CDF) of the IMF. 

206 """ 

207 return lambda x: _np.interp(x, self._masses, self._CDF) 

208 

209 @property 

210 def CDF(self): 

211 """ 

212 Cumulative distribution function (CDF) of the IMF. 

213 """ 

214 return self.cdf 

215 

216 @property 

217 def pdf(self): 

218 """ 

219 Probability density function (PDF) of the IMF. 

220 """ 

221 return self._npdf 

222 

223 @property 

224 def PDF(self): 

225 """ 

226 Probability density function (PDF) of the IMF. 

227 """ 

228 return self.pdf 

229 

230 @property 

231 def imf(self): 

232 """ 

233 Probability density function (PDF) of the IMF. 

234 """ 

235 return self._imf 

236 

237 @property 

238 def IMF(self): 

239 """ 

240 Probability density function (PDF) of the IMF. 

241 """ 

242 return self.imf