more functionality in the plt class

ess/TraceWin.py

@@ -155,6 +155,106 @@ class plt:
         else:
             print "No data to plot at element",key
 
+    def calc_s(self):
+        '''
+        Generates self.s which holds
+        the position of each element
+        in metres
+        '''
+        import numpy
+
+        self.s=[]
+        for i in self.Nelp:
+            self.s.append(self[i]['Zgen']/100.0)
+        self.s=numpy.array(self.s)
+
+    def calc_rel(self):
+        '''
+        Calculates relativistic gamma/beta
+        at each position, based on 
+        AVERAGE beam energy
+        (NOT necessarily reference)
+        '''
+        import numpy
+
+        if not hasattr(self,'avg'):
+            self.calc_avg()
+        self.gamma=[]
+        self.beta=[]
+        for i,j in zip(self.Nelp,xrange(len(self.Nelp))):
+            Eavg=self.avg['E'][j]
+            self.gamma.append((self.mc2+Eavg)/self.mc2)
+            self.beta.append(numpy.sqrt(1.-1./self.gamma[-1]**2))
+
+    def calc_avg(self):
+        '''
+        Calculates averages of 6D coordinates at each
+        element, such that e.g.
+        self.avg["x"] gives average X at each location.
+
+        Units: cm
+        '''
+        import numpy
+
+
+        self.avg=dict(x=[], xp=[], y=[], yp=[], E=[], phi=[])
+
+        vals=['x', 'xp', 'y', 'yp', 'E', 'phi']
+
+        for i in self.Nelp:
+            data=self[i]
+            for v in vals:
+                self.avg[v].append(numpy.average(data[v]))
+
+        for v in vals:
+            self.avg[v]=numpy.array(self.avg[v])
+
+    def calc_minmax(self,pmin=5,pmax=95):
+        '''
+        Calculates min/max values of beam coordinates
+        in percentile, pmin is lower and pmax upper.
+
+        Units: cm
+        '''
+        import numpy
+
+
+        self.min=dict(x=[], xp=[], y=[], yp=[], E=[])
+        self.max=dict(x=[], xp=[], y=[], yp=[], E=[])
+
+        for i in self.Nelp:
+            data=self[i]
+            for v in self.min.keys():
+                self.min[v].append(numpy.percentile(data[v],pmin))
+                self.max[v].append(numpy.percentile(data[v],pmax))
+
+        for v in self.min.keys():
+            self.min[v]=numpy.array(self.min[v])
+            self.max[v]=numpy.array(self.max[v])
+
+    def calc_std(self):
+        '''
+        Calculates the sigma of each coordinate
+        such that e.g. self.sigma['x'] is the sigma_x
+
+        Note: SigmaX etc from partran1.out are in mm, not cm
+
+        Units: cm
+        '''
+        import numpy
+
+
+        self.sigma=dict(x=[], xp=[], y=[], yp=[])
+
+        for i in self.Nelp:
+            data=self[i]
+            for key in self.sigma.keys():
+                self.sigma[key].append(data[key].std())
+
+        for key in self.sigma.keys():
+            self.sigma[key]=numpy.array(self.sigma[key])
+
+
 class density_file:
     '''
     Simple class to read a TraceWin density file