Created
June 20, 2016 18:07
-
-
Save bencbartlett/a3b9ddf9287d12da5fdb6891cccd6b84 to your computer and use it in GitHub Desktop.
Code to make 2D evolving multiplot histogram with side plots, like this one: https://gfycat.com/PersonalGraciousAfricanharrierhawk
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# Importations | |
import os, shutil | |
import numpy as np | |
import matplotlib | |
import matplotlib.pyplot as plt | |
import matplotlib.mlab as mlab | |
import pylab | |
import warnings | |
from mpl_toolkits.mplot3d import Axes3D, art3d | |
from matplotlib.offsetbox import TextArea, VPacker, AnnotationBbox | |
from Libraries.FastProgressBar import progressbar | |
from scipy.stats import norm | |
from ExternalLibraries.blessings import Terminal | |
term = Terminal() | |
# Constants | |
ECalRadius = 129.0 #|Radius of ECal in cm | |
ECalZ = 317.0 | |
def layerVarianceFrame(enWeightedEtaPhi, etaPhi, frameNo, center, rng=None, maxE=1.0, | |
xbins=30, ybins=30, numEvents=1, saveAs=None): | |
'''Plots individual frames of eta-phi 2D histogram showing energy variance''' | |
from matplotlib.ticker import NullFormatter, MaxNLocator | |
from matplotlib.colors import LogNorm | |
def ellipse(ra,rb,ang,x0,y0,Nb=100): | |
'''Plots 1,2,3 sigma rings''' | |
xpos,ypos = x0,y0 | |
radm,radn = ra,rb | |
an = ang | |
co,si = np.cos(an),np.sin(an) | |
the = np.linspace(0,2*np.pi,Nb) | |
X = radm*np.cos(the)*co - si*radn*np.sin(the) + xpos | |
Y = radm*np.cos(the)*si + co*radn*np.sin(the) + ypos | |
return X,Y | |
# Unpack data | |
x, y = enWeightedEtaPhi | |
eta, phi = etaPhi | |
if rng != None: | |
xlim, ylim = rng | |
else: | |
xlim = [np.min(x), np.max(x)] | |
ylim = [np.min(y), np.max(y)] | |
# Format axes | |
nullfmt = NullFormatter() | |
left, width = 0.12, 0.54 #|These determine where the subplots go | |
bottom, height = 0.12, 0.54 | |
bottom_h = left_h = left+width+0.02 | |
rectscatter = [left, bottom, width, height] | |
recthistx = [left, bottom_h, width, 0.2] | |
recthisty = [left_h, bottom, 0.2, height] | |
rectcbar = [left_h+0.22, bottom, 0.05, height] | |
# Set up figure and axes | |
plt.figure(1, figsize=(8,8)) | |
axHist2D = plt.axes(rectscatter) | |
axHistx = plt.axes(recthistx) | |
axHisty = plt.axes(recthisty) | |
axCbar = plt.axes(rectcbar) | |
# Remove labels from supplementary plots | |
axHistx.xaxis.set_major_formatter(nullfmt) | |
axHisty.yaxis.set_major_formatter(nullfmt) | |
axCbar.xaxis.set_major_formatter(nullfmt) | |
axCbar.yaxis.set_major_formatter(nullfmt) | |
axCbar.set_visible(False) | |
# Configure and plot center and sigma bounds | |
contourcolor = 'black' | |
xcenter = np.mean(x) | |
ycenter = np.mean(y) | |
ra = np.std(x) | |
rb = np.std(y) | |
ang = 0 | |
axHist2D.plot(center[0], center[1], '+w', markersize=50, mew=3.0) | |
X,Y = ellipse(ra,rb,ang,xcenter,ycenter) | |
axHist2D.plot(X,Y,"k:",ms=1,linewidth=2.0) | |
axHist2D.annotate('$1\\sigma$', xy=(X[15], Y[15]), xycoords='data',xytext=(10, 10), | |
textcoords='offset points', horizontalalignment='right', | |
verticalalignment='bottom',fontsize=25) | |
X,Y = ellipse(2*ra,2*rb,ang,xcenter,ycenter) | |
axHist2D.plot(X,Y,"k:",color = contourcolor,ms=1,linewidth=2.0) | |
axHist2D.annotate('$2\\sigma$', xy=(X[15], Y[15]), xycoords='data',xytext=(10, 10), | |
textcoords='offset points',horizontalalignment='right', | |
verticalalignment='bottom',fontsize=25, color=contourcolor) | |
X,Y = ellipse(3*ra,3*rb,ang,xcenter,ycenter) | |
axHist2D.plot(X,Y,"k:",color = contourcolor, ms=1,linewidth=2.0) | |
axHist2D.annotate('$3\\sigma$', xy=(X[15], Y[15]), xycoords='data',xytext=(10, 10), | |
textcoords='offset points',horizontalalignment='right', | |
verticalalignment='bottom',fontsize=25, color=contourcolor) | |
# Plot the main 2D histogram | |
H, xedges, yedges, img = axHist2D.hist2d(x, y, bins=(xbins,ybins), range=rng, | |
norm=LogNorm(vmax=numEvents*maxE/50.0)) | |
cbar = plt.colorbar(img, ax=axCbar, fraction=1.0) #|Fraction makes the colorbar fill the entire axis | |
cbar.set_label("$\log$ Energy ($\log$ MIPs) (total for $n$ events)", rotation=90) | |
# Set main limits | |
axHist2D.set_xlim(xlim) | |
axHist2D.set_ylim(ylim) | |
axHist2D.set_xlabel("$\eta$",fontsize=25) | |
axHist2D.set_ylabel("$\phi$",fontsize=25) | |
axHist2D.set_axis_bgcolor(pylab.cm.jet(0.0)) | |
axHist2D.text(1.05, 1.20, "Layer %s/30"%str(frameNo),transform=axHist2D.transAxes,fontsize=25) | |
# Plot two side histograms | |
axHistx.hist(eta, range = rng[0], bins=xbins, color='blue') | |
axHisty.hist(phi, range = rng[1], bins=ybins, orientation='horizontal', color='red') | |
# Set those limits | |
axHistx.set_xlim(axHist2D.get_xlim()) | |
axHistx.set_ylim([0,numEvents*5]) | |
axHisty.set_ylim(axHist2D.get_ylim()) | |
axHisty.set_xlim([0,numEvents*5]) | |
axHistx.set_ylabel("Energy-unweighted hits", rotation=90) | |
axHisty.set_xlabel("Energy-unweighted hits") | |
axHistx.yaxis.set_major_locator(MaxNLocator(4)) | |
axHisty.xaxis.set_major_locator(MaxNLocator(4)) | |
# Finish everything up | |
if saveAs: | |
plt.savefig(saveAs) | |
else: | |
plt.show() | |
plt.clf() | |
def layerVarianceAnalysis(data, eventNo, title, mode=None, clusterID=0, rng=None, endLoop=0, | |
numEvents=1, delete=False, cumulative=False): | |
''' | |
Plots an evolving 2D histogram by layer of a hit. | |
Warning: running this on cumulative mode for very large datasets can make you run out of RAM. | |
''' | |
# Parse data for proper cluster ID and existence of timing data | |
if not cumulative: | |
hits = np.extract(data[eventNo][0]['clusterID']==clusterID, data[eventNo][0]) #|Select only the rechits corresponding to the given cluster | |
# Further data parsing | |
xh,yh,zh,Eh = hits['x'], hits['y'], hits['z'], hits['en'] #|Separating hits | |
center = [data[eventNo][1]['eta'][clusterID], data[eventNo][1]['phi'][clusterID]] | |
eta, phi = XYZtoEtaPhi([xh,yh,zh]) | |
layers = hits['layerID'] | |
if mode == "log": | |
repeats = np.ceil(np.log(Eh)).astype(int) | |
etaw = np.repeat(eta, repeats) #|Energy-weighted eta and phi values | |
phiw = np.repeat(phi, repeats) | |
layersw = np.repeat(layers, repeats) | |
elif mode == "linear": | |
repeats = np.ceil(Eh).astype(int) | |
etaw = np.repeat(eta, repeats) #|Energy-weighted eta and phi values | |
phiw = np.repeat(phi, repeats) | |
layersw = np.repeat(layers, repeats) | |
elif mode == "flat": | |
etaw, phiw, layersw = eta, phi, layers | |
else: | |
center = [0,0] | |
eta = np.array([]) | |
phi = np.array([]) | |
etaw = np.array([]) | |
phiw = np.array([]) | |
layers = np.array([]) | |
layersw = np.array([]) | |
Eh = np.array([]) | |
k = 0 | |
pbar = progressbar("Processing &count& events:", len(data)+1) | |
pbar.start() | |
for event in data: | |
hits = np.extract(event[0]['clusterID']==clusterID, event[0]) #|Select only the rechits corresponding to the given cluster | |
# Further data parsing | |
xh,yh,zh,EhE = hits['x'], hits['y'], hits['z'], hits['en'] #|Separating hits | |
etaE, phiE = XYZtoEtaPhi([xh,yh,zh]) | |
layersE = hits['layerID'] | |
for i in range(len(hits)): | |
clusterID = hits[i]['clusterID'] | |
etaE[i] -= event[1]['eta'][clusterID] | |
phiE[i] -= event[1]['phi'][clusterID] | |
if phiE[i] > 3.0: | |
phiE[i] -= 2*np.pi #|This fixes some modular issues we were having, with -epsilon = 2pi-epsilon | |
elif phiE[i] < -3.0: | |
phiE[i] += 2*np.pi | |
if mode == "log": | |
repeats = np.ceil(np.log(EhE)).astype(int) | |
etawE = np.repeat(etaE, repeats) #|Energy-weighted eta and phi values | |
phiwE = np.repeat(phiE, repeats) | |
layerswE = np.repeat(layersE, repeats) | |
elif mode == "linear": | |
repeats = np.ceil(EhE).astype(int) | |
etawE = np.repeat(etaE, repeats) #|Energy-weighted eta and phi values | |
phiwE = np.repeat(phiE, repeats) | |
layerswE = np.repeat(layersE, repeats) | |
elif mode == "flat": | |
etawE, phiwE, layerswE = etaE, phiE, layersE | |
eta = np.append(eta, etaE) | |
phi = np.append(phi, phiE) | |
etaw = np.append(etaw, etawE) | |
phiw = np.append(phiw, phiwE) | |
layers = np.append(layers, layersE) | |
layersw = np.append(layersw, layerswE) | |
Eh = np.append(Eh, EhE) | |
k += 1 | |
pbar.update(k) | |
pbar.finish() | |
# print "Saving array..." | |
# np.save("Data/etaPhiProcessed.npy", [eta, phi, etaw, phiw, layers, layersw, Eh]) | |
# eta, phi, etaw, phiw, layers, layersw, Eh = np.load("Data/etaPhiProcessed.npy") | |
# center = [0,0] | |
# Set plot ranges | |
if rng == None: | |
pltrange = np.array([(np.min(eta), np.max(eta)), (np.min(phi), np.max(phi))]) | |
else: | |
pltrange = rng | |
# Clear out existing crap | |
path = "Plots/LayerHitAnimations/"+title+"/Event"+str(eventNo) | |
if os.path.exists(path): #|Remove previous crap | |
shutil.rmtree(path) | |
os.makedirs(path+"/frames") | |
else: | |
os.makedirs(path+"/frames") | |
# Set up animation | |
minlayer = 1 #|Minimum layer | |
maxlayer = 30 #|Maximum layer | |
count = minlayer | |
pbar = progressbar("Rendering &count& frames:", maxlayer-minlayer+1) | |
pbar.start() | |
# Render frames | |
while count <= maxlayer: | |
etapltw = np.extract(layersw == count, etaw) | |
phipltw = np.extract(layersw == count, phiw) | |
etaplt = np.extract(layers == count, eta) | |
phiplt = np.extract(layers == count, phi) | |
# Eplt = np.extract(layers == count, Eh) | |
filename = path + "/frames/"+str(count).zfill(3)+".gif" | |
if len(etaplt) > 0: | |
layerVarianceFrame([etapltw, phipltw], [etaplt, phiplt], count, center, rng=pltrange, | |
numEvents=numEvents, maxE=np.max(Eh), xbins=50, ybins=50, | |
saveAs=filename) | |
pbar.update(count) | |
count += 1 | |
pbar.finish() | |
# Render tail if needed: | |
if endLoop: print "Copying tail frames..." | |
for i in xrange(endLoop): | |
shutil.copyfile(filename, filename[:-7]+str(count).zfill(3)+".gif") | |
count += 1 | |
# Combine frames to animated gif | |
print "Combining frames..." | |
import subprocess | |
args = (['convert', '-delay', '25', '-loop', '0', path+"/frames/*.gif", path+"/Shower.gif"]) #|This part requires ImageMagick to function | |
print "Saved to path "+str(os.path.abspath(path))+"/Shower.gif" | |
subprocess.check_call(args) | |
# Delete frames if told to do so | |
if delete: | |
shutil.rmtree(path+"/frames") | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment