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
« 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
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.
12 Parameters:
13 - x: Mass value (float).
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))
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.
25 Parameters:
26 - x: Mass value (float).
28 Returns:
29 - Probability density at the given mass value.
30 """
31 return A * x ** -2.3
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.
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 """
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
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
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
62 # Recalculate the IMF properties based on the updated parameters
63 self._recalculate()
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]
75 # Create a normalized PDF
77 npdf = lambda x: pdf(x) / normalization_factor
78 self._npdf = _np.vectorize(npdf)
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)
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)
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)
91 def random_mass(self, size=1, **kwargs):
92 """
93 Generates random mass values from the IMF.
95 Parameters:
96 - size: Number of random mass values to generate (default: 1).
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]
108 def masses(self, number_samples, endpoint=True):
109 """
110 Generates an array of mass values logarithmically spanning the IMF range.
112 Parameters:
113 - number_samples: Number of mass values to generate.
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)
120 @property
121 def min_mass(self):
122 """
123 Minimum mass value of the IMF range.
124 """
125 return self._min_mass
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.
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()
141 @property
142 def max_mass(self):
143 """
144 Maximum mass value of the IMF range.
145 """
146 return self._max_mass
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.
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()
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
167 @property
168 def seed(self):
169 """
170 Seed for the random number generator.
171 """
172 return self._seed
174 @seed.setter
175 def seed(self, value):
176 """
177 Setter for the seed for the random number generator
179 Parameters:
180 - value: New seed for the random number generator (int, or None to turn off)
181 """
182 self._seed = value
184 @property
185 def number_samples(self):
186 """
187 Seed for the random number generator.
188 """
189 return self._number_samples
191 @number_samples.setter
192 def number_samples(self, value):
193 """
194 Setter for the seed for the random number generator
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()
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)
209 @property
210 def CDF(self):
211 """
212 Cumulative distribution function (CDF) of the IMF.
213 """
214 return self.cdf
216 @property
217 def pdf(self):
218 """
219 Probability density function (PDF) of the IMF.
220 """
221 return self._npdf
223 @property
224 def PDF(self):
225 """
226 Probability density function (PDF) of the IMF.
227 """
228 return self.pdf
230 @property
231 def imf(self):
232 """
233 Probability density function (PDF) of the IMF.
234 """
235 return self._imf
237 @property
238 def IMF(self):
239 """
240 Probability density function (PDF) of the IMF.
241 """
242 return self.imf