From ec600721b069b1a924e335a4d34915bd217fbddb Mon Sep 17 00:00:00 2001
From: Yngve Inntjore Levinsen <Yngve.Levinsen@esss.se>
Date: Wed, 20 Apr 2016 14:10:10 +0200
Subject: [PATCH] added function ess.TraceWin.dst.subplot resembles TraceWin
distribution plot
---
ess/TraceWin.py | 52 +++++++++++++++++++++++++++++++++++++++++++++++++
1 file changed, 52 insertions(+)
diff --git a/ess/TraceWin.py b/ess/TraceWin.py
index 54a909e..a304694 100644
--- a/ess/TraceWin.py
+++ b/ess/TraceWin.py
@@ -121,6 +121,58 @@ class dst:
print >>fout, out
fout.close()
+ def subplot(self, index, x, y, nb=100, mask=None):
+ '''
+ Create a subplot histogram similar to TraceWin.
+
+ Example::
+ import numpy as np
+ from ess import TraceWin
+ from matplotlib import pyplot as plt
+ data=TraceWin.dst('part_dtl1.dst')
+ m=np.where(data['E']>3.5)
+ data.subplot(221,'x','xp',mask=m)
+ data.subplot(222,'y','yp',mask=m)
+ data.subplot(223,'phi','E',mask=m)
+ data.subplot(224,'x','y',mask=m)
+ plt.show()
+
+ '''
+ from matplotlib.colors import LogNorm
+ import matplotlib.pyplot as plt
+ import numpy as np
+
+ units={ 'x': 'mm', 'y': 'mm',
+ 'xp': 'mrad', 'yp': 'mrad',
+ 'E': 'MeV', 'phi': 'deg'
+ }
+ # get X and Y data
+ dx=self[x]
+ dy=self[y]
+ if mask!=None:
+ dx=dx[mask]
+ dy=dy[mask]
+
+ if x in ['x','y','xp','yp']:
+ dx*=1e3
+ if y in ['x','y','xp','yp']:
+ dy*=1e3
+ if x in ['phi']:
+ dx-=np.average(dx)
+ dx*=180/np.pi
+ if y in ['phi']:
+ dy-=np.average(dy)
+ dy*=180/np.pi
+
+ plt.subplot(index)
+ plt.hist2d(dx, dy, bins=nb, norm=LogNorm())
+ plt.title('{} [{}] - {} [{}]'.format(x,units[x],y,units[y]))
+ hist,bin_edges=np.histogram(dx,bins=nb)
+ b=bin_edges[:-1]+0.5*(bin_edges[1]-bin_edges[0])
+ plt.plot(b,hist*0.2*(max(dy)-min(dy))/max(hist)+min(dy),'k',drawstyle='steps')
+ hist,bin_edges=np.histogram(dy,bins=nb)
+ b=bin_edges[:-1]+0.0*(bin_edges[1]-bin_edges[0])
+ plt.plot(hist*0.2*(max(dx)-min(dx))/max(hist)+min(dx),b,'k',drawstyle='steps')
class plt:
'''
--
GitLab