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