added function ess.TraceWin.dst.subplot

resembles TraceWin distribution plot

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:
     '''