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

1 import numpy as _np

2 from scipy.integrate import quad as _quad

3 from . import units as u

4 from . import tools

5

6

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

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

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