def legend(self, ax, xlabel='', ylabel=''):
"""Draw a legend showing the colormap on the given axes.
"""
# Find extrema from normalization
xlim = [self.norm[0].vmin, self.norm[0].vmax]
ylim = [self.norm[1].vmin, self.norm[1].vmax]
# Create spacings and mesh
xx = zmath.spacing(xlim, num=self._LEGEND_DIV, scale=self.xscale)
yy = zmath.spacing(ylim, num=self._LEGEND_DIV, scale=self.yscale)
xx, yy = np.meshgrid(xx, yy)
# Use the colormap to convert to RGB values
cols = self.to_rgba(xx, yy)
# Reshape the colors array appropriately ``(N,3)``
cols = cols.reshape([cols.shape[0] * cols.shape[1], cols.shape[2]])
# `pcolormesh` uses y-positions as rows, and x-pos as columns so use transpose
im = ax.pcolormesh(xx.T, color=cols)
# Ticks
# -----
# This method of using `pcolormesh` loses the x and y value data
# fake the limits/scaling by placing custom tick marks
# Set x-ticks
xticks = np.linspace(0, self._LEGEND_DIV - 1, self._LEGEND_NTICKS)
xticklabels = zmath.spacing(xlim,
num=self._LEGEND_NTICKS,
scale=self.xscale)
if self.xscale.startswith('log'):
xticklabels = [zplot.strSciNot(tt, 0, 0) for tt in xticklabels]
else:
xticklabels = ["{:.2f}".format(tt) for tt in xticklabels]
# Set y-ticks
yticks = np.linspace(0, self._LEGEND_DIV - 1, self._LEGEND_NTICKS)
yticklabels = zmath.spacing(ylim,
num=self._LEGEND_NTICKS,
scale=self.yscale)
if self.yscale.startswith('log'):
yticklabels = [zplot.strSciNot(tt, 0, 0) for tt in yticklabels]
else:
yticklabels = ["{:.2f}".format(tt) for tt in yticklabels]
ax.set(xlabel=xlabel,
ylabel=ylabel,
xticks=xticks,
yticks=yticks,
xticklabels=xticklabels,
yticklabels=yticklabels)
return im
def plot_hour_vs_day(vals,
station=None,
clabel='',
ax=None,
levels=None,
smooth=3,
upsample=2,
cmap=None):
if ax is None:
fig, ax = plt.subplots(figsize=[12, 4])
if station is not None:
_draw_station_title(ax, station)
cmap = _def_temp_cmap(cmap)
vave = np.mean(vals, axis=0)
edges = [
np.arange(1, 365 + 1).astype(float),
np.arange(0, 24).astype(float)
]
if levels is None:
levels = zmath.spacing(vave / 5, integers=True, scale='lin') * 5
mesh = np.meshgrid(*edges, indexing='ij')
# pcm = ax.pcolormesh(vave.T)
pcm, data = _draw_hist2d(ax,
mesh,
vals,
smooth=smooth,
upsample=upsample,
cmap=cmap)
cbar = plt.colorbar(pcm, label=clabel)
# ax.axhline(12, color='0.5', ls='--')
ax.set_yticks(np.arange(0, 25, 6))
*_, cc = kale.plot.draw_contour2d(ax,
edges,
vave,
cmap=pcm.cmap.reversed(),
levels=levels,
smooth=smooth * 2,
upsample=upsample / 4,
pad=1,
cbar=cbar)
ax.set(xlim=[0, 365],
xlabel='Day of Year',
ylim=[0, 24],
ylabel='Hour of Day')
_draw_twiny_months(ax)
return ax, data, cc
def rad_proj_from_2d(path,
snap_num,
scatter_param='Density',
hist_param='Masses',
hist_dens=True,
extr=None,
nbins=20):
bin_sigma = 3
_params = ['Coordinates', 'Masses']
params = [pp for pp in _params]
if (scatter_param is not None) and (scatter_param not in _params):
params.append(scatter_param)
if (hist_param is not None) and (hist_param not in _params):
params.append(hist_param)
snap = lysis.readio.snapshot(path, snap_num=snap_num, pars=params)
pos = snap['Coordinates']
mass = snap['Masses']
scat = None
if scatter_param is not None:
scat = snap[scatter_param]
hist = None
if hist_param is not None:
hist = snap[hist_param]
com = np.sum(pos * mass[:, np.newaxis], axis=0) / np.sum(mass)
rads = np.linalg.norm(pos - com, axis=-1)
if extr is None:
extr = zmath.quantiles(rads, sigmas=[-bin_sigma, bin_sigma])
edges = None
if hist is not None:
stat = 'sum' if hist_dens else 'mean'
edges = zmath.spacing(extr, 'log', nbins + 1)
hist, *_ = sp.stats.binned_statistic(rads,
hist,
statistic=stat,
bins=edges)
if hist_dens:
da = np.diff(np.pi * edges**2)
hist = hist / da
return rads, scat, hist, edges
def main():
# Initialize MPI Parameters
# -------------------------
from mpi4py import MPI
from illpy_lib.illcosmo import Illustris_Cosmology_TOS
COSMO = Illustris_Cosmology_TOS()
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size
# name = MPI.Get_processor_name()
# stat = MPI.Status()
if rank == 0:
NAME = sys.argv[0]
print(("\n%s\n%s\n%s" % (NAME, '=' * len(NAME), str(datetime.now()))))
# Parse Arguments
# ---------------
args = _parseArguments()
RUN = args.RUN
VERBOSE = args.verbose
if args.check: CHECK_EXISTS = True
elif args.nocheck: CHECK_EXISTS = False
# Create Radial Bins (in simulation units)
radExtrema = np.array(
RAD_EXTREMA) / COSMO.CONV_ILL_TO_SOL.DIST.value # [pc] ==> [ill]
radBins = zmath.spacing(radExtrema, num=RAD_BINS)
# Master Process
# --------------
if rank == 0:
print(("RUN = %d " % (RUN)))
print(("VERSION = %.2f" % (_VERSION)))
print(("MPI COMM SIZE = %d " % (size)))
print("")
print(("VERBOSE = %s " % (str(VERBOSE))))
print(("CHECK_EXISTS = %s " % (str(CHECK_EXISTS))))
print("")
print(("RAD_BINS = %d " % (RAD_BINS)))
print(("RAD_EXTREMA = [%.2e, %.2e] [pc]" %
(RAD_EXTREMA[0], RAD_EXTREMA[1])))
print((" = [%.2e, %.2e] [sim]" %
(radExtrema[0], radExtrema[1])))
beg_all = datetime.now()
try:
_runMaster(RUN, comm)
except Exception as err:
_mpiError(comm, err)
# Check subhalo files to see if/what is missing
checkSubhaloFiles(RUN, verbose=VERBOSE, version=_VERSION)
end_all = datetime.now()
print((" - - Total Duration '%s'" % (str(end_all - beg_all))))
# Slave Processes
# ---------------
else:
try:
_runSlave(RUN, comm, radBins, verbose=True)
except Exception as err:
_mpiError(comm, err)
return
def subhaloRadialProfiles(run, snapNum, subhalo, radBins=None, nbins=NUM_RAD_BINS,
mostBound=None, verbose=True):
"""
Construct binned, radial profiles of density for each particle species.
Profiles for the velocity dispersion and gravitational potential are also constructed for
all particle types together.
Arguments
---------
run <int> : illustris simulation run number {1, 3}
snapNum <int> : illustris simulation snapshot number {1, 135}
subhalo <int> : subhalo index number for target snapshot
radBins <flt>[N] : optional, right-edges of radial bins in simulation units
nbins <int> : optional, numbers of bins to create if ``radBins`` is `None`
mostBound <int> : optional, ID number of the most-bound particle for this subhalo
verbose <bool> : optional, print verbose output
Returns
-------
radBins <flt>[N] : coordinates of right-edges of ``N`` radial bins
posRef <flt>[3] : coordinates in simulation box of most-bound particle (used as C.O.M.)
partTypes <int>[M] : particle type numbers for ``M`` types, (``illpy_lib.constants.PARTICLE``)
partNames <str>[M] : particle type strings for each type
numsBins <int>[M, N] : binned number of particles for ``M`` particle types, ``N`` bins each
massBins <flt>[M, N] : binned radial mass profile
densBins <flt>[M, N] : binned mass density profile
potsBins <flt>[N] : binned gravitational potential energy profile for all particles
dispBins <flt>[N] : binned velocity dispersion profile for all particles
"""
if verbose: print(" - - Profiler.subhaloRadialProfiles()")
if verbose: print(" - - - Loading subhalo partile data")
# Redirect output during this call
with zio.StreamCapture():
partData, partTypes = Subhalo.importSubhaloParticles(run, snapNum, subhalo, verbose=False)
partNums = [pd['count'] for pd in partData]
partNames = [PARTICLE.NAMES(pt) for pt in partTypes]
numPartTypes = len(partNums)
# Find the most-bound Particle
# ----------------------------
posRef = None
# If no particle ID is given, find it
if (mostBound is None):
# Get group catalog
mostBound = Subhalo.importGroupCatalogData(
run, snapNum, subhalos=subhalo, fields=[SUBHALO.MOST_BOUND])
if (mostBound is None):
warnStr = "Could not find mostBound particle ID Number!"
warnStr += "Run %d, Snap %d, Subhalo %d" % (run, snapNum, subhalo)
warnings.warn(warnStr, RuntimeWarning)
return None
thisStr = "Run %d, Snap %d, Subhalo %d, Bound ID %d" % (run, snapNum, subhalo, mostBound)
if verbose: print((" - - - - {:s} : Loaded {:s} particles".format(thisStr, str(partNums))))
# Find the most-bound particle, store its position
for pdat, pname in zip(partData, partNames):
# Skip, if no particles of this type
if (pdat['count'] == 0): continue
inds = np.where(pdat[SNAPSHOT.IDS] == mostBound)[0]
if (len(inds) == 1):
if verbose: print((" - - - Found Most Bound Particle in '{:s}'".format(pname)))
posRef = pdat[SNAPSHOT.POS][inds[0]]
break
# } pdat, pname
# Set warning and return ``None`` if most-bound particle is not found
if (posRef is None):
warnStr = "Could not find most bound particle in snapshot! %s" % (thisStr)
warnings.warn(warnStr, RuntimeWarning)
return None
mass = np.zeros(numPartTypes, dtype=object)
rads = np.zeros(numPartTypes, dtype=object)
pots = np.zeros(numPartTypes, dtype=object)
disp = np.zeros(numPartTypes, dtype=object)
radExtrema = None
# Iterate over all particle types and their data
# ==============================================
if verbose: print(" - - - Extracting and processing particle properties")
for ii, (data, ptype) in enumerate(zip(partData, partTypes)):
# Make sure the expected number of particles are found
if (data['count'] != partNums[ii]):
warnStr = "%s" % (thisStr)
warnStr += "Type '%s' count mismatch after loading!! " % (partNames[ii])
warnStr += "Expecting %d, Retrieved %d" % (partNums[ii], data['count'])
warnings.warn(warnStr, RuntimeWarning)
return None
# Skip if this particle type has no elements
# use empty lists so that call to ``np.concatenate`` below works (ignored)
if (data['count'] == 0):
mass[ii] = []
rads[ii] = []
pots[ii] = []
disp[ii] = []
continue
# Extract positions from snapshot, make sure reflections are nearest most-bound particle
posn = reflectPos(data[SNAPSHOT.POS], center=posRef)
# DarkMatter Particles all have the same mass, store that single value
if (ptype == PARTICLE.DM): mass[ii] = [GET_ILLUSTRIS_DM_MASS(run)]
else: mass[ii] = data[SNAPSHOT.MASS]
# Convert positions to radii from ``posRef`` (most-bound particle), and find radial extrema
rads[ii] = zmath.dist(posn, posRef)
pots[ii] = data[SNAPSHOT.POT]
disp[ii] = data[SNAPSHOT.SUBF_VDISP]
radExtrema = zmath.minmax(rads[ii], prev=radExtrema, nonzero=True)
# Create Radial Bins
# ------------------
# Create radial bin spacings, these are the upper-bound radii
if (radBins is None):
radExtrema[0] = radExtrema[0]*0.99
radExtrema[1] = radExtrema[1]*1.01
radBins = zmath.spacing(radExtrema, scale='log', num=nbins)
# Find average bin positions, and radial bin (shell) volumes
numBins = len(radBins)
binVols = np.zeros(numBins)
for ii in range(len(radBins)):
if (ii == 0): binVols[ii] = np.power(radBins[ii], 3.0)
else: binVols[ii] = np.power(radBins[ii], 3.0) - np.power(radBins[ii-1], 3.0)
# Bin Properties for all Particle Types
# -------------------------------------
densBins = np.zeros([numPartTypes, numBins], dtype=DTYPE.SCALAR) # Density
massBins = np.zeros([numPartTypes, numBins], dtype=DTYPE.SCALAR) # Mass
numsBins = np.zeros([numPartTypes, numBins], dtype=DTYPE.INDEX) # Count of particles
# second dimension to store averages [0] and standard-deviations [1]
potsBins = np.zeros([numBins, 2], dtype=DTYPE.SCALAR) # Grav Potential Energy
dispBins = np.zeros([numBins, 2], dtype=DTYPE.SCALAR) # Velocity dispersion
# Iterate over particle types
if verbose: print(" - - - Binning properties by radii")
for ii, (data, ptype) in enumerate(zip(partData, partTypes)):
# Skip if this particle type has no elements
if (data['count'] == 0): continue
# Get the total mass in each bin
numsBins[ii, :], massBins[ii, :] = zmath.histogram(rads[ii], radBins, weights=mass[ii],
edges='right', func='sum', stdev=False)
# Divide by volume to get density
densBins[ii, :] = massBins[ii, :]/binVols
if verbose: print((" - - - - Binned {:s} particles".format(str(np.sum(numsBins, axis=1)))))
# Consistency check on numbers of particles
# -----------------------------------------
# The total number of particles ``numTot`` shouldn't necessarily be in bins.
# The expected number of particles ``numExp`` are those that are within the bounds of bins
for ii in range(numPartTypes):
numExp = np.size(np.where(rads[ii] <= radBins[-1])[0])
numAct = np.sum(numsBins[ii])
numTot = np.size(rads[ii])
# If there is a discrepancy return ``None`` for error
if (numExp != numAct):
warnStr = "%s\nType '%s' count mismatch after binning!" % (thisStr, partNames[ii])
warnStr += "\nExpecting %d, Retrieved %d" % (numExp, numAct)
warnings.warn(warnStr, RuntimeWarning)
return None
# If a noticeable number of particles are not binned, warn, but still continue
elif (numAct < numTot-10 and numAct < 0.9*numTot):
warnStr = "%s : Type %s" % (thisStr, partNames[ii])
warnStr += "\nTotal = %d, Expected = %d, Binned = %d" % (numTot, numExp, numAct)
warnStr += "\nBin Extrema = %s" % (str(zmath.minmax(radBins)))
warnStr += "\nRads = %s" % (str(rads[ii]))
warnings.warn(warnStr, RuntimeWarning)
raise RuntimeError("")
# Convert list of arrays into 1D arrays of all elements
rads = np.concatenate(rads)
pots = np.concatenate(pots)
disp = np.concatenate(disp)
# Bin Grav Potentials
counts, aves, stds = zmath.histogram(rads, radBins, weights=pots,
edges='right', func='ave', stdev=True)
potsBins[:, 0] = aves
potsBins[:, 1] = stds
# Bin Velocity Dispersion
counts, aves, stds = zmath.histogram(rads, radBins, weights=disp,
edges='right', func='ave', stdev=True)
dispBins[:, 0] = aves
dispBins[:, 1] = stds
return radBins, posRef, mostBound, partTypes, partNames, \
numsBins, massBins, densBins, potsBins, dispBins
def plotScattering(sample, snap, mbhb, log, plotNames):
"""Illustrate the scattering calculation for a single, sample system.
Performs calculation by calling the 'scattering()' method, just like in
"MBHBinaryEvolution.py".
Arguments
---------
sample : int
Target galaxy/merger number to examine (this is the number out of *all* binaries,
not just the valid ones [included in ``mbhb.evolution``]).
snap : int
Illustris snapshot number {1, 135}.
mbhb : `Binaries.MBHBinaries` object
log : `logging.Logger` object
plotNames : str
Base name of plots.
Returns
-------
plotNames : list of str
Filenames of the plots created.
"""
log.debug("plotScattering()")
PLOT_A10 = True # LC Occupancy
PLOT_A11 = True # Flux
PLOT_A15 = True # Flux vs. Separation
PLOT_A02 = False # Model Galaxy
PLOT_A08 = False # Dist Func
from . import GM_Figures
figNames = []
# Initialization
# --------------
radialRange = np.array(mbhb.sets.RADIAL_RANGE_MODELS) * PC
numSeps = mbhb.sets.PLOT_SCATTERING_SAMPLE_SEPS
# Convert from `sample` number, of all binaries to index for valid (``mbhb.evolution``) ones
val_inds = np.where(mbhb.valid)[0]
valid_sample = np.where(val_inds == sample)[0]
if valid_sample.size == 1:
valid_sample = valid_sample[0]
else:
raise ValueError("`sample` '{}' returned '{}' from `val_inds`".format(sample, valid_sample))
mstar = mbhb.galaxies.mstar
log.debug(" - Sample subhalo %d, Snapshot %d" % (sample, snap))
# Binary Properties (snapshot dependent)
# `evolution` class includes only valid binaries, so use `valid_sample`
# m1 = np.max(mbhb.evolution.masses[valid_sample, snap])
# m2 = np.min(mbhb.evolution.masses[valid_sample, snap])
m1 = np.max(mbhb.initMasses[sample])
m2 = np.min(mbhb.initMasses[sample])
# Galaxy properties (snapshot independent)
# `galaxies` class includes *all* binaries, so use `sample` itself
gals = mbhb.galaxies
eps = gals.eps[sample]
# functions of energy
rads = gals.rads
periods = gals.perOrb[sample]
dist_func = gals.dist_func[sample]
diffCoef = gals.diffCoef[sample]
j2Circs = gals.j2Circ[sample]
dnStarsAll = gals.dnStarsAll[sample]
ndensStars = gals.densStars[sample]/mstar
radHard = gals.rads_hard[sample]
numStarsAll = sp.integrate.cumtrapz(dnStarsAll[::-1], eps[::-1], initial=0.0)[::-1]
loss_cone = Loss_Cone_Explicit(mbhb.sets, log)
# Wrap the Scattering function (for convenience)
def evalScattering(binSep):
# retvals = scattering(m1, m2, binSep, rads, eps, periods, j2Circs, diffCoef, dist_func,
# [sample], radHard, mbhb.sets)
retvals = loss_cone.harden(
m1, m2, binSep, rads, eps, periods, j2Circs, diffCoef, dist_func,
[sample], radHard, mbhb.sets)
radLC, j2LC, enrLC, dnStarsFLC, numStarsFLC, numStarsSSLC, \
dfStarsFLC, dfStarsSSLC, fluxStarsFLC, fluxStarsSSLC, flux, dadt_lc = retvals
return dnStarsFLC, numStarsFLC, dfStarsFLC, fluxStarsFLC, dfStarsSSLC, fluxStarsSSLC
num_flc = np.zeros(numSeps)
flx_flc = np.zeros(numSeps)
flx_sslc = np.zeros(numSeps)
dadt_sslc = np.zeros(numSeps)
dadt_flc = np.zeros(numSeps)
subPlotNames = zio.modify_filename(plotNames, prepend='stellar_scattering/')
zio.check_path(subPlotNames)
# Iterate over range of binary separations, plot profiles for each
log.debug(" - Calculting scattering for %d binary separations" % (numSeps))
flx_seps = zmath.spacing(radialRange, scale='log', num=numSeps)
for ii, binSep in enumerate(tqdm.tqdm(flx_seps, desc="Calculating scattering")):
dnStarsFLC, numStarsFLC, dfStarsFLC, fluxStarsFLC, dfStarsSSLC, fluxStarsSSLC = \
evalScattering(binSep)
hard_sslc = dadt_scattering(m1+m2, binSep, np.max(fluxStarsSSLC), mstar)
hard_flc = dadt_scattering(m1+m2, binSep, np.max(fluxStarsFLC), mstar)
eps_rad = zmath.spline(gals.rads, gals.eps[sample], log=True, pos=True, extrap=True)
sepEner = eps_rad(binSep)
# Plot Fig A10 - Loss-Cone Occupancy
if PLOT_A10:
fig = GM_Figures.figa10_lc_occupancy(
gals.rads, eps, dnStarsAll, dnStarsFLC, numStarsAll, numStarsFLC,
binr=binSep, bine=sepEner)
fname1 = subPlotNames + "lcPipe_figa10-occ_%02d.png" % (ii)
zplot.saveFigure(fig, fname1, verbose=False) # , log=log)
plt.close(fig)
# Plot Fig A11 - loss-cone fluxes
if PLOT_A11:
fig = GM_Figures.figa11_lc_flux(
eps, gals.rads, dfStarsFLC, dfStarsSSLC, fluxStarsFLC, fluxStarsSSLC,
binr=binSep, bine=sepEner)
fname2 = subPlotNames + "lcPipe_figa11-flx_%02d.png" % (ii)
zplot.saveFigure(fig, fname2, verbose=False) # , log=log)
plt.close(fig)
# Store values to arrays
num_flc[ii] = np.max(numStarsFLC)
flx_flc[ii] = np.max(fluxStarsFLC)
flx_sslc[ii] = np.max(fluxStarsSSLC)
dadt_sslc[ii] = hard_sslc
dadt_flc[ii] = hard_flc
# pbar.update(ii)
# pbar.finish()
if PLOT_A10:
log.debug(" - - Saved FigA10 to (e.g.) '%s'" % (fname1))
if PLOT_A11:
log.debug(" - - Saved FigA11 to (e.g.) '%s'" % (fname2))
# Plot Figs A15 - Scattering vs. Separations
if PLOT_A15:
log.debug(" - Plotting Fig15")
fig = GM_Figures.figA15_flux_sep(flx_seps, flx_flc, flx_sslc, dadt_flc, dadt_sslc,
sample, snap)
fname = plotNames + "lcPipe_figa15-sep.png"
zplot.saveFigure(fig, fname, verbose=False, log=log)
plt.close(fig)
figNames.append(fname)
# Plot Figure A02 - Density/Mass Profiles
if PLOT_A02:
log.info(" - Plotting FigA02")
dens = [gals.densStars[sample], gals.densDM[sample]]
mass = [gals.massStars[sample], gals.massDM[sample], gals.massTot[sample]]
names = ['Stars', 'DM', 'Total']
vels = [gals.vdisp_stars[sample]]
fig = GM_Figures.figA02_ModelGalaxy(gals.rads, dens, mass, vels, names=names)
fname = plotNames + "lcPipe_figa02-dens.png"
zplot.saveFigure(fig, fname, verbose=False, log=log)
plt.close(fig)
figNames.append(fname)
# Plot Figure A08 - Reconstructed Number Density
if PLOT_A08:
log.debug(" - Plotting FigA08")
df_pot = zmath.spline(eps, dist_func, mono=True, pos=True, sort=True, log=True)
dfErrs = np.zeros(np.size(eps))
fig = GM_Figures.figA08_distFunc(eps, gals.rads, ndensStars, dist_func, dfErrs, df_pot, log,
sample=sample)
fname = plotNames + "lcPipe_figa08-distfunc.png"
zplot.saveFigure(fig, fname, verbose=False, log=log)
plt.close(fig)
figNames.append(fname)
return figNames
def plot_grid(grid, grid_names, temps, valid, interp=None):
import matplotlib.pyplot as plt
import zcode.plot as zplot
extr = zmath.minmax(temps, filter='>')
smap = zplot.colormap(extr, 'viridis')
# bads = valid & np.isclose(temps, 0.0)
num = len(grid)
fig, axes = plt.subplots(figsize=[14, 14], nrows=num, ncols=num)
plt.subplots_adjust(hspace=0.4, wspace=0.4)
def_idx = [-4, -4, 4, -4]
for (ii, jj), ax in np.ndenumerate(axes):
if ii < jj:
ax.set_visible(False)
continue
ax.set(xscale='log', yscale='log')
xx = grid[jj]
if ii == jj:
# print(grid_names[ii], zmath.minmax(grid[ii], filter='>'))
# idx = list(range(num))
# idx.pop(ii)
# idx = tuple(idx)
# vals = np.mean(temps, axis=idx)
idx = [
slice(None) if aa == ii else def_idx[aa] for aa in range(num)
]
vals = temps[tuple(idx)]
ax.plot(xx, vals, 'k-')
if interp is not None:
num_test = 10
test = [
np.ones(num_test) * grid[aa][def_idx[aa]]
for aa in range(num)
]
test[ii] = zmath.spacing(grid[ii], 'log', num_test)
test_vals = [interp(tt) for tt in np.array(test).T]
ax.plot(test[ii], test_vals, 'r--')
# bad_vals = np.count_nonzero(bads, axis=idx)
# tw = ax.twinx()
# tw.plot(xx, bad_vals, 'r--')
else:
# print(ii, jj)
# print("\t", ii, grid_names[ii], zmath.minmax(grid[ii], filter='>'))
# print("\t", jj, grid_names[jj], zmath.minmax(grid[jj], filter='>'))
# idx = [0, 1, 2, 3]
# idx.pop(np.max([ii, jj]))
# idx.pop(np.min([ii, jj]))
# vals = np.mean(temps, axis=tuple(idx))
# idx = [slice(None) if aa in [ii, jj] else num//2 for aa in range(num)]
idx = [
slice(None) if aa in [ii, jj] else def_idx[aa]
for aa in range(num)
]
vals = temps[tuple(idx)]
if len(vals) == 0:
continue
yy = grid[ii]
xx, yy = np.meshgrid(xx, yy, indexing='ij')
ax.pcolor(xx, yy, vals, cmap=smap.cmap, norm=smap.norm)
if np.count_nonzero(vals > 0.0) == 0:
continue
tit = "{:.1e}, {:.1e}".format(*zmath.minmax(vals, filter='>'))
ax.set_title(tit, size=10)
# bad_vals = np.count_nonzero(bads, axis=tuple(idx))
# idx = (bad_vals > 0.0)
# aa = xx[idx]
# bb = yy[idx]
# cc = bad_vals[idx]
# ax.scatter(aa, bb, s=2*cc**2, color='0.5', alpha=0.5)
# ax.scatter(aa, bb, s=cc**2, color='r')
if interp is not None:
for kk in range(10):
idx = (vals > 0.0)
x0 = 10**np.random.uniform(
*zmath.minmax(np.log10(xx[idx])))
y0 = 10**np.random.uniform(
*zmath.minmax(np.log10(yy[idx])))
# y0 = np.random.choice(yy[idx])
temp = [grid[ll][def_idx[ll]] for ll in range(num)]
temp[ii] = y0
temp[jj] = x0
if temp[2] >= temp[3]:
temp[2] = 3.1
iv = interp(temp)
if not np.isfinite(iv) or np.isclose(iv, 0.0):
print("\nBAD")
print(temp)
print(iv)
for kk in range(num):
if def_idx[kk] == 0:
temp[kk] = temp[kk] * 1.11
elif def_idx[kk] == -1:
temp[kk] = 0.99 * temp[kk]
iv = interp(temp)
print("\t", temp)
print("\t", iv)
cc = smap.to_rgba(iv)
ss = 20
ax.scatter(temp[jj], temp[ii], color='0.5', s=2 * ss)
ax.scatter(temp[jj], temp[ii], color=cc, s=ss)
if ii == num - 1:
ax.set_xlabel(grid_names[jj])
if jj == 0 and ii != 0:
ax.set_ylabel(grid_names[ii])
return fig
def plot2DHist(ax, xvals, yvals, hist,
cax=None, cbax=None, cscale='log', cmap=None, smap=None, extrema=None,
contours=None, clabel={}, ccolors=None, clw=2.5, fs=None, rasterized=True,
scale='log', csmooth=None, cbg=True,
title=None, labels=None, overlay=None, overlay_fmt="",
cbar_kwargs={}, **kwargs):
"""Plot the given 2D histogram of data.
Use with (e.g.) ``numpy.histogram2d``,
Arguments
---------
ax : ``matplotlib.axes.Axes`` object
Axes object on which to plot.
xvals : (N,) array of scalars
Positions (edges) of x values, assumed to be the same for all rows of
the input data `hist`.
yvals : (M,) array of scalars
Positions (edges) of y values, assumed to be the same for all columns of
the input data `hist`.
hist : (N,M) ndarray of scalars
Grid of data-points to plot.
cax : `matplotlib.axes.Axes` object
Axes object on which to add a colorbar.
See the `cax` parameter of `plt.colorbar`.
cbax : `matplotlib.axes.Axes` object(s)
Axes object from which to make space for a colorbar axis.
See the `ax` parameter of `plt.colorbar`.
cscale : str
Scale to use for the colormap {'linear', 'log'}.
Overridden if `smap` is provided.
cmap : ``matplotlib.colors.Colormap`` object
Matplotlib colormap to use for coloring histogram.
Overridden if `smap` is provided.
fs : int
Fontsize specification.
title : str or `None`
Title to add to top of axes.
smap : `matplotlib.cm.ScalarMappable` object or `None`
A scalar-mappable object to use for colormaps, or `None` for one to be created.
extrema : (2,) array_like of scalars
Minimum and maximum values for colormap scaling.
contours : (L,) array_like of scalar or `None`
Histogram values at which to draw contours using the `plt.contour` `levels` argument.
clabel : dict or `None`
If `None`, no contours labels are drawn, otherwise labels are drawn on the contours,
where additional labeling parameters can be passed in the `clabel` dictionary.
labels : (2,) or (3,) array_like of strings
The first two string are the 'x' and 'y' axis labels respectively. If a third string is
provided it is used as the colorbar label.
overlay : (N,M) ndarray of int or `None`
Number of elements in each bin if overlaid-text is desired.
overlay_fmt : str
Format specification on overlayed values, e.g. "02d" (no colon or brackets).
Returns
-------
pcm : `matplotlib.collections.QuadMesh` object
The resulting plotted object, returned by ``ax.pcolormesh``.
smap : `matplotlib.cm.ScalarMappable` object
Colormap and color-scaling information. See: ``zcode.plot.plot_core.colormap``.
cbar : colorbar or `None`
cs : contours or `None`
"""
cblab = 'Counts'
xgrid, ygrid = np.meshgrid(xvals, yvals)
hist = np.asarray(hist)
if plot_core._scale_to_log_flag(cscale):
filter = 'g'
else:
filter = None
extrema = _set_extrema(extrema, hist, filter=filter)
# Make sure the given `scale` is valid
if np.size(scale) == 1:
scale = [scale, scale]
elif np.size(scale) != 2:
raise ValueError("`scale` must be one or two scaling specifications!")
if labels is not None:
if np.size(labels) >= 2:
ax.set_xlabel(labels[0], size=fs)
ax.set_ylabel(labels[1], size=fs)
if np.size(labels) > 2:
cblab = labels[2]
# Create scalar-mappable if needed
if smap is None:
smap = plot_core.colormap(extrema, cmap=cmap, scale=cscale)
# Plot
pcm = ax.pcolormesh(xgrid, ygrid, hist.T, norm=smap.norm, cmap=smap.cmap, linewidth=0,
rasterized=rasterized, vmin=smap.norm.vmin, vmax=smap.norm.vmax, **kwargs)
pcm.set_edgecolor('face')
# Add color bar
# -------------
cbar = None
if cbax is not None or cax is not None:
if cbax is not None:
cbar = plt.colorbar(smap, cax=cbax, **cbar_kwargs)
else:
cbar = plt.colorbar(smap, ax=cax, **cbar_kwargs)
if fs is not None:
cbar.ax.tick_params(labelsize=fs)
cbar.set_label(cblab, fontsize=fs)
else:
cbar.set_label(cblab)
ticks = [smap.norm.vmin, smap.norm.vmax]
ticks = zmath.spacing(ticks, cscale, integers=True)
cbar.ax.yaxis.set_ticks(smap.norm(ticks), minor=True)
if fs is not None:
ax.tick_params(labelsize=fs)
if title is not None:
ax.set_title(title, size=fs)
# Add overlay
# -----------
if overlay is not None:
form = "{:%s}" % (overlay_fmt)
overlay = np.asarray(overlay) # .astype(int)
# Make sure sizes are correct
if overlay.shape != hist.shape:
raise ValueError("shape of `overlay` ({}) must match `hist` ({})".format(
overlay.shape, hist.shape))
# Remember these are transposes
for ii in range(xgrid.shape[1] - 1):
for jj in range(xgrid.shape[0] - 1):
if scale[0].startswith('log'):
xx = np.sqrt(xgrid[jj, ii] * xgrid[jj, ii+1])
else:
xx = np.average([xgrid[jj, ii], xgrid[jj, ii+1]])
if scale[1].startswith('log'):
yy = np.sqrt(ygrid[jj, ii] * ygrid[jj+1, ii])
else:
yy = np.average([ygrid[jj, ii], ygrid[jj+1, ii]])
ax.text(xx, yy, form.format(overlay.T[jj, ii]), ha='center', va='center',
fontsize=8, bbox=dict(facecolor='white', alpha=0.2, edgecolor='none'))
# Add contour lines
# -----------------
cs = None
if contours is not None:
if isinstance(contours, bool) and contours:
levels = None
else:
levels = np.array(contours)
if csmooth is not None:
data = sp.ndimage.filters.gaussian_filter(hist.T, csmooth)
else:
data = hist.T
if cbg:
ax.contour(xgrid, ygrid, data, colors='0.50', norm=smap.norm,
levels=levels, linewidths=2*clw, antialiased=True, zorder=10, alpha=0.4)
if ccolors is None:
_kw = {'cmap': smap.cmap, 'norm': smap.norm}
else:
_kw = {'colors': ccolors}
cs = ax.contour(xgrid, ygrid, data, **_kw,
levels=levels, linewidths=clw, antialiased=True, zorder=11, alpha=0.8)
if levels is not None and clabel is not None:
clabel.setdefault('inline', True)
if fs is not None:
clabel.setdefault('fontsize', fs)
plt.clabel(cs, **clabel, zorder=50)
plot_core.set_lim(ax, 'x', data=xvals)
plot_core.set_lim(ax, 'y', data=yvals)
return pcm, smap, cbar, cs
def plot2DHistProj(xvals, yvals, weights=None, statistic=None, bins=10, filter=None, extrema=None,
cumulative=None,
fig=None, xproj=True, yproj=True, hratio=0.7, wratio=0.7, pad=0.0, alpha=1.0,
cmap=None, smap=None, type='hist', scale_to_cbar=True,
fs=12, scale='log', histScale='log', labels=None, cbar=True,
overlay=None, overlay_fmt=None,
left=_LEFT, bottom=_BOTTOM, right=_RIGHT, top=_TOP, lo=None, hi=None,
overall=False, overall_bins=20, overall_wide=False, overall_cumulative=False):
"""Plot a 2D histogram with projections of one or both axes.
Arguments
---------
xvals : (N,) array_like,
Values corresponding to the x-points of the given data
yvals : (N,) array_like,
Values corresponding to the y-points of the given data
weights : (N,) array_like or `None`,
Weights used to create histograms. If `None`, then counts are used.
statistic : str or `None`,
Type of statistic to be calculated, passed to ``scipy.stats.binned_statistic``.
e.g. {'count', 'sum', 'mean'}.
If `None`, then either 'sum' or 'count' is used depending on if `weights` are
provieded or not.
bins : int or [int, int] or array_like or [array, array],
Specification for bin sizes. integer values are treated as the number of bins to use,
while arrays are used as the bin edges themselves. If a tuple of two values is given, it
is assumed that the first is for the x-axis and the second for the y-axis.
filter : str or `None`, or [2,] tuple of str or `None`, or [3,] tubple of str or `None`
String specifying how to filter the input `data` relative to zero.
If this is a single value, it is applies to both `xvals` and `yvals`.
If this is a tuple/list of two values, they correspond to `xvals` and `yvals` respectively.
If `weights` are provided, the tuple/list should have three values.
extrema :
cumulative :
fig : ``matplotlib.figure.figure``,
Figure instance to which axes are added for plotting. One is created if not given.
xproj : bool,
Whether to also plot the projection of the x-axis (i.e. histogram ignoring y-values).
yproj : bool,
Whether to also plot the projection of the y-axis (i.e. histogram ignoring x-values).
hratio : float,
Fraction of the total available height-space to use for the primary axes object (2D hist)
wratio : float,
Fraction of the total available width-space to use for the primary axes object (2D hist)
pad : float,
Padding between central axis and the projected ones.
cmap : ``matplotlib.colors.Colormap`` object
Matplotlib colormap to use for coloring histogram.
Overridden if `smap` is provided.
smap : `matplotlib.cm.ScalarMappable` object or `None`
A scalar-mappable object to use for colormaps, or `None` for one to be created.
type : str, {'hist', 'scatter'}
What type of plot should be in the center, a 2D Histogram or a scatter-plot.
scale_to_cbar :
fs : int,
Font-size
scale : str or [str, str],
Specification for the axes scaling {'log','lin'}. If two values are given, the first is
used for the x-axis and the second for the y-axis.
histScale : str,
Scaling to use for the histograms {'log','lin'}-- the color scale on the 2D histogram,
or the Counts axis on the 1D histograms.
labels : (2,) str
cbar : bool,
Add a colorbar.
overlay : str or 'None', if str {'counts', 'values'}
Print a str on each bin writing,
'counts' - the number of values included in that bin, or
'values' - the value of the bin itself.
overlay_fmt : str or 'None'
Format specification on overlayed values, e.g. "02d" (no colon or brackets).
left : float {0.0, 1.0}
Location of the left edge of axes relative to the figure.
bottom : float {0.0, 1.0}
Location of the bottom edge of axes relative to the figure.
right : float {0.0, 1.0}
Location of the right edge of axes relative to the figure.
top : float {0.0, 1.0}
Location of the top edge of axes relative to the figure.
lo : scalar or 'None'
When autocalculating `extrema`, ignore histogram entries below this value.
hi : scalar or 'None'
When autocalculating `extrema`, ignore histogram entries above this value.
overall :
overall_bins
overall_wide
overall_cumulative
Returns
-------
fig : matplotlib.figure.Figure
Figure object containing plots.
"""
# Make sure shapes of input arrays are valid
if np.shape(xvals) != np.shape(yvals):
raise ValueError("Shape of `xvals` ({}) must match `yvals` ({}).".format(
np.shape(xvals), np.shape(yvals)))
if weights is not None and np.shape(weights) != np.shape(xvals):
raise ValueError("Shape of `weights` ({}) must match `xvals` and `yvals` ({}).".format(
np.shape(weights), np.shape(xvals)))
if overlay is not None:
if not (overlay.startswith('val') or overlay.startswith('count')):
raise ValueError("`overlay` = '{}', must be {'values', 'count'}".format(overlay))
# Make sure the given `scale` is valid
if np.size(scale) == 1:
scale = [scale, scale]
elif np.size(scale) != 2:
raise ValueError("`scale` must be one or two scaling specifications!")
# Check the `labels`
if labels is None:
labels = ['', '', '']
elif np.size(labels) == 2:
labels = [labels[0], labels[1], '']
if np.size(labels) != 3:
raise ValueError("`labels` = '{}' is invalid.".format(labels))
# Make sure scale strings are matplotlib-compliant
scale = [plot_core._clean_scale(sc) for sc in scale]
# Determine type of central plot
if type.startswith('hist'):
type_hist = True
elif type.startswith('scat'):
type_hist = False
cblabel = str(labels[2])
labels[2] = 'Count'
else:
raise ValueError("`type` = '{}', must be either 'hist', or 'scatter'.".format(type))
# Infer default statistic
if statistic is None:
if weights is None: statistic = 'count'
else: statistic = 'sum'
if filter is None and histScale.startswith('log'):
filter = 'g'
# Filter input data
if filter is not None:
# Make sure `filter` is an iterable pair
# if weights is None:
# num = 2
# else:
# num = 3
if not np.iterable(filter):
filter = 3*[filter]
elif len(filter) == 1:
filter = 3*[filter[0]]
# if len(filter) != num:
# raise ValueError("If `weights` are provided, number of `filter` values must match.")
# Filter `xvals`
if filter[0] is not None:
inds = zmath.comparison_filter(xvals, filter[0], inds=True)
xvals = xvals[inds]
yvals = yvals[inds]
if weights is not None:
weights = weights[inds]
# Filter `yvals`
if filter[1] is not None:
inds = zmath.comparison_filter(yvals, filter[1], inds=True)
xvals = xvals[inds]
yvals = yvals[inds]
if weights is not None:
weights = weights[inds]
if weights is not None and filter[2] is not None:
inds = zmath.comparison_filter(yvals, filter[2], inds=True)
xvals = xvals[inds]
yvals = yvals[inds]
weights = weights[inds]
# Create and initializae figure and axes
fig, prax, xpax, ypax, cbax, ovax = _constructFigure(
fig, xproj, yproj, overall, overall_wide, hratio, wratio, pad,
scale, histScale, labels, cbar,
left, bottom, right, top, fs=fs)
# Create bins
# -----------
# `bins` is a single scalar value -- apply to both
if np.isscalar(bins):
xbins = bins
ybins = bins
else:
# `bins` is a pair of bin specifications, separate and apply
if len(bins) == 2:
xbins = bins[0]
ybins = bins[1]
# `bins` is a single array -- apply to both
elif len(bins) > 2:
xbins = bins
ybins = bins
# unrecognized option -- error
else:
raise ValueError("Unrecognized shape of ``bins`` = %s" % (str(np.shape(bins))))
# If a number of bins is given, create an appropriate spacing
if np.ndim(xbins) == 0:
xbins = zmath.spacing(xvals, num=xbins+1, scale=scale[0])
if np.ndim(ybins) == 0:
ybins = zmath.spacing(yvals, num=ybins+1, scale=scale[1])
# Make sure bins look okay
for arr, name in zip([xbins, ybins], ['xbins', 'ybins']):
delta = np.diff(arr)
if np.any(~np.isfinite(delta) | (delta == 0.0)):
raise ValueError("Error constructing `{}` = {}, delta = {}".format(name, arr, delta))
# Calculate Histograms
# --------------------
# 2D
try:
hist_2d, xedges_2d, yedges_2d, binnums_2d = sp.stats.binned_statistic_2d(
xvals, yvals, weights, statistic=statistic, bins=[xbins, ybins], expand_binnumbers=True)
hist_2d = np.nan_to_num(hist_2d)
# X-projection (ignore Y)
hist_xp, edges_xp, bins_xp = sp.stats.binned_statistic(
xvals, weights, statistic=statistic, bins=xbins)
# Y-projection (ignore X)
hist_yp, edges_yp, bins_yp = sp.stats.binned_statistic(
yvals, weights, statistic=statistic, bins=ybins)
except:
hist_2d, xedges_2d, yedges_2d, binnums_2d = sp.stats.binned_statistic_2d(
xvals, yvals, weights, statistic=statistic, bins=[xbins, ybins])
hist_2d = np.nan_to_num(hist_2d)
# X-projection (ignore Y)
hist_xp, edges_xp, bins_xp = sp.stats.binned_statistic(
xvals, weights, statistic=statistic, bins=xbins)
# Y-projection (ignore X)
hist_yp, edges_yp, bins_yp = sp.stats.binned_statistic(
yvals, weights, statistic=statistic, bins=ybins)
if cumulative is not None:
hist_2d = _cumulative_stat2d(
weights, hist_2d.shape, binnums_2d, statistic, cumulative)
hist_xp = _cumulative_stat1d(
weights, hist_xp.size, bins_xp, statistic, cumulative[0])
hist_yp = _cumulative_stat1d(
weights, hist_yp.size, bins_yp, statistic, cumulative[1])
# Calculate Extrema - Preserve input extrema if given, otherwise calculate
extrema = _set_extrema(extrema, [hist_2d, hist_xp, hist_yp], filter=filter[2], lo=lo, hi=hi)
# Create scalar-mappable if needed
if smap is None:
smap = plot_core.colormap(extrema, cmap=cmap, scale=histScale)
# Plot Histograms and Projections
# -------------------------------
# Plot 2D Histogram
if type_hist:
overlay_values = None
# If we should overlay strings labeling the num values in each bin, calculate those `counts`
if overlay is not None:
# Overlay the values themselves
if overlay.startswith('val'):
overlay_values = hist_2d
if overlay_fmt is None:
overlay_fmt = ''
# Get the 'counts' to overlay on plot
else:
if overlay_fmt is None:
overlay_fmt = 'd'
try:
overlay_values, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
xvals, yvals, weights, statistic='count', bins=[xbins, ybins],
expand_binnumbers=True)
except:
overlay_values, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
xvals, yvals, weights, statistic='count', bins=[xbins, ybins])
if cumulative is not None:
overlay_values = _cumulative_stat2d(
np.ones_like(xvals), overlay_values.shape, binnums, 'count', cumulative)
overlay_values = overlay_values.astype(int)
pcm, smap, cbar, cs = plot2DHist(prax, xedges_2d, yedges_2d, hist_2d, cscale=histScale,
cbax=cbax, labels=labels, cmap=cmap, smap=smap,
extrema=extrema, fs=fs, scale=scale,
overlay=overlay_values, overlay_fmt=overlay_fmt)
# Colors
# X-projection
if xpax:
colhist_xp = np.array(hist_xp)
# Enforce positive values for colors in log-plots.
if smap.log:
tmin, tmax = zmath.minmax(colhist_xp, filter='g')
colhist_xp = np.maximum(colhist_xp, tmin)
colors_xp = smap.to_rgba(colhist_xp)
if ypax:
colhist_yp = np.array(hist_yp)
# Enforce positive values for colors in log-plots.
if smap.log:
tmin, tmax = zmath.minmax(colhist_yp, filter='g')
colhist_xp = np.maximum(colhist_yp, tmin)
colors_yp = smap.to_rgba(colhist_yp)
# colors_yp = smap.to_rgba(hist_yp)
# Scatter Plot
else:
colors = smap.to_rgba(weights)
prax.scatter(xvals, yvals, c=colors, alpha=alpha)
if cbar:
cbar = plt.colorbar(smap, cax=cbax)
cbar.set_label(cblabel, fontsize=fs)
cbar.ax.tick_params(labelsize=fs)
# Make projection colors all grey
colors_xp = '0.8'
colors_yp = '0.8'
hist_log = plot_core._scale_to_log_flag(histScale)
# Plot projection of the x-axis (i.e. ignore 'y')
if xpax:
islog = scale[0].startswith('log')
xpax.bar(edges_xp[:-1], hist_xp, color=colors_xp, log=islog, width=np.diff(edges_xp),
alpha=_BAR_ALPHA)
# set tick-labels to the top
plt.setp(xpax.get_yticklabels(), fontsize=fs)
xpax.xaxis.tick_top()
# set bounds to bin edges
plot_core.set_lim(xpax, 'x', data=xedges_2d)
# Set axes limits to match those of colorbar
if scale_to_cbar:
# extrema_y = [zmath.floor_log(extrema[0]), zmath.ceil_log(extrema[1])]
round = 0
# if hist_log: round = -1
extrema_y = zmath.minmax(extrema, round=round)
xpax.set_ylim(extrema_y)
# Plot projection of the y-axis (i.e. ignore 'x')
if ypax:
ypax.barh(edges_yp[:-1], hist_yp, color=colors_yp, log=hist_log, height=np.diff(edges_yp),
alpha=_BAR_ALPHA)
# set tick-labels to the top
plt.setp(ypax.get_yticklabels(), fontsize=fs, rotation=90)
ypax.yaxis.tick_right()
# set bounds to bin edges
plot_core.set_lim(ypax, 'y', data=yedges_2d)
try:
ypax.locator_params(axis='x', tight=True, nbins=4)
except:
ypax.locator_params(axis='x', tight=True)
# Set axes limits to match those of colorbar
if scale_to_cbar:
round = 0
# if hist_log: round = -1
extrema_x = zmath.minmax(extrema, round=round)
ypax.set_xlim(extrema_x)
# Plot Overall Histogram of values
if overall:
ov_bins = zmath.spacing(weights, num=overall_bins)
bin_centers = zmath.midpoints(ov_bins, log=hist_log)
nums, bins, patches = ovax.hist(weights, ov_bins, log=hist_log, facecolor='0.5', edgecolor='k')
for pp, cent in zip(patches, bin_centers):
pp.set_facecolor(smap.to_rgba(cent))
# Add cumulative distribution
if overall_cumulative:
cum_sum = np.cumsum(nums)
ovax.plot(bin_centers, cum_sum, 'k--')
prax.set(xlim=zmath.minmax(xedges_2d), ylim=zmath.minmax(yedges_2d))
return fig