def sfh(sim, filename=None, massform=True, clear=False, legend=False,
subplot=False, trange=False, bins=100, **kwargs):
'''
star formation history
**Optional keyword arguments:**
*trange*: list, array, or tuple
size(t_range) must be 2. Specifies the time range.
*bins*: int
number of bins to use for the SFH
*massform*: bool
decides whether to use original star mass (massform) or final star mass
*subplot*: subplot object
where to plot SFH
*legend*: boolean
whether to draw a legend or not
*clear*: boolean
if False (default), plot on the current axes. Otherwise, clear the figure first.
By default, sfh will use the formation mass of the star. In tipsy, this will be
taken from the starlog file. Set massform=False if you want the final (observed)
star formation history
**Usage:**
>>> import pynbody.plot as pp
>>> pp.sfh(s,linestyle='dashed',color='k')
'''
if subplot:
plt = subplot
else:
import matplotlib.pyplot as plt
if "nbins" in kwargs:
bins = kwargs['nbins']
if 'nbins' in kwargs:
bins = kwargs['nbins']
del kwargs['nbins']
if ((len(sim.g)>0) | (len(sim.d)>0)): simstars = sim.star
else: simstars = sim
if trange:
assert len(trange) == 2
else:
trange = [simstars['tform'].in_units(
"Gyr").min(), simstars['tform'].in_units("Gyr").max()]
binnorm = 1e-9 * bins / (trange[1] - trange[0])
trangefilt = filt.And(filt.HighPass('tform', str(trange[0]) + ' Gyr'),
filt.LowPass('tform', str(trange[1]) + ' Gyr'))
tforms = simstars[trangefilt]['tform'].in_units('Gyr')
if massform:
try:
weight = simstars[trangefilt][
'massform'].in_units('Msol') * binnorm
except (KeyError, units.UnitsException):
warnings.warn(
"Could not load massform array -- falling back to current stellar masses", RuntimeWarning)
weight = simstars[trangefilt]['mass'].in_units('Msol') * binnorm
else:
weight = simstars[trangefilt]['mass'].in_units('Msol') * binnorm
if clear:
plt.clf()
sfhist, thebins, patches = plt.hist(tforms, weights=weight, bins=bins,
histtype='step', **kwargs)
if not subplot:
# don't set the limits
#plt.ylim(0.0, 1.2 * np.max(sfhist))
plt.xlabel('Time [Gyr]', fontsize='large')
plt.ylabel('SFR [M$_\odot$ yr$^{-1}$]', fontsize='large')
else:
plt.set_ylim(0.0, 1.2 * np.max(sfhist))
# Make both axes have the same start and end point.
if subplot:
x0, x1 = plt.get_xlim()
else:
x0, x1 = plt.gca().get_xlim()
# add a z axis on top if it has not been already done by an earlier plot:
from pynbody.analysis import pkdgrav_cosmo as cosmo
c = cosmo.Cosmology(sim=sim)
old_axis = plt.gca()
if len(plt.gcf().axes)<2:
pz = plt.twiny()
labelzs = np.arange(5, int(sim.properties['z']) - 1, -1)
times = [13.7 * c.Exp2Time(1.0 / (1 + z)) / c.Exp2Time(1) for z in labelzs]
pz.set_xticks(times)
pz.set_xticklabels([str(x) for x in labelzs])
pz.set_xlim(x0, x1)
pz.set_xlabel('$z$')
plt.sca(old_axis)
if legend:
plt.legend(loc=1)
if filename:
logger.info("Saving %s", filename)
plt.savefig(filename)
return array.SimArray(sfhist, "Msol yr**-1"), array.SimArray(thebins, "Gyr")
def sfh(sim, filename=None, massform=True, clear=False, legend=False,
subplot=False, trange=False, bins=100, **kwargs):
'''
star formation history
**Optional keyword arguments:**
*trange*: list, array, or tuple
size(t_range) must be 2. Specifies the time range.
*bins*: int
number of bins to use for the SFH
*massform*: bool
decides whether to use original star mass (massform) or final star mass
*subplot*: subplot object
where to plot SFH
*legend*: boolean
whether to draw a legend or not
*clear*: boolean
if False (default), plot on the current axes. Otherwise, clear the figure first.
By default, sfh will use the formation mass of the star. In tipsy, this will be
taken from the starlog file. Set massform=False if you want the final (observed)
star formation history
**Usage:**
>>> import pynbody.plot as pp
>>> pp.sfh(s,linestyle='dashed',color='k')
'''
import matplotlib.pyplot as pyplot
if subplot:
plt = subplot
else:
plt = pyplot
if "nbins" in kwargs:
bins = kwargs['nbins']
if 'nbins' in kwargs:
bins = kwargs['nbins']
del kwargs['nbins']
if ((len(sim.g)>0) | (len(sim.d)>0)): simstars = sim.star
else: simstars = sim
if trange:
assert len(trange) == 2
else:
trange = [simstars['tform'].in_units(
"Gyr").min(), simstars['tform'].in_units("Gyr").max()]
binnorm = 1e-9 * bins / (trange[1] - trange[0])
trangefilt = filt.And(filt.HighPass('tform', str(trange[0]) + ' Gyr'),
filt.LowPass('tform', str(trange[1]) + ' Gyr'))
tforms = simstars[trangefilt]['tform'].in_units('Gyr')
if massform:
try:
weight = simstars[trangefilt][
'massform'].in_units('Msol') * binnorm
except (KeyError, units.UnitsException):
warnings.warn(
"Could not load massform array -- falling back to current stellar masses", RuntimeWarning)
weight = simstars[trangefilt]['mass'].in_units('Msol') * binnorm
else:
weight = simstars[trangefilt]['mass'].in_units('Msol') * binnorm
if clear:
plt.clf()
sfhist, thebins, patches = plt.hist(tforms, weights=weight, bins=bins,
histtype='step', **kwargs)
if not subplot:
# don't set the limits
#plt.ylim(0.0, 1.2 * np.max(sfhist))
plt.xlabel('Time [Gyr]', fontsize='large')
plt.ylabel('SFR [M$_\odot$ yr$^{-1}$]', fontsize='large')
else:
plt.set_ylim(0.0, 1.2 * np.max(sfhist))
# Make both axes have the same start and end point.
if subplot:
x0, x1 = plt.get_xlim()
else:
x0, x1 = plt.gca().get_xlim()
# add a z axis on top if it has not been already done by an earlier plot:
from pynbody.analysis import pkdgrav_cosmo as cosmo
c = cosmo.Cosmology(sim=sim)
old_axis = pyplot.gca()
pz = plt.twiny()
labelzs = np.arange(5, int(sim.properties['z']) - 1, -1)
times = [13.7 * c.Exp2Time(1.0 / (1 + z)) / c.Exp2Time(1) for z in labelzs]
pz.set_xticks(times)
pz.set_xticklabels([str(x) for x in labelzs])
pz.set_xlim(x0, x1)
pz.set_xlabel('$z$')
pyplot.sca(old_axis)
if legend:
plt.legend(loc=1)
if filename:
logger.info("Saving %s", filename)
plt.savefig(filename)
return array.SimArray(sfhist, "Msol yr**-1"), array.SimArray(thebins, "Gyr")
def make_figure_markers(means_mo, stddevs_mo, \
means_ma, stddevs_ma, \
means_w0, stddevs_w0, \
means_w1, stddevs_w1, \
means_w2, stddevs_w2, \
means_w3, stddevs_w3, \
means_w4, stddevs_w4, \
means_w5, stddevs_w5, \
title, xlabel, ylabel, ylim):
fig1 = plt.figure(figsize=(10, 10))
#plt.title(title)
plt.ylabel(xlabel)
plt.xlabel(ylabel)
plt.ylim(ylim)
start, end = plt.get_xlim()
plt.set_xticks(np.arange(start, end, 0.1))
plt.errorbar(range(len(means_mo)),
means_mo,
stddevs_mo,
capsize=5,
errorevery=2,
label='no attenuation')
plt.errorbar(range(len(means_ma)),
means_ma,
stddevs_ma,
capsize=5,
errorevery=3,
label='model')
plt.errorbar(range(len(means_w0)),
means_w0,
stddevs_w0,
capsize=5,
errorevery=4,
label='w0 set')
plt.errorbar(range(len(means_w1)),
means_w1,
stddevs_w1,
capsize=5,
errorevery=5,
label='w1 set')
plt.errorbar(range(len(means_w2)),
means_w2,
stddevs_w2,
capsize=5,
errorevery=6,
label='w2 set')
plt.errorbar(range(len(means_w3)),
means_w3,
stddevs_w3,
capsize=5,
errorevery=7,
label='w3 set')
plt.errorbar(range(len(means_w4)),
means_w4,
stddevs_w4,
capsize=5,
errorevery=8,
label='w4 set')
plt.errorbar(range(len(means_w5)),
means_w5,
stddevs_w5,
capsize=5,
errorevery=9,
label='w5 set')
plt.legend(ncol=2, loc='lower right')
filename = title + '.jpg'
plt.savefig(filename)
plt.close()
_table_stats = []
_table_p = []
_line_s = []
_line_p = []
_line_s.append('MO')
_line_p.append('MO')
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_mo, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_mo, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('MA')
_line_p.append('MA')
stat, p = stats.ttest_ind(means_ma, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# ma ma
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_ma, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_ma, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_ma, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_ma, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_ma, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_ma, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W0')
_line_p.append('W0')
stat, p = stats.ttest_ind(means_w0, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w0, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w0 w0
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_w0, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w0, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w0, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w0, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w0, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W1')
_line_p.append('W1')
stat, p = stats.ttest_ind(means_w1, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w1, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w1, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w1 w1
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_w1, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w1, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w1, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w1, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W2')
_line_p.append('W2')
stat, p = stats.ttest_ind(means_w2, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w2, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w2, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w2, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w2 w2
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_w2, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w2, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w2, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W3')
_line_p.append('W3')
stat, p = stats.ttest_ind(means_w3, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w3, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w3, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w3, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w3, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w3 w3
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_w3, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w3, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W4')
_line_p.append('W4')
stat, p = stats.ttest_ind(means_w4, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w4, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w4, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w4, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w4, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w4, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w4 w4
_line_s.append('-')
_line_p.append('-')
stat, p = stats.ttest_ind(means_w4, means_w5)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
_line_s.append('W5')
_line_p.append('W5')
stat, p = stats.ttest_ind(means_w5, means_mo)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_ma)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_w0)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_w1)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_w2)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_w3)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
stat, p = stats.ttest_ind(means_w5, means_w4)
_line_s.append(round(stat, 1))
_line_p.append(round(p, 4))
# w5 w5
_line_s.append('-')
_line_p.append('-')
_table_stats.append(_line_s)
_table_p.append(_line_p)
_line_s = []
_line_p = []
# print('Statistics=%.3f, p=%.3f' % (stat, p))
headers = ['', 'MO', 'MA', 'W0', 'W1', 'W2', 'W3', 'W4', 'W5']
np.savetxt(
title + '_stat.txt',
["%s" % tabulate.tabulate(_table_stats, headers, tablefmt="latex")],
fmt='%s')
np.savetxt(title + '_p.txt',
["%s" % tabulate.tabulate(_table_p, headers, tablefmt="latex")],
fmt='%s')
def sfh(sim,filename=None,massform=True,initstarmass=False,makeplot=True,
subplot=False, trange=False, bins=100, binsize=False, zmin=False,overplot=False,linestyle='-',color='k',linewidth=2,label=None,dored=True,**kwargs):
'''
star formation history
**Optional keyword arguments:**
*trange*: list, array, or tuple
size(t_range) must be 2. Specifies the time range.
*bins*: int
number of bins to use for the SFH
*label*: string
label line for legend
*zmin*: float
set min z to plot on second axes
*overplot*: bool
set to True if you are plotting a line on an already existing plot
*massform*: bool
decides whether to use original star mass (massform) or final star mass
*subplot*: subplot object
where to plot SFH
*legend*: boolean
whether to draw a legend or not
By default, sfh will use the formation mass of the star. In tipsy, this will be
taken from the starlog file. Set massform=False if you want the final (observed)
star formation history
**Usage:**
>>> import pynbody.plot as pp
>>> pp.sfh(s,linestyle='dashed',color='k')
'''
if subplot:
plt = subplot
else:
import matplotlib.pyplot as plt
if 'nbins' in kwargs:
bins=kwargs['nbins']
del kwargs['nbins']
if trange:
assert len(trange) == 2
else:
trange = [0,sim.star['tform'].in_units("Gyr").max()]
if binsize:
bins = int((trange[1] - trange[0])/binsize)
binnorm = 1./(1e9*binsize)
else:
binnorm = 1e-9*(bins / (trange[1] - trange[0]))
trangefilt = filt.And(filt.HighPass('tform',str(trange[0])+' Gyr'),
filt.LowPass('tform',str(trange[1])+' Gyr'))
tforms = sim.star[trangefilt]['tform'].in_units('Gyr')
if massform and not initstarmass:
try:
weight = sim.star[trangefilt]['massform'].in_units('Msol') * binnorm
except (KeyError, units.UnitsException) :
warnings.warn("Could not load massform array -- falling back to current stellar masses", RuntimeWarning)
weight = sim.star[trangefilt]['mass'].in_units('Msol') * binnorm
if initstarmass:
weight = np.zeros(np.size(tforms))
weight[:] = initstarmass*binnorm
if not initstarmass and not massform:
weight = sim.star[trangefilt]['mass'].in_units('Msol') * binnorm
if not makeplot:
sfhist, thebins = np.histogram(tforms, weights=weight, range=trange,bins=bins)
if makeplot:
sfhist, thebins, patches = plt.hist(tforms, weights=weight, range=trange,bins=bins,histtype='step',linestyle=linestyle,color=color,linewidth=linewidth,label=label)
plt.legend(loc='upper left',fontsize=20)
if not overplot:
if not subplot:
plt.ylim(0.0,1.2*np.max(sfhist))
plt.xlim(trange)
plt.xlabel('Time [Gyr]',fontsize=30)
plt.ylabel('SFR [M$_\odot$ yr$^{-1}$]',fontsize=30)
else:
plt.set_ylim(0.0,1.2*np.max(sfhist))
# Make both axes have the same start and end point.
if subplot: x0,x1 = plt.get_xlim()
else: x0,x1 = plt.gca().get_xlim()
from pynbody.analysis import pkdgrav_cosmo as cosmo
c = cosmo.Cosmology(sim=sim)
if dored:
pz = plt.twiny()
if not zmin:
labelzs = np.arange(10,int(sim.properties['z'])-1,-1)
else:
labelzs = np.arange(10,int(sim.properties['z'])-1,-1)
times = [13.7*c.Exp2Time(1.0 / (1+z))/c.Exp2Time(1) for z in labelzs]
pz.set_xticks(times)
pz.set_xticklabels([str(x) for x in labelzs])
pz.set_xlim(x0, x1)
pz.set_xlabel('Redshift',fontsize=30)
if (filename):
if config['verbose']: print "Saving "+filename
plt.savefig(filename)
return array.SimArray(sfhist, "Msol yr**-1"), array.SimArray(thebins, "Gyr")
def make_contour_plot(arr, xs, ys, x_range=None, y_range=None, nlevels = 20,
axbgclr = 'w', axfgclr = 'k',
logscale=True, xlogrange=False, ylogrange=False,
subplot=False, colorbar=False, ret_im=False, cmap=None,
clear=True,legend=False, scalemin = None, clip=False,
scalemax = None, filename = None, annotate = False, **kwargs) :
"""
Plot a contour plot of grid *arr* corresponding to bin centers
specified by *xs* and *ys*. Labels the axes and colobar with
proper units taken from x
Called by :func:`~pynbody.plot.generic.hist2d` and
:func:`~pynbody.plot.generic.gauss_density`.
**Input**:
*arr*: 2D array to plot
*xs*: x-coordinates of bins
*ys*: y-coordinates of bins
**Optional Keywords**:
*x_range*: list, array, or tuple (default = None)
size(x_range) must be 2. Specifies the X range.
*y_range*: tuple (default = None)
size(y_range) must be 2. Specifies the Y range.
*xlogrange*: boolean (default = False)
whether the x-axis should have a log scale
*ylogrange*: boolean (default = False)
whether the y-axis should have a log scale
*nlevels*: int (default = 20)
number of levels to use for the contours
*logscale*: boolean (default = True)
whether to use log or linear spaced contours
*colorbar*: boolean (default = False)
draw a colorbar
*scalemin*: float (default = arr.min())
minimum value to use for the color scale
*scalemax*: float (default = arr.max())
maximum value to use for the color scale
"""
from matplotlib import ticker, colors
arrows=kwargs.get('arrows',[])
if not subplot :
import matplotlib.pyplot as plt
else :
plt = subplot
if x_range is None:
if subplot: x_range=plt.get_xlim()
else : x_range = (np.min(xs),np.max(xs))
if y_range is None:
if subplot: y_range=plt.get_ylim()
else : y_range = (np.min(ys),np.max(ys))
if scalemin is None: scalemin = np.min(arr[arr>0])
if scalemax is None: scalemax = np.max(arr[arr>0])
if 'norm' in kwargs: del(kwargs['norm'])
# Adjust colormap
if cmap:
from mmlutils.mmlplot import get_cmaparray
cmap_colors=get_cmaparray(cmap)
axbgclr=cmap_colors[0]
axfgclr=cmap_colors[-1]
if logscale:
cmap=colors.ListedColormap(cmap_colors[int(0.01*len(cmap_colors)):])
cmap.set_under(color=axbgclr)
cmap.set_over(color=axfgclr)
cmap.set_bad(color=axbgclr) #no effect
levelmin=scalemin/100.
levelmax=scalemax*100.
if logscale:
try:
levels = np.logspace(np.log10(scalemin),np.log10(scalemax),nlevels)
cont_color=colors.LogNorm(clip=False)
except ValueError:
raise ValueError('crazy contour levels -- try specifying the *levels* keyword or use a linear scale')
return
if arr.units != NoUnit() and arr.units != 1 :
cb_label = '$log_{10}('+arr.units.latex()+')$'
else :
cb_label = '$log_{10}(N)$'
else:
levels = np.linspace(scalemin,scalemax,nlevels)
cont_color=None
if arr.units != NoUnit() and arr.units != 1 :
cb_label = '$'+arr.units.latex()+'$'
else :
cb_label = '$N$'
if not subplot and clear : plt.clf()
if logscale: arr=np.ma.masked_less_equal(arr,0).filled(scalemin/100.)
if ret_im :
if logscale: arr = np.log10(arr)
return plt.imshow(arr, origin='down', vmin=scalemin, vmax=scalemax,
aspect = 'auto',cmap=cmap,
#aspect = np.diff(x_range)/np.diff(y_range),cmap=cmap,
extent=[x_range[0],x_range[1],y_range[0],y_range[1]])
# Adjust to log then plot
if logscale: linarr=np.log10(arr) ; linlvl=np.log10(levels)
else : linarr=arr ; linlvl=levels
cs = plt.contourf(xs,ys,linarr, linlvl, norm=None,cmap=cmap,extend='both',**kwargs)
#cs = plt.contourf(xs,ys,arr, levels, norm=cont_color,cmap=cmap,**kwargs)
plt.gca().set_axis_bgcolor(axbgclr)
if kwargs.has_key('xlabel'):
xlabel = kwargs['xlabel']
else :
try:
if xlogrange: xlabel=r''+'$log_{10}('+xs.units.latex()+')$'
else : xlabel = r''+'$x/' + xs.units.latex() +'$'
except AttributeError:
xlabel = None
if xlabel :
try:
if subplot :
plt.set_xlabel(xlabel)
else:
plt.xlabel(xlabel)
except:
pass
if kwargs.has_key('ylabel'):
ylabel = kwargs['ylabel']
else :
try:
if ylogrange: ylabel='$log_{10}('+ys.units.latex()+')$'
else : ylabel = r''+'$y/' + ys.units.latex() +'$'
except AttributeError:
ylabel=None
if ylabel :
try:
if subplot :
plt.set_ylabel(ylabel)
else:
plt.ylabel(ylabel)
except:
pass
# if not subplot:
# plt.xlim((x_range[0],x_range[1]))
# plt.ylim((y_range[0],y_range[1]))
if colorbar :
cbfmt="%.2f" if logscale else "%.2e"
cb = plt.colorbar(cs, format = cbfmt).set_label(r''+cb_label)
#cb = plt.colorbar(cs, format = "%.2e").set_label(r''+cb_label)
if legend : plt.legend(loc=2)
# Annotate
if annotate:
imgsiz = (x_range[1]-x_range[0],y_range[1]-y_range[0])
# Arrow
from mmlutils.mmlmath import polar_square
dx=imgsiz[0]/40.
dy=imgsiz[1]/40.
for iarrow in arrows:
iphi=np.arctan((iarrow[1]/imgsiz[1])/(iarrow[0]/imgsiz[0]))
idir=polar_square(iphi)
ix=(idir[0]+0.5)*(imgsiz[0]-dx)+x_range[0]
iy=(idir[1]+0.5)*(imgsiz[1]-dy)+y_range[0]
plt.scatter(ix,iy,c=axfgclr,marker=(3,0,-90+iphi*180./np.pi),s=arr.shape[0])
plt.xlim((x_range[0],x_range[1]))
plt.ylim((y_range[0],y_range[1]))
# Strings
# imgsiz = (array.SimArray(x_range[1]-x_range[0],units=xs.units),
# array.SimArray(y_range[1]-y_range[0],units=ys.units))
xlocpad=imgsiz[0]/20.
ylocpad=imgsiz[1]/20.
for locstr in ['ul','ur','dl','dr']:
iloc=[0.,0.]
if locstr[1]=='l': iloc[0]=x_range[0]+xlocpad ; ha='left'
elif locstr[1]=='r': iloc[0]=x_range[1]-xlocpad ; ha='right'
if locstr[0]=='u': iloc[1]=y_range[1]-ylocpad ; va='top'
elif locstr[0]=='d': iloc[1]=y_range[0]+ylocpad ; va='bottom'
plt.text(iloc[0],iloc[1],kwargs.get(locstr+'str',''),color=axfgclr,
horizontalalignment=ha, verticalalignment=va)
# Save
if (filename):
if config['verbose']: print "Saving "+filename
plt.savefig(filename)
