From 0e504916e00b148440550f0672ad134470854976 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Sun, 3 May 2026 20:26:10 -0300 Subject: [PATCH 1/7] Change all retention functions to use initial mass This is the source of a major bug that has been in pretty much every version of SSPtools. The fallback fraction is a function of the *initial* progenitor mass, not the final BH mass, as we have been using. This will, unfortunately, not be a simple fix, and will require rethinking how natal kicks are implemented entirely. --- ssptools/kicks.py | 63 ++++++++++++++++++++++++++++++++++++++-------- ssptools/masses.py | 4 +++ 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/ssptools/kicks.py b/ssptools/kicks.py index 640f3ad..b5b2623 100644 --- a/ssptools/kicks.py +++ b/ssptools/kicks.py @@ -10,7 +10,7 @@ import scipy.interpolate as interp -__all__ = ["natal_kicks", "KickStats", "maxwellian_kick_v"] +__all__ = ["kick_retention_fraction", "KickStats", "maxwellian_kick_v"] @dataclasses.dataclass(eq=False, frozen=True) @@ -53,7 +53,7 @@ def _maxwellian_retention_frac(m, vesc, FeH, vdisp=265., *, SNe_method='rapid'): Parameters ---------- m : float - The mean mass of a BH bin. + The mean initial mass of the progenitor star which will make a BH. vesc : float The initial escape velocity of the cluster. @@ -73,7 +73,7 @@ def _maxwellian_retention_frac(m, vesc, FeH, vdisp=265., *, SNe_method='rapid'): Returns ------- float - The retention fraction of BHs of this mass. + The retention fraction of BHs created by stars of this mass. ''' @@ -137,11 +137,10 @@ def _F12_fallback_frac(FeH, *, model='uSSE', SNe_method='rapid'): # feh_path = get_data(f"sse/MP_FEH{FeH:+.2f}.dat") # .2f snaps to the grid feh_path = get_data(f"ifmr/{model}_{SNe_method}/IFMR_FEH{FeH:+.2f}.dat") - # load in the data (only final remnant mass and fbac) - fb_grid = np.loadtxt(feh_path, usecols=(1, 3), unpack=True) + # load in the data (only initial star mass and fbac) + fb_grid = np.loadtxt(feh_path, usecols=(0, 3), unpack=True) - # TODO this is incorrect, and should be a function of mi - # Interpolate the mr-fb grid + # Interpolate the mi-fb grid return interp.interp1d(fb_grid[0], fb_grid[1], kind="linear", bounds_error=False, fill_value=(0.0, 1.0)) @@ -176,7 +175,7 @@ def _sigmoid_retention_frac(m, slope, scale): Parameters ---------- m : float - The mean mass of a BH bin. + The mean initial mass of the progenitor star which will make a BH. slope : float The "slope" of the sigmoid function, defining the "sharpness" of the @@ -192,7 +191,7 @@ def _sigmoid_retention_frac(m, slope, scale): Returns ------- float - The retention fraction of BHs of this mass. + The retention fraction of BHs created by stars of this mass. ''' return erf(np.exp(slope * (m - scale))) @@ -213,7 +212,7 @@ def _tanh_retention_frac(m, slope, scale): Parameters ---------- m : float - The mean mass of a BH bin. + The mean initial mass of the progenitor star which will make a BH. slope : float The "slope" of the sigmoid function, defining the "sharpness" of the @@ -228,12 +227,47 @@ def _tanh_retention_frac(m, slope, scale): Returns ------- float - The retention fraction of BHs of this mass. + The retention fraction of BHs created by stars of this mass. ''' return 0.5 * (np.tanh(slope * (m - scale)) + 1) # return np.tanh(np.exp(slope * (m - scale))) # alternative +def kick_retention_fraction(m_ini, f_kick=None, method='fryer2012', **ret_kwargs): + ''' + + m_ini : float + The initial (ZAMS) mass of the progenitor star which will form a BH. + ''' + + f_ret = _get_kick_method(method) + + # If no given total kick fraction, use old-style of directly using f_ret + if f_kick is None: + + return f_ret(m_ini, **ret_kwargs) + + # return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, **ret_kwargs) + + else: + raise NotImplementedError('TODO') + + # get the BHMF using the new fast InitialBHPopulation and then do this + + # # Fit for the desired kick scale + # try: + # slp, scl = _determine_kick_params(Mr_BH, Nr_BH, f_ret, f_kick, + # **ret_kwargs) + # except TypeError as err: + # mssg = (f"Can only use `f_kick` with methods that use a `scale` and" + # f" `slope` parameter, not '{method}'.") + # raise ValueError(mssg) from err + + # # Compute the natal kicks based on fit scale + return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, slope=slp, scale=scl) + + + # -------------------------------------------------------------------------- # Performing natal kicks on BH mass functions # -------------------------------------------------------------------------- @@ -391,6 +425,13 @@ def natal_kicks(Mr_BH, Nr_BH, f_kick=None, method='fryer2012', **ret_kwargs): stats : KickStats Statistics on how and how many BHs are kicked. + Notes + ----- + This was the old method of applying natal kicks, when the ejections of BHs + were only considered after generating all of the masses. This relied on + faulty logic w.r.t. the fallback fractions and are no longer used, but left + here for posterity. This method is still valid for non-Maxwellian kicks. + See Also -------- _maxwellian_retention_frac : Maxwellian retention fraction algorithm. diff --git a/ssptools/masses.py b/ssptools/masses.py index 2af1fe8..d71e475 100644 --- a/ssptools/masses.py +++ b/ssptools/masses.py @@ -808,6 +808,10 @@ def determine_index(self, mass, massbins, *, allow_overflow=False): if massbins in star_classes._fields: massbins = getattr(self.bins, massbins) + # TODO I think you can replace this entire thing with just + # ind = np.digitize(mass, bin_edges, right=False) - 1 + # That is vectorized and clearer. + # Since mass bins always increasing, can look at only the lower bound try: ind = np.flatnonzero(massbins.lower <= mass)[-1] From dcf70fc2b8a7f4992d753a220d875914dd14b536 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Sun, 3 May 2026 20:51:17 -0300 Subject: [PATCH 2/7] Completely rewrite `InitialBHPopulation` While working on kick stuff, where this class became necessary in order to work out some statistics, I realized that this entire thing, which before consisted of the same ODE as the rest of the classes, but stripped of anything not BH related, could be accomplished with some simple math, given that we know the IMF and the IFMR. This class should give (nearly) the same results as the previous, but in a fraction of the time (~1 sec to ~0.02 sec). This allows it to be used in other places without much worry. Note that this *does* some of the init signatures, so it is not entirely backwards compatible --- ssptools/evolve_mf.py | 268 ++++++++++++++++++------------------------ ssptools/masses.py | 11 +- 2 files changed, 124 insertions(+), 155 deletions(-) diff --git a/ssptools/evolve_mf.py b/ssptools/evolve_mf.py index de8dae5..849516e 100644 --- a/ssptools/evolve_mf.py +++ b/ssptools/evolve_mf.py @@ -1362,41 +1362,9 @@ def f_BH(self, M_cluster): '''The total mass fraction in BHs, given a total population mass.''' return self.Mtot / M_cluster - def __init__(self, M_BH, N_BH, BH_bins, FeH, *, - natal_kicks=False, kick_method='maxwellian', f_kick=None, - SNe_method='rapid', vesc=90, kick_vdisp=265., - kick_slope=1, kick_scale=20): + def __init__(self, M_BH, N_BH, BH_bins): '''Should not init from this, use provided classmethods.''' - # ------------------------------------------------------------------ - # Optionally perform all natal kicks on the input mass arrays - # ------------------------------------------------------------------ - - if natal_kicks: - - self.natal_kicks = natal_kicks - - match kick_method.casefold(): - - case 'maxwellian' | 'f12' | 'fryer2012': - - from .ifmr import _check_IFMR_FeH_bounds - FeH_BH = _check_IFMR_FeH_bounds(FeH) - - kick_kw = dict(method=kick_method, f_kick=f_kick, vesc=vesc, - FeH=FeH_BH, vdisp=kick_vdisp, - SNe_method=SNe_method) - - case 'sigmoid' | 'tanh': - kick_kw = dict(method=kick_method, f_kick=f_kick, - slope=kick_slope, scale=kick_scale) - - case _: - mssg = f"Invalid natal kick algorithm '{kick_method=}'" - raise ValueError(mssg) - - *_, self._kick_stats = kicks.natal_kicks(M_BH, N_BH, **kick_kw) - # ------------------------------------------------------------------ # Compute and store all final BH mass and number arrays and values # ------------------------------------------------------------------ @@ -1419,10 +1387,14 @@ def __init__(self, M_BH, N_BH, BH_bins, FeH, *, # ---------------------------------------------------------------------- @classmethod - def from_IMF(cls, IMF, nbins, FeH, N0=5e5, *, natal_kicks=True, - binning_breaks=None, binning_method='default', + def from_IMF(cls, IMF, FeH, *, natal_kicks=True, + kick_method='maxwellian', f_kick=None, + SNe_method='rapid', vesc=90, kick_vdisp=265., + kick_slope=1, kick_scale=20, BH_IFMR_method='banerjee20', WD_IFMR_method='mist18', - BH_IFMR_kwargs=None, WD_IFMR_kwargs=None, **kwargs): + BH_IFMR_kwargs=None, WD_IFMR_kwargs=None, + bins=None, nbins=[5, 5, 20], + binning_breaks=None, binning_method='default', **kwargs): '''Initialize a BH population by evolving from an IMF. Based on a given IMF, sharing many arguments with `EvolvedMF`, generate @@ -1448,18 +1420,36 @@ def from_IMF(cls, IMF, nbins, FeH, N0=5e5, *, natal_kicks=True, FeH : float Metallicity, in solar fraction [Fe/H]. - N0 : int, optional - Total initial number of stars, over all bins. - If None, uses the N0 of the given IMF class. Defaults to 5e5 stars. - natal_kicks : bool, optional Whether to account for natal kicks in the BH dynamical retention. Defaults to True (note this difference from `EvolvedMF`). + kick_method : {'maxwellian', 'tanh', 'sigmoid'}, optional + The BH natal kick algorithm to use, if `natal_kicks=True`. + See `kicks` for more details. + vesc : float, optional Initial cluster escape velocity, in km/s, for use in the computation of BH natal kick effects. Defaults to 90 km/s. + kick_vdisp : float, optional + The dispersion of the Maxwellian natal kick velocity distribution. + Only used if `kick_method='maxwellian'`. + See `kicks._maxwellian_retention_frac` for more information. + + f_kick : float, optional + + kick_slope : float + The "slope" of the sigmoid retention function, defining the + "sharpness" of the increase. Only used if `kick_method='sigmoid'`. + See `kicks._sigmoid_retention_frac` for more information. + + kick_scale : float + The scale-mass of the sigmoid retention function, defining the + approximate mass of the turn-over from 0 to 1. + Only used if `kick_method='sigmoid'`. + See `kicks._sigmoid_retention_frac` for more information. + binning_breaks : list of float, optional The binning break masses to use when constructing the mass bins, including outer edges. See `masses.MassBins` for more information. @@ -1490,149 +1480,123 @@ def from_IMF(cls, IMF, nbins, FeH, N0=5e5, *, natal_kicks=True, All other arguments are passed to the `InitialBHPopulation`. ''' - def compute_tms(mi): - a = tms_constants - return a[0] * np.exp(a[1] * mi ** a[2]) + # ------------------------------------------------------------------ + # Initialise the initial mass function and mass bins given the + # power-law IMF slopes and bins + # ------------------------------------------------------------------ - def compute_mto(t): - a0, a1, a2 = tms_constants - if t > a0: - return (np.log(t / a0) / a1) ** (1 / a2) - else: - return np.inf + _ifmr = IFMR(FeH, WD_method=WD_IFMR_method, WD_kwargs=WD_IFMR_kwargs, + BH_method=BH_IFMR_method, BH_kwargs=BH_IFMR_kwargs) + + binning_breaks = IMF.mb if binning_breaks is None else binning_breaks - def _derivs_BHs(t, y): - '''Derivatives relevant to mass changes due to stellar evolution. - This is a massively simplified case of the DEs from `EvolvedMF`, - and is only valid up to the age all BHs are formed. - `y` is [MS N array, BH N array, BH M array]. - ''' + if bins is None: + # TODO it now makes more sense to have nbins be just the number of + # BH bins, should maybe remove some options + massbins = MassBins(binning_breaks, nbins, IMF, _ifmr, + binning_method=binning_method) - # Setup derivative bins - Ns = y[:nbin_MS] # only Ns is relevant right now - dNs = np.zeros(nbin_MS) - dNr, dMr = np.zeros(nbin_BH), np.zeros(nbin_BH) + bins = massbins.bins.BH - frem = 1.0 # Retain all BHs + nbin_BH = bins.lower.size - # Apply only if this time is atleast later than the earliest tms - if t > tms_u[-1]: + # ------------------------------------------------------------------ + # Compute the initial amounts of BHs expected from the IMF and IFMR + # ------------------------------------------------------------------ - # Find out which mass bin is the current turn-off bin - isev = np.where(t > tms_u)[0][0] + # Get domain of progenitor masses that make BHs + mi = np.linspace(*_ifmr.BH_mi, 50_000) + dm = mi[1] - mi[0] - mto, m1 = compute_mto(t), massbins.bins.MS.lower[isev] + # Determine the final BH masses + mf = _ifmr.predict(mi) - # Avoid "hitting" the bin edge - if mto > m1 and (Nj := Ns[isev]) > 0.1: + # Determine which BH bin each BH will fall into + # TODO this may be behaving weird when out of bounds? (!) + inds = np.digitize(mf, np.r_[bins.lower, bins.upper[-1]]) - 1 - # The normalization constant - # TODO deal with the constant divide-by-zero warning here - Aj = Nj / Pk(IMF.a[-1], 1, m1, mto) + # Optionally perform natal kicks + if natal_kicks: - # Get the number of turn-off stars per unit of mass - dNdm = Aj * mto**IMF.a[-1] + match kick_method.casefold(): - else: - dNdm = 0 - # TODO just break??? + case 'maxwellian' | 'f12' | 'fryer2012': - # Compute the full dN/dt = dN/dm * dm/dt - a = tms_constants - dmdt = abs((1.0 / (a[1] * a[2] * t)) - * (np.log(t / a[0]) / a[1]) ** (1 / a[2] - 1)) + from .ifmr import _check_IFMR_FeH_bounds + FeH_BH = _check_IFMR_FeH_bounds(FeH) - dNdt = -dNdm * dmdt + kick_kw = dict(method=kick_method, f_kick=f_kick, vesc=vesc, + FeH=FeH_BH, vdisp=kick_vdisp, + SNe_method=SNe_method) - # Fill in star derivatives - dNs[isev] = dNdt + case 'sigmoid' | 'tanh': + kick_kw = dict(method=kick_method, f_kick=f_kick, + slope=kick_slope, scale=kick_scale) - # Skip 0-mass remnants - if t <= final_age and (m_rem := _ifmr.predict(mto)) > 0: + case _: + mssg = f"Invalid natal kick algorithm '{kick_method=}'" + raise ValueError(mssg) - # If this happens its because IFMR is (numerically) broken! - if (cls_rem := _ifmr.predict_type(mto)) != 'BH': - mssg = f"Initial mass {mto} made a {cls_rem}, not a BH." - raise RuntimeError(mssg) + frem = kicks.kick_retention_fraction(mi, **kick_kw) - # Find bin based on lower bin edge (must be careful later) + else: + frem = 1.0 - irem = massbins.determine_index(m_rem, cls_rem) + # Compute the masses and numbers of BHs which will form - # Fill in remnant derivatives - dNr[irem] = -dNdt * frem - dMr[irem] = -m_rem * dNdt * frem + M_BH = np.bincount(inds, weights=frem * mf * IMF.N(mi) * dm, + minlength=nbin_BH) - return np.r_[dNs, dNr, dMr] + N_BH = np.bincount(inds, weights=frem * IMF.N(mi) * dm, + minlength=nbin_BH) # ------------------------------------------------------------------ - # Initialise the initial mass function and mass bins given the - # power-law IMF slopes and bins + # Create the final class instance, and populate some extra values + # specific to this init method, related to the evolution # ------------------------------------------------------------------ - _ifmr = IFMR(FeH, WD_method=WD_IFMR_method, WD_kwargs=WD_IFMR_kwargs, - BH_method=BH_IFMR_method, BH_kwargs=BH_IFMR_kwargs) + out = cls(M_BH, N_BH, bins) - binning_breaks = IMF.mb if binning_breaks is None else binning_breaks - - massbins = MassBins(binning_breaks, nbins, IMF, _ifmr, - binning_method=binning_method) - - nbin_MS, nbin_BH = massbins.nbin.MS, massbins.nbin.BH + # Helpful classes - # Load a_i coefficients derived from interpolated Dartmouth models - mstogrid = np.loadtxt(get_data("sevtables/msto.dat")) - nearest_FeH = np.argmin(np.abs(mstogrid[:, 0] - FeH)) - tms_constants = mstogrid[nearest_FeH, 1:] + out.IMF = IMF + out.IFMR = _ifmr - # Compute t_ms for all bin edges - tms_u = compute_tms(massbins.bins.MS.upper) + # Natal kicks - init_N, init_M = massbins.initial_values(packed=False, N0=N0) + out.natal_kicks = natal_kicks - # ------------------------------------------------------------------ - # Integrate and solve the derivatives (evolve) - # ------------------------------------------------------------------ - - y0 = np.r_[init_N.MS, init_N.BH, init_M.BH] - - sol = ode(_derivs_BHs) - sol.set_integrator("dopri5", max_step=1e12, atol=1e-5, rtol=1e-5) - sol.set_initial_value(y0, t=0.0) + if natal_kicks: - final_age = compute_tms(_ifmr.BH_mi.lower + 0.1) # fltpnt bound errors + # Kinda crazy but I think we have to recompute this entirely + # without the kicks to get the difference. - # Integrate the solver at each bin bound, to match EvolvedMF. - # The ODE solvers are suspiciously sensitive to this, and while this - # may not be the most sane solution, this will recreate EvolvedMF. - for ti in np.sort(tms_u[tms_u < final_age]): - sol.integrate(ti) + M_BH_nokick = np.bincount(inds, weights=mf * IMF.N(mi) * dm, + minlength=nbin_BH) - sol.integrate(final_age) + N_BH_nokick = np.bincount(inds, weights=IMF.N(mi) * dm, + minlength=nbin_BH) - N_BH, M_BH = sol.y[nbin_MS:nbin_MS + nbin_BH], sol.y[nbin_MS + nbin_BH:] + ibh_nokick = cls(M_BH_nokick, N_BH_nokick, bins) - # ------------------------------------------------------------------ - # Create the final class instance, and populate some extra values - # specific to this init method, related to the evolution - # ------------------------------------------------------------------ + out._kick_stats = kicks.KickStats.from_final(M_BH, N_BH, ibh_nokick, + parameters=kick_kw) - out = cls(M_BH, N_BH, massbins.bins.BH, - FeH=FeH, natal_kicks=natal_kicks, **kwargs) - - out.IMF = IMF + else: - out.age = final_age + out._kick_stats = kicks.KickStats.no_kicks(nbin_BH) - Ns = sol.y[:nbin_MS] + # Compute the final age (i.e. compute_tms()) - alphas = np.repeat(IMF.a, massbins._nbin_MS_each) - As = Ns / Pk(alphas, 1, *massbins.bins.MS) - Ms = As * Pk(alphas, 2, *massbins.bins.MS) + mstogrid = np.loadtxt(get_data("sevtables/msto.dat")) + nearest_FeH = np.argmin(np.abs(mstogrid[:, 0] - FeH)) + a = mstogrid[nearest_FeH, 1:] + out.age = a[0] * np.exp(a[1] * _ifmr.BH_mi.lower ** a[2]) # Stellar losses - out.Ns_lost = init_N.MS.sum() - Ns.sum() - out.Ms_lost = init_M.MS.sum() - Ms.sum() + + out.Ns_lost = IMF.integrate_N(*_ifmr.BH_mi) # number of stars now BHs + out.Ms_lost = IMF.integrate_M(*_ifmr.BH_mi) # mass of stars lost to BHs return out @@ -1700,17 +1664,19 @@ def from_powerlaw(cls, m_breaks, a_slopes, nbins, FeH, N0=5e5, *, @classmethod def from_BHMF(cls, m_breaks, a_slopes, nbins, FeH, N0=1000, *, - natal_kicks=True, binning_method='default', + binning_method='default', BH_IFMR_method='banerjee20', WD_IFMR_method='mist18', - BH_IFMR_kwargs=None, WD_IFMR_kwargs=None, **kwargs): + BH_IFMR_kwargs=None, WD_IFMR_kwargs=None): '''Initialize a BH population based on a given power-law mass function. Based on an explicit given power-law parametrization of a BH mass function, generate a corresponding population of BH mass bins. This class *does not* simulate any evolution, and is not based on any sort of stellar mass function, but simply creates bins of mass - (normalized to N0) for a given BH MF, and optionally provides - natal kicks. + (normalized to N0) for a given BH MF. + + The IFMR kwargs are solely used to determine min and max BH masses for + creating the mass bins (depending on the binning methods used). Parameters ---------- @@ -1730,10 +1696,6 @@ def from_BHMF(cls, m_breaks, a_slopes, nbins, FeH, N0=1000, *, N0 : int, optional Total initial number of BHs, over all bins. Defaults to 1000 BHs. - natal_kicks : bool, optional - Whether to account for natal kicks in the BH dynamical retention. - Defaults to True (note this difference from `EvolvedMF`). - vesc : float, optional Initial cluster escape velocity, in km/s, for use in the computation of BH natal kick effects. Defaults to 90 km/s. @@ -1772,5 +1734,5 @@ def from_BHMF(cls, m_breaks, a_slopes, nbins, FeH, N0=1000, *, N_BH, M_BH, _ = MF.binned_eval(bins) - return cls(M_BH, N_BH, bins, FeH=FeH, - natal_kicks=natal_kicks, **kwargs) + return cls(M_BH, N_BH, bins) + diff --git a/ssptools/masses.py b/ssptools/masses.py index d71e475..e1a0973 100644 --- a/ssptools/masses.py +++ b/ssptools/masses.py @@ -117,8 +117,7 @@ def mmean(self): def Mtot(self): '''Total mass of system under this IMF (assuming `self.N0` stars).''' # TODO integrating this every time is wasteful - from scipy.integrate import quad - return quad(self.M, self.mb[0], self.mb[-1])[0] + return self.integrate_M(self.mb[0], self.mb[-1]) @classmethod def from_M0(cls, m_break, a, M0, *, ext='zeros'): @@ -303,6 +302,14 @@ def binned_M(self, bins, *, N=None): '''Return the total mass of stars within given mass bins.''' return self.binned_eval(bins=bins, N=N)[1] + def integrate_N(self, ml, mu): + from scipy.integrate import quad + return quad(self.N, ml, mu)[0] + + def integrate_M(self, ml, mu): + from scipy.integrate import quad + return quad(self.M, ml, mu)[0] + # -------------------------------------------------------------------------- # Mass Bins From 5a9cb7f0c329967455efaab84c7cf61e29ad94b7 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Mon, 4 May 2026 12:17:21 -0300 Subject: [PATCH 3/7] remake `KickStats` for new kick design The kick statistics cannot be easily computed now that the kicks are done during the ODE solving directly, not on the BH bins after the fact. However, these are still useful to have. Therefore this commit changes `KickStats` to try to recreate this logic as best as possible. Note that these will be approximate at best, but it is less important now, as they will not be used directly for anything crucial. --- ssptools/kicks.py | 62 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/ssptools/kicks.py b/ssptools/kicks.py index b5b2623..40b53bf 100644 --- a/ssptools/kicks.py +++ b/ssptools/kicks.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- from .ifmr import get_data +from .evolve_mf import InitialBHPopulation import dataclasses @@ -17,20 +18,81 @@ class KickStats: retention: np.ndarray mass_kicked: np.ndarray + num_kicked: np.ndarray parameters: dict @property def total_kicked(self) -> float: + '''The total amount of BH mass kicked.''' return self.mass_kicked.sum() @classmethod def no_kicks(cls, nmbin): + '''Kick stats when no kicks are applied.''' return cls( retention=np.ones(nmbin), mass_kicked=np.zeros(nmbin), + num_kicked=np.zeros(nmbin), parameters=dict() ) + @classmethod + def from_final(cls, Mr_BH, Nr_BH, ibh: InitialBHPopulation, parameters): + '''Compute the kick stats based on final BH arrays. + + This constructor (the main one, likely) uses the final amounts of BHs, + after all kicks have been applied (but no other ejections have taken + place, e.g. no dynamical ejections), to compute the statistics based on + a corresponding `InitialBHPopulation` where the kicks have *not* been + applied. + + This is done by simply assuming that all mass that is in the initial + BH population and not the given BH bins must have been natally kicked. + + This will, at best, provide an approximately correct view of the kicks. + Any differences attributable to the binning effects in `Mr_BH` compared + to `InitialBHPopulation` will be especially notable, and may lead to + cases where, e.g., the kick amounts look very slightly negative. + There is unfortunately no better way to determine the effects of + natal kicks on the final BH bins themselves. + + Also note that this *will not* be valid for cases where the BH bins + are created *before* the full amounts of BHs have been formed (i.e. + younger than `ibh.age`). + + Parameters + ---------- + Mr_BH : ndarray + Array[nbin] of the total final masses of black holes in each + BH mass bin, after natal kicks but before any other ejections. + + Nr_BH : ndarray + Array[nbin] of the total final numbers of black holes in each + BH mass bin, after natal kicks but before any other ejections. + + ibh : InitialBHPopulation + The `InitialBHPopulation` instance representing the expected + complete initial BH mass function. This must align exactly with + the BH bins, and must have been created with `natal_kicks=False`. + + parameters : dict + The kick retention function parameters used, for convenient storage. + ''' + + retention = np.ones_like(Mr_BH) + c = ibh.M > 0 + retention[c] = Mr_BH[c] / ibh.M[c] + + M_kicked = ibh.M - Mr_BH + N_kicked = ibh.N - Nr_BH + + return cls( + retention=retention, + mass_kicked=M_kicked, + num_kicked=N_kicked, + parameters=parameters + ) + # -------------------------------------------------------------------------- # Retention fraction functions From 1e8d5dd0663df71487d48a9f9a9057c2ab166bd7 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Mon, 4 May 2026 12:29:59 -0300 Subject: [PATCH 4/7] Completely redo kick logic in `evolve_mf` This commit changes entirely how the BH natal kicks are handled and applied to the `evolve_mf` classes. Namely, this is the fix from the old incorrect computations based on the final mass, to the correct version using progenitor masses. This removes all kick logic from the after-the-fact BH ejection stage, and instead replaces the largely unused `_frem` map with a new function which calls the relevant retention functions from `kicks`. This means that the kicks are now applied during the creation of the BHs themselves, removing some fraction over time based on the progenitor star masses. A lot of other logic had to change to account for these changes. Notably, the kick stats are no longer as important, and have become more approximate. Relatedly, the meaning of `BH_ret_dyn` has now changed, in a *not backwards compatible* manner, and now represents only the dynamical ejection retention fraction. Also changed is the handling of `f_kick`, which now has the parameters determined once at the start, and the actual used kick params stored. --- ssptools/evolve_mf.py | 220 +++++++++++++++++++++++++++++------------ ssptools/ifmr.py | 2 + ssptools/kicks.py | 222 +++++++++++------------------------------- 3 files changed, 217 insertions(+), 227 deletions(-) diff --git a/ssptools/evolve_mf.py b/ssptools/evolve_mf.py index 849516e..48a83f1 100644 --- a/ssptools/evolve_mf.py +++ b/ssptools/evolve_mf.py @@ -23,8 +23,6 @@ class EvolvedMF: to a binned present-day mass function (PDMF) at a given set of ages, and computes the numbers and masses of stars and remnants in each mass bin. - # TODO add more in-depth explanation or references to actual algorithm here - Parameters ---------- IMF : PowerLawIMF @@ -64,12 +62,10 @@ class EvolvedMF: NS_ret : float, optional Neutron star retention fraction (0 to 1). Defaults to 0.1 (10%). - BH_ret_int : float, optional - Initial black hole retention fraction (0 to 1). Defaults to 1 (100%). - BH_ret_dyn : float, optional - Dynamical black hole retention fraction (0 to 1), including both - dynamical ejections and natal kicks. Defaults to 1 (100%). + Dynamical black hole retention fraction (0 to 1). Does not include the + natal kicks, only the a posteriori dynamical ejections. + Defaults to 1 (100%). natal_kicks : bool, optional Whether to account for natal kicks in the BH dynamical retention. @@ -221,7 +217,7 @@ class EvolvedMF: Notes ----- - The BH kicks and ejections are computed post facto on the BH mass + The dynamical ejections are computed after the fact on the BH mass bins for each requested output age (`tout`). As such, parameters such as `BH_ret_dyn` are applied equally at each age, and not accounted for during the evolution, which may not be entirely realistic. @@ -289,7 +285,7 @@ def nmr(self): return (np.c_[self.Nr][-1] > 10 * self.Nmin).sum() def __init__(self, IMF, nbins, FeH, tout, esc_rate, N0=5e5, - tcc=0.0, NS_ret=0.1, BH_ret_int=1.0, BH_ret_dyn=1.0, *, + tcc=0.0, NS_ret=0.1, BH_ret_dyn=1.0, *, natal_kicks=False, kick_method='maxwellian', SNe_method='rapid', vesc=90, kick_vdisp=265., f_kick=None, kick_slope=1, kick_scale=20, @@ -320,9 +316,7 @@ def __init__(self, IMF, nbins, FeH, tout, esc_rate, N0=5e5, self._stellar_ev = stellar_evolution self.NS_ret = NS_ret - self.BH_ret_int = BH_ret_int self.BH_ret_dyn = BH_ret_dyn - self._frem = {'WD': 1., 'NS': NS_ret, 'BH': BH_ret_int} self.FeH = FeH @@ -349,6 +343,16 @@ def __init__(self, IMF, nbins, FeH, tout, esc_rate, N0=5e5, self.massbins = MassBins(binning_breaks, nbins, self.IMF, self.IFMR, binning_method=binning_method) + # ------------------------------------------------------------------ + # Compute the full initial BH mass function + # ------------------------------------------------------------------ + + self.ibh = InitialBHPopulation.from_IMF( + IMF, FeH=FeH, N0=N0, bins=self.massbins.bins.BH, natal_kicks=False, + WD_IFMR_method=WD_IFMR_method, WD_IFMR_kwargs=WD_IFMR_kwargs, + BH_IFMR_method=BH_IFMR_method, BH_IFMR_kwargs=BH_IFMR_kwargs, + ) + # ------------------------------------------------------------------ # Setup lifetime approximations and compute t_ms of all bin edges # ------------------------------------------------------------------ @@ -379,6 +383,19 @@ def __init__(self, IMF, nbins, FeH, tout, esc_rate, N0=5e5, self.natal_kicks = natal_kicks + # If f_kick, try to determine kick scale automatically + if f_kick is not None: + try: + kick_slope, kick_scale = kicks._determine_kick_params( + kick_method, f_kick, self.IMF, self.IFMR, + kick_slope, kick_scale + ) + + except TypeError as err: + mssg = (f"Can only use `f_kick` with methods that use " + f"`scale` and `slope` params, not '{kick_method}'.") + raise ValueError(mssg) from err + match kick_method.casefold(): case 'maxwellian' | 'f12' | 'fryer2012': @@ -387,14 +404,18 @@ def __init__(self, IMF, nbins, FeH, tout, esc_rate, N0=5e5, FeH_BH = _check_IFMR_FeH_bounds(FeH) self._kick_kw = dict( - method=kick_method, f_kick=f_kick, vesc=vesc, + method=kick_method, vesc=vesc, FeH=FeH_BH, vdisp=kick_vdisp, SNe_method=SNe_method ) case 'sigmoid' | 'tanh': - self._kick_kw = dict(method=kick_method, f_kick=f_kick, + + self._kick_kw = dict(method=kick_method, slope=kick_slope, scale=kick_scale) + case 'full' | 'everything' | 'all' | 'none': + self._kick_kw = dict(method=kick_method) + case _: mssg = f"Invalid natal kick algorithm '{kick_method=}'" raise ValueError(mssg) @@ -482,6 +503,31 @@ def compute_mto(self, t): return out + def _frem(self, mi, rem_type): + '''The *initial* retention fraction of a given remnant type.''' + + match rem_type: + + # White dwarfs, assume always 100% retained + + case 'WD': + return 1.0 + + # Neutron stars, input `NS_ret`, default to constant 10% + + case 'NS': + return self.NS_ret + + # Black Holes, compute full natal kicks + + case 'BH': + + if self.natal_kicks: + return kicks.kick_retention_fraction(mi, **self._kick_kw) + + else: + return 1.0 + def _derivs(self, t, y): '''Main function for computing derivatives relevant to mass evolution @@ -579,7 +625,7 @@ def _derivs_sev(self, t, y): irem = self.massbins.determine_index(m_rem, cls_rem) # Compute Remnant retention fractions based on remnant type - frem = self._frem[cls_rem] + frem = self._frem(mto, cls_rem) # Fill in remnant derivatives getattr(dNr, cls_rem)[irem] = -dNdt * frem @@ -683,6 +729,7 @@ def _derivs_esc(self, t, y): dNs[depl_mask] += B * Is + # TODO is B correct, given alpha is only for the stars?? dalpha[depl_mask] += ( B * ((bins_MS.lower[depl_mask] / md) ** 0.5 - (bins_MS.upper[depl_mask] / md) ** 0.5) @@ -794,10 +841,6 @@ def _evolve(self): self.rem_types = np.repeat(self.massbins.nbin._fields[1:], self.massbins.nbin[1:]) - # To save some repetition, just note these stats here - if not self.natal_kicks: - self._kick_stats = kicks.KickStats.no_kicks(self.massbins.nbin.BH) - # ------------------------------------------------------------------ # Initialise ODE solver # ------------------------------------------------------------------ @@ -843,53 +886,61 @@ def _evolve(self): Ms = As * Pk(alpha, 2, *bins_MS) # ---------------------------------------------------------- - # Eject BHs, first through natal kicks, then dynamically + # Eject BHs # ---------------------------------------------------------- # TODO really feels like this should be done during the - # evolution/derivs? At least for the natal kicks. + # evolution/derivs? # TODO due to rem_classes tuple, ejections done in place, which # is not ideal # Check if any BH have been created if ti > self.compute_tms(self.IFMR.BH_mi.upper): + # ------------------------------------------------------ + # Get BH natal kick statistics. + # These are somewhat approximate, and also must be done + # before the dynamical ejections so we can try to work out + # how much mass was kicked before them. + # ------------------------------------------------------ + + if self.natal_kicks: + + # if not all BHs have formed yet, kick stats is invalid + # this is just due to how we now compute them using ibh + if ((ti == self.tout[-1]) and (ti < self.ibh.age)): + + mssg = ("_kick_stats will not be valid for ages " + f"before all BHs form ({self.ibh.age} Myr)") + warnings.warn(mssg) + + self._kick_stats = kicks.KickStats.from_final( + Mr.BH, Nr.BH, self.ibh, parameters=self._kick_kw + ) + + else: + + self._kick_stats = kicks.KickStats.no_kicks( + self.massbins.nbin.BH + ) + + # ------------------------------------------------------ + # Eject BHs dynamically, based on desired BH_ret_dyn + # ------------------------------------------------------ + # calculate total mass we want to eject M_eject = Mr.BH.sum() * (1.0 - self.BH_ret_dyn) + M_ret = Mr.BH.sum() - M_eject # If kicking basically all, skip ahead if 0. <= M_ret / (Mr.BH[0] / Nr.BH[0]) < self.Nmin: - # Compute the natal kicks, just to store the stats - if self.natal_kicks: - - *_, self._kick_stats = kicks.natal_kicks( - Mr.BH.copy(), Nr.BH.copy(), **self._kick_kw - ) - # Remove all BHs and skip ahead Mr.BH[:] = 0 Nr.BH[:] = 0 else: - # First remove mass from all bins by natal kicks - if self.natal_kicks: - - *_, self._kick_stats = kicks.natal_kicks( - Mr.BH, Nr.BH, **self._kick_kw - ) - M_eject -= self._kick_stats.total_kicked - - if M_eject < 0: - mssg = ( - f"Natal kicks already removed {-M_eject} Msun " - "more than total ejections desired by " - f"'BH_ret_dyn={self.BH_ret_dyn}'. " - "Increase BH_ret_dyn or alter natal kicks." - ) - raise ValueError(mssg) - # Remove dynamical BH ejections self._dyn_eject_BH(Mr.BH, Nr.BH, M_eject=M_eject) @@ -919,8 +970,13 @@ def _evolve(self): self.mr[c][iout, :] = mr + # ------------------------------------------------------------------ + # Get some final helpful quantities + # ------------------------------------------------------------------ + Mtot = self.Ms.sum(axis=1) + np.c_[self.Mr].sum(axis=1) Ntot = self.Ns.sum(axis=1) + np.c_[self.Nr].sum(axis=1) + # TODO sometimes results in nan, because Ntot has nans? self.mmean = Mtot / Ntot self.converged = sol.successful() @@ -928,6 +984,8 @@ def _evolve(self): mssg = "ODE solver has *not* converged, this MF will not be valid." warnings.warn(mssg) + self.sol = sol + class EvolvedMFWithBH(EvolvedMF): r'''Evolve an IMF to a present-day mass function at a given age and f_BH. @@ -1012,7 +1070,7 @@ def __init__(self, IMF, nbins, FeH, tout, esc_rate, f_BH, self.strict_BH_target = strict_BH_target - # leave BH_ret_dyn as default, will be ignored. BH_ret_int is fine + # leave BH_ret_dyn as default, will be ignored. super().__init__(IMF, nbins, FeH, tout, esc_rate, *args, **kwargs) @@ -1225,41 +1283,68 @@ def _evolve(self): Ms = As * Pk(alpha, 2, *bins_MS) # ---------------------------------------------------------- - # Eject BHs, first through natal kicks, then dynamically + # Eject BHs # ---------------------------------------------------------- # Check if any BH have been created if ti > self.compute_tms(self.IFMR.BH_mi.upper): - fBH_target = self._fBH_target[iout] + # ------------------------------------------------------ + # Get BH natal kick statistics. + # These are somewhat approximate, and also must be done + # before the dynamical ejections so we can try to work out + # how much mass was kicked before them. + # ------------------------------------------------------ - # First remove mass from all bins by natal kicks, if desired - # Do it first because we dont control the exact amount so - # this could make the target fBH invalid afterwards if self.natal_kicks: - *_, self._kick_stats = kicks.natal_kicks( - Mr.BH, Nr.BH, **self._kick_kw - ) + + # if not all BHs have formed yet, kick stats is invalid + # this is just due to how we now compute them using ibh + if ((ti == self.tout[-1]) and (ti < self.ibh.age)): + + mssg = ("_kick_stats will not be valid for ages " + f"before all BHs form ({self.ibh.age} Myr)") + warnings.warn(mssg) + + self._kick_stats = kicks.KickStats.from_final( + Mr.BH, Nr.BH, self.ibh, parameters=self._kick_kw + ) else: + self._kick_stats = kicks.KickStats.no_kicks( self.massbins.nbin.BH ) + # ------------------------------------------------------ + # Eject BHs dynamically, based on desired f_BH target + # ------------------------------------------------------ + + fBH_target = self._fBH_target[iout] + Mtot = np.r_[Mr].sum() + Ms.sum() Mbhtot = Mr.BH.sum() fBH_current = Mbhtot / Mtot # can only make fBH go down by removing BHs if fBH_target > fBH_current: + + if self.natal_kicks: + nkwmssgs = ("; after natal kicks", + " or turn off / alter natal kicks") + else: + nkwmssgs = ("", "") + mssg = ( f"Target `f_BH` ({fBH_target}) is greater than " - f"f_BH formed at t={ti:.1f} ({fBH_current:.3f}; " - f"after natal kicks). Reduce `f_BH`, alter IMF " - f"slopes or turn off / alter natal kicks." + f"f_BH formed at t={ti:.1f} ({fBH_current:.3f}" + f"{nkwmssgs[0]}). Reduce `f_BH` or alter IMF " + f"slopes{nkwmssgs[1]}." ) + if self.strict_BH_target: raise ValueError(mssg) + else: # Will pass harmlessly through _dyn_eject_BH mssg += (" Proceeding anyways due to " @@ -1518,6 +1603,18 @@ def from_IMF(cls, IMF, FeH, *, natal_kicks=True, # Optionally perform natal kicks if natal_kicks: + # If f_kick, try to determine kick scale automatically + if f_kick is not None: + try: + kick_slope, kick_scale = kicks._determine_kick_params( + kick_method, f_kick, IMF, _ifmr, + kick_slope, kick_scale + ) + except TypeError as err: + mssg = (f"Can only use `f_kick` with methods that use " + f"`scale` and `slope` params, not '{kick_method}'.") + raise ValueError(mssg) from err + match kick_method.casefold(): case 'maxwellian' | 'f12' | 'fryer2012': @@ -1525,14 +1622,18 @@ def from_IMF(cls, IMF, FeH, *, natal_kicks=True, from .ifmr import _check_IFMR_FeH_bounds FeH_BH = _check_IFMR_FeH_bounds(FeH) - kick_kw = dict(method=kick_method, f_kick=f_kick, vesc=vesc, + kick_kw = dict(method=kick_method, vesc=vesc, FeH=FeH_BH, vdisp=kick_vdisp, SNe_method=SNe_method) case 'sigmoid' | 'tanh': - kick_kw = dict(method=kick_method, f_kick=f_kick, + + kick_kw = dict(method=kick_method, slope=kick_slope, scale=kick_scale) + case 'full' | 'everything' | 'all' | 'none': + kick_kw = dict() + case _: mssg = f"Invalid natal kick algorithm '{kick_method=}'" raise ValueError(mssg) @@ -1658,7 +1759,7 @@ def from_powerlaw(cls, m_breaks, a_slopes, nbins, FeH, N0=5e5, *, imf = PowerLawIMF(m_break=m_breaks, a=a_slopes, N0=N0, ext='zeros') - return cls.from_IMF(imf, nbins, FeH, N0=N0, natal_kicks=natal_kicks, + return cls.from_IMF(imf, FeH, nbins=nbins, natal_kicks=natal_kicks, binning_breaks=binning_breaks, binning_method=binning_method, **kwargs) @@ -1735,4 +1836,3 @@ def from_BHMF(cls, m_breaks, a_slopes, nbins, FeH, N0=1000, *, N_BH, M_BH, _ = MF.binned_eval(bins) return cls(M_BH, N_BH, bins) - diff --git a/ssptools/ifmr.py b/ssptools/ifmr.py index 3aebf61..5485f3a 100644 --- a/ssptools/ifmr.py +++ b/ssptools/ifmr.py @@ -676,6 +676,8 @@ def __init__(self, FeH, *, NS_mass=1.4, def predict_type(self, m_in): '''Predict the remnant type (WD, NS, BH) given the initial mass(es)''' + # TODO note that m_in=nan will always return a WD, should maybe error? + rem_type = np.where( m_in >= self.BH_mi[0], 'BH', np.where( diff --git a/ssptools/kicks.py b/ssptools/kicks.py index 40b53bf..05d44a8 100644 --- a/ssptools/kicks.py +++ b/ssptools/kicks.py @@ -107,7 +107,7 @@ def _maxwellian_retention_frac(m, vesc, FeH, vdisp=265., *, SNe_method='rapid'): drawn from a Maxwellian distribution with a certain kick dispersion scaled down by a fallback fraction, as described by Fryer et al. (2012). - The fraction of black holes retained in each mass bin is then found by + The fraction of black holes retained is then found by integrating the kick velocity distribution from 0 to the estimated initial system escape velocity. In other words, by evaluating the CDF at the escape velocity. @@ -116,6 +116,7 @@ def _maxwellian_retention_frac(m, vesc, FeH, vdisp=265., *, SNe_method='rapid'): ---------- m : float The mean initial mass of the progenitor star which will make a BH. + Used to determine the fallback fraction. vesc : float The initial escape velocity of the cluster. @@ -126,7 +127,7 @@ def _maxwellian_retention_frac(m, vesc, FeH, vdisp=265., *, SNe_method='rapid'): SNe_method : {'rapid', 'delayed', 'NS', 'none'}, optional Which method to use to determine the fallback fraction as a function of - the black hole mass, which scales the dispersion as σ(1-fb). + the initial mass, which scales the dispersion as σ(1-fb). Available methods include the "rapid" (default) or "delayed" supernovae prescriptions described by Fryer+2012, or the ratio of the neutron star to black hole mass. @@ -226,7 +227,7 @@ def _sigmoid_retention_frac(m, slope, scale): r'''Retention fraction alg. based on a paramatrized sigmoid function. This method is based on a simple parametrization of the relationship - between the retention fraction and the BH mass as a sigmoid function, + between the retention fraction and the initial mass as a sigmoid function, increasing smoothly between 0 and 1 around a scale mass. .. math:: @@ -243,7 +244,8 @@ def _sigmoid_retention_frac(m, slope, scale): The "slope" of the sigmoid function, defining the "sharpness" of the increase. A value of 0 is completely flat (at fret=erf(1)~0.85), an increasingly positive value approaches a step function at the scale - mass, and a negative value will retain more low-mass bins than high. + mass, and a negative value will retain more low-mass progenitor BHs + than high. scale : float The scale-mass of the sigmoid function, defining the approximate mass @@ -262,7 +264,7 @@ def _tanh_retention_frac(m, slope, scale): r'''Retention fraction alg. based on a paramatrized function of tanh. This method is based on a simple parametrization of the relationship - between the retention fraction and the BH mass as a sigmoid function, + between the retention fraction and the initial mass as a sigmoid function, namely the hyperbolic tangent, increasing smoothly between 0 and 1 and reaching 50% at the given scale mass. @@ -280,7 +282,8 @@ def _tanh_retention_frac(m, slope, scale): The "slope" of the sigmoid function, defining the "sharpness" of the increase. A value of 0 is completely flat (at 50% for all masses), an increasingly positive value approaches a step function at the scale - mass, and a negative value will retain more low-mass bins than high. + mass, and a negative value will retain more low-mass progenitor BHs + than high. scale : float The scale-mass of the sigmoid function, defining the approximate mass @@ -295,99 +298,79 @@ def _tanh_retention_frac(m, slope, scale): # return np.tanh(np.exp(slope * (m - scale))) # alternative -def kick_retention_fraction(m_ini, f_kick=None, method='fryer2012', **ret_kwargs): - ''' - - m_ini : float - The initial (ZAMS) mass of the progenitor star which will form a BH. - ''' - - f_ret = _get_kick_method(method) - - # If no given total kick fraction, use old-style of directly using f_ret - if f_kick is None: - - return f_ret(m_ini, **ret_kwargs) - - # return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, **ret_kwargs) +def kick_retention_fraction(m_ini, method='fryer2012', **ret_kwargs): + r'''Compute the retention fraction of BH natal-kicks from progenitor masses - else: - raise NotImplementedError('TODO') - - # get the BHMF using the new fast InitialBHPopulation and then do this - - # # Fit for the desired kick scale - # try: - # slp, scl = _determine_kick_params(Mr_BH, Nr_BH, f_ret, f_kick, - # **ret_kwargs) - # except TypeError as err: - # mssg = (f"Can only use `f_kick` with methods that use a `scale` and" - # f" `slope` parameter, not '{method}'.") - # raise ValueError(mssg) from err - - # # Compute the natal kicks based on fit scale - return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, slope=slp, scale=scl) - - - -# -------------------------------------------------------------------------- -# Performing natal kicks on BH mass functions -# -------------------------------------------------------------------------- + Determines the effects of BH natal-kicks, in the form of a retention + fraction to be used when creating the BH masses, based on the given natal + kick algorithm. + Various natal kick algorithms are currently available. All methods require + different arguments, which can be passed to `ret_kwargs`. + See the respective retention functions for details on these arguments and + the algorithms themselves. -def _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, **ret_kwargs): + Parameters + ---------- + m_ini : float or ndarray[float] + The initial (ZAMS) mass of the progenitor star which will form a BH. + Is passed to the requested retention fraction function. - c = Nr_BH > 0.1 - mr_BH = Mr_BH[c] / Nr_BH[c] - natal_ejecta = np.zeros_like(Mr_BH) - retention = np.full_like(Mr_BH, np.nan) + method : {'sigmoid', 'tanh', 'maxwellian'}, optional + Natal kick algorithm to use, defining the retention fraction as a + function of mean bin mass. Defaults to the Maxwellian method. - retention[c] = f_ret(mr_BH, **ret_kwargs) + **ret_kwargs : dict, optional + All other arguments are passed to the retention fraction function. - # keep track of how much we eject - natal_ejecta[c] = Mr_BH[c] * (1 - retention[c]) + Returns + ------- + f_rem : float + The retention fraction of the BH created by this input progenitor mass. - Mr_BH[c] *= retention[c] - Nr_BH[c] *= retention[c] + See Also + -------- + _maxwellian_retention_frac : Maxwellian retention fraction algorithm. + _sigmoid_retention_frac : Sigmoid retention fraction algorithm. + _tanh_retention_frac : Hyperbolic tangent retention fraction algorithm. + ''' - stats = KickStats( - retention=retention, mass_kicked=natal_ejecta, parameters=ret_kwargs - ) + f_ret = _get_kick_method(method) - return Mr_BH, Nr_BH, stats + return f_ret(m_ini, **ret_kwargs) -def _determine_kick_params(Mr_BH, Nr_BH, f_ret, f_target, slope, scale=10.): - ''' - Use a root finding algorithm to determine the value of `scale` needed in - order to eject a total fraction of the given BHs `f_target`, using the - given `f_ret` function to distribute the kicks among mass bins. - Note that while it says "params", right now can only compute the scale. - ''' +def _determine_kick_params(method, f_target, IMF, IFMR, slope, scale=10.): import scipy.optimize as opt - c = Nr_BH > 0.1 + f_ret = _get_kick_method(method) - m_BH = Mr_BH[c] / Nr_BH[c] - M_BH_tot = Mr_BH.sum() + # Get domain of progenitor masses that make BHs + mi = np.linspace(*IFMR.BH_mi, 50_000) + dm = mi[1] - mi[0] - f_ini = Mr_BH[c] / M_BH_tot + # Determine the final BH masses + mf = IFMR.predict(mi) - # f_ret = _tanh_retention_frac + M_BH_tot_ini = np.sum(mf * IMF.N(mi) * dm) # Keep first guess on optionally given scale scale = scale if scale is not None else 10.0 def target_fret(scl): - retention = f_ret(m_BH, scale=scl, slope=slope) + retention = f_ret(mi, scale=scl, slope=slope) - f_BH_final = (f_ini * (1 - retention)).sum(axis=0) + M_BH_tot = np.sum(retention * mf * IMF.N(mi) * dm) + + f_BH_final = 1 - (M_BH_tot / M_BH_tot_ini) return f_target - f_BH_final try: - sol = opt.root_scalar(target_fret, x0=scale, bracket=(-25, 75)) + sol = opt.root_scalar(target_fret, x0=scale, + bracket=(-25, IFMR.BH_mi.upper + 25)) + except ValueError as err: mssg = ("Root finder failed to find scale parameter matching target " f"{f_target}. 'f_target' or 'slope' need to be adjusted.") @@ -427,101 +410,6 @@ def _get_kick_method(method): return f_ret -def natal_kicks(Mr_BH, Nr_BH, f_kick=None, method='fryer2012', **ret_kwargs): - r'''Computes the effects of BH natal-kicks on the mass and number of BHs - - Determines the effects of BH natal-kicks, and distributes said kicks - throughout the different BH mass bins, based on the given natal kick - algorithm. In general, BHs are preferentially lost in the low mass - bins, with the lowest masses being entirely kicked, and the highest masses - being entirely retained. - - Two natal kick algorithms are currently available. Both methods require - different arguments, which can be passed to `ret_kwargs`. - See the respective retention functions for details on these arguments. - - The first is based on the assumption that the kick velocity is drawn from - a Maxwellian distribution with a certain kick dispersion (scaled down by - a “fallback fraction” interpolated from a grid of SSE models). The - fraction of black holes retained in each mass bin is then found by - integrating the kick velocity distribution from 0 to the estimated initial - system escape velocity. See Fryer et al. (2012) for more information. - - A second, more directly flexible method is not based on any modelled BH - physics, but simply determines the retention fraction of BHs in each bin - based on the simple sigmoid function - :math:`f_{ret}(m)=\operatorname{erf}\left(e^{a(m-b)}\right)`. - - Both input BH arrays are modified *in place*, as well as returned. - - Parameters - ---------- - Mr_BH : ndarray - Array[nbin] of the total initial masses of black holes in each - BH mass bin. - - Nr_BH : ndarray - Array[nbin] of the total initial numbers of black holes in each - BH mass bin. - - f_kick : float, optional - Unused. - - method : {'sigmoid', 'maxwellian'}, optional - Natal kick algorithm to use, defining the retention fraction as a - function of mean bin mass. Defaults to the Maxwellian method. - - **ret_kwargs : dict, optional - All other arguments are passed to the retention fraction function. - - Returns - ------- - Mr_BH : ndarray - Array[nbin] of the total final masses of black holes in each - BH mass bin, after natal kicks. - - Nr_BH : ndarray - Array[nbin] of the total final numbers of black holes in each - BH mass bin, after natal kicks. - - stats : KickStats - Statistics on how and how many BHs are kicked. - - Notes - ----- - This was the old method of applying natal kicks, when the ejections of BHs - were only considered after generating all of the masses. This relied on - faulty logic w.r.t. the fallback fractions and are no longer used, but left - here for posterity. This method is still valid for non-Maxwellian kicks. - - See Also - -------- - _maxwellian_retention_frac : Maxwellian retention fraction algorithm. - _sigmoid_retention_frac : Sigmoid retention fraction algorithm. - _tanh_retention_frac : Hyperbolic tangent retention fraction algorithm. - ''' - - f_ret = _get_kick_method(method) - - # If no given total kick fraction, use old-style of directly using f_ret - if f_kick is None: - return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, **ret_kwargs) - - else: - - # Fit for the desired kick scale - try: - slp, scl = _determine_kick_params(Mr_BH, Nr_BH, f_ret, f_kick, - **ret_kwargs) - except TypeError as err: - mssg = (f"Can only use `f_kick` with methods that use a `scale` and" - f" `slope` parameter, not '{method}'.") - raise ValueError(mssg) from err - - # Compute the natal kicks based on fit scale - return _unbound_natal_kicks(Mr_BH, Nr_BH, f_ret, slope=slp, scale=scl) - - # -------------------------------------------------------------------------- # Computing kick velocities directly # -------------------------------------------------------------------------- From 808c1bfc39e9dbc7dd3330ad8055b5fad90c6eee Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Mon, 4 May 2026 12:41:21 -0300 Subject: [PATCH 5/7] Fix bug in circular imports Also bump version a major release, due to backwards-incompatibility. --- README.md | 4 ++-- pyproject.toml | 4 ++-- ssptools/__init__.py | 2 +- ssptools/kicks.py | 3 +-- 4 files changed, 6 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 4792661..3f4d33f 100755 --- a/README.md +++ b/README.md @@ -38,14 +38,14 @@ import ssptools m_break, a_slopes, nbins = [0.08, 0.5, 150.], [+1.3, -2.3], [5, 30] pdmf = ssptools.EvolvedMF.from_powerlaw(m_break, a_slopes, nbins, - FeH=-1.5, tout=13000, Ndot=0, N0=1e6) + FeH=-1.5, tout=13000, esc_rate=0, N0=1e6) ``` Alternatively, an IMF class can be instantiated and used directly: ```python imf = ssptools.masses.PowerLawIMF(m_break, a_slopes, N0=1e6) -pdmf = ssptools.EvolvedMF(imf, nbins, FeH=-1.5, tout=[0, 1000, 13000], Ndot=0) +pdmf = ssptools.EvolvedMF(imf, nbins, FeH=-1.5, tout=[0, 1000, 13000], esc_rate=0) ``` See the documentation of each class for more details on all possible parameters. diff --git a/pyproject.toml b/pyproject.toml index 876df20..2b12c67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "astro-ssptools" -version = "2.1.4" +version = "3.0.0" description = "Simple Stellar Population Tools" authors = [ {name = "Eduardo Balbinot", email = "eduardo.balbinot@gmail.com"}, @@ -20,7 +20,7 @@ dependencies = [ ] classifiers = [ - "Development Status :: 4 - Beta", + "Development Status :: 5 - Production/Stable", "Programming Language :: Python", "Intended Audience :: Science/Research", "Topic :: Scientific/Engineering :: Astronomy", diff --git a/ssptools/__init__.py b/ssptools/__init__.py index 84c5886..94b0095 100755 --- a/ssptools/__init__.py +++ b/ssptools/__init__.py @@ -4,7 +4,7 @@ __author__ = """Eduardo Balbinot""" __email__ = "eduardo.balbinot@gmail.com" -__version__ = "2.1.4" +__version__ = "3.0.0" from .evolve_mf import * from .sev import * diff --git a/ssptools/kicks.py b/ssptools/kicks.py index 05d44a8..6e686b1 100644 --- a/ssptools/kicks.py +++ b/ssptools/kicks.py @@ -2,7 +2,6 @@ # -*- coding: utf-8 -*- from .ifmr import get_data -from .evolve_mf import InitialBHPopulation import dataclasses @@ -37,7 +36,7 @@ def no_kicks(cls, nmbin): ) @classmethod - def from_final(cls, Mr_BH, Nr_BH, ibh: InitialBHPopulation, parameters): + def from_final(cls, Mr_BH, Nr_BH, ibh, parameters): '''Compute the kick stats based on final BH arrays. This constructor (the main one, likely) uses the final amounts of BHs, From 4cdd3a25957ac1497f59147417aaa4b73094d085 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Tue, 5 May 2026 11:36:17 -0300 Subject: [PATCH 6/7] fix bug with 0-mass and very high mass BHs In IFMRS with a wider range of BH masses (like SEVN, with PISN), we both need to be sure we ignore the destroyed m=0 BHs, and need to make sure any bins constructed for the BHs can hold BHs of mass greater than `m_breaks[-1]`. --- ssptools/evolve_mf.py | 5 +++++ ssptools/masses.py | 1 + 2 files changed, 6 insertions(+) diff --git a/ssptools/evolve_mf.py b/ssptools/evolve_mf.py index 48a83f1..b94b107 100644 --- a/ssptools/evolve_mf.py +++ b/ssptools/evolve_mf.py @@ -1590,12 +1590,17 @@ def from_IMF(cls, IMF, FeH, *, natal_kicks=True, # ------------------------------------------------------------------ # Get domain of progenitor masses that make BHs + # TODO I could potentially get away with less point if logscaled mi mi = np.linspace(*_ifmr.BH_mi, 50_000) dm = mi[1] - mi[0] # Determine the final BH masses mf = _ifmr.predict(mi) + # Skip 0-mass remnants + mi = mi[mf > 0.0] + mf = mf[mf > 0.0] + # Determine which BH bin each BH will fall into # TODO this may be behaving weird when out of bounds? (!) inds = np.digitize(mf, np.r_[bins.lower, bins.upper[-1]]) - 1 diff --git a/ssptools/masses.py b/ssptools/masses.py index e1a0973..7f3de3a 100644 --- a/ssptools/masses.py +++ b/ssptools/masses.py @@ -547,6 +547,7 @@ def __init__(self, m_break, nbins, imf, ifmr, *, bins_BH = mbin(bins_MS.lower[BH_mask].copy(), bins_MS.upper[BH_mask].copy()) bins_BH.lower[0] = ifmr.BH_mf.lower + bins_BH.upper[-1] = ifmr.BH_mf.upper # Neutron Stars From 7bc5807284bffee433fc8816f6fae628a01fcee5 Mon Sep 17 00:00:00 2001 From: Nolan Dickson Date: Tue, 5 May 2026 17:31:51 -0300 Subject: [PATCH 7/7] update unit tests for new kicks, ibh --- ssptools/evolve_mf.py | 2 +- tests/test_evolve_mf.py | 107 +++++++++++++++++++- tests/test_kicks.py | 219 +++++++++------------------------------- 3 files changed, 152 insertions(+), 176 deletions(-) diff --git a/ssptools/evolve_mf.py b/ssptools/evolve_mf.py index b94b107..688cf6c 100644 --- a/ssptools/evolve_mf.py +++ b/ssptools/evolve_mf.py @@ -1583,7 +1583,7 @@ def from_IMF(cls, IMF, FeH, *, natal_kicks=True, bins = massbins.bins.BH - nbin_BH = bins.lower.size + nbin_BH = len(bins.lower) # ------------------------------------------------------------------ # Compute the initial amounts of BHs expected from the IMF and IFMR diff --git a/tests/test_evolve_mf.py b/tests/test_evolve_mf.py index 6ff77d9..2b1b86c 100644 --- a/tests/test_evolve_mf.py +++ b/tests/test_evolve_mf.py @@ -19,7 +19,7 @@ # TODO need a test for when multiple tout, as that can be unstable. DEFAULT_KWARGS = dict( IMF=DEFAULT_IMF, nbins=[5, 5, 20], FeH=-1.00, tout=[12_000], esc_rate=0., - N0=5e5, tcc=0.0, NS_ret=0.1, BH_ret_int=1.0, BH_ret_dyn=1.0, + N0=5e5, tcc=0.0, NS_ret=0.1, BH_ret_dyn=1.0, natal_kicks=False, vesc=90 ) @@ -376,3 +376,108 @@ def M_t(t): ydot = emf._derivs_esc(t, y) assert ydot == pytest.approx(expected) + + + +class TestInitialBHPopulation: + + bins = masses.mbin(lower=[0.1, 5.0, 10.0, 20.0, 50.], + upper=[5.0, 10.0, 20.0, 50., 150]) + + ibh_kw = dict(IMF=DEFAULT_IMF, natal_kicks=False, bins=bins) + + # ---------------------------------------------------------------------- + # Initialization routines + # ---------------------------------------------------------------------- + + # def test_from_bhmf + + @pytest.mark.parametrize('FeH', [-2, -1.5, -1, -0.5, 0.0]) + def test_from_imf(self, FeH): + + # Initialize Initial BH Population + + ibh = evolve_mf.InitialBHPopulation.from_IMF(FeH=FeH, **self.ibh_kw) + + # Recompute the amounts of BHs + + _ifmr = ifmr.IFMR(FeH=FeH) + + mi = np.linspace(*_ifmr.BH_mi, 500_000) # use high resolution + dm = mi[1] - mi[0] + + # Determine the final BH masses + mf = _ifmr.predict(mi) + + # Skip 0-mass remnants + mi = mi[mf > 0.0] + mf = mf[mf > 0.0] + + Nbins = len(self.bins.lower) + + inds = np.digitize(mf, np.r_[self.bins.lower, self.bins.upper[-1]]) - 1 + + # natal kicks off, no frem term + expected_MBH = np.bincount(inds, weights=mf * DEFAULT_IMF.N(mi) * dm, + minlength=Nbins) + + expected_NBH = np.bincount(inds, weights=DEFAULT_IMF.N(mi) * dm, + minlength=Nbins) + + assert ibh.M == pytest.approx(expected_MBH, rel=0.01) + assert ibh.N == pytest.approx(expected_NBH, rel=0.01) + + # ---------------------------------------------------------------------- + # Properties and attributes + # ---------------------------------------------------------------------- + + @pytest.mark.parametrize( + 'FeH, expected_Mtot', + [ + (-2, 21843.38486), + (-1.5, 20827.24313), + (-1, 17298.61685), + (-0.5, 14599.08499), + (0.0, 7491.10354) + ], + ) + def test_Mtot(self, FeH, expected_Mtot): + ibh = evolve_mf.InitialBHPopulation.from_IMF(FeH=FeH, **self.ibh_kw) + + assert ibh.Mtot == pytest.approx(expected_Mtot, rel=0.001) + + @pytest.mark.parametrize( + 'FeH, expected_age', + [ + (-2, 8.85291), + (-1.5, 8.67542), + (-1, 8.46274), + (-0.5, 8.10353), + (0.0, 7.55464) + ], + ) + def test_age(self, FeH, expected_age): + ibh = evolve_mf.InitialBHPopulation.from_IMF(FeH=FeH, **self.ibh_kw) + + assert ibh.age == pytest.approx(expected_age, rel=0.0001) + + @pytest.mark.parametrize( + 'FeH, expected_Ms_lost', + [ + (-2, 44383.71036), + (-1.5, 43786.47009), + (-1, 43010.53556), + (-0.5, 42071.81981), + (0.0, 39949.14502) + ], + ) + def test_Ms_lost(self, FeH, expected_Ms_lost): + ibh = evolve_mf.InitialBHPopulation.from_IMF(FeH=FeH, **self.ibh_kw) + + assert ibh.Ms_lost == pytest.approx(expected_Ms_lost, rel=0.0001) + + + + + + diff --git a/tests/test_kicks.py b/tests/test_kicks.py index be7af09..4503585 100644 --- a/tests/test_kicks.py +++ b/tests/test_kicks.py @@ -1,11 +1,22 @@ #!/usr/bin/env python +import dataclasses + import pytest import numpy as np import scipy.special import scipy.integrate as integ -from ssptools import kicks +from ssptools import kicks, evolve_mf, masses, ifmr + + +DEFAULT_M_BREAK = [0.1, 0.5, 1.0, 100] + +DEFAULT_IMF = masses.PowerLawIMF( + m_break=DEFAULT_M_BREAK, a=[-0.5, -1.3, -2.5], N0=5e5 +) + +DEFAULT_IFMR = ifmr.IFMR(FeH=-1) class TestRetentionAlgorithms: @@ -83,192 +94,52 @@ def test_flat_retention_frac(self, value): assert fret == pytest.approx(value) - -class TestNatalKicks: - - # TODO these aren't really representative m_BH (m=M/N=.5,1,2) - @pytest.fixture() - def Mi(self): - return np.array([10., 10., 10.]) - - @pytest.fixture() - def Ni(self): - return np.array([20., 10., 5.]) - - # ---------------------------------------------------------------------- - # Test Maxwellian / Fryer+2012 natal kick algorithms - # ---------------------------------------------------------------------- - @pytest.mark.parametrize( - 'FeH, vesc, expected', + 'mi, rem_type, expected_frem', [ - (-1., 25., np.stack((np.array([0.002227, 0.002227, 0.002810]), - np.array([0.004454, 0.002227, 0.001405])))), - (-1., 100., np.stack((np.array([0.136963, 0.136963, 0.171672]), - np.array([0.273926, 0.136963, 0.085836])))), - (-1., 200., np.stack((np.array([0.966443, 0.966443, 1.186696]), - np.array([1.932887, 0.966443, 0.593348])))), - (0.3, 25., np.stack((np.array([0.002227, 0.002227, 0.002727]), - np.array([0.004454, 0.002227, 0.001363])))), - (0.3, 100., np.stack((np.array([0.136963, 0.136963, 0.166788]), - np.array([0.273926, 0.136963, 0.08339])))), - (0.3, 200., np.stack((np.array([0.966443, 0.966443, 1.156175]), - np.array([1.932887, 0.966443, 0.578087])))) + (0.5, 'WD', 1.0), + (1.0, 'WD', 1.0), + (1.4, 'NS', 0.1), + (5.0, 'BH', 0.01006), + (10.0, 'BH', 0.01006,), + (20.0, 'BH', 0.01229), + (50.0, 'BH', 1.0), + (100.0, 'BH', 1.0) ], ) - def test_F12_kicks_quantities(self, Mi, Ni, FeH, vesc, expected): + def test_frem(self, mi, rem_type, expected_frem): - Mf, Nf, _ = kicks.natal_kicks(Mi, Ni, method='F12', FeH=FeH, vesc=vesc) + @dataclasses.dataclass + class SpoofEMF: + NS_ret = 0.1 + natal_kicks = True + _kick_kw = dict(FeH=-1, method='fryer2012', vesc=90) - assert np.stack((Mf, Nf)) == pytest.approx(expected, rel=1e-3) + frem = evolve_mf.EvolvedMF._frem(SpoofEMF(), mi, rem_type) - @pytest.mark.parametrize( - 'FeH, vesc, expected', - [ - (-1., 25., 29.992735), - (-1., 100., 29.554400), - (-1., 200., 26.880415), - (0.3, 25., 29.992818), - (0.3, 100., 29.559284), - (0.3, 200., 26.910936), - ], - ) - def test_F12_kicks_total(self, Mi, Ni, FeH, vesc, expected): - - _, _, ks = kicks.natal_kicks(Mi, Ni, method='F12', FeH=FeH, vesc=vesc) - - assert ks.total_kicked == pytest.approx(expected, rel=1e-3) - - # ---------------------------------------------------------------------- - # Test sigmoid natal kick algorithm - # ---------------------------------------------------------------------- + assert frem == pytest.approx(expected_frem, abs=0.0001) @pytest.mark.parametrize( - 'slope, scale, expected', + 'f_target, expected_scale', [ - (0.5, 0., np.stack((np.array([9.306121, 9.802805, 9.998790]), - np.array([18.612243, 9.802805, 4.999395])))), - (0.5, 20., np.stack((np.array([0.000657, 0.000844, 0.001392]), - np.array([0.001315, 0.000844, 0.000696])))), - (0.5, 50., np.stack((np.array([0.0, 0.0, 0.0]), - np.array([0.0, 0.0, 0.0])))), - (10, 0., np.stack((np.array([10., 10., 10.]), - np.array([20., 10., 5.])))), - (10, 20., np.stack((np.array([0.0, 0.0, 0.0]), - np.array([0.0, 0.0, 0.0])))), - (10, 50., np.stack((np.array([0.0, 0.0, 0.0]), - np.array([0.0, 0.0, 0.0])))), + (0.001, 16.24579), + (0.1, 21.66143), + (0.5, 35.07088), + (0.8, 62.84709), + (0.99, 97.70552) ], ) - def test_sigmoid_kicks_quantities(self, Mi, Ni, slope, scale, expected): - - Mf, Nf, _ = kicks.natal_kicks(Mi, Ni, method='sigmoid', - slope=slope, scale=scale) - - assert np.stack((Mf, Nf)) == pytest.approx(expected, abs=1e-3) - - @pytest.mark.parametrize( - 'slope, scale, expected', - [ - (0.5, 0, 0.892281), - (0.5, 20, 29.997105), - (0.5, 50, 29.999999), - (10, 0, 0.0), - (10, 20, 30.0), - (10, 50, 30.0) - ], - ) - def test_sigmoid_kicks_total(self, Mi, Ni, slope, scale, expected): - - _, _, ks = kicks.natal_kicks(Mi, Ni, method='sigmoid', - slope=slope, scale=scale) - - assert ks.total_kicked == pytest.approx(expected, rel=1e-3) - - # ---------------------------------------------------------------------- - # Test tanh natal kick algorithm - # ---------------------------------------------------------------------- - - @pytest.mark.parametrize( - 'slope, scale, expected', - [ - (0.5, 0., np.stack((np.array([6.224593, 7.310585, 8.807970]), - np.array([12.449186, 7.310585, 4.403985])))), - (0.5, 20., np.stack((np.array([3.3982e-08, 5.6027e-08, 1.5229e-07]), - np.array([6.7965e-08, 5.603e-08, 7.615e-08])))), - (0.5, 50., np.stack((np.array([0., 0., 0.]), - np.array([0., 0., 0.])))), - (10., 0., np.stack((np.array([9.999546, 9.999999, 10.]), - np.array([19.999092, 9.999999, 5.])))), - (10., 20., np.stack((np.array([0., 0., 0.]), - np.array([0., 0., 0.])))), - (10., 50., np.stack((np.array([0., 0., 0.]), - np.array([0., 0., 0.])))) - ], - ) - def test_tanh_kicks_quantities(self, Mi, Ni, slope, scale, expected): - - Mf, Nf, _ = kicks.natal_kicks(Mi, Ni, method='tanh', - slope=slope, scale=scale) - - assert np.stack((Mf, Nf)) == pytest.approx(expected, abs=1e-3) - - @pytest.mark.parametrize( - 'slope, scale, expected', - [ - (0.5, 0, 7.65685), - (0.5, 20, 29.99999), - (0.5, 50, 30.0), - (10, 0, 0.000454), - (10, 20, 30.0), - (10, 50, 30.0) - ], - ) - def test_tanh_kicks_total(self, Mi, Ni, slope, scale, expected): - - _, _, ks = kicks.natal_kicks(Mi, Ni, method='tanh', - slope=slope, scale=scale) - - assert ks.total_kicked == pytest.approx(expected, rel=1e-3) - - # ---------------------------------------------------------------------- - # Test natal kick algorithm when explicitly specifying f_kick - # ---------------------------------------------------------------------- - - @pytest.mark.parametrize( - 'slope, f_kick, expected', - [ - (0.5, 0.1, np.stack((np.array([8.421878, 8.979450, 9.598670]), - np.array([16.843757, 8.979450, 4.799335])))), - (0.5, 0.5, np.stack((np.array([3.40972, 4.603431, 6.986839]), - np.array([6.819459, 4.603431, 3.493419])))), - (0.5, 0.9, np.stack((np.array([0.464521, 0.743462, 1.792015]), - np.array([0.929043, 0.743462, 0.896007])))), - (10., 0.1, np.stack((np.array([7.000194, 9.999805, 10.]), - np.array([14.00038, 9.999805, 5.])))), - (10., 0.5, np.stack((np.array([4.5389e-04, 4.999546, 9.999999]), - np.array([9.0779e-04, 4.999546, 4.999999])))), - (10., 0.9, np.stack((np.array([4.0134e-13, 8.8335e-09, 2.999999]), - np.array([8.0269e-13, 8.8335e-09, 1.500000])))) - ], - ) - def test_expl_fkick_kicks_quantities(self, Mi, Ni, slope, f_kick, expected): - - Mf, Nf, _ = kicks.natal_kicks(Mi, Ni, method='tanh', f_kick=f_kick, - slope=slope) - - assert np.stack((Mf, Nf)) == pytest.approx(expected, abs=1e-3) - - @pytest.mark.parametrize('slope', [0.5, 10.]) - @pytest.mark.parametrize('f_kick', [0.1, 0.5, 0.9]) - def test_expl_fkick_kicks_total(self, Mi, Ni, slope, f_kick): - - expected = f_kick * Mi.sum() - - _, _, ks = kicks.natal_kicks(Mi, Ni, method='tanh', f_kick=f_kick, - slope=slope) + def test_determine_params(self, f_target, expected_scale): + + _, scale = kicks._determine_kick_params( + method='tanh', + f_target=f_target, + IMF=DEFAULT_IMF, + IFMR=DEFAULT_IFMR, + slope=0.5 + ) - assert ks.total_kicked == pytest.approx(expected, rel=1e-3) + assert scale == pytest.approx(expected_scale, abs=0.0001) class TestKickVelocities: