def load_ic_mstar(self):
'''
separate C and M stars, sets their indicies as attributes: icstar and
imstar, will include artificial star tests (if there are any).
Trilegal simulation must have been done with -l and -a flags.
This is done using self.rec meaning use should be e.g.:
self.ast_mag2[self.rec][self.icstar]
Hard coded:
M star: C/O <= 1, LogL >= 3.3, and TPAGB flag
C star: C/O >= 1 and TPAGB flag
'''
if not hasattr(self, 'rec'):
self.rec = np.arange(len(self.data['CO']))
try:
co = self.data['CO'][self.rec]
except KeyError as e:
logger.error('No AGB stars. Trilegal must be run with -a')
raise e
logl = self.data['logL'][self.rec]
itpagb = trilegal.get_stage_label('TPAGB')
stage = self.data['stage'][self.rec]
self.imstar, = np.nonzero((co <= 1) & (logl >= 3.3) & (stage == itpagb))
self.icstar, = np.nonzero((co >= 1) & (stage == itpagb))
def stage_inds(self, name):
'''
indices where self.data['stage'] is [name] evolutionary stage
see trilegal.get_stage_label
'''
assert 'stage' in self.data.keys(), 'no stages marked in this file'
inds, = np.nonzero(self.data['stage'] == trilegal.get_stage_label(name))
return inds
def all_stages(self, *stages):
'''
add the indices of evolutionary stage(s) as an attribute i[stage]
'''
if stages is ():
stages = trilegal.get_stage_label()
for stage in stages:
i = self.stage_inds(stage)
self.__setattr__('i%s' % stage.lower(), i)
return
def scatter_hist(starpop, xdata, ydata, coldata='stage', xbins=50, ybins=50,
slice_inds=None, xlim=None, ylim=None, clim=None,
discrete=False, scatter_kw={}):
""" """
import palettable
import matplotlib.gridspec as gridspec
def side_hist(ax, data, cdata, bins=50, cols=None, binsx=True,
discrete=False, cbins=10):
"""
Add a histogram of data to ax made from inds of unique values of cdata
"""
if discrete:
if type(cbins) == int:
hist, cbins = np.histogram(cdata, bins=cbins)
dinds = np.digitize(cdata, cbins)
uinds = range(len(cbins))
else:
uinds = np.unique(cdata)
dinds = np.digitize(cdata, uinds)
cols = palettable.get_map('Spectral', 'Diverging',
len(uinds)).mpl_colors
for j, i in enumerate(uinds):
hist, hbins = np.histogram(data[dinds == i], bins=bins)
if binsx:
x = hbins[1:]
y = hist
else:
y = hbins[1:]
x = hist
ax.plot(x, y, linestyle='steps-pre', color=cols[j], zorder=100-j)
return ax
fig = plt.figure(figsize=(11, 7))
gs = gridspec.GridSpec(3, 3)
axt = plt.subplot(gs[0,:-1])
axr = plt.subplot(gs[1:,-1])
ax = plt.subplot(gs[1:,:-1])
if type(xdata) is str:
ax.set_xlabel(xdata.replace('_', '\_'))
xdata = starpop.data[xdata]
if type(ydata) is str:
ax.set_ylabel(ydata.replace('_', '\_'))
ydata = starpop.data[ydata]
collabel = None
if type(coldata) is str:
collabel = coldata.replace('_', '\_')
coldata = starpop.data[coldata]
if slice_inds is not None:
xdata = xdata[slice_inds]
ydata = ydata[slice_inds]
coldata = coldata[slice_inds]
scatter_kw = dict({'edgecolors': 'none', 'cmap': plt.cm.Spectral}.items() \
+ scatter_kw.items())
l = ax.scatter(xdata, ydata, c=coldata, **scatter_kw)
axt = side_hist(axt, xdata, coldata, bins=xbins, discrete=discrete)
axr = side_hist(axr, ydata, coldata, bins=ybins, binsx=False,
discrete=discrete)
axt.set_yscale('log')
axr.set_xscale('log')
if xlim is not None:
axt.set_xlim(xlim)
ax.set_xlim(axt.get_xlim())
if ylim is not None:
axr.set_ylim(ylim)
ax.set_ylim(axr.get_ylim())
if clim is not None:
l.set_clim(clim)
axt.set_ylabel('$\#$')
axr.set_xlabel('$\#$')
axt.tick_params(labelbottom=False, labeltop=True, right=False)
axr.tick_params(labelright=True, labelleft=False, top=False)
gs.update(left=0.1, right=0.9, wspace=0, hspace=0)
if collabel == 'stage':
axtr = plt.subplot(gs[0, -1])
cols = palettable.get_map('Spectral', 'Diverging',
9).mpl_colors
[axtr.plot(0, 0, label=get_stage_label()[i], color=cols[i])
for i in range(len(cols))]
axtr.tick_params(labelbottom=False, labelleft=False)
axtr.grid() # should turn it off!!
plt.legend(mode='expand', ncol=2, frameon=False)
axs = [axt, ax, axr, axtr]
else:
axs = [axt, ax, axr]
return axs
