xboa
_bunch.py
Go to the documentation of this file.
1 #This file is a part of xboa
2 #
3 #xboa is free software: you can redistribute it and/or modify
4 #it under the terms of the GNU General Public License as published by
5 #the Free Software Foundation, either version 3 of the License, or
6 #(at your option) any later version.
7 #
8 #xboa is distributed in the hope that it will be useful,
9 #but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11 #GNU General Public License for more details.
12 #
13 #You should have received a copy of the GNU General Public License
14 #along with xboa in the doc folder. If not, see
15 #<http://www.gnu.org/licenses/>.
16 #
17 
18 import xboa.common as Common
19 import xboa.hit.factory
20 from xboa.hit import Hit
21 from xboa.hit import BadEventError
22 from xboa.core import Hitcore
23 from xboa.core import Bunchcore
24 from xboa.common import rg as rg
25 
26 import string
27 import copy
28 import gzip
29 import StringIO
30 import operator
31 import math
32 import sys
33 import bisect
34 
35 try:
36  import json
37 except ImportError:
38  pass
39 
40 try:
41  import ROOT
42 except ImportError:
43  pass
44 try:
45  import numpy
46  from numpy import matrix
47  from numpy import linalg
48 except ImportError:
49  pass
50 
51 class Bunch:
52  """
53  Represents a bunch of particles. Methods to:
54 
55  - Read/write from a file
56  - Calculate arbitrary moments
57  - Apply cuts
58  - Calculate emittances, beta functions etc
59  - Extract hit data as float or arrays of floats
60  - Plots (with ROOT imported)
61  """
62 
63  def __init__(self):
64  """Initialise to an empty bunch. Alternatively use static initialisers defined below - I prefer static initialisers"""
65  self.__hits = []
66  self.__bunchcore = Bunchcore()
67  self.__covs = None
68  self.__means = {}
69 
70  def __str__(self):
71  """Return an abbreviated string like <Bunch of n Hits>"""
72  return '<Bunch of '+str(len(self.__hits))+' Hits>'
73 
74  def __repr__(self):
75  """Return a long string that contains data in every hit"""
76  covs = 'None'
77  if self.__covs != None:
78  covs = 'numpy.'+repr(self.__covs)
79  out = 'Bunch.new_from_hits( ['
80  for hit in self.__hits:
81  out += str(hit)+',\n'
82  out += '],'+covs+','+repr(self.__means)+')'
83  return out
84 
85  def __copy__(self):
86  """Make a copy but use references to self's data"""
87  copy = self
88  return copy
89 
90  def copy(self):
91  return self.__copy__()
92 
93  def __deepcopy__(self, target):
94  """Make a copy and copy self's data"""
95  target = eval(self.__repr__())
96  target.__covs = copy.deepcopy(self.__covs)
97  target.__means = copy.deepcopy(self.__means)
98  return target
99 
100  def deepcopy(self):
101  """Make a copy and copy self's data"""
102  target = Bunch()
103  return self.__deepcopy__(target)
104 
105  def __len__(self):
106  """Number of hits"""
107  return len(self.__hits)
108 
109  def __getitem__(self, key):
110  """Called by subscript operator, returns the key^th hit in the bunch"""
111  return self.__hits[key]
112 
113  def __setitem__(self, key, value):
114  """Called by subscript operator, sets the key^th hit in the bunch"""
115  self.__hits[key] = value
116  self.__bunchcore.set_item(value._Hit__hitcore, key)
117 
118  def __delitem__(self, key):
119  """Called by remove method, deletes the key^th hit in the bunch"""
120  self.__hits.__delitem__(key)
121 
122  def __del__(self):
123  """Called by del"""
124  for hit in self.__hits:
125  del(hit)
126 
127  def __eq__(self, target, float_tolerance=Common.float_tolerance):
128  """Return true if all hits in self are the same as all hits in target, and set_covariance_matrix data are the same"""
129  if type(target) != type(self): return False
130  if len(self.__hits) != len(target.__hits): return False
131  for key,value in self.__means.iteritems():
132  if not key in target.__means: return False
133  if abs(self.__means[key] - target.__means[key]) > float_tolerance: return False
134  try:
135  if not all(self.__covs == target.__covs): return False
136  except:
137  try:
138  if self.__covs != target.__covs: return False
139  except:
140  return True
141  for index in range(len(self.__hits)):
142  if self.__hits[index].__ne__( target.__hits[index], float_tolerance): return False
143  return True
144 
145  def __ne__(self, target, float_tolerance=Common.float_tolerance):
146  """Return true if any hits in self are not the same as the corresponding hit in target"""
147  return not self.__eq__(target, float_tolerance)
148 
149  def append(self, hit):
150  """Append a hit to Bunch self"""
151  self.__bunchcore.set_item(hit._Hit__hitcore, len(self.__hits))
152  self.__hits.append(hit)
153 
154  def new_from_hits(hits_list, covs=None, means={},weights=[]):
155  """
156  Return a bunch from a list of hits, making a deepcopy of the hits (so allocating new memory for each hit)
157 
158  - hits_list = list of Hit objects
159 
160  e.g. myBunch = new_from_hits([hit1, hit2, hit3]) will return a new bunch containing hit1, hit2, hit3
161  """
162  bunch = Bunch()
163  for hit in hits_list:
164  new_hit = hit.deepcopy()
165  bunch.append( new_hit )
166  bunch.__covs = covs
167  bunch.__means = means
168  return bunch
169  new_from_hits = staticmethod(new_from_hits)
170 
171  def new_dict_from_read_builtin(file_type_string, file_name, indexing_variable='station'):
172  """
173  Return a dict of all bunches in a file using a built-in format
174 
175  - file_type_string = string from Hit.file_types() that defines the formatting of the file
176  - file_name = name of the file that contains hit data
177  - indexing_variable = variable that will be used to define a bunch
178 
179  e.g. bunch_dict = new_dict_from_read_builtin('icool_for003', 'for003.dat', 'pid') will return a new dict of bunches
180  loaded from for003.dat in icool_for003 format, where each entry in the dict will be a reference from a pid value to a bunch
181 
182  e.g. bunch_dict = new_dict_from_read_builtin('icool_for009', 'for009.dat', 'station') will return a new dict of bunches
183  loaded from for009.dat in icool_for009 format, where each entry in the dict will be a reference from a station value to a bunch
184  """
185  # can we just use bunch.new_from_read_builtin to generate bunch?
186  if not file_type_string in Hit.file_types(): raise IOError('Attempt to load file of unknown file type '+str(file_type_string)+' - try one of '+str(Hit.file_types()))
187  if file_type_string.find('maus_root') > -1:
188  hit_list = Bunch.read_maus_root_file(file_name, list_of_maus_types=[file_type_string])
189  bunch = Bunch.new_from_hits(hit_list)
190  elif file_type_string.find('maus_') > -1:
191  hit_list = Bunch.read_maus_json_file(file_name, list_of_maus_types=[file_type_string])
192  bunch = Bunch.new_from_hits(hit_list)
193  else:
194  filehandle = Bunch.setup_file(file_type_string, file_name)
195  bunch = Bunch.new_from_read_builtin(file_type_string, file_name)
196  filehandle.close()
197  bunch_dict = {}
198  for hit in bunch:
199  if not hit.check():
200  pass
201  if not int(hit.get(indexing_variable)) in bunch_dict:
202  bunch_dict[ int(hit.get(indexing_variable)) ] = Bunch()
203  bunch_dict[ int(hit.get(indexing_variable)) ].append(hit)
204  return bunch_dict
205  new_dict_from_read_builtin = staticmethod(new_dict_from_read_builtin)
206 
207  def new_list_from_read_builtin(file_type_string, file_name, sort_variable='station'):
208  """
209  Return a sorted list of all bunches in a file using a built-in format
210 
211  - file_type_string = string from Hit.file_types() that defines the formatting of the file
212  - file_name = name of the file that contains hit data
213  - sort_variable = variable that will be used to define a bunch and used for sort order
214 
215  e.g. bunch_list = new_list_from_read_builtin('icool_for003', 'for003.dat', 'pid') will return a new list of bunches
216  loaded from for003.dat in icool_for003 format, where each entry in the list will contain only one pid value with first entry
217  having lowest pid
218 
219  e.g. bunch_list = new_list_from_read_builtin('icool_for009', 'for009.dat', 'station') will return a new list of bunches
220  loaded from for009.dat in icool_for009 format, where each entry in the list will contain only one station value with first entry
221  having lowest station
222  """
223  bunch_dict = Bunch.new_dict_from_read_builtin(file_type_string, file_name, sort_variable)
224  key_list = []
225  bunch_list = []
226  for key in bunch_dict:
227  key_list.append(key)
228  key_list.sort()
229  for key in key_list:
230  bunch_list.append(bunch_dict[key])
231  return bunch_list
232  new_list_from_read_builtin = staticmethod(new_list_from_read_builtin)
233 
234  def new_from_read_builtin(file_type_string, file_name, test_function=None, number_of_hits=-1):
235  """
236  Initialise a bunch from a file using a built in format
237 
238  - file_type_string = string from Hit.file_types() that defines the file format
239  - file_name = string that defines the file_name to be used
240  - test_function = Hits with test_function(hit) == False will be ignored, unless test_function==None
241  - number_of_hits = only loads the first number_of_hits Hits
242 
243  e.g. myBunch = Bunch.new_from_read_builtin('icool_for009', for009.dat, lambda hit: hit['station'] == 1, 1000)
244  will return a bunch containing first 1000 hits from for009.dat with stationNumber=1
245  """
246  if not file_type_string in Hit.file_types():
247  raise IOError('Attempt to load file of unknown file type '+str(file_type_string)+\
248  ' - try one of '+str(Hit.file_types()))
249  elif file_type_string.find('maus_root') > -1:
250  hit_list = Bunch.read_maus_root_file(file_name, number_of_hits, list_of_maus_types=[file_type_string])
251  print "LEN HITS", len(hit_list)
252  return Bunch.new_from_hits(hit_list)
253  elif file_type_string.find('maus_') > -1:
254  hit_list = Bunch.read_maus_json_file(file_name, number_of_hits, list_of_maus_types=[file_type_string])
255  return Bunch.new_from_hits(hit_list)
256  filehandle = Bunch.setup_file(file_type_string, file_name)
257  bunch = Bunch()
258  while(len(bunch) < number_of_hits or number_of_hits < 0):
259  try: hit = Hit.new_from_read_builtin(file_type_string, filehandle)
260  except(EOFError): break
261  except(BadEventError): pass
262  if not hit.check():
263  pass
264  elif test_function == None: bunch.append(hit)
265  elif test_function(hit): bunch.append(hit)
266  filehandle.close()
267  return bunch
268  new_from_read_builtin = staticmethod(new_from_read_builtin)
269 
270  def new_from_read_user(format_list, format_units_dict, filehandle, number_of_skip_lines, test_function=None, number_of_hits=-1):
271  """
272  Initialise a bunch from a file using a user defined format
273 
274  - format_list = ordered list of variables from get_variables() that contains the particle variables on each line of your input file
275  - format_units_dict = dict of variables:units that are used to transform to internal system of units
276  - filehandle = file handle, created using e.g. filehandle = open('for009.dat')
277  - number_of_skip_lines = integer; skip this many lines at the start of the file
278  - test_function = Hits with test_function(hit) == False will be ignored, unless test_function==None
279  - number_of_hits = only loads the first number_of_hits Hits. Will read to end of file if this is negative
280 
281  e.g. myBunch = Bunch.new_from_read_user(['x','y','z','px','py','pz'], {'x':'mm', 'y':'mm', 'z':'mm', 'px':'MeV/c', 'py':'MeV/c', 'pz':'MeV/c'}, my_input_file, 3)
282  will skip the first 3 lines of the file and then read in all events from my_input_file, assuming the formatting listed
283  """
284  for dummy in range(number_of_skip_lines):
285  filehandle.readline()
286  bunch = Bunch()
287  while(len(bunch.__hits) < number_of_hits or number_of_hits < 0):
288  try: hit = Hit.new_from_read_user(format_list, format_units_dict, filehandle)
289  except(EOFError): break
290  if test_function == None: bunch.append(hit)
291  elif test_function(hit): bunch.append(hit)
292  filehandle.close()
293  return bunch
294  new_from_read_user = staticmethod(new_from_read_user)
295 
296  def new_hit_shell(n_per_dimension, ellipse, set_variable_list, mass_shell_variable, defaults={}):
297  """
298  Create a set of particles that sit on the shell of an ellipse in some arbitrary dimensional space
299 
300  - n_per_dimension = number of particles in each dimension. Total number of particles is n_per_dimension^dimension
301  - ellipse = numpy.matrix that defines the ellipse of particles
302  - set_variable_list = list of set variables that defines what each dimension in the ellipse corresponds to
303  - mass_shell_variable = variable that will be used to set the mass shell condition for each hit
304  - defaults = dict of default set variables that each hit will get.
305 
306  e.g. new_bunch_as_hit_shell(2, numpy.matrix([[1,0],[0,1]]),['x','y'],'energy',{'pid':-13,'mass':Common.pdg_pid_to_mass[13],'pz':200.})
307  will make a set of muons on a circle in x-y space with 200. MeV/c momentum. Event number is automatically allocated
308  incrementing from 0.
309  """
310  grid = Common.make_shell(n_per_dimension, ellipse)
311  bunch = Bunch()
312  for ev,vector in enumerate(grid):
313  bunch.append(Hit.new_from_dict(defaults))
314  for i,var in enumerate(set_variable_list):
315  bunch[-1][var] = vector[0,i]
316  bunch[-1]['event_number'] = ev
317  bunch[-1].mass_shell_condition(mass_shell_variable)
318  return bunch
319  new_hit_shell = staticmethod(new_hit_shell)
320 
321  def read_maus_root_file(file_name, number_of_hits=-1, list_of_maus_types=['maus_root_virtual_hit', 'maus_root_primary']):
322  """
323  Read hits from a maus ROOT file
324 
325  - file_name = name of the file to read
326  - number_of_hits = integer maximum number of hits to read in. If negative,
327  this variable is ignored. Note that XBOA will read in a complete spill,
328  check if it has reached the maximum, then reject any excess.
329  - list_of_maus_types = only make hits from these types
331  Returns a list of hits
332  """
333  Common.has_maus()
334  root_file = ROOT.TFile(file_name, "READ")
335  root_tree = root_file.Get("Spill")
336  hits = []
337  for maus_type in list_of_maus_types:
338  fac = xboa.hit.factory.MausRootHitFactory(root_tree, maus_type)
339  while number_of_hits < 0 or len(hits) < number_of_hits:
340  try:
341  hits.append(fac.make_hit())
342  except BadEventError:
343  break
344  root_file.Close()
345  return hits
346  read_maus_root_file = staticmethod(read_maus_root_file)
347 
348  def read_maus_json_file(file_name, number_of_hits=-1,
349  list_of_maus_types=['maus_json_virtual_hit',
350  'maus_json_primary']):
351  """
352  Initialise a bunch from a MAUS file.
353 
354  - file_name = string that defines the file_name to be used
355  - list_of_maus_types = loads all specified types from the MAUS file
356  """
357  Common.has_json()
358  json_file = open(file_name, 'r')
359  hits = []
360  for maus_type in list_of_maus_types:
361  fac = xboa.hit.factory.MausJsonHitFactory(json_file, maus_type)
362  while number_of_hits < 0 or len(hits) < number_of_hits:
363  try:
364  hits.append(fac.make_hit())
365  except BadEventError:
366  break
367  json_file.close()
368  return hits
369  read_maus_json_file = staticmethod(read_maus_json_file)
370 
371  def setup_file(file_format_type_string, file_name):
372  """Returns a file handle with special phrases and characters stripped. Returned file_handle contains only hit data"""
373  filehandle = open(file_name, 'r')
374  if not file_format_type_string in Hit.file_types():
375  raise KeyError('Could not find filetype '+file_format_type_string+' - Options are '+str(Hit.file_types()))
376  for dummy in range(Bunch.__number_of_header_lines[file_format_type_string]):
377  filehandle.readline()
378  if file_format_type_string == 'g4beamline_bl_track_file':
379  open_file = filehandle
380  filehandle = StringIO.StringIO()
381  try:
382  line = open_file.readline()
383  while( len(line) > 0 ):
384  if(string.find(line, 'cm cm cm')>-1 and line[0] == '#') : Hit.set_g4bl_unit('cm')
385  elif(string.find(line, 'mm mm mm')>-1 and line[0] == '#') : Hit.set_g4bl_unit('mm')
386  elif(line[0] == '#' or line[0] == '\n'): pass
387  else: filehandle.write(line)
388  line = open_file.readline()
389  except IOError: pass
390  open_file.close()
391  filehandle.seek(0)
392  return filehandle
393  setup_file = staticmethod(setup_file)
394 
395  def translate(self, dict_variables_values, mass_shell_variable=''):
396  """
397  Translate all events in the bunch by dict_variables_values
398 
399  - dict_variables_values = dict of variable names from Hit.set_variables() to variable values
400  - mass_shell_variable = set mass shell variable so that E^2=p^2+m^2 after translation.
401  Does nothing if equal to ''
402 
403  e.g. bunch.translate({'pz':-10,'z':100}, 'energy') reduces particle pz by 10 MeV/c and increases
404  E by 100 mm then sets E so that mass shell condition is still obeyed
405  """
406  for hit in self.__hits:
407  hit.translate(dict_variables_values, mass_shell_variable)
408  ##
409  def abelian_transformation(self, rotation_list, rotation_matrix, translation_dict={}, origin_dict={}, mass_shell_variable=''):
410  """
411  Perform an Abelian Transformation on all variables in the bunch so that V_{out} - O = R*(V_{in}-O)+T
412 
413  - rotation_list = list of Hit.set_variables() strings that defines the vector V
414  - rotation_matrix = matrix R
415  - translation_dict = dict of Hit.set_variables() strings to values that defines the vector T
416  - mass_shell_variable = set mass shell variable so that E^2=p^2+m^2 after transformation.
417  Does nothing if equal to ''
418  - origin_dict = dict of Hit.set_variables() strings to values that defines the vector O. Defaults to
419  bunch mean.
420 
421  e.g. bunch.abelian_transformation(['x','px','z'], [[1,0.1,0],[0,1,0],[0,0,1]], {'z':1000}, 'pz') has
422  V_{in} = ('x','px','z')
423  R = 1 0.1 0
424  0 1 0
425  0 0 1
426  T = (0,0,1000)
427  O = bunch mean
428  (This is like a drift). Note that 'z' has to be explicitly included in the rotation to be present in T
429  """
430  for hit in self.__hits:
431  hit.abelian_transformation(rotation_list, rotation_matrix, translation_dict, origin_dict, mass_shell_variable)
432 
433  def transform_to(self, transform_list, target_covariances, translation_dict={}, origin_dict={}, mass_shell_variable=''):
434  """
435  Perform a linear transformation on all particles in the bunch to get a bunch with covariance matrix
436  similar to target_covariances. The transformation always has determinant 1, and so is emittance conserving.
437  Then performs a translation using translation_dict
438 
439  - transform_list = list of Hit.set_variables() strings that defines the vector V
440  - target_covariances = target covariance matrix. Must be n*n dimensions where n is the length of transform_list
441  - translation_dict = dict of Hit.set_variables() strings to values that defines a translation
442  - origin_dict = dict of Hit.set_variables() strings to values that defines the vector O. Defaults to
443  bunch mean.
444  - mass_shell_variable = set mass shell variable so that E^2=p^2+m^2 after transformation.
445  Does nothing if equal to \'\'
446 
447  e.g. bunch.transform_to(['x','px'], [[1,0],[0,1]], {'pz':10}, 'energy') will apply a transformation to all hits in the
448  bunch that diagonalises the bunch and makes Var(x,x) equal to Var(px,px), but conserving emittance; then will add
449  10 to pz for all tracks; then adjust energy to keep mass shell condition.
450  """
451  ### transformation to turn covariance matrix into unity matrix
452  Common.has_numpy()
453  (self_evals, self_evecs) = linalg.eigh(self.covariance_matrix(transform_list, origin_dict))
454  unite = matrix(numpy.zeros((len(transform_list),len(transform_list))))
455  for i in range(len(transform_list)): unite[i,i] = self_evals[i]**0.5
456  unite = linalg.inv(self_evecs*unite)
457  ### transformation to turn unity matrix into target_covariances
458  (target_evals, target_evecs) = linalg.eigh(target_covariances)
459  ununite = matrix(numpy.zeros((len(transform_list),len(transform_list))))
460  for i in range(len(transform_list)): ununite[i,i] = target_evals[i]**0.5
461  ununite = target_evecs*ununite
462  ### force determinant to 1 =>emittance conservation
463  transform = ununite*unite
464  power = 1./float(len(transform_list))
465  transform = transform/(abs(linalg.det(transform))**power)
466  self.abelian_transformation(transform_list, transform, translation_dict, origin_dict, mass_shell_variable)
467 
468  def period_transformation(self, offset, frequency, variable='t'):
469  """
470  Transform all hit's variable from a linear space to a periodic space with some frequency and offset
471 
472  - offset = float that defines the offset
473  - frequency = float that defines the frequency of the periodic space
474  - variable = string from hit_set_variables. variable for which the periodicity is applied
475 
476  e.g. bunch.period_transformation( dt, rf_freq) would transform all times into a range between
477  (-0.5/rf_freq, 0.5/rf_freq). Hits with 't' = n*dt would transform to t=0 (n arbitrary integer)
478  Used to transform a set of many rf buckets into one rf bucket, for calculating longitudinal emittance etc
479  """
480  for hit in self:
481  hit[variable] -= offset #all points relative to offset
482  number_pulses = hit[variable]*frequency+0.5
483  # if number_pulses < 0.: number_pulses -= 1. #this is weird!
484  hit[variable] -= math.floor(number_pulses)/frequency
485 
486  def hits(self):
487  """Returns the list of all hits in the bunch"""
488  return self.__hits
489 
490  def get_hits(self, variable, value, comparator=operator.eq):
491  """
492  Returns a list of all hits in the bunch with comparator(variable, value) == True
493  By default comparator is equality operator, so returns a list of all hits in the bunch with variable == value
494 
495  - variable = string from Hit.get_variables()
496  - value = value for comparison
497  - comparator = function that operates on a hit and returns a boolean
498 
499  e.g. get_hits('eventNumber', 50) returns a list of all hits in the bunch with hitNumber 50
500  e.g. get_hits('eventNumber', 50, operator.lt) returns a list of all hits in the bunch with hitNumber less than 50
501  """
502  some_hits = []
503  for key in self.__hits:
504  if comparator(self.get_hit_variable(key, variable), value):
505  some_hits.append(key)
506  return some_hits
508  def hit_equality(self, target, float_tolerance=Common.float_tolerance):
509  """
510  Return True if all hits in self are the same as all hits in target
511 
512  - target = bunch to test against
513  """
514  if len(self) != len(target): return False
515  for i in range( len(self) ):
516  if self[i] != target[i]: return False
517  return True
518 
519  def conditional_remove(self, variable_value_dict, comparator, value_is_nsigma_bool = False):
520  """
521  For when a cut is not good enough, remove hits from the bunch altogether if comparator(value, get([variable])) is False
522 
523  - variable_value_dict = dict of Bunch.get_variables() to the cut value; cut if the hit variable has value > value
524  - value_is_nsigma_bool = boolean; if True, value is the number of standard deviations
525  - global_cut = boolean; if True, apply cut to global weights; else apply to local weights
526 
527  e.g. bunch.cut({'energy':300,}, operator.ge) removes particles if they have
528  ge(hit.get('energy'),300) here operator.ge is a functional representation of the >= operator
529  ge(hit.get('energy'),300) is the same as (hit.get('energy') >= 300)
530 
531  A whole load of useful comparators can be found in the operator module. e.g. ge is >=; le is <=; etc
532  See also python built-in function "filter".
533  """
534  delete_items = []
535  for variable, cut_value in variable_value_dict.iteritems():
536  if(value_is_nsigma_bool):
537  cut_value *= self.standard_deviation(variable, self.mean([variable]))**0.5
538  #optimisation for amplitude cut; calculate covariance matrix first
539  hit_value_list = self.list_get_hit_variable([variable])[0]
540  for i in range ( len(hit_value_list) ):
541  if comparator(hit_value_list[i], cut_value):
542  delete_items.append(self.__hits[i])
543  for item in delete_items:
544  self.__hits.remove(item)
545 
546  def standard_deviation(self, variable, variable_mean={}):
547  """
548  Return a standard deviation, given by sum( (u-u_mean)**2 )**0.5 where u is the variable specified
549 
550  - variable = a variable from Hit.get_variables(). Can be a list, in which
551  case only the first hit is taken
552  - variable_mean = use this as u_mean rather than calculating a priori. Can
553  be a dict, in which case if variable is in the keys will
554  use that value.
555 
556  e.g. bunch.standard_deviation('x') returns the standard deviation of x
557 
558  e.g. bunch.standard_deviation('x', 1.) returns the standard deviation of x
559  about a mean of 1.
560 
561  e.g. bunch.standard_deviation(['x'], {'x':1.}) returns the standard
562  deviation of x about a mean of 1.
563  """
564  a_variable = variable
565  if type(variable) == type([]):
566  a_variable = variable[0]
567  a_variable_mean = variable_mean
568  if type(variable_mean) != type({}):
569  a_variable_mean = {a_variable:variable_mean}
570  return self.moment(['x', 'x'], a_variable_mean)**0.5
571 
572 
573  def moment(self, variable_list, variable_mean_dict={}):
574  """
575  Return a moment, sum( (u1 - u1_mean)*(u2-u2_mean)*(etc)) where u is a set of Hit get_variables indexed by variable_list
576 
577  - variable_list = list of strings that index the moment
578  - variable_mean_dict = dict of variable to means for moment calculation. Assume bunch mean if no value is specified
579 
580  e.g. bunch.moment(['x','x'], {'x':0.1}) returns the variance of x about a mean of 0.1
581 
582  e.g. bunch.moment(['x','px']) returns the covariance of x,px about the bunch mean
583  """
584  bunch_weight = self.bunch_weight()
585  if abs(bunch_weight) < 1e-9: raise ZeroDivisionError('Trying to find moment of bunch with 0 weight')
586  mean_list = []
587  if self.__covs != None and len(variable_list) == 2:
588  try:
589  i1 = Bunch.__axis_list.index(variable_list[0])
590  i2 = Bunch.__axis_list.index(variable_list[1])
591  return float(self.__covs[i1,i2]) # not numpy.float64
592  except ValueError:
593  pass
594  if variable_mean_dict == {}:
595  variable_mean_dict = self.mean(variable_list)
596  for var in variable_list:
597  mean_list.append(variable_mean_dict[var])
598  try:
599  moment = self.__bunchcore.moment(variable_list, mean_list)
600  return moment
601  except: #variables not in bunchcore.get_variable list
602  moment = 0
603  for hit in self.__hits:
604  change = 1.
605  for index in range(len(variable_list)):
606  change *= self.get_hit_variable(hit, variable_list[index]) - mean_list[index]
607  moment += change*hit.get('weight')
608  return moment/self.bunch_weight()
609 
610  def mean(self, variable_list):
611  """
612  Return dict of variables to means
613 
614  - variable_list = list of strings to calculate means
615 
616  e.g. bunch.mean(['x','px']) returns a dict like {'x':x_mean, 'px':px_mean}
617  """
618  variable_mean_dict = {}
619  for var in variable_list:
620  if not var in variable_mean_dict:
621  if not var in self.__means:
622  variable_mean_dict[var] = self.moment([var], {var:0.0})
623  else: variable_mean_dict[var] = self.__means[var]
624  return variable_mean_dict
625 
626  def covariance_matrix(self, get_variable_list, origin_dict={}):
627  """
628  Return covariance matrix of variables
629 
630  - variable_list = list of strings to calculate covariance
631  - origin_dict = dict of strings to values for origin about which covariances are calculated. Defaults to mean
632 
633  e.g. bunch.covariance_matrix(['x','px'], {'x':0, 'px':0}) returns a matrix like [[Var(x,x), Var(x,px)], [Var(x,px), Var(px,px)]
634  """
635  Common.has_numpy()
636  dim = len(get_variable_list)
637  cov = matrix(numpy.zeros((dim,dim)))
638  origin_dict1 = copy.deepcopy(origin_dict)
639  origin_list = []
640  for var in get_variable_list:
641  if not var in origin_dict1:
642  origin_dict1[var] = self.mean([var])[var]
643  origin_list.append(origin_dict1[var])
644  try:
645  covariances = self.__bunchcore.covariance_matrix(get_variable_list, origin_list)
646  for i1 in range(dim):
647  for i2 in range(i1, dim):
648  cov[i1,i2] = covariance_list[i1][i2]
649  return cov
650  except:
651  pass
652  for i1 in range(dim):
653  for i2 in range(i1, dim):
654  covariance_list = [get_variable_list[i1], get_variable_list[i2]]
655  cov[i1,i2] = self.moment(covariance_list, origin_dict1)
656  cov[i2,i1] = cov[i1,i2]
657  return cov
658 
659  def set_covariance_matrix(self, use_internal_covariance_matrix=True, covariance_matrix=None, mean_dict={}):
660  """
661  Choose whether to use an internal covariance matrix for calculations
662 
663  - use_internal_covariance_matrix = boolean. If set to True, x-boa will recalculate its internal covariance matrix and use this
664  for some calculations. If set to False, x-boa will calculate the covariance matrix each time.
665  - covariance_matrix = NumPy matrix. x-boa will use this matrix as the internal covariance matrix; should be a 10x10 matrix ordered like
666  (x, px, y, py, z, pz, t, enegry, ct, energy). Ignored if equal to None
667  - mean_dict = dict of variables to means. x-boa will use this dict to index means; should contain means of
668  (x, px, y, py, z, pz, t, enegry, ct, energy). Ignored if equal to None
669 
670  As a speed optimisation, x-boa can calculate a covariance matrix and use this for all calculations involving covariances, i.e. Twiss parameters, emittances, amplitudes etc. Otherwise x-boa will re-calculate this each time, which can be slow. Be careful though - x-boa does not automatically detect for hits being added or removed from the bunch, etc. The user must call this function each time the bunch changes (events added, weightings changed, etc) to update the internal covariance matrix
671  """
672  if use_internal_covariance_matrix:
673  if covariance_matrix == None:
674  self.__covs = self.covariance_matrix(Bunch.axis_list_to_covariance_list(Bunch.__axis_list))
675  else: self.__covs = covariance_matrix
676  if mean_dict == {}:
677  self.__means = self.mean(Bunch.axis_list_to_covariance_list(Bunch.__axis_list))
678  else: self.__means = mean_dict
679  else:
680  self.__covs = None
681  self.__means = {}
682 
683  def covariances_set(self):
684  """If internal covariances are set by set_covariance_matrix, return True; else return False"""
685  return self.__covs != None
686 
687 
688  def means_set(self):
689  """If means are set by set_covariance_matrix, return True; else return False"""
690  return self.__means != {}
691 
692  def bunch_weight(self):
693  """Return statistical weight of all hits in the bunch"""
694  weight = 0
695  for key in self.__hits:
696  weight += key.get('weight')
697  return weight
698 
699  def clear_local_weights(self):
700  """Set local_weight of all hits in the bunch to 1"""
701  for key in self.__hits:
702  key.set('local_weight', 1)
703 
704  def clear_global_weights():
705  """Set global_weight of all hits in the bunch to 1"""
706  Hit.clear_global_weights()
707  clear_global_weights = staticmethod(clear_global_weights)
709  def clear_weights(self):
710  """Set global_weight and local_weight of all hits in the bunch to 1"""
711  self.clear_local_weights()
712  self.clear_global_weights()
713 
714  def cut(self, variable_value_dict, comparator, value_is_nsigma_bool = False, global_cut = False):
715  """
716  Set weight of hits in the bunch to 0 if comparator(value, get([variable])) is False
717 
718  - variable_value_dict = dict of Bunch.get_variables() to the cut value; cut if the hit variable has value > value
719  - comparator = callable; comparator(value, cut_value) returns true, cut the variable (set weight to 0.)
720  - value_is_nsigma_bool = boolean; if True, value is the number of standard deviations
721  - global_cut = boolean; if True, apply cut to global weights; else apply to local weights
723  e.g. bunch.cut({'energy':300,}, operator.ge) sets weight of particles to zero if they have
724  ge(hit.get('energy'),300) here operator.ge is a functional representation of the >= operator
725  ge(hit.get('energy'),300) is the same as (hit.get('energy') >= 300)
726 
727  A whole load of useful comparators can be found in the operator module. e.g. ge is >=; le is <=; etc
728  """
729  if global_cut: set_var = 'global_weight'
730  else: set_var = 'local_weight'
731  for variable, cut_value in variable_value_dict.iteritems():
732  if(value_is_nsigma_bool):
733  cut_value *= self.moment([variable, variable], self.mean([variable]))**0.5
734  # check for non-int get variables
735  if variable in Hitcore.get_variables() and type(Hit()[variable]) == type(0.):
736  self.__bunchcore.cut_double(variable, comparator, cut_value, global_cut)
737  else:
738  self._cut(variable, comparator, cut_value, set_var)
739 
740  def _cut(self, variable, comparator, cut_value, set_var):
741  hit_value_list = self.list_get_hit_variable([variable])[0]
742  for i in range ( len(hit_value_list) ):
743  if comparator(hit_value_list[i], cut_value):
744  self.__hits[i].set(set_var, 0)
745 
746  def transmission_cut(self, test_bunch, global_cut=False, test_variable=['spill', 'event_number', 'particle_number'], float_tolerance=Common.float_tolerance):
747  """
748  Set weight of hits in the bunch to 0 if no events can be found in test_bunch with the
749  same test_variable
750 
751  - test_bunch = bunch to test against
752  - global_cut = if True, apply cut to global weights
753  - test_variable = cut a hit in self if no hits are found in test_bunch with
754  the same test_variable. If test_variable is a list, then
755  cut from self if no hits are found in test_bunch with all
756  test_variables the same.
757 
758  E.g. bunch.transmission_cut( some_other_bunch, True ) will apply a global
759  cut if a hit with the same [spill, event_number, particle_number] is
760  not in some_other_bunch (for all hits)
761  """
762  if global_cut: wt = 'global_weight'
763  else: wt = 'local_weight'
764  hit_list = [[] for i in range(len(test_bunch))]
765  if type(test_variable) == type(''):
766  test_variable = [test_variable]
767  for i in range( len(test_bunch) ):
768  for var in test_variable:
769  hit_list[i].append(test_bunch[i][var])
770  hit_list.sort()
771  for hit in self:
772  weight = 0
773  hit_value = [hit[var] for var in test_variable]
774  next_value = bisect.bisect_left(hit_list, hit_value)
775  if next_value > -1 and next_value < len(hit_list):
776  diff = [abs(hit_list[next_value][i]-hit_value[i]) for i in range(len(hit_value))]
777  diff = sum(diff)
778  if diff < float_tolerance: weight = hit[wt]
779  if next_value+1 < len(hit_list):
780  diff = [abs(hit_list[next_value+1][i]-hit_value[i]) for i in range(len(hit_value))]
781  diff = sum(diff)
782  if diff < float_tolerance: weight = hit[wt]
783  hit.set(wt, weight)
785  def get_beta (self, axis_list, geometric=None):
786  """
787  Return the bunch beta function, defined by
788 
789  - For normalised beta = Sum_{i} Var(x_i)*p/(emittance*mass*n)
790  - For geometric beta = Sum_{i} Var(x_i)/(emittance*n)
791 
792  where x_i is a position variable from axis_list.
793 
794  - axis_list = list of axes strings
795 
796  and n is the length of axis_list. E.g.
797  ~~~~~~~~~~~~~~~{.py}
798  get_beta(['x','y'])
799  ~~~~~~~~~~~~~~~
800  will return ( Var(x,x)+Var(y,y) )/(2*pz*emittance(['x','y')*mass)
801  """
802  if geometric == None: geometric = Bunch.__geometric_momentum
803  beta = 0.0
804  emittance = self.get_emittance(axis_list, None, geometric)
805  for axis in axis_list:
806  beta += self.moment([axis,axis], self.mean([axis]))
807  if not geometric:
808  beta *= self.mean(['p'])['p']/(emittance*self.hits()[0].get('mass')*float(len(axis_list)))
809  else:
810  beta /= emittance*float(len(axis_list))
811  return beta
812 
813  def get_gamma(self, axis_list, geometric=None):
814  """
815  Return the bunch gamma function, defined by
816 
817  - For normalised gamma = Sum_{i} Var(p_i)/(pz*emittance*mass*n)
818  - For geometric gamma = Sum_{i} Var(p_i)/(emittance*n)
819 
820  where p_i is a momentum variable from axis_list.
821 
822  - axis_list = list of axes strings
823 
824  and n is the length of axis_list. E.g.
825  ~~~~~~~~~~~~~~~{.py}
826  get_gamma(['x','y'])
827  ~~~~~~~~~~~~~~~
828  will return ( Var(p_x)+Var(p_y) )/(2*pz*emittance(['x','y')*mass)
829  """
830  if geometric == None: geometric = Bunch.__geometric_momentum
831  gamma = 0.0
832  emittance = self.get_emittance(axis_list)
833  for axis in axis_list:
834  mom = self.momentum_variable(axis, geometric)
835  gamma += self.moment([mom,mom], self.mean([mom]))
836  if not geometric:
837  gamma /= self.mean(['p'])['p']*self.hits()[0].get('mass')
838  gamma /= (emittance*len(axis_list))
839  return gamma
840 
841  def get_alpha(self, axis_list, geometric=None):
842  """
843  Return the bunch alpha function, defined by
844 
845  - For normalised alpha = Sum_{i} Cov(x_i, p_i)/(emittance*mass*n)
846  - For geometric alpha = Sum_{i} Cov(x_i, p_i)/(emittance*n)
848  where p_i is a momentum variable from axis_list.
849 
850  - axis_list = list of axes strings
851 
852  and n is the length of axis_list. E.g.
853  ~~~~~~~~~~~~~~~{.py}
854  get_alpha(['x','y'])
855  ~~~~~~~~~~~~~~~
856  will return ( Var(p_x)+Var(p_y) )/(2*emittance(['x','y')*mass)
857  """
858  if geometric == None: geometric = Bunch.__geometric_momentum
859  alpha = 0.0
860  emittance = self.get_emittance(axis_list)
861  for axis in axis_list:
862  mom = self.momentum_variable(axis, geometric)
863  alpha += self.moment([axis,mom], self.mean([axis, mom]))
864  if not geometric:
865  alpha /= (emittance*self.hits()[0].get('mass')*len(axis_list))
866  else:
867  alpha /= emittance*len(axis_list)
868  return -alpha
869 
870  def get_dispersion(self, axis):
871  """
872  Return the bunch dispersion D in axis q, defined by
873 
874  - D = cov(q,E)*mean(E)/var(E)
875  - axis = string from axis_list that defines the axis q
877  E.g.
878  ~~~~~~~~~~~~~~~{.py}
879  my_bunch.get_dispersion('y')
880  ~~~~~~~~~~~~~~~
881  would return the dispersion in the y-direction
882  """
883  return self.moment([axis,'energy'])*self.mean(['energy'])['energy']/self.moment(['energy','energy'])
884 
885  def get_dispersion_rsquared(self):
886  """
887  Return a special bunch dispersion defined by
888 
889  - D = (<r2,E>-<r2>*<E>)*mean(E)/var(E)
890 
891  where <> are raw moments (not central moments)
892  """
893  mean_dict = {'r_squared':0., 'energy':0.}
894  return (self.moment(['r_squared','energy'],mean_dict)-(self.moment(['x','x'])+self.moment(['y','y']))*self.mean(['energy'])['energy'])\
895  *self.mean(['energy'])['energy']/self.moment(['energy','energy'])
896 
897 
898  def get_dispersion_prime(self, axis, geometric = None):
899  """
900  Return the bunch dispersion prime D in axis q, defined by
901 
902  - D' = cov(p_q,E)*mean(E)/var(E)
903  - axis = string from axis_list that defines the axis q
904 
905  E.g.
906  ~~~~~~~~~~~~~~~{.py}
907  my_bunch.get_dispersion('y')
908  ~~~~~~~~~~~~~~~
909  would return the dispersion in the y-direction
910  """
911  mom = Bunch.momentum_variable(axis, geometric)
912  return self.moment([mom,'energy'])*self.mean(['energy'])['energy']/self.moment(['energy','energy'])
913 
914  def get_decoupled_emittance(self, coupled_axis_list, emittance_axis_list, covariance_matrix=None, geometric=None):
915  """Reserved for future development"""
916  # the problem here is the sort order of the eigenvalues, which doesn't necessarily reflect what was in the
917  # original array
918  if geometric == None: geometric = Bunch.__geometric_momentum
919  coupled_axis_list2 = []
920  emit_var = []
921  for ax in Bunch.__axis_list:
922  if ax in emittance_axis_list:
923  emit_var.append(len(coupled_axis_list2)) #mark position of emittance variables
924  if ax in coupled_axis_list or ax in emittance_axis_list:
925  coupled_axis_list2.append(ax) #build sorted list of coupled variables
926  cov_var_list = Bunch.axis_list_to_covariance_list(coupled_axis_list2)
927  my_cov = copy.deepcopy(covariance_matrix)
928  if my_cov == None:
929  my_cov = self.__cov_mat_picker(coupled_axis_list2)
930  if my_cov == None:
931  my_cov = self.covariance_matrix(cov_var_list)
932  (evalues, ematrix) = numpy.linalg.eig(my_cov)
933  ematrix = matrix(ematrix)
934  my_cov = matrix(my_cov)
935  array_out = []
936  for var in emit_var:
937  array_out = array_out+[float(evalues[2*var]), float(evalues[2*var+1])]
938  print array_out
939  print 'Undiagonalisation:'
940  print ematrix*numpy.diagflat(evalues)*ematrix.T
941  print ematrix
942  print numpy.diagflat(evalues)
943  cov_out = numpy.diagflat(array_out)
944  return self.get_emittance(emittance_axis_list, covariance_matrix=cov_out, geometric=geometric)
945  """
946  Return the decoupled n dimensional normalised RMS emittance of the bunch
947 
948  - emittance = (|V|**(1/2n))/mass
949 
950  where n is the number of elements in axis list and V is a covariance matrix
951  with elements V_{ij} = cov(u_i, u_j) where u is a vector of elements in the
952  axis_list and their momentum conjugates, u = (q_i, p_i).
953  The slight subtlety is that the emittance is calculated with any coupling between
954  axis_list and the elements in coupled_axis_list removed by a transformation.
955  In particular, I calculate the eigenvalues of the matrix generated from
956  emittance_axis_list+coupled_axis_list (duplicated elements removed) and return the
957  product of eigenvaues pertaining to emittance_axis_list
958 
959  - coupled_axis_list = list of axes from Bunch.get_axes().
960  - emittance_axis_list = list of axes from Bunch.get_axes().
961  - covariance_matrix = if specified, use this covariance matrix rather than
962  calculating a new one each time.
963  - geometric = if specified, use geometric variables (dx/dz, dy/dz, etc) rather
964  than normalised variables (px, py, pz).
965 
966  E.g.
967  ~~~~~~~~~~~~~~~{.py}
968  get_decoupled_emittance(['x','y'], ['x'])
969  ~~~~~~~~~~~~~~~
970  will return
971  |V|**(1/2)
972  where V is the matrix with elements
973  V = Var(x',x') Cov(x',px')
974  Cov(x',px') Var(px',px')
975  and |V| is the determinant. Here x' indicates x following a decoupling
976  transformation (to remove, for example, rotation due to solenoid field).
977  """
978 
979  def get_emittance(self, axis_list, covariance_matrix=None, geometric = None):
980  """
981  Return the n dimensional normalised emittance of the bunch
982 
983  - emittance = (|V|**(1/2n))/mass
984 
985  where n is the number of elements in axis list and V is a covariance matrix
986  with elements V_{ij} = cov(u_i, u_j) where u is a vector of elements in the
987  axis_list and their momentum conjugates, u = (q_i, p_i)
988 
989  - axis_list = list of axes from Bunch.get_axes().
990  - covariance_matrix = if specified, use this covariance matrix rather than
991  calculating a new one each time.
992  - geometric = if specified, use geometric variables (dx/dz, dy/dz, etc) rather
993  than normalised variables (px, py, pz).
994 
995  E.g.
996  ~~~~~~~~~~~~~~~{.py}
997  get_emittance(['x'])
998  ~~~~~~~~~~~~~~~
999  will return
1000  |V|**(1/2)
1001  where V is the matrix with elements
1002  V = Var(x,x) Cov(x,px)
1003  Cov(x,px) Var(px,px)
1004  and |V| is the determinant.
1005  """
1006  if geometric == None: geometric = Bunch.__geometric_momentum
1007  cov_list = Bunch.axis_list_to_covariance_list(axis_list)
1008  my_cov = copy.deepcopy(covariance_matrix)
1009  if my_cov == None:
1010  my_cov = self.__cov_mat_picker(axis_list)
1011  if my_cov == None:
1012  my_cov = self.covariance_matrix(cov_list)
1013  emittance = linalg.det(my_cov)**(1./len(cov_list))
1014  if not Bunch.__geometric_momentum:
1015  emittance /= self.__hits[0].get('mass')
1016  return float(emittance)
1017 
1018  def get_kinetic_angular_momentum(self, rotation_axis_dict={'x':0,'y':0,'px':0,'py':0,'x\'':0,'y\'':0}):
1019  """
1020  Return the bunch kinetic angular momentum about some arbitrary axis, defined by
1021 
1022  - L = <x py> - <y px> (momentum variables)
1023  - L = <p>*(<x y'> - <y x'>) (geometric variables)
1024 
1025  - rotation_axis_dict = dict that defines the axis of rotation
1026  """
1027  if not Bunch.__geometric_momentum:
1028  return self.moment(['x','py'], rotation_axis_dict) - self.moment(['y','px'], rotation_axis_dict)
1029  else:
1030  return (self.moment(['x','y\''], rotation_axis_dict) - self.moment(['y','x\''], rotation_axis_dict)) * self.mean(['pz'])['pz']
1031 
1032  def get_canonical_angular_momentum(self, bz=None, field_axis_dict={'x':0, 'y':0}, rotation_axis_dict={'x':0,'y':0,'px':0,'py':0}):
1033  """
1034  Return the bunch canonical angular momentum about some arbitrary axis, defined by
1035 
1036  - L = <x py> - <y px> + e Bz (<x^2>+y^2) (momentum variables)
1037  - L = <p>*(<x y'> - <y x'>) + e Bz (<x^2>+y^2) (geometric variables)
1038 
1039  - bz = nominal on-axis magnetic field - where axis is as; defaults to the field of the 1st particle in the bunch
1040  - field_axis_dict = dict that defines the nominal solenoid axis
1041  - rotation_axis_dict = dict that defines the axis of rotation
1042  """
1043  if bz==None:
1044  bz = (self[0]['bz'])
1045  return self.get_kinetic_angular_momentum(rotation_axis_dict) + \
1046  Common.constants['c_light']*bz*(self.moment(['x','x'], field_axis_dict) + self.moment(['y','y'], field_axis_dict))/2.
1047 
1048  def momentum_variable(axis_variable, geometric_momentum):
1049  """
1050  Return the momentum conjugate for axis_variable
1051  """
1052  if not axis_variable in Bunch.__axis_list:
1053  raise IndexError(str(axis_variable)+' does not appear to be an axis variable. Options are '+str(Bunch.get_axes()))
1054  if type(axis_variable) == int:
1055  if geometric_momentum:
1056  return Hit.get_variables()[axis_variable]+'\''
1057  else:
1058  return Hit.get_variables()[axis_variable+4]
1059  if geometric_momentum:
1060  return axis_variable+'\''
1061  elif axis_variable == 't' or axis_variable == 'ct' or axis_variable==3:
1062  return 'energy'
1063  else:
1064  return 'p'+axis_variable
1065  momentum_variable = staticmethod(momentum_variable)
1066 
1067  def get_amplitude (bunch, hit, axis_list, covariance_matrix=None, mean_dict={}, geometric=None):
1068  """
1069  Return the particle amplitude for a hit relative to a bunch, defined by
1070 
1071  - amplitude = emittance*x^T.V^-1.x
1072 
1073  where x is a vector of particle coordinates and V is a covariance matrix
1074 
1075  - bunch = bunch from which the covariance matrix is calculated
1076  - hit = hit from which the particle vector is taken
1077  - axis_list = list of axes that defines the covariance matrix and particle vector
1078  - covariance_matrix = if this is not set to None, will use this covariance_matrix
1079  for the calculation rather than taking one from bunch. Should be a numpy matrix
1080  like "x","px","y","py"
1081  E.g.
1082  ~~~~~~~{.py}
1083  Bunch.get_amplitude(my_bunch, my_hit, ['x','y'])
1084  ~~~~~~~
1085  will return
1086  emittance*(x,px,y,py)^T*V^{-1}(x,px,y,py)*(x,px,y,py)
1087  """
1088  if geometric == None: geometric = Bunch.__geometric_momentum
1089  cov_list = Bunch.axis_list_to_covariance_list(axis_list, geometric)
1090  if mean_dict == {}:
1091  mean_dict = bunch.__mean_picker(axis_list)
1092  if mean_dict == {}:
1093  mean_dict = bunch.mean(cov_list)
1094  my_vector = hit.get_vector(cov_list, mean_dict)
1095  my_cov = copy.deepcopy( covariance_matrix )
1096  if my_cov == None:
1097  my_cov = bunch.__cov_mat_picker(axis_list)
1098  if my_cov == None:
1099  my_cov = bunch.covariance_matrix(cov_list)
1100  my_amp = my_vector*linalg.inv(my_cov)*my_vector.T
1101  return float(my_amp[0,0])*bunch.get_emittance(axis_list, my_cov)
1102  get_amplitude = staticmethod(get_amplitude)
1103 
1104  def axis_list_to_covariance_list(axis_list, geometric=None):
1105  """
1106  Convert from a list of position variables to a list of position
1107  and conjugate momentum variables
1108 
1109  - axis_list = list of strings from get_axes()
1110  """
1111  if geometric == None: geometric = Bunch.__geometric_momentum
1112  cov_list = []
1113  for axis in axis_list:
1114  cov_list.append(axis)
1115  cov_list.append(Bunch.momentum_variable(axis, geometric))
1116  return cov_list
1117  axis_list_to_covariance_list = staticmethod(axis_list_to_covariance_list)
1118 
1119  def convert_string_to_axis_list(axis_string):
1120  """
1121  Return a list of axes from a string.
1122 
1123  - axis_string = string that contains a list of axis. Either format
1124  as axis_strings separated by white space or as a python list
1125  of axis_strings
1126  """
1127  axis_list = []
1128  try:
1129  axis_list = eval(axis_string) #try to evaluate it as a list
1130  except:
1131  axis_list = axis_string.split(' ') #assume it is a white space separated dataset instead
1132  for key in axis_list:
1133  if not key in Bunch.__axis_list:
1134  raise KeyError('Could not resolve '+str(axis_string)+' into a list of axes.')
1135  return axis_list
1136  convert_string_to_axis_list = staticmethod(convert_string_to_axis_list)
1137 
1138  def set_geometric_momentum(new_value_bool):
1139  """
1140  Set the default for conjugate momenta; either geometric (x', y', etc) or kinetic (px, py, etc)
1141 
1142  - new_value_bool = if True, use geometric momenta. If False, use kinetic momenta (default)
1143  """
1144  Bunch.__geometric_momentum = new_value_bool
1145  set_geometric_momentum = staticmethod(set_geometric_momentum)
1146 
1147  def get_geometric_momentum():
1148  """
1149  Set the default for conjugate momenta; either geometric (x', y', etc) or kinetic (px, py, etc)
1150 
1151  - new_value_bool = if True, use geometric momenta. If False, use kinetic momenta (default)
1152  """
1153  return Bunch.__geometric_momentum
1154  get_geometric_momentum = staticmethod(get_geometric_momentum)
1155 
1156  def get_axes():
1157  """Return list of axis variables (position-type variables)"""
1158  return Bunch.__axis_list
1159  get_axes = staticmethod(get_axes)
1160 
1161  def get_hit_variable(self, hit, variable_name, covariance_matrix=None, mean_dict={}):
1162  """
1163  Return axis variable for hit. Special power is that it can calculate amplitude, unlike hit.get(blah)
1164 
1165  - hit = a Hit object
1166  - variable_name = either a variable from Hit.get_variables() or an amplitude variable.
1167  amplitude variables should be formatted like 'amplitude x y' or 'amplitude x y t'
1168  - covariance_matrix = use a pre-calculated covariance matrix, or calculate if set to None
1169  """
1170  if type(variable_name) == str:
1171  if variable_name.find('amplitude') > -1:
1172  axis_list = Bunch.convert_string_to_axis_list(variable_name[10:len(variable_name)])
1173  return Bunch.get_amplitude(self, hit, axis_list, covariance_matrix, mean_dict)
1174  return hit.get(variable_name)
1176  def list_get_hit_variable(self, list_of_variables, list_of_units=[], covariance_matrix = None, mean_dict = {}):
1177  """
1178  Return a list of get_hit_variable results, one element for each hit in the bunch
1179 
1180  - variable_name = either a variable from Hit.get_variables() or an amplitude variable.
1181  amplitude variables should be formatted like 'amplitude x y' or 'amplitude x y t'
1182  i.e. amplitude followed by a list of axes separated by white space
1183  - covariance_matrix = use a pre-calculated covariance matrix, or calculate if set to None
1184  """
1185  values = []
1186  for dummy in list_of_variables:
1187  values.append([])
1188  for i in range(len(list_of_variables)):
1189  var = list_of_variables[i]
1190  if type(var) is str:
1191  if(var.find('amplitude') > -1):
1192  axis_list = Bunch.convert_string_to_axis_list(var[10:len(var)])
1193  if covariance_matrix == None:
1194  covariance_matrix = self.covariance_matrix( Bunch.axis_list_to_covariance_list(axis_list) )
1195  if mean_dict == {}:
1196  mean_dict = self.mean(Bunch.axis_list_to_covariance_list(axis_list))
1197  for hit in self.__hits:
1198  values[i].append( self.get_hit_variable(hit, var, covariance_matrix, mean_dict) )
1199  if not list_of_units == []:
1200  values[i][-1] /= Common.units[list_of_units[i]]
1201  return values
1202 
1203  def get(self, variable_string, variable_list):
1204  """
1205  Return a bunch variable taken from the list Bunch.get_variables()
1206 
1207  - variable_string = string variable from Bunch.get_variables()
1208  - variable_list = list of variables (see below, a bit complicated)
1209 
1210  For 'bunch_weight', axis_list is ignored
1211 
1212  For 'dispersion', 'dispersion_prime' the first value of variable_list is used. It should be taken from the list Bunch.get_axes()
1214  For 'mean', the first value in variable_list is used. It should be taken from the list Bunch.hit_get_variables()
1216  For others, variable_list should be a list taken from Bunch.get_axes()
1217 
1218  E.g. bunch.get('emittance', ['x','y']) would get emittance x y
1220  E.g. bunch.get('mean', ['x','px','z']) would return mean of x; px and z are ignored
1221  """
1222  if not variable_string in Bunch.__get_dict:
1223  raise KeyError(variable_string+' not available. Options are '+str(Bunch.get_variables()))
1224  else:
1225  return self.__get_dict[variable_string](self, variable_list)
1226 
1227  def list_get(dict_of_bunches, list_of_variables, list_of_axis_lists):
1228  """
1229  Get a list of lists of variables from each bunch in the dict_of_bunches
1230 
1231  - dict_of_bunches = dict of bunches from which the lists will be taken
1232  - list_of_variables = list of variables that will be used in the calculation
1233  - list_of_axis_lists = axis_list that will be used in the calculation of the
1234  corresponding variable
1235 
1236  E.g. list_get(my_dict_of_bunches, ['mean', 'emittance'], [['z'],['x','y']])
1237  would calculate a list of two arrays: (i) mean z; (ii) emittance x y
1238 
1239  Output is sorted by the first variable in the list. This might be useful to interface
1240  with E.g. a plotting package
1241  """
1242  values = []
1243  for dummy in list_of_variables:
1244  values.append([])
1245  for key, bunch in dict_of_bunches.iteritems():
1246  for i in range(len(list_of_variables)):
1247  values[i].append(bunch.get(list_of_variables[i], list_of_axis_lists[i]))
1248  Common.multisort(values)
1249  return values
1250  list_get = staticmethod(list_get)
1251 
1252  def get_variables():
1253  """Return a list of variables suitable for calls to Bunch.get"""
1254  return Bunch.__get_list
1255  get_variables = staticmethod(get_variables)
1256 
1257  def hit_get_variables():
1258  """Return a list of variables suitable for calls to Bunch.get"""
1259  return Hit.get_variables()+['amplitude x', 'amplitude y', 'amplitude x y', 'amplitude ct', 'amplitude x y ct']
1260  hit_get_variables = staticmethod(hit_get_variables)
1261 
1262  def hit_write_builtin(self, file_type_string, file_name, user_comment=None):
1263  """
1264  Write the hits in the bunch using some built-in format to the file file_name
1265 
1266  - file_type_string = string that controls which predefined file type will be used
1267  - file_name = string that defines the file name
1268  - user_comment = comment to be included in file header (e.g. problem title)
1269  """
1270  if not file_type_string in Hit.file_types(): raise IOError('Attempt to write file of unknown file type '+str(file_type_string)+' - try one of '+str(Hit.file_types()))
1271  Hit.write_list_builtin_formatted(self.__hits, file_type_string, file_name, user_comment)
1272  return
1273 
1274  def hit_write_builtin_from_dict(dict_of_bunches, file_type_string, file_name, user_comment=None):
1275  """
1276  Write the hits in the dict of bunches using some built-in format to the file file_name
1277 
1278  - dict_of_bunches = list of hits
1279  - file_type_string = string that controls which predefined file type will be used
1280  - file_name = string that defines the file name
1281  - user_comment = comment to be included in file header (e.g. problem title)
1282  """
1283  if not file_type_string in Hit.file_types(): raise IOError('Attempt to write file of unknown file type '+str(file_type_string)+' - try one of '+str(Hit.file_types()))
1284  all_hits = []
1285  for key,value in dict_of_bunches.iteritems():
1286  all_hits += value.hits()
1287  Hit.write_list_builtin_formatted(all_hits, file_type_string, file_name)
1288  hit_write_builtin_from_dict = staticmethod(hit_write_builtin_from_dict)
1289 
1290  def hit_write_user(self, format_list, format_units_dict, file_handle, separator=' ', comparator=None):
1291  """
1292  Write the hits in the bunch in some user defined format to file_handle, in an order defined by comparator
1293 
1294  - format_list = ordered list of strings from get_variables()
1295  - format_units_dict = dict of formats from format_list to units
1296  - file_handle = file handle made using e.g. open() command to which hits are written
1297  - separator = separator that is used to separate hits
1298  - comparator = comparator that is used for sorting prior to writing the file; if left as None, no sorting is performed
1299 
1300  e.g. aBunch.hit_write_user(['x','px','y','py'], ['x':'m','y':'m','px':'MeV/c','py':'MeV/c'], some_file, '@') would make output like
1301  0.001@0.002@0.001@0.002
1302  0.003@0.004@0.003@0.004
1303  in some_file
1304  """
1305  if not comparator==None:
1306  self.__hits.sort(comparator)
1307  for hit in self.__hits:
1308  hit.write_user_formatted(format_list, format_units_dict, file_handle, separator)
1309 
1310  def build_ellipse_2d(beta, alpha, emittance, p, mass, geometric):
1311  """
1312  Build a 2x2 ellipse matrix, given by
1313  For geometric = true
1314 
1315  - var(x, x) = beta*emittance
1316  - var(x, x') = alpha*emittance
1317  - var(x',x') = gamma*emittance
1318 
1319  For geometric = false
1321  - var(x, x ) = beta*emittance*mass/p
1322  - var(x, px) = alpha*emittance*mass
1323  - var(px,px) = gamma*emittance*mass*p
1324  """
1325  Common.has_numpy()
1326  cov = matrix(numpy.zeros((2,2)))
1327  gamma = (1.+alpha**2)/beta
1328  if beta <= 0.:
1329  raise ValueError("Beta "+str(beta)+" must be positive")
1330  if emittance <= 0.:
1331  raise ValueError("Emittance "+str(emittance)+" must be positive")
1332  if geometric:
1333  cov[0,0] = beta *emittance
1334  cov[1,0] = -alpha*emittance
1335  cov[0,1] = cov[1,0]
1336  cov[1,1] = gamma*emittance
1337  if not geometric:
1338  cov[0,0] = beta *emittance*mass/p
1339  cov[1,0] = -alpha*emittance*mass
1340  cov[0,1] = cov[1,0]
1341  cov[1,1] = gamma*emittance*mass*p
1342  return cov
1343  build_ellipse_2d = staticmethod(build_ellipse_2d)
1344 
1345  def build_MR_ellipse():
1346  """
1347  Not implemented
1348  """
1349 # Build a 6x6 ellipse using the Mais-Ripken parameterisation - not implemented
1350  raise NotImplementedError('Mais-Ripken methods will be released later')
1351  build_MR_ellipse = staticmethod(build_MR_ellipse)
1353  def build_ET_ellipse():
1354  """
1355  Not implemented
1356  """
1357 # Build a 6x6 ellipse using the Edwards-Teng parameterisation
1358  raise NotImplementedError('Edwards-Teng methods will be released later')
1359  build_ET_ellipse = staticmethod(build_ET_ellipse)
1360 
1362  """
1363  Not implemented
1364  """
1365 # Build an ellipse using the transfer matrix M such that M.T() * ellipse * M = ellipse
1366  raise NotImplementedError('This function not ready yet')
1367  build_ellipse_from_transfer_matrix = staticmethod(build_ellipse_from_transfer_matrix)
1368 
1369  def build_penn_ellipse(emittance_t, mass, beta_t, alpha_t, p, Ltwiddle_t, bz, q):
1370  """
1371  Build an ellipse using Penn formalism for solenoids.
1372 
1373  - emittance_t = nominal beam emittance
1374  - mass = particle mass
1375  - beta_t = transverse beta function
1376  - alpha_t = transverse alpha function
1377  - p = reference momentum
1378  - Ltwiddle_t = normalised canonical angular momentum
1379  - bz = magnetic field strength
1380  - q = particle charge
1381 
1382  Output ellipse is a 4*4 matrix with elements v_ij where i,j indexes the vector (x,px,y,py)
1383  """
1384  if beta_t <= 0.:
1385  raise ValueError("Beta "+str(beta_t)+" must be positive")
1386  if emittance_t <= 0.:
1387  raise ValueError("Emittance "+str(emittance_t)+" must be positive")
1388  if mass <= 0.:
1389  raise ValueError("Mass "+str(mass)+" must be positive")
1390 
1391  k = Common.constants['c_light']*bz/2./p
1392  gamma_t = (1+alpha_t*alpha_t+(beta_t*k - Ltwiddle_t)*(beta_t*k-Ltwiddle_t))/beta_t
1393  s_xx = emittance_t*mass*beta_t/p
1394  s_xpx = -emittance_t*mass*alpha_t
1395  s_pxpx = emittance_t*mass*p*gamma_t
1396  s_xpy = -emittance_t*mass*(beta_t*k-Ltwiddle_t)*q
1397  covariance_matrix = numpy.matrix([
1398  [s_xx, s_xpx, 0., s_xpy],
1399  [s_xpx, s_pxpx,-s_xpy, 0],
1400  [ 0., -s_xpy, s_xx, s_xpx],
1401  [s_xpy, 0., s_xpx, s_pxpx],
1402  ])
1403  return covariance_matrix
1404  build_penn_ellipse = staticmethod(build_penn_ellipse)
1405 
1406  def histogram(self, x_axis_string, x_axis_units='', y_axis_string='', y_axis_units='', nx_bins=None, ny_bins=None, xmin=None, xmax=None, ymin=None, ymax=None):
1407  """
1408  Returns a binned 1D or 2D histogram of hits in the bunch (but doesnt draw anything or call any plotting package).
1409 
1410  - x_axis_string = string for call to get_hit_variables()
1411  - x_axis_units = units for x axis
1412  - y_axis_string = string for call to get_hit_variables(). If string is empty, makes a 1D histogram
1413  - y_axis_units = units for y axis
1414  - nx_bins = force histogram to use this number of bins for the x-axis rather than choosing number of bins itself
1415  - ny_bins = force histogram to use this number of bins for the y-axis rather than choosing number of bins itself
1416  - xmin = float that overrides auto-detection of minimum x-axis value
1417  - xmax = float that overrides auto-detection of maximum x-axis value
1418  - ymin = float that overrides auto-detection of minimum y-axis value
1419  - ymax = float that overrides auto-detection of maximum y-axis value
1420 
1421  Return value is a list of bin weights for a 1D histogram or a list of lists of bin weights for a 2D histogram
1422  """
1423  Common.has_numpy()
1424  weights = self.list_get_hit_variable(['weight'], [''])[0]
1425  x_points = self.list_get_hit_variable([x_axis_string], [x_axis_units])[0]
1426  y_bins = None
1427  if xmin!=None: xmin*=Common.units[x_axis_units]
1428  if xmax!=None: xmax*=Common.units[x_axis_units]
1429  if ymin!=None: ymin*=Common.units[y_axis_units]
1430  if ymax!=None: ymax*=Common.units[y_axis_units]
1431  if y_axis_string=='':
1432  (nx_bins, ny_bins, dummy) = Common.n_bins(len(x_points), nx_bins, 1)
1433  else:
1434  y_points = self.list_get_hit_variable([y_axis_string], [y_axis_units])[0]
1435  (nx_bins, ny_bins, dummy) = Common.n_bins(len(y_points), nx_bins, ny_bins, 2)
1436  y_bins = Common.get_bin_edges(y_points, ny_bins, ymin, ymax)
1437  x_bins = Common.get_bin_edges(x_points, nx_bins, xmin, xmax)
1438  return self.histogram_var_bins(x_axis_string, x_bins, x_axis_units, y_axis_string, y_bins, y_axis_units)
1440  def histogram_var_bins(self, x_axis_string, x_bins, x_axis_units='', y_axis_string='', y_bins=None, y_axis_units=''):
1441  """
1442  Returns a binned histogram of hits in the bunch.
1443 
1444  Requires numpy
1445 
1446  - x_axis_string = string for call to get_hit_variables()
1447  - x_bins = list of bin edges used to generate the histogram
1448  - x_axis_units = units for x axis
1449  - y_axis_string = string for call to get_hit_variables()
1450  - y_bins = list of bin edges used to generate the histogram
1451  - y_axis_units = units for y axis
1452 
1453  Return value is a tuple containing ((bin weights array),x_bins,y_bins). The bin weights array is always nx*ny, with ny defaulting to 1 in the case of a 1d histogram.
1454  """
1455  Common.has_numpy()
1456  weights = self.list_get_hit_variable(['weight'], [''])[0]
1457  x_points = self.list_get_hit_variable([x_axis_string], [x_axis_units])[0]
1458  if y_axis_string == '':
1459  return Common.histogram(x_points, x_bins, weights=weights)
1460  else:
1461  y_points = self.list_get_hit_variable([y_axis_string], [y_axis_units])[0]
1462  return Common.histogram(x_points, x_bins, y_points, y_bins, weights=weights)
1463 
1464  def root_histogram(self, x_axis_string, x_axis_units='', y_axis_string='', y_axis_units='', nx_bins=None, ny_bins=None, canvas=None, xmin=None, xmax=None, ymin=None, ymax=None,
1465  line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.fill_color, stats=rg.stats, hist_title_string='', draw_option=''):
1466  """
1467  Prints a 1D or 2D histogram of hits in the bunch. Axes are automatically detected
1468  to encapsulate all data.
1469 
1470  Needs correct root installation.
1471 
1472  - x_axis_string = string for call to get_hit_variables()
1473  - x_axis_units = units for x axis
1474  - y_axis_string = string for call to get_hit_variables(). If string is empty, makes a 1D histogram
1475  - y_axis_units = units for y axis
1476  - nx_bins = force root_histogram to use this number of bins for the x-axis rather than choosing number of bins itself
1477  - ny_bins = force root_histogram to use this number of bins for the y-axis rather than choosing number of bins itself
1478  - canvas = if canvas=None, root_histogram will create a new tcanvas and plot histograms on that
1479  else will plot on the existing canvas, attempting to plot on existing axes
1480  - xmin = float that overrides auto-detection of minimum x-axis value
1481  - xmax = float that overrides auto-detection of maximum x-axis value
1482  - ymin = float that overrides auto-detection of minimum y-axis value
1483  - ymax = float that overrides auto-detection of maximum y-axis value
1484  - line_color = int that sets the line colour of the histogram
1485  - line_style = int that sets the line style of the histogram
1486  - line_width = int that sets the line width of the histogram
1487  - fill_color = int that sets the fill color of the histogram
1488  - stats = set to True to plot a stats box on the histogram
1489  - hist_title_string = specify the string that will appear as a title on the canvas
1490 
1491  e.g. bunch.root_histogram('x', 'cm', 'py', 'GeV/c') will histogram x vs py.
1492  Returns a tuple of the (canvas, histogram)
1493  """
1494  weights = []
1495  num_events = float(len(self))
1496  x_points = self.list_get_hit_variable([x_axis_string], [x_axis_units])[0]
1497  if y_axis_string == '' and nx_bins==None:
1498  nx_bins = int(num_events/10.)
1499  if not y_axis_string == '' and draw_option=='':
1500  draw_option = 'COL'
1501  y_points = self.list_get_hit_variable([y_axis_string], [y_axis_units])[0]
1502  if nx_bins == None: nx_bins = int(num_events**0.7/10.)
1503  if ny_bins == None: ny_bins = int(num_events**0.7/10.)
1504  else: y_points = []
1505  for key in self.__hits:
1506  weights.append(key.get('weight'))
1507  if not x_axis_units == '': x_axis_units = " ["+x_axis_units+"]"
1508  if not y_axis_units == '': y_axis_units = " ["+y_axis_units+"]"
1509  x_name = x_axis_string+x_axis_units
1510  y_name = y_axis_string+y_axis_units
1511  num_evts = 0.
1512  for hit in self.__hits:
1513  if abs(hit.get('weight')) > 0.: num_evts += 1.
1514  name = x_axis_string
1515  if not y_axis_string == '': name += ":"+y_axis_string
1516  if canvas == '' or canvas == None: canvas = Common.make_root_canvas(name)
1517  else:
1518  draw_option += 'same'
1519  canvas.cd()
1520  hist = Common.make_root_histogram(name, x_points, x_name, nx_bins, y_points, y_name, ny_bins, weights, xmin, xmax, ymin, ymax,
1521  line_color, line_style, line_width, fill_color, stats, hist_title_string)
1522  hist.SetStats(False)
1523  hist.Draw(draw_option)
1524  canvas.Update()
1525  return (canvas, hist)
1526 
1527  def root_scatter_graph(self, x_axis_string, y_axis_string, x_axis_units='', y_axis_units='', include_weightless=True, canvas=None, xmin=None, xmax=None, ymin=None, ymax=None,
1528  line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.graph_fill_color, hist_title_string=''):
1529  """
1530  Prints a 2d scatter plot of Hit.get(...) data over a dict of bunches.
1531 
1532  Needs correct root installation.
1533 
1534  - x_axis_string = string for call to Bunch.get_hit_variable() used to calculate x-values
1535  - x_axis_units = units for x axis
1536  - y_axis_string = string for call to Bunch.get_hit_variable() used to calculate y-values
1537  - y_axis_units = units for y axis
1538  - include_weightless = set to True to plot hits with <= 0. weight
1539  - canvas = if canvas=None, root_graph will create a new tcanvas and plot graphs on that
1540  else will plot on the existing canvas, attempting to plot on existing axes
1541  - xmin = float that overrides auto-detection of minimum x-axis value
1542  - xmax = float that overrides auto-detection of maximum x-axis value
1543  - ymin = float that overrides auto-detection of minimum y-axis value
1544  - ymax = float that overrides auto-detection of maximum y-axis value
1545  - line_color = int that sets the line colour of the histogram
1546  - line_style = int that sets the line style of the histogram
1547  - line_width = int that sets the line width of the histogram
1548  - fill_color = int that sets the fill color of the histogram
1549  - hist_title_string = specify the string that will appear as a title on the canvas
1550  - fit_ellipse = set to True to draw a fitted ellipse corresponding to the relevant covariances matrix
1551 
1552  e.g. bunch.root_scatter_graph('x', 'p', 'cm', 'MeV/c') will graph x vs p.
1553  Returns a tuple of (canvas, histogram, graph) where the histogram contains
1554  the axes.
1555  """
1556  fit_ellipse=False # hard coded - I am not sure this is the right way
1557  x_points = []
1558  y_points = []
1559  for hit in self:
1560  if abs(hit['weight'])>Common.float_tolerance or include_weightless:
1561  x_points.append(self.get_hit_variable(hit, x_axis_string)/Common.units[x_axis_units])
1562  y_points.append(self.get_hit_variable(hit, y_axis_string)/Common.units[y_axis_units])
1563  name = x_axis_string+":"+y_axis_string
1564  if not x_axis_units == '': x_axis_string += " ["+x_axis_units+"]"
1565  if not y_axis_units == '': y_axis_string += " ["+y_axis_units+"]"
1566  graph = Common.make_root_graph(name, x_points, x_axis_string, y_points, y_axis_string, True, xmin, xmax, ymin, ymax,
1567  line_color, line_style, line_width, fill_color, hist_title_string)
1568  if canvas == '' or canvas==None:
1569  canvas = Common.make_root_canvas(name)
1570  graph[0].SetStats(False)
1571  graph[0].Draw()
1572  else:
1573  canvas.cd()
1574  graph[1].Draw('p')
1575  if fit_ellipse:
1576  points = [(x_points[i], y_points[i]) for i in range(len(x_points))]
1577  mean, cov = Common.fit_ellipse(points, -1., self.list_get_hit_variable(['weight'])[0], 1, True)
1578  func = Common.make_root_ellipse_function(mean, cov, [1.],
1579  graph[0].GetXaxis().GetXmin(),
1580  graph[0].GetXaxis().GetXmax(),
1581  graph[0].GetYaxis().GetXmin(),
1582  graph[0].GetYaxis().GetXmax())
1583  func.Draw('SAME')
1584  _function_persistent.append(fit_func)
1585  canvas.Update()
1586  return (canvas, graph[0], graph[1])
1587 
1588  def root_graph(bunches, x_axis_string, x_axis_list, y_axis_string, y_axis_list, x_axis_units='', y_axis_units='', canvas='', comparator=None, xmin=None, xmax=None, ymin=None, ymax=None,
1589  line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.graph_fill_color, hist_title_string=''):
1590  """
1591  Prints a graph of Bunch.get(...) data over a set of bunches and returns the root canvas.
1592 
1593  Needs correct root installation.
1594 
1595  - bunches = either a dict or a list of bunches
1596  - x_axis_string = string for call to Bunch.get() used to calculate x axis variables
1597  - x_axis_list = list of variables for call to Bunch.get() used to calculate x axis variables
1598  - x_axis_units = units for x axis
1599  - y_axis_string = string for call to Bunch.get() used to calculate y axis variables
1600  - y_axis_list = list of variables for call to Bunch.get() used to calculate y axis variables
1601  - y_axis_units = units for y axis
1602  - canvas = if canvas=None, root_graph will create a new tcanvas and plot graphs on that
1603  else will plot on the existing canvas, attempting to plot on existing axes
1604  - comparator = comparator of bunch1, bunch2 - if none given, will sort by x-axis value
1605  - xmin = float that overrides auto-detection of minimum x-axis value
1606  - xmax = float that overrides auto-detection of maximum x-axis value
1607  - ymin = float that overrides auto-detection of minimum y-axis value
1608  - ymax = float that overrides auto-detection of maximum y-axis value
1609  - line_color = int that sets the line colour of the histogram
1610  - line_style = int that sets the line style of the histogram
1611  - line_width = int that sets the line width of the histogram
1612  - fill_color = int that sets the fill color of the histogram
1613  - hist_title_string = specify the string that will appear as a title on the canvas
1614 
1615  e.g. root_graph(bunch_dict, 'mean', ['x'], 'emittance', ['x','y'], 'cm', 'mm') will graph mean x vs emittance x y.
1616  Returns a tuple of (canvas, histogram, graph) where the histogram contains the axes
1617  """
1618  x_points = []
1619  y_points = []
1620  list_of_bunches = []
1621  try: list_of_bunches = bunches.values()
1622  except: list_of_bunches = bunches
1623  if comparator != None:
1624  list_of_bunches.sort(comparator)
1625  for i, bunch in enumerate(list_of_bunches):
1626  try:
1627  x_points.append(bunch.get(x_axis_string, x_axis_list)/Common.units[x_axis_units])
1628  y_points.append(bunch.get(y_axis_string, y_axis_list)/Common.units[y_axis_units])
1629  except Exception:
1630  if len(x_points) > len(y_points):
1631  x_points.pop(-1)
1632  sys.excepthook(*sys.exc_info())
1633  print 'Failed to get data from bunch', i
1634  x_axis_string += "( "
1635  y_axis_string += "( "
1636  for value in x_axis_list: x_axis_string += value+" "
1637  for value in y_axis_list: y_axis_string += value+" "
1638  x_axis_string += ")"
1639  y_axis_string += ")"
1640  name = x_axis_string+":"+y_axis_string
1641  if not x_axis_units == '': x_axis_string += " ["+x_axis_units+"]"
1642  if not y_axis_units == '': y_axis_string += " ["+y_axis_units+"]"
1643  graph = Common.make_root_graph(name, x_points, x_axis_string, y_points, y_axis_string, comparator==None, xmin, xmax, ymin, ymax,
1644  line_color, line_style, line_width, fill_color, hist_title_string)
1645  if canvas == '' or canvas == None:
1646  canvas = Common.make_root_canvas(name)
1647  graph[0].SetStats(False)
1648  graph[0].Draw()
1649  else:
1650  canvas.cd()
1651  graph[1].Draw('l')
1652  canvas.Update()
1653  return (canvas, graph[0], graph[1])
1654  root_graph = staticmethod(root_graph)
1655 
1656  def matplot_histogram(self, x_axis_string, x_axis_units='', y_axis_string='', y_axis_units='', canvas=None, comparator=None):
1657  """
1658  Prints a 1D or 2D histogram of hits in the bunch. Axes are automatically detected
1659  to encapsulate all data. Use Bunch.cut(...) to shrink the axes.
1660 
1661  Needs correct matplot installation.
1662 
1663  - x_axis_string = string for call to get_hit_variables()
1664  - x_axis_units = units for x axis
1665  - y_axis_string = string for call to get_hit_variables(). If string is empty, makes a 1D histogram
1666  - y_axis_units = units for y axis
1667  - comparator = comparator of bunch1, bunch2 - if none given, will sort by x-axis value
1668 
1669  e.g. bunch.matplot_histogram('x', 'cm', 'py', 'GeV/c') will histogram x vs py.
1670 
1671  To display plots, call Common.wait_for_matplot() - script waits until all matplot windows are closed
1672  """
1673  weights = []
1674 
1675  num_evts = 0.
1676  for hit in self.__hits:
1677  if abs(hit.get('weight')) > 0.: num_evts += 1.
1678  n_bins = num_evts/10.+1.;
1679 
1680  x_points = self.list_get_hit_variable([x_axis_string], [x_axis_units])[0]
1681  if not y_axis_string == '':
1682  y_points = self.list_get_hit_variable([y_axis_string], [y_axis_units])[0]
1683  n_bins = num_evts**0.5/10.+1.;
1684  else: y_points = []
1685  for key in self.__hits:
1686  weights.append(key.get('weight'))
1687  if not x_axis_units == '': x_axis_units = " ["+x_axis_units+"]"
1688  if not y_axis_units == '': y_axis_units = " ["+y_axis_units+"]"
1689  x_name = x_axis_string+x_axis_units
1690  y_name = y_axis_string+y_axis_units
1691  sort = True
1692  if comparator != None:
1693  sort = False
1694  Common.make_matplot_histogram( x_points, x_name, int(n_bins), y_points, y_name, int(n_bins), weight_list=weights)
1695 
1696  def matplot_scatter_graph(self, x_axis_string, y_axis_string, x_axis_units='', y_axis_units='', include_weightless=True):
1697  """
1698  Prints a 2d scatter plot of Hit.get(...) data over a dict of bunches.
1699 
1700  Needs correct matplot installation.
1701 
1702  - x_axis_string = string for call to Bunch.get_hit_variable() used to calculate x-values
1703  - x_axis_units = units for x axis
1704  - y_axis_string = string for call to Bunch.get_hit_variable() used to calculate y-values
1705  - y_axis_units = units for y axis
1706  - include_weightless = set to True to plot hits with <= 0. weight
1707 
1708  e.g. bunch.matplot_scatter_graph('x', 'cm', 'p', 'MeV/c') will graph x vs p.
1709 
1710  To display plots, call Common.wait_for_matplot() - script waits until all matplot windows are closed
1711  """
1712  x_points = []
1713  y_points = []
1714  for hit in self:
1715  if abs(hit['weight'])>Common.float_tolerance or include_weightless:
1716  x_points.append(self.get_hit_variable(hit, x_axis_string)/Common.units[x_axis_units])
1717  y_points.append(self.get_hit_variable(hit, y_axis_string)/Common.units[y_axis_units])
1718  if not x_axis_units == '': x_axis_string += " ["+x_axis_units+"] "
1719  if not y_axis_units == '': y_axis_string += " ["+y_axis_units+"] "
1720  graph = Common.make_matplot_scatter(x_points, x_axis_string, y_points, y_axis_string)
1721 
1722  def matplot_graph(bunches, x_axis_string, x_axis_list, y_axis_string, y_axis_list, x_axis_units='', y_axis_units='', comparator=None):
1723  """
1724  Prints a graph of Bunch.get(...) data over a dict of bunches.
1725 
1726  Needs correct matplot installation.
1727 
1728  - x_axis_string = string for call to Bunch.get() used to calculate x axis variables
1729  - x_axis_list = list of variables for call to Bunch.get() used to calculate x axis variables
1730  - y_axis_string = string for call to Bunch.get() used to calculate y axis variables
1731  - y_axis_list = list of variables for call to Bunch.get() used to calculate y axis variables
1732  - x_axis_units = units for x axis
1733  - y_axis_units = units for y axis
1734  - comparator = comparator of bunch1, bunch2 - if none given, will sort by x-axis value
1736  e.g. bunch.matplot_graph('mean', ['x'], 'cm', 'emittance', ['x','y'], 'mm') will graph mean x vs emittance x y.
1737 
1738  To display plots, call Common.wait_for_matplot() - script waits until all matplot windows are closed
1739  """
1740  x_points = []
1741  y_points = []
1742  try: list_of_bunches = bunches.values()
1743  except: list_of_bunches = bunches
1744  if comparator != None:
1745  list_of_bunches.sort(comparator)
1746  for bunch in list_of_bunches:
1747  x_points.append(bunch.get(x_axis_string, x_axis_list)/Common.units[x_axis_units])
1748  y_points.append(bunch.get(y_axis_string, y_axis_list)/Common.units[y_axis_units])
1749  x_axis_string += "( "
1750  y_axis_string += "( "
1751  for value in x_axis_list: x_axis_string += value+" "
1752  for value in y_axis_list: y_axis_string += value+" "
1753  x_axis_string += ") "
1754  y_axis_string += ") "
1755  if not x_axis_units == '': x_axis_string += " ["+x_axis_units+"] "
1756  if not y_axis_units == '': y_axis_string += " ["+y_axis_units+"] "
1757  graph = Common.make_matplot_graph(x_points, x_axis_string, y_points, y_axis_string)
1758  matplot_graph = staticmethod(matplot_graph)
1759 
1760 
1761 #privates
1762  def __mean_for_get(self, variable_list):
1763  return self.mean(variable_list)[variable_list[0]]
1764 
1765  def __weight_for_get(self, variable_list):
1766  return self.bunch_weight()
1767 
1768  def __ang_mom_for_get(self, variable_list):
1769  return self.get_kinetic_angular_momentum()
1770 
1771  def __dispersion_for_get(self, variable_list):
1772  return self.get_dispersion(variable_list[0])
1773 
1774  def __dispersion_prime_for_get(self, variable_list):
1775  return self.get_dispersion_prime(variable_list[0])
1776 
1777  def __cov_mat_picker(self, axis_list):
1778  Common.has_numpy()
1779  if self.__covs == None: return None
1780  covs = numpy.zeros( (2*len(axis_list),2*len(axis_list) ))
1781  ind = []
1782  for axis in axis_list:
1783  ind.append( Bunch.__axis_list.index( axis ) )
1784  for i in range( len(ind) ):
1785  for j in range( len(ind) ):
1786  covs[2*i, 2*j] = self.__covs[2*ind[i], 2*ind[j]]
1787  covs[2*i+1,2*j] = self.__covs[2*ind[i]+1, 2*ind[j]]
1788  covs[2*i+1,2*j+1] = self.__covs[2*ind[i]+1, 2*ind[j]+1]
1789  covs[2*i, 2*j+1] = self.__covs[2*ind[i], 2*ind[j]+1]
1790  return covs
1791 
1792  def __mean_picker(self, var_list):
1793  if self.__means == None or self.__means == {}: return {}
1794  means = {}
1795  all_list = Bunch.axis_list_to_covariance_list(var_list)
1796  try:
1797  for var in all_list: means[var] = self.__means[var]
1798  except:
1799  return {}
1800  return means
1801 
1802  __number_of_header_lines = {'icool_for009':3, 'icool_for003':2, 'g4beamline_bl_track_file':0,'zgoubi':0, 'turtle':0, 'madx':0,'mars_1':0, 'maus_json_virtual_hit':0, 'maus_json_primary':0, 'opal_loss':1}
1803  __axis_list = ['x','y','z','t', 'ct']
1804  __get_dict = {'angular_momentum':__ang_mom_for_get, 'emittance':get_emittance, 'dispersion':__dispersion_for_get,
1805  'dispersion_prime':__dispersion_prime_for_get, 'beta':get_beta, 'alpha':get_alpha, 'gamma':get_gamma,
1806  'moment':moment, 'mean':__mean_for_get, 'bunch_weight':__weight_for_get, 'standard_deviation':standard_deviation}
1807  __geometric_momentum = False
1808 
1809  __get_list = []
1810  for key,value in __get_dict.iteritems():
1811  __get_list.append(key)
1812 
1813  def bunch_overview_doc(verbose = False):
1814  """Creates some summary documentation for the Bunch class. If verbose is True then will also print any functions or data not included in summary"""
1815  name_list = ['initialise', 'transforms', 'hit', 'moments', 'weights', 'twiss', 'twiss_help', 'io', 'ellipse', 'root', 'matplotlib', 'generic graphics']
1816  function_list = {
1817  'initialise' : ['new_dict_from_read_builtin', 'new_from_hits', 'new_from_read_builtin', 'new_from_read_user', 'new_list_from_read_builtin', 'new_hit_shell', 'copy', 'deepcopy'],
1818  'transforms' : ['abelian_transformation', 'period_transformation', 'transform_to', 'translate'],
1819  'hit' : ['get', 'get_hit_variable', 'get_hits', 'hit_equality', 'list_get', 'list_get_hit_variable', 'append', 'hits', 'hit_get_variables', 'get_variables', 'get_amplitude'],
1820  'moments' : ['mean', 'moment', 'covariance_matrix'],
1821  'weights' : ['bunch_weight', 'clear_global_weights', 'clear_local_weights', 'clear_weights', 'cut', 'transmission_cut', 'conditional_remove'],
1822  'twiss' : ['get_emittance', 'get_beta', 'get_alpha', 'get_gamma', 'get_emittance', 'get_canonical_angular_momentum', 'get_dispersion', 'get_dispersion_prime','get_dispersion_rsquared', 'get_kinetic_angular_momentum'],
1823  'twiss_help' : ['convert_string_to_axis_list', 'covariances_set', 'means_set', 'momentum_variable', 'set_geometric_momentum', 'set_covariance_matrix', 'get_axes', 'get_geometric_momentum', 'axis_list_to_covariance_list'],
1824  'io' : ['hit_write_builtin', 'hit_write_builtin_from_dict', 'hit_write_user', 'setup_file', 'read_maus_file'],
1825  'ellipse' : ['build_ellipse_2d', 'build_ellipse_from_transfer_matrix', 'build_penn_ellipse'], #, 'build_ET_ellipse', 'build_MR_ellipse'
1826  'root' : ['root_graph', 'root_histogram', 'root_scatter_graph'],
1827  'matplotlib' : ['matplot_graph', 'matplot_histogram', 'matplot_scatter_graph'],
1828  'generic graphics': ['histogram', 'histogram_var_bins'],
1829  }
1830  function_doc = {
1831  'initialise':'Functions that can be used to initialise a Bunch in various different ways:',
1832  'transforms':'Functions that transform all data in the bunch in some way:',
1833  'hit':'Functions to examine individual hits within the bunch:',
1834  'moments':'Functions to calculate beam moments:',
1835  'weights':'Functions to access and modify statistical weights:',
1836  'twiss':'Functions to calculate Twiss parameters, etc, in the bunch:',
1837  'twiss_help':'Helper functions to aid Twiss parameter calculation:',
1838  'io':'Output and some input helper functions:',
1839  'ellipse':'Functions to calculate beam ellipses based on Twiss parameters etc:',
1840  'root':'Interfaces to root plotting library:',
1841  'matplotlib':'Interfaces to matplotlib plotting library:',
1842  'generic graphics':'General algorithms to support graphics functions:',
1843  }
1844  bunch_doc = """
1845 Bunch class is a container class for a set of hits. The main purpose is to store a set of hits that are recorded at an output station in some tracking code. Functions are provided to calculate Twiss parameters, RMS emittances etc for this bunch. List of functions:
1846 """
1847  dir_bunch = dir(Bunch)
1848  if verbose:
1849  print 'The following functions and data are in Bunch but not in overview_doc:'
1850  for func in dir_bunch:
1851  found = False
1852  for func_sublist in function_list.values():
1853  if func in func_sublist: found = True
1854  if not found: print func,
1855  print '\n'
1856 
1857  print 'The following functions and data are in bunch_overview_doc but not in Bunch:'
1858  for func_sublist in function_list.values():
1859  for func in func_sublist:
1860  if func not in dir_bunch:
1861  print func,
1862  print
1863 
1864  doc = bunch_doc
1865  for key in name_list:
1866  doc = doc + function_doc[key]+'\n'
1867  for item in function_list[key]:
1868  doc = doc+' '+item+'\n'
1869  __doc__ = doc
1870  return doc
1871  bunch_overview_doc = staticmethod(bunch_overview_doc)
1872 
1873 #summary documentation
1874 __doc__ = Bunch.bunch_overview_doc()
1875 
def __init__
Initialise to an empty bunch.
Definition: _bunch.py:66
MausJsonHitFactory strips hits of a specified type from a json_document and returns to the user...
def standard_deviation
Return a standard deviation, given by sum( (u-u_mean)**2 )**0.5 where u is the variable specified...
Definition: _bunch.py:596
def build_ET_ellipse
Not implemented.
Definition: _bunch.py:1430
The hit factory module defines a number of factory classes used for generating hit objects...
Definition: __init__.py:1
def get_variables
Return a list of variables suitable for calls to Bunch.get.
Definition: _bunch.py:1320
def __setitem__
Called by subscript operator, sets the key^th hit in the bunch.
Definition: _bunch.py:124
tuple new_from_read_builtin
Definition: _bunch.py:287
tuple convert_string_to_axis_list
Definition: _bunch.py:1195
def period_transformation
Transform all hit's variable from a linear space to a periodic space with some frequency and offset...
Definition: _bunch.py:507
def root_scatter_graph
Prints a 2d scatter plot of Hit.get(...) data over a dict of bunches.
Definition: _bunch.py:1635
def get_kinetic_angular_momentum
Return the bunch kinetic angular momentum about some arbitrary axis, defined by.
Definition: _bunch.py:1080
def cut
Set weight of hits in the bunch to 0 if comparator(value, get([variable])) is False.
Definition: _bunch.py:772
tuple hit_write_builtin_from_dict
Definition: _bunch.py:1358
MausRootHitFactory reads hits of a specified type from a root spill.
def clear_local_weights
Set local_weight of all hits in the bunch to 1.
Definition: _bunch.py:741
def set_covariance_matrix
Choose whether to use an internal covariance matrix for calculations.
Definition: _bunch.py:708
def deepcopy
Make a copy and copy self's data.
Definition: _bunch.py:108
def __deepcopy__
Make a copy and copy self's data.
Definition: _bunch.py:100
def __del__
Called by del.
Definition: _bunch.py:135
def clear_weights
Set global_weight and local_weight of all hits in the bunch to 1.
Definition: _bunch.py:753
def __getitem__
Called by subscript operator, returns the key^th hit in the bunch.
Definition: _bunch.py:119
def mean
Return dict of variables to means.
Definition: _bunch.py:652
dictionary __get_dict
Definition: _bunch.py:1888
def transform_to
Perform a linear transformation on all particles in the bunch to get a bunch with covariance matrix s...
Definition: _bunch.py:477
def __delitem__
Called by remove method, deletes the key^th hit in the bunch.
Definition: _bunch.py:130
def covariances_set
If internal covariances are set by set_covariance_matrix, return True; else return False...
Definition: _bunch.py:722
def get_emittance
Return the n dimensional normalised emittance of the bunch.
Definition: _bunch.py:1058
def means_set
If means are set by set_covariance_matrix, return True; else return False.
Definition: _bunch.py:728
def get_decoupled_emittance
Reserved for future development.
Definition: _bunch.py:967
def get_canonical_angular_momentum
Return the bunch canonical angular momentum about some arbitrary axis, defined by.
Definition: _bunch.py:1097
def conditional_remove
For when a cut is not good enough, remove hits from the bunch altogether if comparator(value, get([variable])) is False.
Definition: _bunch.py:565
common module defines common utility data and functions that are used elsewhere
Definition: __init__.py:1
def root_histogram
Prints a 1D or 2D histogram of hits in the bunch.
Definition: _bunch.py:1572
tuple set_geometric_momentum
Definition: _bunch.py:1205
def get_gamma
Return the bunch gamma function, defined by.
Definition: _bunch.py:876
def get_hits
Returns a list of all hits in the bunch with comparator(variable, value) == True By default comparato...
Definition: _bunch.py:531
tuple clear_global_weights
Definition: _bunch.py:749
Represents a bunch of particles.
Definition: _bunch.py:62
def get_beta
Return the bunch beta function, defined by.
Definition: _bunch.py:847
def __str__
Return an abbreviated string like <Bunch of="" n="" hits>="">
Definition: _bunch.py:74
def append
Append a hit to Bunch self.
Definition: _bunch.py:165
def __repr__
Return a long string that contains data in every hit.
Definition: _bunch.py:79
def matplot_histogram
Prints a 1D or 2D histogram of hits in the bunch.
Definition: _bunch.py:1754
tuple axis_list_to_covariance_list
Definition: _bunch.py:1175
def get_dispersion_rsquared
Return a special bunch dispersion defined by.
Definition: _bunch.py:942
def get_dispersion
Return the bunch dispersion D in axis q, defined by.
Definition: _bunch.py:931
def get_axes
Return list of axis variables (position-type variables)
Definition: _bunch.py:1219
def build_MR_ellipse
Not implemented.
Definition: _bunch.py:1421
def moment
Return a moment, sum( (u1 - u1_mean)*(u2-u2_mean)*(etc)) where u is a set of Hit get_variables indexe...
Definition: _bunch.py:617
def hit_write_builtin
Write the hits in the bunch using some built-in format to the file file_name.
Definition: _bunch.py:1338
def hit_write_user
Write the hits in the bunch in some user defined format to file_handle, in an order defined by compar...
Definition: _bunch.py:1375
def histogram_var_bins
Returns a binned histogram of hits in the bunch.
Definition: _bunch.py:1532
def hits
Returns the list of all hits in the bunch.
Definition: _bunch.py:516
Implemented within this module:
Definition: __init__.py:1
def __len__
Number of hits.
Definition: _bunch.py:114
list all
Definition: __init__.py:34
def get_dispersion_prime
Return the bunch dispersion prime D in axis q, defined by.
Definition: _bunch.py:961
def translate
Translate all events in the bunch by dict_variables_values.
Definition: _bunch.py:430
def hit_get_variables
Return a list of variables suitable for calls to Bunch.get.
Definition: _bunch.py:1326
def __eq__
Return true if all hits in self are the same as all hits in target, and set_covariance_matrix data ar...
Definition: _bunch.py:141
def transmission_cut
Set weight of hits in the bunch to 0 if no events can be found in test_bunch with the same test_varia...
Definition: _bunch.py:806
tuple new_list_from_read_builtin
Definition: _bunch.py:250
def get
Return a bunch variable taken from the list Bunch.get_variables()
Definition: _bunch.py:1286
def get_hit_variable
Return axis variable for hit.
Definition: _bunch.py:1232
def histogram
Returns a binned 1D or 2D histogram of hits in the bunch (but doesnt draw anything or call any plotti...
Definition: _bunch.py:1499
def matplot_scatter_graph
Prints a 2d scatter plot of Hit.get(...) data over a dict of bunches.
Definition: _bunch.py:1794
def __copy__
Make a copy but use references to self's data.
Definition: _bunch.py:91
def clear_global_weights
Set global_weight of all hits in the bunch to 1.
Definition: _bunch.py:747
def __ne__
Return true if any hits in self are not the same as the corresponding hit in target.
Definition: _bunch.py:160
def get_geometric_momentum
Set the default for conjugate momenta; either geometric (x', y', etc) or kinetic (px, py, etc)
Definition: _bunch.py:1213
tuple new_dict_from_read_builtin
Definition: _bunch.py:222
def get_alpha
Return the bunch alpha function, defined by.
Definition: _bunch.py:905
def hit_equality
Return True if all hits in self are the same as all hits in target.
Definition: _bunch.py:544
def bunch_weight
Return statistical weight of all hits in the bunch.
Definition: _bunch.py:733
def list_get_hit_variable
Return a list of get_hit_variable results, one element for each hit in the bunch. ...
Definition: _bunch.py:1248
tuple build_ellipse_from_transfer_matrix
Definition: _bunch.py:1442
def covariance_matrix
Return covariance matrix of variables.
Definition: _bunch.py:670