alpha = {self.alpha}
" + f"m_min = {self.m_min}
" + f"m_max = {self.m_max}
") + + def imf(self, m): + '''Computes the IMF value for a given mass or array of masses 'm'. + Raises a ValueError if any value in 'm' is less than or equal to zero. + + Parameters + ---------- + m : float or array_like + Stellar mass or array of stellar masses [Msun]. + + Returns + ------- + float or ndarray + The IMF value for the given mass or masses. + ''' + m = np.asarray(m) + if np.any(m <= 0): + raise ValueError("Mass must be positive.") + valid = self._check_valid(m) + return m ** (-self.alpha) + + def rvs(self, size=1, rng=None): + """Draw random samples from the Kroupa1993 IMF. + + Uses analytical inverse transform sampling for efficiency. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random mass samples in solar masses. + """ + if rng is None: + rng = np.random.default_rng() + + # Analytical inverse transform sampling + # For power law IMF: m^(-alpha), the inverse CDF is: + # m = (u * (m_max^(1-alpha) - m_min^(1-alpha)) + m_min^(1-alpha))^(1/(1-alpha)) + normalization_constant = (1.0 - self.alpha) / ( + self.m_max**(1.0 - self.alpha) - self.m_min**(1.0 - self.alpha) + ) + u = rng.uniform(size=size) + masses = (u * (1.0 - self.alpha) / normalization_constant + + self.m_min**(1.0 - self.alpha))**(1.0 / (1.0 - self.alpha)) + + return masses + + class Kroupa2001(IMFBase): """Initial Mass Function based on Kroupa (2001), which is defined as a broken power-law: @@ -275,6 +412,39 @@ def imf(self, m): out[mask3] = const2 * (m[mask3] / self.m2break) ** (-self.alpha3) return out + def rvs(self, size=1, rng=None): + """Draw random samples from the Kroupa2001 IMF. + + Uses inverse transform sampling with discretized PDF for the + broken power-law distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random mass samples in solar masses. + """ + if rng is None: + rng = np.random.default_rng() + + # Create discretized PDF for inverse sampling + # Use more points near the breaks for better accuracy + n_points = 2000 + m_grid = np.linspace(self.m_min, self.m_max, n_points) + pdf_values = self.imf(m_grid) + + # Sample using inverse transform method + masses = inverse_sampler(m_grid, pdf_values, size=size, rng=rng) + + return masses + + class Chabrier2003(IMFBase): """Chabrier2003 Initial Mass Function (IMF), which is defined as a lognormal distribution for low-mass stars and a power-law distribution @@ -362,3 +532,35 @@ def imf(self, m): C = (1.0 / (self.m_break * sqrt_2pi_sigma)) * np.exp(-log_term_break) powerlaw = C * (m / self.m_break) ** (-self.alpha) return np.where(m < self.m_break, lognormal, powerlaw) + + def rvs(self, size=1, rng=None): + """Draw random samples from the Chabrier2003 IMF. + + Uses inverse transform sampling with discretized PDF for the + combined lognormal and power-law distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random mass samples in solar masses. + """ + if rng is None: + rng = np.random.default_rng() + + # Create discretized PDF for inverse sampling + # Use more points near the transition for better accuracy + n_points = 2000 + m_grid = np.linspace(self.m_min, self.m_max, n_points) + pdf_values = self.imf(m_grid) + + # Sample using inverse transform method + masses = inverse_sampler(m_grid, pdf_values, size=size, rng=rng) + + return masses diff --git a/posydon/popsyn/distributions.py b/posydon/popsyn/distributions.py index 36247403ab..8ac96ab764 100644 --- a/posydon/popsyn/distributions.py +++ b/posydon/popsyn/distributions.py @@ -5,6 +5,7 @@ ] import numpy as np from scipy.integrate import quad +from scipy.stats import truncnorm class FlatMassRatio: @@ -12,12 +13,12 @@ class FlatMassRatio: A uniform distribution for mass ratios q = m2/m1 within specified bounds. This distribution assigns equal probability to all mass ratios within the - given range [q_min, q_max]. + given range (q_min, q_max], exclusive bottom, inclusive top. Parameters ---------- q_min : float, optional - Minimum mass ratio (default: 0.05). Must be in (0, 1]. + Minimum mass ratio (default: 0.05). Must be in [0, 1). q_max : float, optional Maximum mass ratio (default: 1.0). Must be in (0, 1]. @@ -38,19 +39,19 @@ def __init__(self, q_min=0.05, q_max=1): Parameters ---------- q_min : float, optional - Minimum mass ratio (default: 0.05). Must be in (0, 1]. + Minimum mass ratio (default: 0.05). Must be in [0, 1). q_max : float, optional Maximum mass ratio (default: 1.0). Must be in (0, 1]. Raises ------ ValueError - If q_min or q_max are not in (0, 1], or if q_min >= q_max. + If q_min or q_max are not in valid range, or if q_min >= q_max. """ - if not (0 < q_min <= 1): - raise ValueError("q_min must be in (0, 1)") + if not (0 <= q_min < 1): + raise ValueError("q_min must be in [0, 1)") if not (0 < q_max <= 1): - raise ValueError("q_max must be in (0, 1)") + raise ValueError("q_max must be in (0, 1]") if q_min >= q_max: raise ValueError("q_min must be less than q_max") @@ -124,11 +125,31 @@ def pdf(self, q): Probability density at mass ratio q. """ q = np.asarray(q) - valid = (q >= self.q_min) & (q <= self.q_max) + valid = (q > self.q_min) & (q <= self.q_max) pdf_values = np.zeros_like(q, dtype=float) pdf_values[valid] = self.flat_mass_ratio(q[valid]) * self.norm return pdf_values + def rvs(self, size=1, rng=None): + """Draw random samples from the flat mass ratio distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + float or ndarray + Random samples from the distribution. + """ + if rng is None: + rng = np.random.default_rng() + + return rng.uniform(self.q_min, self.q_max, size=size) + class Sana12Period(): """Period distribution from Sana et al. (2012). @@ -310,6 +331,73 @@ def pdf(self, p, m1): * self.norm(m1[valid])) return pdf_values + def rvs(self, size=1, m1=None, rng=None): + """Draw random samples from the Sana12 period distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + m1 : float or array_like, optional + Primary mass(es). If array, must have length equal to size. + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random period samples in days. + + Raises + ------ + ValueError + If m1 is None or has incorrect size. + """ + if rng is None: + rng = np.random.default_rng() + + if m1 is None: + raise ValueError("m1 (primary mass) must be provided for Sana12Period sampling") + + m1 = np.atleast_1d(m1) + if m1.size == 1: + m1 = np.full(size, m1[0]) + elif m1.size != size: + raise ValueError(f"m1 must be a single value or have size={size}" + f"\n m1 = {m1}\n m1.size = {m1.size}") + + # Import here to avoid circular dependency + from posydon.utils.common_functions import rejection_sampler + + periods = np.zeros(size) + + # For low mass stars (m1 <= 15), use log-uniform distribution + low_mass_mask = m1 <= self.mbreak + n_low = np.sum(low_mass_mask) + if n_low > 0: + periods[low_mass_mask] = 10**rng.uniform( + np.log10(self.p_min), + np.log10(self.p_max), + size=n_low + ) + + # For high mass stars (m1 > 15), use rejection sampling + high_mass_mask = ~low_mass_mask + n_high = np.sum(high_mass_mask) + if n_high > 0: + # Create PDF function for rejection sampler + def pdf_high_mass(logp): + return self.sana12_period(logp, self.mbreak + 1) + + periods[high_mass_mask] = 10**rejection_sampler( + size=n_high, + x_lim=[np.log10(self.p_min), np.log10(self.p_max)], + pdf=pdf_high_mass, + rng=rng + ) + + return periods + class PowerLawPeriod(): '''Power law period distribution with slope pi and boundaries m_min, m_max. @@ -438,6 +526,38 @@ def pdf(self, p): return pdf_values + def rvs(self, size=1, rng=None): + """Draw random samples from the power law period distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random period samples in days. + """ + if rng is None: + rng = np.random.default_rng() + + # Import here to avoid circular dependency + from posydon.utils.common_functions import inverse_sampler + + # Create discretized PDF for inverse sampling + n_points = 1000 + logp_grid = np.linspace(np.log10(self.p_min), np.log10(self.p_max), n_points) + pdf_values = self.power_law_period(logp_grid) + + # Sample in log space + logp_samples = inverse_sampler(logp_grid, pdf_values, size=size, rng=rng) + + # Convert back to linear space + return 10**logp_samples + class LogUniform(): """Log-uniform distribution between specified minimum and maximum values. @@ -468,6 +588,28 @@ def __init__(self, min=5.0, max=1e5): self.norm = self._calculate_normalization() + def __repr__(self): + """Return string representation of the distribution. + + Returns + ------- + str + String representation showing the distribution parameters. + """ + return f"LogUniform(min={self.min}, max={self.max})" + + def _repr_html_(self): + """Return HTML representation for Jupyter notebooks. + + Returns + ------- + str + HTML string for rich display in notebooks. + """ + return (f"min = {self.min}
" + f"max = {self.max}
") + def _calculate_normalization(self): """ Calculate the normalization constant for the log-uniform distribution. @@ -497,6 +639,465 @@ def pdf(self, x): valid = (x > 0) & (x >= self.min) & (x <= self.max) pdf_values = np.zeros_like(x, dtype=float) + pdf_values[valid] = self.norm / x[valid] # PDF is constant in log space, so divide by x for linear space + + return pdf_values + + def rvs(self, size=1, rng=None): + """Draw random samples from the log-uniform distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random samples from the distribution. + """ + if rng is None: + rng = np.random.default_rng() + + # Sample uniformly in log space + log_samples = rng.uniform(np.log10(self.min), np.log10(self.max), size=size) + + # Convert back to linear space + return 10**log_samples + + +class ThermalEccentricity: + """Thermal eccentricity distribution for binary star systems. + + The thermal distribution follows pdf(e) = 2*e, which is the + distribution expected for binaries that have undergone significant + dynamical interactions or thermal relaxation. + + Parameters + ---------- + e_min : float, optional + Minimum eccentricity (default: 0.0). Must be in [0, 1). + e_max : float, optional + Maximum eccentricity (default: 1.0). Must be in (0, 1]. + + Raises + ------ + ValueError + If e_min or e_max are not in [0, 1], or if e_min >= e_max. + + Examples + -------- + >>> dist = ThermalEccentricity() + >>> e_samples = dist.rvs(size=1000) + """ + + def __init__(self, e_min=0.0, e_max=1.0): + """Initialize the thermal eccentricity distribution. + + Parameters + ---------- + e_min : float, optional + Minimum eccentricity (default: 0.0). Must be in [0, 1). + e_max : float, optional + Maximum eccentricity (default: 1.0). Must be in (0, 1]. + + Raises + ------ + ValueError + If e_min or e_max are not in [0, 1], or if e_min >= e_max. + """ + if not (0 <= e_min < 1): + raise ValueError("e_min must be in [0, 1)") + if not (0 < e_max <= 1): + raise ValueError("e_max must be in (0, 1]") + if e_min >= e_max: + raise ValueError("e_min must be less than e_max") + + self.e_min = e_min + self.e_max = e_max + self.norm = self._calculate_normalization() + + def __repr__(self): + """Return string representation of the distribution.""" + return f"ThermalEccentricity(e_min={self.e_min}, e_max={self.e_max})" + + def _repr_html_(self): + """Return HTML representation for Jupyter notebooks.""" + return (f"e_min = {self.e_min}
" + f"e_max = {self.e_max}
") + + def _calculate_normalization(self): + """Calculate the normalization constant for the thermal eccentricity + distribution. + + Returns + ------- + float + The normalization constant ensuring the PDF integrates to 1. + """ + # Integral of 2*e from e_min to e_max is e_max^2 - e_min^2 + integral = self.e_max**2 - self.e_min**2 + if integral == 0: # pragma: no cover + raise ValueError("Cannot normalize distribution: e_min == e_max") + return 1.0 / integral + + def thermal_eccentricity(self, e): + """Compute the thermal eccentricity distribution value. + + Parameters + ---------- + e : float or array_like + Eccentricity value(s). + + Returns + ------- + float or ndarray + Distribution value (2*e). + """ + return 2.0 * np.asarray(e) + + def pdf(self, e): + """Probability density function of the thermal eccentricity distribution. + + Parameters + ---------- + e : float or array_like + Eccentricity value(s). + + Returns + ------- + float or ndarray + Probability density at eccentricity e. + """ + e = np.asarray(e) + valid = (e >= self.e_min) & (e <= self.e_max) + pdf_values = np.zeros_like(e, dtype=float) + pdf_values[valid] = self.thermal_eccentricity(e[valid]) * self.norm + return pdf_values + + def rvs(self, size=1, rng=None): + """Draw random samples from the thermal eccentricity distribution. + + Uses the analytical inverse CDF: e = sqrt(u * (e_max^2 - e_min^2) + e_min^2) + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random eccentricity samples in [e_min, e_max]. + """ + if rng is None: + rng = np.random.default_rng() + + # Inverse CDF: e = sqrt(u * (e_max^2 - e_min^2) + e_min^2) + u = rng.uniform(size=size) + return np.sqrt(u * (self.e_max**2 - self.e_min**2) + self.e_min**2) + + +class UniformEccentricity: + """Uniform eccentricity distribution for binary star systems. + + A flat distribution over eccentricities between e_min and e_max. + + Parameters + ---------- + e_min : float, optional + Minimum eccentricity (default: 0.0). Must be in [0, 1). + e_max : float, optional + Maximum eccentricity (default: 1.0). Must be in (0, 1]. + + Raises + ------ + ValueError + If e_min or e_max are not in [0, 1], or if e_min >= e_max. + + Examples + -------- + >>> dist = UniformEccentricity() + >>> e_samples = dist.rvs(size=1000) + """ + + def __init__(self, e_min=0.0, e_max=1.0): + """Initialize the uniform eccentricity distribution. + + Parameters + ---------- + e_min : float, optional + Minimum eccentricity (default: 0.0). Must be in [0, 1). + e_max : float, optional + Maximum eccentricity (default: 1.0). Must be in (0, 1]. + + Raises + ------ + ValueError + If e_min or e_max are not in [0, 1], or if e_min >= e_max. + """ + if not (0 <= e_min < 1): + raise ValueError("e_min must be in [0, 1)") + if not (0 < e_max <= 1): + raise ValueError("e_max must be in (0, 1]") + if e_min >= e_max: + raise ValueError("e_min must be less than e_max") + + self.e_min = e_min + self.e_max = e_max + self.norm = 1.0 / (e_max - e_min) + + def __repr__(self): + """Return string representation of the distribution.""" + return f"UniformEccentricity(e_min={self.e_min}, e_max={self.e_max})" + + def _repr_html_(self): + """Return HTML representation for Jupyter notebooks.""" + return (f"e_min = {self.e_min}
" + f"e_max = {self.e_max}
") + + def pdf(self, e): + """Probability density function of the uniform eccentricity distribution. + + Parameters + ---------- + e : float or array_like + Eccentricity value(s). + + Returns + ------- + float or ndarray + Probability density at eccentricity e. + """ + e = np.asarray(e) + valid = (e >= self.e_min) & (e <= self.e_max) + pdf_values = np.zeros_like(e, dtype=float) pdf_values[valid] = self.norm + return pdf_values + + def rvs(self, size=1, rng=None): + """Draw random samples from the uniform eccentricity distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random eccentricity samples in [e_min, e_max]. + """ + if rng is None: + rng = np.random.default_rng() + + return rng.uniform(self.e_min, self.e_max, size=size) + + +class ZeroEccentricity: + """Zero eccentricity distribution for circular binary orbits. + + All samples are exactly zero (circular orbits). The PDF is a Dirac delta + at e=0. + + Examples + -------- + >>> dist = ZeroEccentricity() + >>> e_samples = dist.rvs(size=1000) # All zeros + """ + + def __init__(self): + """Initialize the zero eccentricity distribution.""" + pass + + def __repr__(self): + """Return string representation of the distribution.""" + return "ZeroEccentricity()" + + def _repr_html_(self): + """Return HTML representation for Jupyter notebooks.""" + return "e = 0 (circular orbits)
" + + def pdf(self, e): + """Probability density function of the zero eccentricity distribution. + + This is formally a Dirac delta at e=0. Returns 1.0 at e=0, 0.0 elsewhere. + + Parameters + ---------- + e : float or array_like + Eccentricity value(s). + + Returns + ------- + float or ndarray + 1.0 where e==0, 0.0 elsewhere. + """ + e = np.asarray(e) + return np.where(e == 0, 1.0, 0.0) + + def rvs(self, size=1, rng=None): + """Draw random samples from the zero eccentricity distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator (unused, included for API consistency). + + Returns + ------- + ndarray + Array of zeros with shape (size,). + """ + return np.zeros(size) + + +class LogNormalSeparation: + """Log-normal orbital separation distribution for binary star systems. + + Orbital separations are drawn from a log-normal distribution in log10 space, + truncated between specified minimum and maximum values. + + Parameters + ---------- + mean : float, optional + Mean of the log10 distribution (default: 0.85, corresponding to ~7.08 Rsun). + sigma : float, optional + Standard deviation of the log10 distribution (default: 0.37). + min : float, optional + Minimum orbital separation in solar radii (default: 5.0). + max : float, optional + Maximum orbital separation in solar radii (default: 1e5). + + Raises + ------ + ValueError + If min is not positive, max <= min, or sigma <= 0. + + Examples + -------- + >>> dist = LogNormalSeparation(mean=3.0, sigma=1.5, min=10, max=1e6) + >>> separations = dist.rvs(size=1000) + + Notes + ----- + Uses scipy.stats.truncnorm for efficient sampling from the truncated + normal distribution in log10 space. + """ + + def __init__(self, mean=0.85, sigma=0.37, min=5.0, max=1e5): + """Initialize the log-normal separation distribution. + + Parameters + ---------- + mean : float, optional + Mean of the log10 distribution (default: 0.85). + sigma : float, optional + Standard deviation of the log10 distribution (default: 0.37). + min : float, optional + Minimum orbital separation in solar radii (default: 5.0). + max : float, optional + Maximum orbital separation in solar radii (default: 1e5). + + Raises + ------ + ValueError + If min is not positive, max <= min, or sigma <= 0. + """ + if min <= 0: + raise ValueError("min must be positive") + if max <= min: + raise ValueError("max must be greater than min") + if sigma <= 0: + raise ValueError("sigma must be positive") + + self.mean = mean + self.sigma = sigma + self.min = min + self.max = max + + # Compute truncation bounds for scipy.stats.truncnorm + self.a_low = (np.log10(min) - mean) / sigma + self.a_high = (np.log10(max) - mean) / sigma + + def __repr__(self): + """Return string representation of the distribution.""" + return (f"LogNormalSeparation(mean={self.mean}, sigma={self.sigma}, " + f"min={self.min}, max={self.max})") + + def _repr_html_(self): + """Return HTML representation for Jupyter notebooks.""" + return (f"mean (log10) = {self.mean}
" + f"sigma (log10) = {self.sigma}
" + f"min = {self.min} R☉
" + f"max = {self.max} R☉
") + + def pdf(self, a): + """Probability density function of the log-normal separation distribution. + + Parameters + ---------- + a : float or array_like + Orbital separation(s) in solar radii. + + Returns + ------- + float or ndarray + Probability density at separation a. + """ + a = np.asarray(a) + valid = (a > 0) & (a >= self.min) & (a <= self.max) + pdf_values = np.zeros_like(a, dtype=float) + + log_a = np.log10(a[valid]) + # PDF of truncnorm in log space, transformed to linear space + # pdf(a) = pdf_log(log a) * |d(log a)/da| = pdf_log(log a) / (a * ln(10)) + pdf_values[valid] = truncnorm.pdf( + log_a, + self.a_low, self.a_high, + loc=self.mean, + scale=self.sigma + ) / (a[valid] * np.log(10)) return pdf_values + + def rvs(self, size=1, rng=None): + """Draw random samples from the log-normal separation distribution. + + Parameters + ---------- + size : int, optional + Number of samples to draw (default: 1). + rng : numpy.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ndarray + Random orbital separation samples in solar radii. + """ + if rng is None: + rng = np.random.default_rng() + + # Sample from truncated normal in log10 space + log_separations = truncnorm.rvs( + self.a_low, self.a_high, + loc=self.mean, + scale=self.sigma, + size=size, + random_state=rng + ) + + # Convert back to linear space + return 10**log_separations diff --git a/posydon/popsyn/independent_sample.py b/posydon/popsyn/independent_sample.py index 7bb0ee0f11..95d9c3c909 100644 --- a/posydon/popsyn/independent_sample.py +++ b/posydon/popsyn/independent_sample.py @@ -16,6 +16,7 @@ import numpy as np from scipy.stats import truncnorm +from posydon.popsyn import IMFs, distributions from posydon.popsyn.Moes_distributions import Moe_17_PsandQs from posydon.utils.common_functions import rejection_sampler @@ -130,46 +131,11 @@ def generate_orbital_periods(primary_masses, # Sana H., et al., 2012, Science, 337, 444 if orbital_period_scheme == 'Sana+12_period_extended': - # compute periods as if all M1 <= 15Msun (where pi = 0.0) - orbital_periods_M_lt_15 = 10**RNG.uniform( - low=np.log10(orbital_period_min), - high=np.log10(orbital_period_max), - size=number_of_binaries) - - # compute periods as if all M1 > 15Msun - def pdf(logp): - pi = 0.55 - beta = 1 - pi - A = np.log10(10**0.15)**(-pi)*(np.log10(10**0.15) - - np.log10(orbital_period_min)) - B = 1./beta*(np.log10(orbital_period_max)**beta - - np.log10(10**0.15)**beta) - C = 1./(A + B) - pdf = np.zeros(len(logp)) - - for j, logp_j in enumerate(logp): - # for logP<=0.15 days, the pdf is uniform - if np.log10(orbital_period_min) <= logp_j and logp_j < 0.15: - pdf[j] = C*0.15**(-pi) - - # original Sana H., et al., 2012, Science, 337, 444 - elif 0.15 <= logp_j and logp_j < np.log10(orbital_period_max): - pdf[j] = C*logp_j**(-pi) - - else: - pdf[j] = 0. - - return pdf - - orbital_periods_M_gt_15 = 10**(rejection_sampler( - size=number_of_binaries, - x_lim=[np.log10(orbital_period_min), np.log10(orbital_period_max)], - pdf=pdf, - rng=RNG)) - - orbital_periods = np.where(primary_masses <= 15.0, - orbital_periods_M_lt_15, - orbital_periods_M_gt_15) + period_dist = distributions.Sana12Period( + p_min=orbital_period_min, + p_max=orbital_period_max + ) + orbital_periods = period_dist.rvs(size=number_of_binaries, m1=primary_masses, rng=RNG) else: raise ValueError("You must provide an allowed orbital period scheme.") @@ -217,10 +183,11 @@ def generate_orbital_separations(number_of_binaries=1, "orbital separation scheme.") if orbital_separation_scheme == 'log_uniform': - orbital_separations = 10**RNG.uniform( - low=np.log10(orbital_separation_min), - high=np.log10(orbital_separation_max), - size=number_of_binaries) + sep_dist = distributions.LogUniform( + min=orbital_separation_min, + max=orbital_separation_max + ) + orbital_separations = sep_dist.rvs(size=number_of_binaries, rng=RNG) if orbital_separation_max < orbital_separation_min: raise ValueError("`orbital_separation_max` must be " @@ -234,20 +201,13 @@ def generate_orbital_separations(number_of_binaries=1, "`log_orbital_separation_mean`, " "`log_orbital_separation_sigma`.") - # Set limits for truncated normal distribution - a_low = (np.log10(orbital_separation_min) - - log_orbital_separation_mean) / log_orbital_separation_sigma - a_high = (np.log10(orbital_separation_max) - - log_orbital_separation_mean) / log_orbital_separation_sigma - - # generate orbital separations from a truncted normal distribution - log_orbital_separations = truncnorm.rvs( - a_low, a_high, - loc=log_orbital_separation_mean, - scale=log_orbital_separation_sigma, - size=number_of_binaries, - random_state=RNG) - orbital_separations = 10**log_orbital_separations + sep_dist = distributions.LogNormalSeparation( + mean=log_orbital_separation_mean, + sigma=log_orbital_separation_sigma, + min=orbital_separation_min, + max=orbital_separation_max + ) + orbital_separations = sep_dist.rvs(size=number_of_binaries, rng=RNG) else: pass @@ -285,11 +245,14 @@ def generate_eccentricities(number_of_binaries=1, raise ValueError("You must provide an allowed eccentricity scheme.") if eccentricity_scheme == 'thermal': - eccentricities = np.sqrt(RNG.uniform(size=number_of_binaries)) + ecc_dist = distributions.ThermalEccentricity() + eccentricities = ecc_dist.rvs(size=number_of_binaries, rng=RNG) elif eccentricity_scheme == 'uniform': - eccentricities = RNG.uniform(size=number_of_binaries) + ecc_dist = distributions.UniformEccentricity() + eccentricities = ecc_dist.rvs(size=number_of_binaries, rng=RNG) elif eccentricity_scheme == 'zero': - eccentricities = np.zeros(number_of_binaries) + ecc_dist = distributions.ZeroEccentricity() + eccentricities = ecc_dist.rvs(size=number_of_binaries, rng=RNG) else: # This should never be reached pass @@ -332,30 +295,18 @@ def generate_primary_masses(number_of_binaries=1, # Salpeter E. E., 1955, ApJ, 121, 161 if primary_mass_scheme == 'Salpeter': - alpha = 2.35 - normalization_constant = (1.0-alpha) / (primary_mass_max**(1-alpha) - - primary_mass_min**(1-alpha)) - random_variable = RNG.uniform(size=number_of_binaries) - primary_masses = (random_variable*(1.0-alpha)/normalization_constant - + primary_mass_min**(1.0-alpha))**(1.0/(1.0-alpha)) + imf = IMFs.Salpeter(alpha=2.35, m_min=primary_mass_min, m_max=primary_mass_max) + primary_masses = imf.rvs(size=number_of_binaries, rng=RNG) # Kroupa P., Tout C. A., Gilmore G., 1993, MNRAS, 262, 545 elif primary_mass_scheme == 'Kroupa1993': - alpha = 2.7 - normalization_constant = (1.0-alpha) / (primary_mass_max**(1-alpha) - - primary_mass_min**(1-alpha)) - random_variable = RNG.uniform(size=number_of_binaries) - primary_masses = (random_variable*(1.0-alpha)/normalization_constant - + primary_mass_min**(1.0-alpha))**(1.0/(1.0-alpha)) + imf = IMFs.Kroupa1993(alpha=2.7, m_min=primary_mass_min, m_max=primary_mass_max) + primary_masses = imf.rvs(size=number_of_binaries, rng=RNG) # Kroupa P., 2001, MNRAS, 322, 231 elif primary_mass_scheme == 'Kroupa2001': - alpha = 2.3 - normalization_constant = (1.0-alpha) / (primary_mass_max**(1-alpha) - - primary_mass_min**(1-alpha)) - random_variable = RNG.uniform(size=number_of_binaries) - primary_masses = (random_variable*(1.0-alpha)/normalization_constant - + primary_mass_min**(1.0-alpha))**(1.0/(1.0-alpha)) + imf = IMFs.Kroupa2001(m_min=primary_mass_min, m_max=primary_mass_max) + primary_masses = imf.rvs(size=number_of_binaries, rng=RNG) else: pass @@ -405,13 +356,18 @@ def generate_secondary_masses(primary_masses, # Generate secondary masses if secondary_mass_scheme == 'flat_mass_ratio': - mass_ratio_min = np.max([secondary_mass_min / primary_masses, - np.ones(len(primary_masses))*0.05], axis=0) - mass_ratio_max = np.min([secondary_mass_max / primary_masses, - np.ones(len(primary_masses))], axis=0) - secondary_masses = ( - (mass_ratio_max - mass_ratio_min) * RNG.uniform( - size=number_of_binaries) + mass_ratio_min) * primary_masses + # Calculate mass ratio bounds for each primary mass + q_min = np.maximum(secondary_mass_min / primary_masses, 0.05) + q_max = np.minimum(secondary_mass_max / primary_masses, 1.0) + + # Sample mass ratios using the distribution class + # For mass-dependent bounds, we need to sample individually + mass_ratios = np.zeros(number_of_binaries) + for i in range(number_of_binaries): + q_dist = distributions.FlatMassRatio(q_min=q_min[i], q_max=q_max[i]) + mass_ratios[i] = q_dist.rvs(size=1, rng=RNG)[0] + + secondary_masses = primary_masses * mass_ratios if secondary_mass_scheme == 'q=1': secondary_masses = primary_masses diff --git a/posydon/popsyn/norm_pop.py b/posydon/popsyn/norm_pop.py index 4f2946504e..fbaee17c32 100644 --- a/posydon/popsyn/norm_pop.py +++ b/posydon/popsyn/norm_pop.py @@ -100,23 +100,25 @@ def get_pdf_for_m1(m1): [kwargs['secondary_mass_max'] / m1, np.ones(len(m1))], axis=0) - q_dist = lambda q: np.where((q >= minimum) & (q <= maximum), + # Use FlatMassRatio distribution class + q_dist = lambda q: np.where((q > minimum) & (q <= maximum), 1/(maximum - minimum), 0) return q_dist q_pdf = lambda q, m1: get_pdf_for_m1(m1)(q) elif kwargs['secondary_mass_scheme'] == 'flat_mass_ratio': # flat mass ratio, where bounds are given - q_pdf = lambda q, m1=None: np.where( - (q > kwargs['q_min']) & (q <= kwargs['q_max']), - 1/(kwargs['q_max'] - kwargs['q_min']), - 0) + from posydon.popsyn.distributions import FlatMassRatio + q_dist = FlatMassRatio(q_min=kwargs['q_min'], q_max=kwargs['q_max']) + q_pdf = lambda q, m1=None: q_dist.pdf(q) else: # default to a flat distribution Pwarn("The secondary_mass_scheme is not defined use a flat mass ratio " "distribution in (0,1].", "UnsupportedModelWarning") - q_pdf = lambda q, m1=None: np.where((q > 0.0) & (q<=1.0), 1, 0) + from posydon.popsyn.distributions import FlatMassRatio + q_dist = FlatMassRatio(q_min=0.0, q_max=1.0) + q_pdf = lambda q, m1=None: q_dist.pdf(q) return q_pdf def get_binary_fraction_pdf(kwargs): diff --git a/posydon/unit_tests/popsyn/test_IMFs.py b/posydon/unit_tests/popsyn/test_IMFs.py index 40e29c4839..485268bab9 100644 --- a/posydon/unit_tests/popsyn/test_IMFs.py +++ b/posydon/unit_tests/popsyn/test_IMFs.py @@ -158,6 +158,141 @@ def test_imf_outside_range_warning(self, default_imf): # Verify the function still returns values (just with warning) assert len(result) == 1 + def test_rvs_default(self, default_imf): + """Test random sampling from the Salpeter IMF.""" + # Test single sample + rng = np.random.default_rng(42) + sample = default_imf.rvs(size=1, rng=rng) + assert len(sample) == 1 + assert default_imf.m_min <= sample[0] <= default_imf.m_max + + # Test multiple samples + rng = np.random.default_rng(42) + samples = default_imf.rvs(size=1000, rng=rng) + assert len(samples) == 1000 + assert np.all(samples >= default_imf.m_min) + assert np.all(samples <= default_imf.m_max) + + def test_rvs_without_rng(self, default_imf): + """Test random sampling without providing an RNG.""" + samples = default_imf.rvs(size=100) + assert len(samples) == 100 + assert np.all(samples >= default_imf.m_min) + assert np.all(samples <= default_imf.m_max) + +class TestKroupa1993IMF: + @pytest.fixture + def default_kroupa1993(self): + """Fixture for default Kroupa1993 instance.""" + return IMFs.Kroupa1993() + + @pytest.fixture + def custom_kroupa1993(self): + """Fixture for custom Kroupa1993 instance.""" + return IMFs.Kroupa1993(alpha=2.5, m_min=0.05, m_max=150.0) + + def test_initialization_default(self, default_kroupa1993): + """Test default initialization of Kroupa1993 IMF.""" + assert default_kroupa1993.alpha == 2.7 + assert default_kroupa1993.m_min == 0.01 + assert default_kroupa1993.m_max == 200.0 + integral, _ = quad(default_kroupa1993.imf, + default_kroupa1993.m_min, + default_kroupa1993.m_max) + assert np.isclose(integral*default_kroupa1993.norm, 1.0, rtol=1e-5) + + def test_initialization_custom(self, custom_kroupa1993): + """Test custom initialization of Kroupa1993 IMF.""" + assert custom_kroupa1993.alpha == 2.5 + assert custom_kroupa1993.m_min == 0.05 + assert custom_kroupa1993.m_max == 150.0 + integral, _ = quad(custom_kroupa1993.imf, + custom_kroupa1993.m_min, + custom_kroupa1993.m_max) + assert np.isclose(integral*custom_kroupa1993.norm, 1.0, rtol=1e-5) + + def test_repr(self, default_kroupa1993): + """Test string representation.""" + rep_str = default_kroupa1993.__repr__() + assert "Kroupa1993(" in rep_str + assert "alpha=2.7" in rep_str + assert "m_min=0.01" in rep_str + assert "m_max=200.0" in rep_str + + def test_repr_html(self, default_kroupa1993): + """Test HTML representation.""" + html_str = default_kroupa1993._repr_html_() + assert "