Initial Mass Function (IMF)¶

airball provides an initial mass function class IMF for randomly generating stellar masses for different stellar environments. "The initial mass function (IMF) is an empirical function that describes the initial distribution of masses for a population of stars during star formation" (wikipedia).

For this example we'll import some necessary packages and make sure to use a colour blind friendly palette.

In [1]:

from pathlib import Path

from _notebook_utils import save_and_display_figure

import airball
import numpy as np
import time
import joblib

import matplotlib.pyplot as plt
import matplotlib.style as style
import matplotlib

matplotlib.rcParams["text.usetex"] = True
style.use("tableau-colorblind10")
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

images = Path.cwd() / "images/imf"
images.mkdir(parents=True, exist_ok=True)

from pathlib import Path from _notebook_utils import save_and_display_figure import airball import numpy as np import time import joblib import matplotlib.pyplot as plt import matplotlib.style as style import matplotlib matplotlib.rcParams["text.usetex"] = True style.use("tableau-colorblind10") %matplotlib inline %config InlineBackend.figure_format = 'retina' images = Path.cwd() / "images/imf" images.mkdir(parents=True, exist_ok=True)

Basics¶

The default IMF uses a piecewise combination of Salpeter (1955) for $m \geq 1\,M_\odot$ and Chabrier (2003) for single stars for $m < 1\,M_\odot$.

In [2]:

# Create an IMF with a specified mass range and default mass function
imf = airball.IMF(min_mass=0.08, max_mass=8)

# Generate a random mass from the IMF
random_mass = imf.random_mass()
print(random_mass)
# Output: e.g., 0.2319633663071667 solMass

# Create an IMF with a specified mass range and default mass function imf = airball.IMF(min_mass=0.08, max_mass=8) # Generate a random mass from the IMF random_mass = imf.random_mass() print(random_mass) # Output: e.g., 0.2319633663071667 solMass

0.15849485819172077 solMass

In [3]:

# Generate 5 random masses from the IMF
random_mass = imf.random_mass(size=5)
print(random_mass)
# Output: e.g., [1.1002683  0.94606418 0.7699541  1.1230717  0.13752477] solMass

# Generate 5 random masses from the IMF random_mass = imf.random_mass(size=5) print(random_mass) # Output: e.g., [1.1002683 0.94606418 0.7699541 1.1230717 0.13752477] solMass

[0.08096327 0.54059483 0.66969752 0.39351509 0.30088   ] solMass

In [4]:

# Compute the median mass for the IMF
imf.median_mass  # Output: e.g., 0.22984008 solMass

# Compute the median mass for the IMF imf.median_mass # Output: e.g., 0.22984008 solMass

Out[4]:

$0.22889492 \; \mathrm{M_{\odot}}$

Providing your own IMF¶

You can provide your own IMF when initializing an instance of the class by passing a function to the mass_function parameter. The following examples use an IMF using Salpeter (1955) for the entire mass range or have a uniform IMF. The IMF you provide must be a probability density function (PDF). airball will normalize the PDF for you over the given mass range. You can define an IMF as a standard python function or as a lambda function. Let's also set different mass ranges for each mass function.

You can also use any of the provided IMFs, see available IMFs.

In [5]:

# Define an IMF based on Salpeter (1955) and a (unphysical) uniform IMF.
def salpeter55(x):
    return x**-2.3

uniform = lambda x: 1
loguniform = lambda x: 1 / x

# Create an IMF with a specified mass range and default mass function
salpeter55_imf = airball.IMF(min_mass=0.1, max_mass=5, mass_function=salpeter55)
uniform_imf = airball.IMF(min_mass=0.01, max_mass=10, mass_function=uniform)
loguniform_imf = airball.IMF(min_mass=0.03, max_mass=1, mass_function=loguniform)

# Generate some random masses from the IMFs
print(salpeter55_imf.random_mass(5))
# Output: e.g., [2.13009562 0.10366307 0.91977187 0.11607485 0.11049281] solMass
print(uniform_imf.random_mass(5))
# Output: e.g., [4.81040127 4.13340906 3.64267182 3.37830285 3.81484321] solMass
print(loguniform_imf.random_mass(5))
# Output: e.g., [0.59206326 0.03541626 0.44852198 0.0262296  1.35840607] solMass

# Define an IMF based on Salpeter (1955) and a (unphysical) uniform IMF. def salpeter55(x): return x**-2.3 uniform = lambda x: 1 loguniform = lambda x: 1 / x # Create an IMF with a specified mass range and default mass function salpeter55_imf = airball.IMF(min_mass=0.1, max_mass=5, mass_function=salpeter55) uniform_imf = airball.IMF(min_mass=0.01, max_mass=10, mass_function=uniform) loguniform_imf = airball.IMF(min_mass=0.03, max_mass=1, mass_function=loguniform) # Generate some random masses from the IMFs print(salpeter55_imf.random_mass(5)) # Output: e.g., [2.13009562 0.10366307 0.91977187 0.11607485 0.11049281] solMass print(uniform_imf.random_mass(5)) # Output: e.g., [4.81040127 4.13340906 3.64267182 3.37830285 3.81484321] solMass print(loguniform_imf.random_mass(5)) # Output: e.g., [0.59206326 0.03541626 0.44852198 0.0262296 1.35840607] solMass

[1.63041333 1.37794687 0.68008313 0.10839793 0.14810343] solMass
[9.60785907 8.08928965 8.56781722 6.91317571 6.86433749] solMass
[0.13964636 0.06750722 0.07885003 0.43685646 0.13773602] solMass

/Users/zyrxvo/Documents/Research/airball/src/airball/imf.py:600: UserWarning: mass_function has no 'unit' attribute. Assuming IMF unit 'solMass'. Define a unit on your mass function to silence this warning.
  warnings.warn(

In [6]:

# Compute the median masses for the IMFs
print(f"Salpeter (1955) median mass: {salpeter55_imf.median_mass:1.5f}")
print(f"Uniform PDF median mass: {uniform_imf.median_mass:1.5f}")
print(f"Log-Uniform PDF median mass: {loguniform_imf.median_mass:1.5f}")

# Compute the median masses for the IMFs print(f"Salpeter (1955) median mass: {salpeter55_imf.median_mass:1.5f}") print(f"Uniform PDF median mass: {uniform_imf.median_mass:1.5f}") print(f"Log-Uniform PDF median mass: {loguniform_imf.median_mass:1.5f}")

Salpeter (1955) median mass: 0.16963 solMass
Uniform PDF median mass: 5.00500 solMass
Log-Uniform PDF median mass: 0.17321 solMass

We can visually compare the different distributions by plotting their PDFs. We can also see how the generating function corresponds to the PDF. For convenience, if we need to, we can call the IMF.masses function to generate a numpy.logspace array over the IMF mass range. This is especially useful when you want to bin samples over many orders of magnitude. However, airball will also automatically set the regions outside of the mass range to be zero for the PDF (and 0 or 1 for the CDF).

In [7]:

nsamples = 1e6
default_samples = imf.random_mass(nsamples).value
salpeter55_samples = salpeter55_imf.random_mass(nsamples).value
uniform_samples = uniform_imf.random_mass(nsamples).value
loguniform_samples = loguniform_imf.random_mass(nsamples).value

nsamples = 1e6 default_samples = imf.random_mass(nsamples).value salpeter55_samples = salpeter55_imf.random_mass(nsamples).value uniform_samples = uniform_imf.random_mass(nsamples).value loguniform_samples = loguniform_imf.random_mass(nsamples).value

In [8]:

nbins = 100
npoints = 1000
ms = np.geomspace(0.008, 15, npoints)
common = {"density": True, "alpha": 0.4, "log": True}

plt.rcParams.update({"font.size": 16})
fig, ax = plt.subplots(1, 1, figsize=(10, 5))

ax.hist(loguniform_samples, bins=loguniform_imf.masses(nbins), color="C9", **common)
ax.loglog(ms, loguniform_imf.pdf(ms), "C9", label="IMF (Log-Uniform)")

ax.hist(default_samples, bins=imf.masses(nbins), color="C0", **common)
ax.loglog(ms, imf.pdf(ms), "C0", label="IMF (Default)")

ax.hist(salpeter55_samples, bins=salpeter55_imf.masses(nbins), color="C1", **common)
ax.loglog(ms, salpeter55_imf.pdf(ms), "C1", label="IMF (Salpeter55)")

ax.hist(uniform_samples, bins=uniform_imf.masses(nbins), color="C3", **common)
ax.loglog(ms, uniform_imf.pdf(ms), "C3--", label="IMF (Uniform)")

ax.legend(prop={"size": 12})
ax.set_xlabel(f"Stellar Masses [{imf.unit}]")
ax.set_ylabel("IMF (PDF)")

fig.patch.set_facecolor("#E0E0E0")
filepath = images / "imf-pdfs.png"
alt_text = "A log-log figure showing four different IMF PDFs: Uniform, Log-Uniform, Salpeter (1955), and AIRBALL's Default IMF."
save_and_display_figure(filepath, alt_text, width=900)

nbins = 100 npoints = 1000 ms = np.geomspace(0.008, 15, npoints) common = {"density": True, "alpha": 0.4, "log": True} plt.rcParams.update({"font.size": 16}) fig, ax = plt.subplots(1, 1, figsize=(10, 5)) ax.hist(loguniform_samples, bins=loguniform_imf.masses(nbins), color="C9", **common) ax.loglog(ms, loguniform_imf.pdf(ms), "C9", label="IMF (Log-Uniform)") ax.hist(default_samples, bins=imf.masses(nbins), color="C0", **common) ax.loglog(ms, imf.pdf(ms), "C0", label="IMF (Default)") ax.hist(salpeter55_samples, bins=salpeter55_imf.masses(nbins), color="C1", **common) ax.loglog(ms, salpeter55_imf.pdf(ms), "C1", label="IMF (Salpeter55)") ax.hist(uniform_samples, bins=uniform_imf.masses(nbins), color="C3", **common) ax.loglog(ms, uniform_imf.pdf(ms), "C3--", label="IMF (Uniform)") ax.legend(prop={"size": 12}) ax.set_xlabel(f"Stellar Masses [{imf.unit}]") ax.set_ylabel("IMF (PDF)") fig.patch.set_facecolor("#E0E0E0") filepath = images / "imf-pdfs.png" alt_text = "A log-log figure showing four different IMF PDFs: Uniform, Log-Uniform, Salpeter (1955), and AIRBALL's Default IMF." save_and_display_figure(filepath, alt_text, width=900)

Out[8]:

In [9]:

plt.rcParams.update({"font.size": 16})
fig, ax = plt.subplots(1, 1, figsize=(10, 5))
common = {"density": True, "alpha": 0.4, "cumulative": True}

ax.hist(imf.random_mass(nsamples).value, bins=imf.masses(nbins), color="C0", **common)
ax.plot(ms, imf.cdf(ms), "C0", label="IMF CDF (Default)")

ax.hist(loguniform_imf.random_mass(nsamples).value, bins=loguniform_imf.masses(nbins), color="C9", **common)
ax.plot(ms, loguniform_imf.cdf(ms), "C9", label="IMF CDF (LogUniform)")

ax.hist(salpeter55_imf.random_mass(nsamples).value, bins=salpeter55_imf.masses(nbins), color="C1", **common)
ax.plot(ms, salpeter55_imf.cdf(ms), "C1", label="IMF CDF (Salpeter55)")

ax.hist(uniform_imf.random_mass(nsamples).value, bins=uniform_imf.masses(nbins), color="C3", **common)
ax.plot(ms, uniform_imf.cdf(ms), "C3--", label="IMF CDF (Uniform)")

ax.legend(loc=2, prop={"size": 12})
ax.set_xlabel(f"Stellar Masses [{imf.unit}]")
ax.set_ylabel("IMF (CDF)")
ax.set_xscale("log")

fig.patch.set_facecolor("#E0E0E0")
filepath = images / "imf-cdfs.png"
alt_text = "A log-log figure showing four different IMF CDFs: Uniform, Log-Uniform, Salpeter (1955), and AIRBALL's Default IMF."
save_and_display_figure(filepath, alt_text, width=900)

plt.rcParams.update({"font.size": 16}) fig, ax = plt.subplots(1, 1, figsize=(10, 5)) common = {"density": True, "alpha": 0.4, "cumulative": True} ax.hist(imf.random_mass(nsamples).value, bins=imf.masses(nbins), color="C0", **common) ax.plot(ms, imf.cdf(ms), "C0", label="IMF CDF (Default)") ax.hist(loguniform_imf.random_mass(nsamples).value, bins=loguniform_imf.masses(nbins), color="C9", **common) ax.plot(ms, loguniform_imf.cdf(ms), "C9", label="IMF CDF (LogUniform)") ax.hist(salpeter55_imf.random_mass(nsamples).value, bins=salpeter55_imf.masses(nbins), color="C1", **common) ax.plot(ms, salpeter55_imf.cdf(ms), "C1", label="IMF CDF (Salpeter55)") ax.hist(uniform_imf.random_mass(nsamples).value, bins=uniform_imf.masses(nbins), color="C3", **common) ax.plot(ms, uniform_imf.cdf(ms), "C3--", label="IMF CDF (Uniform)") ax.legend(loc=2, prop={"size": 12}) ax.set_xlabel(f"Stellar Masses [{imf.unit}]") ax.set_ylabel("IMF (CDF)") ax.set_xscale("log") fig.patch.set_facecolor("#E0E0E0") filepath = images / "imf-cdfs.png" alt_text = "A log-log figure showing four different IMF CDFs: Uniform, Log-Uniform, Salpeter (1955), and AIRBALL's Default IMF." save_and_display_figure(filepath, alt_text, width=900)

Out[9]:

Notice the difference between the uniform and log-uniform distributions. Consider which one you need for your use case.

Convergence¶

The IMF class uses scipy.interpolate.PchipInterpolator to fit a cubic spline to the PDF in log-mass space, then computes the CDF analytically via the spline's antiderivative. This may cause issues for particularly unwieldly PDFs. Random sampling uses a second PCHIP spline over the inverse CDF, evaluated directly for bulk draws. The default number of samples used to build the spline is $10^5$. This can be changed when initializing the IMF or afterwards by setting the interpolating_points property, which triggers a full recalculation (as does updating min_mass or max_mass). If you find the default precision insufficient, increasing interpolating_points will improve accuracy at the cost of a longer initialization time.

Let's go through an example where we have an exact solution to compare against. Thank you to Aleksey Generozov for making this faster!

In [10]:

npts = 100
pts = np.geomspace(10, 1e7, npts)
avg_rel_err = np.zeros(npts)
max_rel_err = np.zeros(npts)
walltimes = np.zeros(npts)


# Example (UNNORMALIZED) mass function to integrate
def f(m, A=0.5):
    return np.where(m < A, (m / A) ** (-1.3), (m / A) ** (-2.3))
f.unit = airball.units.solMass

# Integral of f above (UNNORMALIZED)
def cdf_exact(m, A=0.5, mmin=0.01):
    return np.where(
        m <= A,
        (-10 * A / 3) * (m / A) ** (-0.3) + (10 * A / 3) * (mmin / A) ** (-0.3),
        (-100 * A / 39 + (10 * A / 3) * (mmin / A) ** (-0.3))
        - (10 * A / 13) * (m / A) ** (-1.3),
    )


def compare(n):
    start_time = time.perf_counter()
    issue = airball.IMF(0.01, 100, mass_function=f, interpolating_points=n)
    walltime = (time.perf_counter() - start_time)  # seconds

    # skip over the first value in the CDF when we compute the relative error because it is identically zero.
    npoints = 1e6
    cdfx = cdf_exact(issue.masses(npoints))
    high_res = (cdfx / cdfx[-1])[1:]
    diff = np.abs((issue.cdf(issue.masses(npoints))[1:] - high_res) / high_res)
    avg_rel_err = np.nanmean(diff)
    max_rel_err = np.nanmax(diff)

    return avg_rel_err, max_rel_err, walltime


results = joblib.Parallel(n_jobs=-1)(joblib.delayed(compare)(n) for n in pts)
for i in range(len(pts)):
    avg_rel_err[i], max_rel_err[i], walltimes[i] = results[i]

npts = 100 pts = np.geomspace(10, 1e7, npts) avg_rel_err = np.zeros(npts) max_rel_err = np.zeros(npts) walltimes = np.zeros(npts) # Example (UNNORMALIZED) mass function to integrate def f(m, A=0.5): return np.where(m < A, (m / A) ** (-1.3), (m / A) ** (-2.3)) f.unit = airball.units.solMass # Integral of f above (UNNORMALIZED) def cdf_exact(m, A=0.5, mmin=0.01): return np.where( m <= A, (-10 * A / 3) * (m / A) ** (-0.3) + (10 * A / 3) * (mmin / A) ** (-0.3), (-100 * A / 39 + (10 * A / 3) * (mmin / A) ** (-0.3)) - (10 * A / 13) * (m / A) ** (-1.3), ) def compare(n): start_time = time.perf_counter() issue = airball.IMF(0.01, 100, mass_function=f, interpolating_points=n) walltime = (time.perf_counter() - start_time) # seconds # skip over the first value in the CDF when we compute the relative error because it is identically zero. npoints = 1e6 cdfx = cdf_exact(issue.masses(npoints)) high_res = (cdfx / cdfx[-1])[1:] diff = np.abs((issue.cdf(issue.masses(npoints))[1:] - high_res) / high_res) avg_rel_err = np.nanmean(diff) max_rel_err = np.nanmax(diff) return avg_rel_err, max_rel_err, walltime results = joblib.Parallel(n_jobs=-1)(joblib.delayed(compare)(n) for n in pts) for i in range(len(pts)): avg_rel_err[i], max_rel_err[i], walltimes[i] = results[i]

In [11]:

plt.rcParams.update({"font.size": 16})
fig, ax = plt.subplots(1, 2, figsize=(12, 5))
ax[0].loglog(pts, avg_rel_err, "+-", label="Average")
ax[0].loglog(pts, max_rel_err, "x-", label="Maximum")
ax[0].loglog(pts, 1 / pts**2, "k--", label=r"$\mathcal{O}(N^{-2})$")
ax[0].set_xlabel("Number of Sample Points")
ax[0].set_ylabel("Relative Error")
ax[0].legend(prop={"size": 12}, loc=3)
ax[0].axvline(1e5, ls=":", c="k")

ax[1].loglog(pts, walltimes, "C0.-", label="Walltime")
ax[1].loglog(pts, (walltimes[-1] / pts[-1]) * pts, "k-.", label=r"$\mathcal{O}(N)$")
ax[1].set_xlabel("Number of Sample Points")
ax[1].set_ylabel("Walltime [s]")
ax[1].set_ylim(np.min(walltimes)/2, np.max(walltimes) * 2)
ax[1].legend(prop={"size": 12})
ax[1].axvline(1e5, ls=":", c="k")

plt.subplots_adjust(wspace=0.25)
fig.patch.set_facecolor("#E0E0E0")
filepath = images / "interpolating-samples.png"
alt_text = "A two-panel log-log figure showing the relative error and walltime for generating the interpolating functions with respect to the number of sample points."
save_and_display_figure(filepath, alt_text, width=1000)

plt.rcParams.update({"font.size": 16}) fig, ax = plt.subplots(1, 2, figsize=(12, 5)) ax[0].loglog(pts, avg_rel_err, "+-", label="Average") ax[0].loglog(pts, max_rel_err, "x-", label="Maximum") ax[0].loglog(pts, 1 / pts**2, "k--", label=r"$\mathcal{O}(N^{-2})$") ax[0].set_xlabel("Number of Sample Points") ax[0].set_ylabel("Relative Error") ax[0].legend(prop={"size": 12}, loc=3) ax[0].axvline(1e5, ls=":", c="k") ax[1].loglog(pts, walltimes, "C0.-", label="Walltime") ax[1].loglog(pts, (walltimes[-1] / pts[-1]) * pts, "k-.", label=r"$\mathcal{O}(N)$") ax[1].set_xlabel("Number of Sample Points") ax[1].set_ylabel("Walltime [s]") ax[1].set_ylim(np.min(walltimes)/2, np.max(walltimes) * 2) ax[1].legend(prop={"size": 12}) ax[1].axvline(1e5, ls=":", c="k") plt.subplots_adjust(wspace=0.25) fig.patch.set_facecolor("#E0E0E0") filepath = images / "interpolating-samples.png" alt_text = "A two-panel log-log figure showing the relative error and walltime for generating the interpolating functions with respect to the number of sample points." save_and_display_figure(filepath, alt_text, width=1000)

Out[11]:

In [ ]: