Skip to content

Clean up stellardist API #159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sybenzvi merged 2 commits into from
Jul 11, 2026
Merged

Clean up stellardist API #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce38111
Ensure monotonicity in CDF; clean up API.
Jul 11, 2026
24715cd
Make test and docs consistent with new API.
Jul 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions docs/source/nb/stellar_distributions.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,10 +61,10 @@
"outputs": [],
"source": [
"fd1 = FixedDistance(10*u.kpc)\n",
"d1 = fd1.distance(10000)\n",
"d1 = fd1.generate_distance(10000)\n",
"\n",
"fd2 = FixedDistance(10*u.kpc, sigma=0.5*u.kpc)\n",
"d2 = fd2.distance(10000)\n",
"d2 = fd2.generate_distance(10000)\n",
"\n",
"fig, axes = plt.subplots(1,2, figsize=(12,4), sharex=True)\n",
"ax1, ax2 = axes\n",
Expand Down Expand Up @@ -114,7 +114,7 @@
"for model in models:\n",
" sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
" sd = StellarDensity(sdfile)\n",
" ax.plot(sd.dist, sd.cdf, lw=3, alpha=0.7, label=sd.name)\n",
" ax.plot(sd.distance, sd.cdf, lw=3, alpha=0.7, label=sd.name)\n",
" \n",
"ax.set(xlabel='distance [kpc]',\n",
" ylabel='probability',\n",
Expand Down Expand Up @@ -145,7 +145,7 @@
"for model in models:\n",
" sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
" sd = StellarDensity(sdfile, add_LMC=True, add_SMC=True)\n",
" ax.plot(sd.dist, sd.cdf, lw=3, alpha=0.7, label=sd.name)\n",
" ax.plot(sd.distance, sd.cdf, lw=3, alpha=0.7, label=sd.name)\n",
"\n",
"ax.set(xlabel='distance [kpc]',\n",
" ylabel='probability',\n",
Expand Down Expand Up @@ -178,7 +178,7 @@
"for i, (model, ax) in enumerate(zip(models, axes)):\n",
" sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
" sd = StellarDensity(sdfile)\n",
" distances = sd.distance(100000)\n",
" distances = sd.generate_distance(100000)\n",
" \n",
" ax.hist(distances.value, bins=np.linspace(0., 30., 61), color='C{}'.format(i),\n",
" alpha=0.7,\n",
Expand Down Expand Up @@ -216,7 +216,7 @@
"for i, (model, ax) in enumerate(zip(models, axes)):\n",
" sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
" sd = StellarDensity(sdfile, add_LMC=True, add_SMC=True)\n",
" distances = sd.distance(100000)\n",
" distances = sd.generate_distance(100000)\n",
" \n",
" ax.hist(distances.value, bins=np.linspace(0, 70, 71), color=f'C{i}',\n",
" alpha=0.7,\n",
Expand Down
103 changes: 63 additions & 40 deletions python/asteria/stellardist.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ def __init__(self):
super().__init__()

@abstractmethod
def distance(self, size=1):
def generate_distance(self, size=1):
"""Generate distance to a progenitor.

Parameters
Expand Down Expand Up @@ -61,7 +61,7 @@ def __init__(self, d, sigma=None):
self.dist = d
self.sigma = sigma

def distance(self, size=1):
def generate_distance(self, size=1):
"""Generate distance to a progenitor.

Parameters
Expand Down Expand Up @@ -98,7 +98,7 @@ def __str__(self):
class StellarDensity(Distance):
"""Generate distances according to a Sun-centric radial mass density."""

def __init__(self, distcdf_file, add_LMC=False, add_SMC=False):
def __init__(self, distcdf_file, add_LMC=False, add_SMC=False, m_MW=2.012e10*u.Msun, m_LMC=2.7e9*u.Msun, m_SMC=3.1e8*u.Msun):
"""Progenitor distance with some uncertainty.

Parameters
Expand All @@ -109,61 +109,61 @@ def __init__(self, distcdf_file, add_LMC=False, add_SMC=False):
If true, add a Gaussian model of the LMC stellar distribution.
add_SMC : bool
If true, add a Gaussian model of the SMC stellar distribution.
m_MW : Quantity
Stellar mass of the Milky Way, excluding remnants. Default value is from Lian+ ApJL 37:990, 2025.
m_LMC : Quantity
Stellar mass of the LMC. Default value from van der Marel+ Proc. IAU Symp 256 4:81, 2009 (arXiv:0809.4268).
m_SMC : Quantity
Stellar mass of the SMC. Default value from van der Marel+ Proc. IAU Symp 256 4:81, 2009 (arXiv:0809.4268).
"""
super().__init__()

hdu = fits.open(distcdf_file)
distunit = u.Unit(hdu['DIST'].header['BUNIT'])
self.dist = hdu['DIST'].data
self.cdf = hdu['CDF'].data
self.name = hdu['CDF'].header['NAME']
self.publication = hdu['CDF'].header['PUB']
self.use_LMC = add_LMC
self.use_SMC = add_SMC

# Virial mass of the MW is ~1.5e12 Msun; see L. Watkins et al.,
# ApJ 873:118 (2019), arXiv:1804.11348 [astro-ph.GA].
# Assume 5% of this is stored in stellar mass (educated guess).
m_MW = 0.05 * 1.5e12
self._dist = hdu['DIST'].data
self._cdf = hdu['CDF'].data
self._name = hdu['CDF'].header['NAME']
self._publication = hdu['CDF'].header['PUB']
self._use_LMC = add_LMC
self._use_SMC = add_SMC

if add_LMC:
# Combined visible mass of LMC disk and neutral gas; see
# Kinematical Structure of the Magellanic System,
# R. P. van der Marel, N. Kallivayalil, G. Besla, Proc. IAU Symp.
# 256, 2009, 81-92, arXiv:0809.4268 [astro-ph].
m_LMC = 3.2e9
#- Treat the LMC as a Gaussian blob centered 50 kpc from Earth
r_LMC = (50*u.kpc).to(distunit).value
sigma_LMC = (2.5*u.kpc).to(distunit).value

mratio = m_LMC/m_MW
mratio = float(m_LMC/m_MW)
dist_LMC = np.linspace(r_LMC-5*sigma_LMC, r_LMC+5*sigma_LMC, 51)
cdf_LMC = mratio * norm.cdf(dist_LMC, r_LMC, sigma_LMC)

self.dist = np.append(self.dist, dist_LMC)
self.cdf = np.append(self.cdf, 1. + cdf_LMC)
self.cdf /= np.max(self.cdf)
self._dist = np.append(self._dist, dist_LMC)
self._cdf = np.append(self._cdf, 1. + cdf_LMC)

#- Cut repeat entries and ensure monotonicity in the CDF
self._dist, idx = np.unique(self._dist[self._dist.argsort(kind='stable')], return_index=True)
self._cdf = self._cdf[idx] / np.max(self._cdf[idx])
if add_SMC:
# Total stellar mass of the SMC is estimated to be 3.1x10^8 Msun;
# R. P. van der Marel, N. Kallivayalil, G. Besla, Proc. IAU Symp.
# 256, 2009, 81-92, arXiv:0809.4268 [astro-ph].
m_SMC = 3.1e8
#- Treat the LMC as a Gaussian blob centered 60 kpc from Earth
r_SMC = (60*u.kpc).to(distunit).value
sigma_SMC = (1.25*u.kpc).to(distunit).value

mratio = m_SMC/m_MW
mratio = float(m_SMC/m_MW)
dist_SMC = np.linspace(r_SMC-5*sigma_SMC, r_SMC+5*sigma_SMC, 51)
cdf_SMC = mratio * norm.cdf(dist_SMC, r_SMC, sigma_SMC)

self.dist = np.append(self.dist, dist_SMC)
self.cdf = np.append(self.cdf, 1. + cdf_SMC)
self.cdf /= np.max(self.cdf)
self._dist = np.append(self._dist, dist_SMC)
self._cdf = np.append(self._cdf, 1. + cdf_SMC)

#- Cut repeat entries and ensure monotonicity in the CDF
self._dist, idx = np.unique(self._dist[self._dist.argsort(kind='stable')], return_index=True)
self._cdf = self._cdf[idx] / np.max(self._cdf[idx])

self._inv_cdf = PchipInterpolator(self.cdf, self.dist)
self._inv_cdf = PchipInterpolator(self._cdf, self._dist)

# Set distance in distance units.
self.dist *= distunit
#- Set distance in distance units
self._dist *= distunit

def distance(self, size=1):
def generate_distance(self, size=1):
"""Generate distance to a progenitor.

Parameters
Expand All @@ -177,7 +177,7 @@ def distance(self, size=1):
Distance(s) to CCSN progenitor.
"""
u = np.random.uniform(0.,1., size)
return self._inv_cdf(u) * self.dist.unit
return self._inv_cdf(u) * self._dist.unit

def __str__(self):
"""Express as a string.
Expand All @@ -188,9 +188,32 @@ def __str__(self):
String representation of StellarDensity.
"""
s = 'Stellar density:\n'
s += '- model = {}\n'.format(self.name)
s += '- paper = {}\n'.format(self.publication)
s += '- use LMC = {}\n'.format(self.use_LMC)
s += '- use SMC = {}'.format(self.use_SMC)
s += '- model = {}\n'.format(self._name)
s += '- paper = {}\n'.format(self._publication)
s += '- use LMC = {}\n'.format(self._use_LMC)
s += '- use SMC = {}'.format(self._use_SMC)
return s

@property
def distance(self):
return self._dist

@property
def cdf(self):
return self._cdf

@property
def name(self):
return self._name

@property
def publiation(self):
return self._publication

@property
def use_LMC(self):
return self._use_LMC

@property
def use_SMC(self):
return self._use_SMC
6 changes: 3 additions & 3 deletions python/asteria/tests/test_stellardist.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,10 @@ class TestStellarDistributions(unittest.TestCase):

def test_fixed_distance(self):
fd = FixedDistance(10*u.kpc)
d = fd.distance()[0]
d = fd.generate_distance()[0]
self.assertTrue(d == 10*u.kpc)

d = fd.distance(size=5)
d = fd.generate_distance(size=5)
self.assertTrue(len(d) == 5)

def test_stellar_density(self):
Expand All @@ -22,5 +22,5 @@ def test_stellar_density(self):

sd = StellarDensity(stellar_file)

d = sd.distance()
d = sd.generate_distance()
self.assertTrue(np.abs(d.to('kpc').value - 8.853) < 1e-3)
Loading