xboa
BunchTest.py
Go to the documentation of this file.
1 import os
2 import operator
3 import bisect
4 import sys
5 
6 try:
7  import numpy
8  from numpy import linalg
9 except ImportError:
10  pass
11 
12 from xboa import *
13 from xboa.core import *
14 from xboa.Hit import Hit
15 from xboa.Bunch import Bunch
16 from xboa.Common import rg
17 
19 run_test = xboa.test.TestTools.run_test
20 run_test_group = xboa.test.TestTools.run_test_group
21 __float_tol = xboa.test.TestTools.__float_tol
22 parse_tests = xboa.test.TestTools.parse_tests
23 __test_root_hist = xboa.test.TestTools.test_root_hist
24 __test_root_canvas = xboa.test.TestTools.test_root_canvas
25 __test_root_graph = xboa.test.TestTools.test_root_graph
26 
27 #=======================================================================================================================#
28 #====================== BUNCH ==============================================================#
29 #=======================================================================================================================#
31  import xboa.core.Bunchcore
32  bc = Bunchcore.Bunchcore(20)
33  assert bc.length() == 0
34  hit = Hit()
35  bc.set_item(hit, hit._Hit__hitcore, 19)
36  bc.set_item(hit, hit._Hit__hitcore, -20)
37  bc.set_item(hit, hit._Hit__hitcore, 20)
38  bc.set_item(hit, hit._Hit__hitcore, -21)
39  try:
40  bc.set_item(hit, hit._Hit__hitcore)
41  print 'Fail at bunchcore.set_item(<bad input>)'
42  return 'fail'
43  except:
44  pass
45  if bc.get_item(19) != (hit,hit._Hit__hitcore) or bc.get_item(-21) != (hit,hit._Hit__hitcore):
46  print 'Fail at bunchcore.get_item(19) && bc.get_item(-20)',bc.get_item(19), bc.get_item(-20)
47  return 'fail'
48  try:
49  bc.get_item(21)
50  print 'Fail at bunchcore.get_item(20)'
51  return 'fail'
52  except:
53  pass
54  try:
55  bc.get_item(-22)
56  print 'Fail at bunchcore.get_item(-21)'
57  return 'fail'
58  except:
59  pass
60  try:
61  bc.get_item()
62  print 'Fail at bunchcore.get_item(<bad input>)'
63  return 'fail'
64  except:
65  pass
66  for i in range(20):
67  hit = Hit()
68  hit['x'] = float(i)
69  hit['y'] = -float(i)
70  hit['eventNumber'] = i
71  bc.set_item(hit, hit._Hit__hitcore, i)
72  mom = bc.moment(['x'], [1.]) #calculation tested at Bunch level
73  a = bc.covariance_matrix(['x', 'y'], [1.,2.])
74  try:
75  bc.moment('x')
76  return 'fail'
77  except:
78  pass
79  try:
80  bc.moment(['garbage'])
81  return 'fail'
82  except:
83  pass
84  test_strings = [['x', 'garbage'], ['x'], ['x','y','z'], 'x']
85  test_floats = [[1.,2.,3.],[1.],[1.,1],1.]
86  for string_list in test_strings:
87  try:
88  bc.covariance_matrix(string_list, [1.,2.])
89  print 'fail on string_list',string_list
90  return 'fail'
91  except:
92  pass
93  for float_list in test_floats:
94  try:
95  bc.covariance_matrix(['x','y'], float_list)
96  print 'fail on float_list',float_list
97  return 'fail'
98  except:
99  pass
100  return 'pass'
101 
102 def bunch_hit_equality_test(bunch1, bunch2, isEqual):
103  if bunch1.hit_equality(bunch2) and isEqual: return 'pass'
104  elif not bunch1.hit_equality(bunch2) and not isEqual: return 'pass'
105  return 'fail'
106 
107 def bunch_delete_test(bunch):
108  bunch1 = bunch.deepcopy()
109  del(bunch1)
110  try:
111  bunch1.bunch_weight()
112  return 'fail'
113  except:
114  return 'pass'
115 
116 def bunch_equality_test(bunch1, bunch2, isEqual):
117  if bunch1 == bunch2 and isEqual: return 'pass'
118  elif bunch1 != bunch2 and not isEqual: return 'pass'
119  return 'fail'
120 
121 def bunch_str_test(bunch):
122  if type(str(bunch)) == type('string'): return 'pass'
123  return 'fail'
124 
125 def bunch_repr_test(bunch):
126  new_bunch = eval(repr(bunch))
127  if new_bunch == bunch: return 'pass'
128  print repr(bunch)
129  return 'fail'
130 
131 def bunch_copy_test(bunch):
132  bunch1 = bunch.copy()
133  if bunch1 is bunch: return 'pass'
134  return 'fail'
135 
137  bunch1 = bunch.deepcopy()
138  if bunch1 == bunch and not bunch1 is bunch: return 'pass'
139  return 'fail'
140 
141 def bunch_len_test(bunch):
142  if len(bunch) == len(bunch.hits()): return 'pass'
143  return 'fail'
144 
146  for index in range(len(bunch.hits())):
147  if bunch[index] != bunch.hits()[index]: return 'fail'
148  return 'pass'
149 
151  bunch1 = bunch.deepcopy()
152  for index in range(len(bunch1)):
153  bunch1[index] = bunch1[0]
154  if bunch1[index] != bunch1[0]: return 'fail'
155  return 'pass'
156 
158  #def conditional_remove(self, variable_value_dict, comparator, value_is_nsigma_bool = False):
159  pid = -13248324
160  bunch1 = bunch.deepcopy()
161  bunch1.append(Hit.new_from_dict({'x':100., 'pid':pid}))
162  bunch1.conditional_remove({'pid':pid}, operator.eq)
163  if len(bunch1.get_hits('pid', pid)) != 0:
164  print 'failed - bunch1 wrong length',bunch1,bunch1.get_hits('pid', pid)
165  return 'fail'
166  bunch1.append(Hit.new_from_dict({'x':100., 'pid':pid}))
167  bunch1.conditional_remove({'pid':pid}, operator.ne)
168  if len(bunch1.get_hits('pid', pid)) != 1:
169  print 'failed - bunch1 wrong length',bunch1,bunch1.get_hits('pid', pid)
170  return 'fail'
171  if len(bunch1) != 1: return 'fail'
172  return 'pass'
173 
175  bunch1 = bunch.deepcopy()
176  translation = {'x':5,'px':2}
177  bunch1.translate({'x':5,'px':2})
178  for index in range(len(bunch)):
179  for key,value in translation.iteritems():
180  if not abs(bunch[index].get(key) + value - bunch1[index].get(key)) < __float_tol: return 'fail'
181  bunch1 = bunch.deepcopy()
182  translation = {'x':5,'px':2}
183  bunch1.translate({'x':5,'px':2}, 'energy')
184  for index in range(len(bunch)):
185  if not bunch1[index].check(): return 'fail'
186  for key,value in translation.iteritems():
187  if not abs(bunch[index].get(key) + value - bunch1[index].get(key)) < __float_tol: return 'fail'
188  return 'pass'
189 
191  bunch1 = bunch.deepcopy()
192  bunch2 = bunch.deepcopy()
193  bunch1.abelian_transformation(['x','px'], numpy.matrix([[1,0],[0,1]]), {'x':0,'px':0}, {'x':0,'px':0}, '')
194  if bunch != bunch1: return 'fail'
195  bunch1.abelian_transformation(['x','px'], numpy.matrix([[1,0],[0,1]]))
196  if bunch != bunch1: return 'fail'
197  R = numpy.matrix([[0.76,0.43],[0.76,0.29]])
198  O = numpy.matrix([[1.3], [-1.7]])
199  T = numpy.matrix([[1.782],[2.35]])
200  bunch1.abelian_transformation(['x','px'], R, {'x':T[0,0],'px':T[1,0]}, {'x':O[0,0],'px':O[1,0]}, 'energy')
201  for key in range( len(bunch1) ):
202  bunch2[key].abelian_transformation(['x','px'], R, {'x':T[0,0],'px':T[1,0]}, {'x':O[0,0],'px':O[1,0]}, 'energy')
203  if bunch2[key] != bunch1[key]: return 'fail'
204  return 'pass'
205 
207  bunch1 = bunch.deepcopy()
208  bunch1.transform_to(['x','px','pz'], bunch.covariance_matrix(['x','px','pz'])) #NOTE: May fail due to covariance_matrix
209  if bunch != bunch1:
210  print bunch.covariance_matrix(['x','px','pz'])
211  print bunch1.covariance_matrix(['x','px','pz'])
212  for i in range( len(bunch) ):
213  print bunch[i]
214  print bunch1[i],'\n'
215  return 'fail'
216  target_matrix = numpy.matrix([[2,-1],[-1,3]])
217  translation = {'x':4.,'px':5.}
218  trans_back = {'x':-4.,'px':-5.}
219  origin = {'x':6.,'px':7.}
220  bunch1.transform_to(['x','px'], target_matrix, translation, origin, 'energy')
221  bunch1.translate(trans_back, 'energy')
222  covariances = bunch1.covariance_matrix(['x','px'], origin)
223  covariances = covariances*(float(linalg.det( target_matrix ))/float( linalg.det(covariances) ))**0.5
224  if abs(covariances[0,0] - target_matrix[0,0]) > __float_tol: return 'fail'
225  if abs(covariances[1,0] - target_matrix[0,1]) > __float_tol: return 'fail'
226  if abs(covariances[1,1] - target_matrix[1,1]) > __float_tol: return 'fail'
227  d1 = linalg.det(bunch1.covariance_matrix(['x','px'], origin))
228  d2 = linalg.det(bunch.covariance_matrix(['x','px'], origin))
229  if abs(d1 - d2) > __float_tol and abs( (d1-d2)/(d1+d2) )>__float_tol:
230  return 'fail'
231  return 'pass'
232 
233 def bunch_hits_test(bunch):
234  if bunch.hits() is bunch._Bunch__hits: return 'pass'
235  return 'fail'
236 
238  hits = bunch.get_hits('station', 13, operator.ne)
239  for hit in bunch.hits():
240  if hit.get('station') != 13:
241  if not hit in hits: return 'fail'
242  if hit.get('station') == 13:
243  if hit in hits: return 'fail'
244  return 'pass'
245 
247  if bunch.standard_deviation('x') != bunch.moment(['x', 'x'])**0.5:
248  return 'fail'
249  if bunch.standard_deviation('x', {'x':10.}) != bunch.moment(['x', 'x'], {'x':10.})**0.5:
250  return 'fail'
251  return 'pass'
252 
253 def bunch_moment_test(bunch):
254  moment_list = ['x','x','px']
255  mean_list = {'x':1., 'px':2.}
256  moment = bunch.moment(moment_list, mean_list)
257  this_moment = 0.
258  for hit in bunch.hits():
259  x = 1.
260  for var in moment_list:
261  x *= (hit.get(var) - mean_list[var])
262  this_moment += x/bunch.bunch_weight()
263  #ran into floating point issues hence weird test condition
264  if abs(this_moment - moment)/abs(this_moment+moment) < __float_tol or abs(this_moment - moment) < __float_tol: return 'pass'
265  return 'fail'
266 
267 def bunch_mean_test(bunch):
268  mean_list = ['x','x','px']
269  mean = bunch.mean(mean_list)
270  this_mean = {'x':0., 'px':0.}
271  for hit in bunch.hits():
272  for var in mean_list:
273  this_mean[var] += hit.get(var)/bunch.bunch_weight()
274  for var in mean_list:
275  if abs(this_mean[var] - mean[var]) < __float_tol: return 'pass'
276  return 'fail'
277 
279  cov_list = ['x','px','pz']
280  origin_dict = {'x':2.}
281  covariances = bunch.covariance_matrix(cov_list, origin_dict)
282  origin_dict['px'] = bunch.mean(['px'])['px']
283  origin_dict['pz'] = bunch.mean(['pz'])['pz']
284  for i in range(len(cov_list)):
285  for j in range(len(cov_list)):
286  if abs(bunch.moment([cov_list[i], cov_list[j]],origin_dict) - covariances[i,j]) > __float_tol:
287  print cov_list[i],cov_list[j],' ',bunch.moment([cov_list[i], cov_list[j]],origin_dict),covariances[i,j]
288  return 'fail'
289  covariances1 = bunch.covariance_matrix(cov_list)
290  covariances2 = bunch.covariance_matrix(cov_list, {})
291  covariances3 = bunch.covariance_matrix(cov_list, bunch.mean(cov_list))
292  for i in range(len(cov_list)):
293  for j in range(len(cov_list)):
294  if abs(covariances1[i,j] - covariances2[i,j]) > __float_tol or abs(covariances1[i,j] - covariances3[i,j]) > __float_tol:
295  print i,j,covariances1[i,j], covariances2[i,j], covariances3[i,j]
296  return 'fail'
297  return 'pass'
298 
300  weight = 0.
301  for key in bunch:
302  weight += key.get('weight')
303  if abs(bunch.bunch_weight() - weight) < __float_tol: return 'pass'
304  return 'fail'
305 
307  bunch1 = bunch.deepcopy()
308  for key in bunch1:
309  key.set('local_weight', 0.5)
310  bunch1.clear_local_weights()
311  if abs(bunch1.bunch_weight() - bunch.bunch_weight()) < __float_tol: return 'pass'
312  return 'fail'
313 
315  bunch1 = bunch.deepcopy()
316  for key in bunch1:
317  key['global_weight'] = 0.5
318  bunch1.clear_global_weights()
319  if abs(bunch1.bunch_weight() - bunch.bunch_weight()) < __float_tol: return 'pass'
320  return 'fail'
321 
323  bunch1 = bunch.deepcopy()
324  for key in bunch1:
325  key.set('local_weight', 0.5)
326  key.set('global_weight', 0.5)
327  bunch1.clear_weights()
328  if abs(bunch1.bunch_weight() - bunch.bunch_weight()) < __float_tol: return 'pass'
329  return 'fail'
330 
331 def bunch_cut_test(bunch):
332  bunch1 = bunch.deepcopy()
333  bunch1.clear_weights()
334  bunch1.cut({'x':0.1, 'px':0.1}, operator.ge)
335  for key in bunch1:
336  if key.get('x') > 0.1 or key.get('px') > 0.1:
337  if abs(key.get('local_weight') ) > __float_tol : return 'fail'
338  else:
339  if abs(key.get('local_weight') -1. ) > __float_tol : return 'fail'
340  bunch1.clear_local_weights()
341  bunch1.cut({'amplitude x':0.}, operator.ge)
342  if bunch1.bunch_weight() > 1e-9:
343  return 'fail'
344  bunch1.clear_local_weights()
345  sigma_x = bunch1.moment(['x','x']) ** 0.5
346  bunch1.cut({'x':1}, operator.ge, True, True)
347  for key in bunch1:
348  if key.get('x') - sigma_x > __float_tol:
349  if abs(key.get('global_weight') ) > __float_tol :
350  bunch1.clear_weights()
351  return 'fail'
352  else:
353  if abs(key.get('global_weight') -1. ) > __float_tol:
354  bunch1.clear_weights()
355  return 'fail'
356  bunch1.clear_weights()
357  return 'pass'
358 
360  test = 'pass'
361  bunch1 = bunch.deepcopy()
362  bunch2 = Bunch()
363  for hit in range(len(bunch1)/2):
364  bunch2.append(bunch1[hit])
365  bunch2[0]['particle_number'] += 1
366  bunch1.transmission_cut(bunch2, global_cut=True, test_variable='event_number')
367  for hit1 in bunch1:
368  cut = True
369  for hit2 in bunch2:
370  if hit1['event_number'] == hit2['event_number']: cut = False
371  if hit1['weight'] == 1. and cut: test = 'fail'
372  if hit1['weight'] == 0. and not cut: test = 'fail'
373  bunch1.clear_weights()
374  bunch1.transmission_cut(bunch2, global_cut=False, test_variable='event_number')
375  for hit1 in bunch1:
376  cut = True
377  for hit2 in bunch2:
378  if hit1['event_number'] == hit2['event_number']: cut = False
379  if hit1['weight'] == 1. and cut: test = 'fail'
380  if hit1['weight'] == 0. and not cut: test = 'fail'
381  bunch1.transmission_cut(bunch2, global_cut=False, test_variable=['event_number', 'particle_number'])
382  for hit1 in bunch1:
383  cut = True
384  for hit2 in bunch2:
385  if hit1['event_number'] == hit2['event_number'] and \
386  hit1['particle_number'] == hit2['particle_number']: cut = False
387  if hit1['weight'] == 1. and cut: test = 'fail'
388  if hit1['weight'] == 0. and not cut: test = 'fail'
389  bunch1.clear_weights()
390  return test
391 
392 # tests set and get
394  geom = Bunch.get_geometric_momentum()
395  Bunch.set_geometric_momentum(geom)
396  if Bunch.get_geometric_momentum( ) != geom: return 'fail'
397  Bunch.set_geometric_momentum(not geom)
398  if Bunch.get_geometric_momentum( ) == geom: return 'fail'
399  Bunch.set_geometric_momentum(geom)
400  return 'pass'
401 
403  geom = Bunch.get_geometric_momentum()
404  Bunch.set_geometric_momentum(True)
405  if Bunch.axis_list_to_covariance_list( ['x','y','z','t', 'ct'] ) != ['x','x\'','y','y\'','z','z\'','t','t\'', 'ct','ct\'']: return 'fail'
406  Bunch.set_geometric_momentum(False)
407  if Bunch.axis_list_to_covariance_list( ['x','y','z','t', 'ct'] ) != ['x','px','y','py','z','pz','t','energy', 'ct','energy']: return 'fail'
408  geom = not Bunch.get_geometric_momentum()
409  return 'pass'
410 
412  if Bunch.convert_string_to_axis_list( 'x y z t ct' ) != ['x','y','z','t','ct']: return 'fail'
413  try:
414  Bunch.convert_string_to_axis_list( 'x y z t ct bob' )
415  return 'fail' #should throw
416  except: pass
417  return 'pass'
418 
420  if Bunch.get_axes() != ['x','y','z','t','ct']: return 'fail'
421  return 'pass'
422 
424  for var in Bunch.hit_get_variables():
425  for hit in bunch.hits():
426  try:
427  value = bunch.get_hit_variable(hit, var)
428  if type(var) is str:
429  if var.find('amplitude') > 0 or var == '': return 'pass'
430  if abs(value - hit.get(var)) > __float_tol: return 'fail'
431  except ZeroDivisionError:
432  pass
433  return 'pass'
434 
436  var_list = Bunch.hit_get_variables()[0:4]
437  value_lists = bunch.list_get_hit_variable(var_list)
438  for i in range(len(bunch.hits())):
439  for j in range(len(var_list)):
440  if type(var_list[j]) is str:
441  if var_list[j].find('amplitude') < 0 and var_list[j] != '':
442  target = bunch[i].get(var_list[j])
443  else:
444  target = bunch[i].get(var_list[j])
445  if abs(value_lists[j][i] - target): return 'fail'
446  return 'pass'
447 
448 def bunch_get_test(bunch):
449  Bunch.set_geometric_momentum(False)
450  cov = bunch.covariance_matrix(['x','px','y','py'])
451  for var in Bunch.get_variables():
452  value = bunch.get(var, ['x','y'])
453  if var == 'angular_momentum':
454  target = bunch.moment(['x','py'],{'x':0,'py':0})-bunch.moment(['y','px'],{'y':0,'px':0})
455  if var == 'emittance':
456  target = linalg.det(cov)**0.25/bunch.hits()[0].get('mass')
457  if var == 'dispersion':
458  target = bunch.moment(['x','energy'])*bunch.mean(['energy'])['energy']/bunch.moment(['energy','energy'])
459  if var == 'dispersion_prime':
460  target = bunch.moment(['px','energy'])*bunch.mean(['energy'])['energy']/bunch.moment(['energy','energy'])
461  if var == 'beta':
462  target = (cov[0,0] + cov[2,2])/2.*bunch.mean(['p'])['p']/bunch.get('emittance', ['x','y'])/bunch.hits()[0].get('mass')
463  if var == 'alpha':
464  target = -(cov[0,1] + cov[2,3])/2./bunch.get('emittance', ['x','y'])/bunch.hits()[0].get('mass')
465  if var == 'gamma':
466  target = (cov[1,1] + cov[3,3])/2./bunch.mean(['p'])['p']/bunch.get('emittance', ['x','y'])/bunch.hits()[0].get('mass')
467  if var == 'moment':
468  target = bunch.moment(['x','y'])
469  if var == 'mean':
470  target = bunch.mean(['x'])['x']
471  if var == 'bunch_weight':
472  target = bunch.bunch_weight()
473  if var == 'standard_deviation':
474  target = bunch.standard_deviation(['x'])
475  if abs(value - target) > abs(target)*__float_tol+__float_tol:
476  print 'bunch_get_test:',var,target,value
477  return 'fail'
478  return 'pass'
479 
480  Bunch.set_geometric_momentum(True)
481  cov = bunch.covariance_matrix(['x','x\'','y','y\''])
482  for var in Bunch.get_variables():
483  value = bunch.get(var, ['x','y'])
484  if var == 'angular_momentum':
485  target = (-cov[1,2]+cov[0,3]) * bunch.mean(['p'])['p']
486  if var == 'emittance':
487  target = linalg.det(cov)**0.25
488  if var == 'dispersion':
489  target = -bunch.moment(['x','energy'])*bunch.mean(['energy'])['energy']/bunch.moment(['energy','energy'])
490  if var == 'dispersion_prime':
491  target = bunch.moment(['px','energy'])*bunch.mean(['energy'])['energy']/bunch.moment(['energy','energy'])
492  if var == 'beta':
493  target = (cov[0,0] + cov[2,2])/2./bunch.get('emittance', ['x','y'])
494  if var == 'alpha':
495  target = (cov[0,1] + cov[2,3])/2./bunch.get('emittance', ['x','y'])
496  if var == 'gamma':
497  target = (cov[1,1] + cov[3,3])/2./bunch.get('emittance', ['x','y'])
498  if var == 'moment':
499  target = bunch.moment(['x','y'])
500  if var == 'mean':
501  target = bunch.mean(['x'])['x']
502  if var == 'bunch_weight':
503  target = bunch.bunch_weight()
504  if abs(value - target) > __float_tol: return 'fail'
505  Bunch.set_geometric_momentum(False)
506  return 'pass'
507 
509  bunch1 = bunch.deepcopy()
510  axis_list = ['x']
511  amp = 0.
512  Bunch.set_geometric_momentum(False)
513  for hit in bunch1.hits():
514  delta = Bunch.get_amplitude(bunch1, hit, axis_list)
515  amp += delta/float(len(bunch1.hits()))
516  target = bunch1.get_emittance(axis_list,bunch1.covariance_matrix(['x','px']))*2.*len(axis_list)
517  if abs(amp - target) > __float_tol:
518  print 'Failed at geometric=False:',amp, target
519  return 'fail'
520  Bunch.set_geometric_momentum(True)
521  if bunch1.covariances_set(): bunch1.set_covariance_matrix()
522  amp = 0.
523  for hit in bunch1.hits():
524  amp += Bunch.get_amplitude(bunch1, hit, axis_list)/float(len(bunch1))
525  target = bunch1.get_emittance(axis_list)*2.*len(axis_list)
526  if abs(amp - target) > __float_tol:
527  print 'Failed at geometric=True:',amp, target
528  return 'fail'
529  Bunch.set_geometric_momentum(False)
530  return 'pass'
531 
533  testpass = True
534  bin_x = range(-100,100)
535  bin_y = [-0.1,0.1,1.0]
536  bin_weights = [[0]*(len(bin_x)-1)]
537  for hit in bunch:
538  if hit['y']> -0.1*Common.units['m'] and hit['y']<0.1*Common.units['m']:
539  i = bisect.bisect_right(bin_x, hit['x'])-1
540  if i>=0 and i<len(bin_weights[0]): bin_weights[0][i] += hit['weight']
541  h1 = bunch.histogram_var_bins('x', bin_x, 'mm')
542  h2 = bunch.histogram_var_bins('x', bin_x, 'mm', 'y', bin_y, 'm')
543 
544  bunch1 = bunch.deepcopy()
545  bunch1.cut({'x':-100.}, operator.le)
546  bunch1.cut({'x': 100.}, operator.ge)
547  testpass &= h1[1] == bin_x and h1[2] ==[]
548  testpass &= abs(h1[0].sum()-bunch1.bunch_weight()) < __float_tol #I assume binning has been done ok - tested in Common.histogram()
549 
550  bunch2 = bunch.deepcopy()
551  bunch2.cut({'x':-100., 'y':-0.1*Common.units['m']}, operator.le)
552  bunch2.cut({'x': 100., 'y': 1.0*Common.units['m']}, operator.ge)
553  testpass &= h2[1] == bin_x and h2[2] == bin_y
554  testpass &= abs(h2[0].sum().sum()-bunch2.bunch_weight()) < __float_tol #I assume binning has been done ok - tested in Common.histogram()
555 
556  if testpass: return 'pass'
557  return 'fail'
558 
560  h1 = bunch.histogram('x')
561  h2 = bunch.histogram('x', 'm', 'y', 'mm', nx_bins=13, ny_bins=11, xmin=-2., xmax=3., ymin=-4., ymax=4.)
562  (xmin,xmax) = (min(bunch.list_get_hit_variable(['x'], ['mm'])[0]),max(bunch.list_get_hit_variable(['x'], ['mm'])[0]) )
563  testpass = True
564  testpass &= (h1[1][0]-xmin)<__float_tol and (h1[1][-1]-xmax)<__float_tol and h1[2]==[]
565  testpass &= (h2[1][0]+2.*Common.units['m'])<__float_tol and (h2[1][13]-3.*Common.units['m'])<__float_tol and len(h2[1])==14 #nbins+1
566  testpass &= (h2[2][0]+4.*Common.units['mm'])<__float_tol and (h2[2][11]-4.*Common.units['mm'])<__float_tol and len(h2[2])==12 #nbins+1
567  if testpass: return 'pass'
568  return 'fail'
569 
570 # 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,
571 # line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.graph_fill_color, hist_title_string=''):
572 
573 def __bunch_cmp(b1, b2):
574  return cmp(b1.get_emittance(['x','y']), b2.get_emittance(['x','y']))
575 
576 def bunch_root_graph_test(bunch_dict):
577  canvas,hist,graph = Bunch.root_graph(bunch_dict, 'mean', ['z'], 'emittance', ['x','y'], 'mm', 'm')
578  canvas,hist,graph = Bunch.root_graph(bunch_dict, 'mean', ['z'], 'bunch_weight', [], 'mm', 'm')
579  canvas = ''
580  canvas,hist,graph = Bunch.root_graph(bunch_dict, 'mean', ['z'], 'emittance', ['x','y'], 'mm', 'm', canvas, __bunch_cmp, -10., 100., -20., 200., 1, 2, 3, 4, 'mrs title')
581  testpass = __test_root_hist(hist,'mean( z ):emittance( x,y )', 'mean( z ) [mm]','emittance( x y ) [m]', -10., 100., -20., 200., rg.line_color, rg.line_style, rg.line_width, rg.fill_color, 'mrs title')
582  try: list_of_bunches = bunch_dict.values()
583  except: list_of_bunches = bunch_dict
584  x_list,y_list = [],[]
585  for b in list_of_bunches:
586  x_list.append(b.mean(['z'])['z'])
587  y_list.append(b.get_emittance(['x','y'])*1e-3)
588  Common.multisort([y_list,x_list])
589  testpass &= __test_root_graph(graph, 'mean( z ):emittance( x,y )', x_list, y_list, 1, 2, 3, 4)
590  if testpass: return 'pass' #need to test new options
591  return 'fail'
592 
594  #check that defaults are set or non-defaults are passed correctly to hist
595  (canvas,hist) = bunch.root_histogram('energy','MeV')
596  (canvas,hist) = bunch.root_histogram('x','mm','px','MeV/c', 3, 7, canvas, -25, 45, -17, 22, 1, 2, 3, 4, True, 'baby title')
597  testpass = __test_root_hist(hist,'mean( z ):emittance( x,y )', 'x [mm]','px [MeV/c]', -25., 45., -17., 22., 1, 2, 3, 4, 'baby title')
598  if testpass: return 'pass' #need to test new options
599  return 'fail'
600 
602  #check that defaults are set or non-defaults are passed correctly to hist
603  (canvas,hist,graph) = bunch.root_scatter_graph('energy','p')
604  (canvas,hist,graph) = bunch.root_scatter_graph('x','px','m','GeV/c', True, canvas,-25,45,-17,22,1,2,3,4,'abc')
605  x_l,y_l = [],[]
606  for hit in bunch:
607  x_l.append(hit['x']*1e-3)
608  y_l.append(hit['px']*1e-3)
609  Common.multisort([x_l,y_l])
610  testpass = __test_root_hist (hist, 'energy:p', 'x [m]','px [GeV/c]', -25., 45., -17., 22., rg.line_color, rg.line_style, rg.line_width, rg.fill_color, 'abc')
611  testpass &= __test_root_graph(graph, 'energy:p', x_l, y_l, 1, 2, 3, 4)
612  if testpass: return 'pass'
613  return 'fail' #need to test new options
614 
615 def bunch_matplot_graph_test(bunch_dict):
616  canvas1 = Bunch.matplot_graph(bunch_dict, 'mean', ['z'], 'bunch_weight', [], 'mm')
617  canvas2 = Bunch.matplot_graph(bunch_dict, 'mean', ['z'], 'emittance', ['x','y'], 'mm', 'm')
618  return 'pass'
619 
621  canvas1 = bunch.matplot_histogram('x','mm','px','MeV/c')
622  canvas2 = bunch.matplot_histogram('energy','MeV')
623  return 'pass'
624 
626  canvas1 = bunch.matplot_scatter_graph('x','px','mm','MeV/c')
627  canvas1 = bunch.matplot_scatter_graph('energy','p')
628  return 'pass'
629 
630 def bunch_list_get_test(bunch_dict):
631  var_list = ['beta', 'mean']
632  axis_list = [['x','y'], ['z']]
633  Bunch.list_get(bunch_dict,var_list, axis_list)
634  return 'pass'
635 
637  ellipse = Bunch.build_penn_ellipse(6*Common.units['mm'], Common.pdg_pid_to_mass[13], 333.,0.,200.,0.,4*Common.units['T'],1.)
638  bunch = Bunch.new_hit_shell(3, ellipse, ['x','px','y','py'], 'energy', defaults={'pid':-13,'mass':Common.pdg_pid_to_mass[13],'pz':200.})
639  bunch_amp = Bunch.get_amplitude(bunch, bunch[0], ['x','y'], ellipse)
640  for hit in bunch:
641  if abs(Bunch.get_amplitude(bunch, hit, ['x','y'], ellipse) - bunch_amp) > Common.float_tolerance:
642  return 'fail'
643  return 'pass'
644 
645 def bunch_maus_root_io_test(_bunch_not_used):
646  out = 'pass'
647  file_types = filter(lambda format: format.find('maus_root') >= 0, Hit.file_types())
648  formats = {'maus_root_primary':6, 'maus_root_virtual_hit':12}
649  for format in file_types:
650  bunch = Bunch.new_from_read_builtin(format, sys.prefix+'/share/xboa/data/maus_test.root')
651  id_dict = {}
652  # FYI 2 spills; 3 events in each spill; each event has 1 particle; with two stations (hits)
653  for hit in bunch:
654  _id = (hit['spill'], hit['event_number'], hit['particle_number'], hit['station'])
655  if _id in id_dict:
656  out = 'fail'
657  else:
658  id_dict[_id] = True
659  if formats[format] != len(bunch):
660  print 'bunch_maus_root_io_test fail on format', format, 'length', len(bunch)
661  out = 'fail'
662  if out == 'fail':
663  print 'bunch_maus_root_io_test', format
664  print 'spill', 'event', 'track', 'statn'
665  for _id in id_dict.keys():
666  for i in _id:
667  print str(i).rjust(5),
668  print
669  return out
670 
672  out = 'pass'
673  file_types = filter(lambda format: format.find('maus_root') < 0, Hit.file_types())
674  for format in file_types:
675  try:
676  bunch.hit_write_builtin(format, 'bunch_builtin_io_test.gz')
677  bunch1 = Bunch.new_from_read_builtin(format, 'bunch_builtin_io_test.gz')
678  if format == 'icool_for003' or format == 'mars_1' or format == 'g4beamline_bl_track_file': #no station in these formats, defaults to 0
679  for i in range( len(bunch1) ): bunch1[i]['station'] = bunch[i]['station']
680  if bunch.covariances_set() and bunch.means_set():
681  bunch1.set_covariance_matrix(use_internal_covariance_matrix=True, covariance_matrix=bunch._Bunch__covs, mean_dict=bunch._Bunch__means)
682  except Exception:
683  if format.find('muon1') == 0:
684  bunch1 = bunch # muon1 tested elsewhere
685  else:
686  sys.excepthook(*sys.exc_info())
687  bunch1 = None #force a fail on exception
688  os.remove('bunch_builtin_io_test.gz')
689  if format.find('maus') >= 0:
690  for i, hit in enumerate(bunch1): # write op does not conserve particle_number
691  bunch1[i]['particle_number'] = bunch[i]['particle_number']
692  bunch1[i]['event_number'] = bunch[i]['event_number']
693  if format.find('primary') >= 0:
694  bunch1[i]['station'] = bunch[i]['station']
695  if bunch != bunch1:
696  print 'Builtin io failed with format',format
697  try:
698  print 'Input length',len(bunch),'Output length',len(bunch1)
699  for i,hit in enumerate(bunch):
700  counter = 0
701  if bunch[i] != bunch1[i] and counter < 5:
702  print bunch[i], '\n', bunch1[i], '\n'
703  counter += 1
704  except:
705  sys.excepthook(*sys.exc_info())
706  out = 'fail'
707  return out
708 
710  muon1_test_input = sys.prefix+'/share/xboa/data/muon1_output.csv'
711  bunch = Bunch.new_from_read_builtin('muon1_csv', muon1_test_input)
712  testpass = len(bunch) == 5
713  testpass = testpass and bunch[0]['pid'] == -13
714  testpass = testpass and abs(bunch[0]['mass']-Common.pdg_pid_to_mass[13]) < 1e-9
715  testpass = testpass and abs(bunch[0]['x']-6.82814508887e-002) < 1e-9
716  testpass = testpass and abs(bunch[0]['y']+8.73903560505e-002) < 1e-9
717  testpass = testpass and abs(bunch[0]['t']-1.86835078986e-1) < 1e-9
718  testpass = testpass and abs(bunch[0]['charge']-1) < 1e-9
719  testpass = testpass and abs(bunch[0]['p']-208.470289725) < 1e-9
720  print bunch[0], len(bunch)
721  if testpass: return 'pass'
722  else: return 'fail'
723 
725  # we aren't interested in hit details, only number of hits in each station
726  # bunch_dict is not sorted by station, so we first count how many hits there are in each station
727  total_hit_count = 0
728  hit_list_per_station = {}
729  for key, bunch in bunch_dict.iteritems():
730  for hit in bunch:
731  station = hit['station']
732  if not station in hit_list_per_station: hit_list_per_station[ station ] = []
733  hit_list_per_station[station].append(hit)
734  total_hit_count += 1
735 
736  for format in filter(lambda format: format.find('maus_root') < 0 and format.find('muon1') < 0, Hit.file_types()):
737  Bunch.hit_write_builtin_from_dict (bunch_dict, format, 'bunch_dict_builtin_io_test.gz')
738  bunch_dict1 = Bunch.new_dict_from_read_builtin(format, 'bunch_dict_builtin_io_test.gz')
739  bunch_list1 = Bunch.new_list_from_read_builtin(format, 'bunch_dict_builtin_io_test.gz')
740  os.remove('bunch_dict_builtin_io_test.gz')
741  if format == 'icool_for003' or format == 'mars_1' or format == 'maus_primary' or format == 'g4beamline_bl_track_file':
742  if len(bunch_dict1) != 1 or len(bunch_dict1[0]) != total_hit_count: return 'fail'
743  if len(bunch_list1) != 1 or len(bunch_list1[0]) != total_hit_count: return 'fail'
744  else:
745  if len(bunch_dict1) != len(hit_list_per_station):
746  print 'len in',len(hit_list_per_station),'len out',len(bunch_dict1),format
747  return 'fail'
748  for key in bunch_dict1.keys():
749  if len(bunch_dict1[key]) != len(hit_list_per_station[key]):
750  for i,hit in enumerate(hit_list_per_station[key]): print 'in_'+str(i)+' ',hit['particle_number'],hit['event_number'],hit['station'],hit in hit_list_per_station[key]
751  for i,hit in enumerate(bunch_dict1[key]): print 'out_'+str(i),hit['particle_number'],hit['event_number'],hit['station'],hit in bunch_dict1[key]
752  print 'station:',key,'n_hits_in:',len(hit_list_per_station[key]),'n_hits_out:',len(bunch_dict1[key]), format
753  return 'fail'
754  return 'pass'
755 
756 def __test_comptor(hit1, hit2):
757  return cmp(hit2['eventNumber'], hit1['eventNumber'])
758 
759 def __test_truth(hit):
760  return hit['eventNumber'] == 3
761 
763  fh = open('bunch_user_io_test', 'w')
764  fh.write('test output\n')
765  format_list = ['eventNumber', 'eventNumber', 'pid', 'status', 't', 'local_weight', 'x', 'y', 'z', 'px', 'py', 'pz', 'sx', 'sy', 'sz', 'station', 'particleNumber', 'charge']
766  units_dict = {'eventNumber':'', 'particleNumber':'', 'pid':'', 'status':'', 't':'s', 'local_weight':'', 'x':'m', 'y':'m', 'z':'m', 'px':'GeV/c', 'py':'GeV/c', 'pz':'GeV/c', 'sx':'', 'sy':'', 'sz':'', 'station':'', 'charge':''}
767 
768  bunch.hit_write_user(format_list, units_dict, fh)
769  fh.close()
770  fh = open('bunch_user_io_test', 'r')
771  bunch1 = Bunch.new_from_read_user(format_list, units_dict, fh, 1)
772  for hit in bunch1: hit.mass_shell_condition('energy')
773  if len(bunch1) != len(bunch):
774  raise RuntimeError('bunches of different length')
775  if not bunch.hit_equality(bunch1):
776  counter = 0
777  for i in range(len(bunch)):
778  if bunch[i] != bunch1[i] and counter < 5:
779  print i, bunch[i], '\n', bunch1[i], '\n'
780  counter += 1
781  raise RuntimeError('bunch not equal to bunch1')
782  fh.close()
783  fh = open('bunch_user_io_test', 'w')
784  bunch.hit_write_user(format_list, units_dict, fh, comparator=__test_comptor)
785  fh.close()
786  fh = open('bunch_user_io_test', 'r')
787  bunch1 = Bunch.new_from_read_user(format_list, units_dict, fh, 0, __test_truth)
788  nhits = 0
789  os.remove('bunch_user_io_test')
790  for hit in bunch:
791  if __test_truth(hit): nhits+=1
792  if len(bunch1) != nhits: raise RuntimeError('bunch1 length '+str(len(bunch1))+' != nhits '+str(nhits))
793  for hit in bunch1:
794  if hit['eventNumber'] != 3: return 'fail'
795  return 'pass'
796 # def new_from_read_user(format_list, format_units_dict, filehandle, number_of_skip_lines, station_number=None, number_of_hits=-1):
797 
799  print 'set covariances test not defined'
800  return 'fail'
801 
803  dr2 = (bunch.moment(['r_squared','energy'],{'r_squared':0.,'energy':0})-(bunch.moment(['x','x'])+bunch.moment(['y','y']))*bunch.mean(['energy'])['energy'])\
804  *bunch.mean(['energy'])['energy']/bunch.moment(['energy','energy'])
805  dr2_test = bunch.get_dispersion_rsquared()
806  if abs(dr2 - dr2_test) < __float_tol: return 'pass'
807  return 'fail'
808 
810  field_axis_dict = {'x':5., 'y':3.}
811  rotation_axis_dict = {'x':1.,'y':4.,'px':2.,'py':-1.}
812  bz = 5.*Common.units['T']
813  l_can = bunch.get_kinetic_angular_momentum(rotation_axis_dict) \
814  + Common.constants['c_light']*bz*(bunch.moment(['x','x'], field_axis_dict) \
815  + bunch.moment(['y','y'], field_axis_dict))/2.
816  test_lcan = bunch.get_canonical_angular_momentum(bz, field_axis_dict, rotation_axis_dict)
817  if abs(l_can - test_lcan) < __float_tol: return 'pass'
818  return 'fail'
819 
821  for frequency in [10., 100., 1000.]:
822  for offset in [0.1,0.2,0.3,0.4,0.5]:
823  bunch.period_transformation(offset, frequency, 'z')
824  for hit in bunch:
825  if abs(hit['z'])-1./frequency > __float_tol:
826  print 'bunch_period_transformation_test failed with frequency',frequency,' offset ',offset,' z out ',hit['z']
827  return 'fail'
828  return 'pass'
829 
831  ell1_test = [[ 1331.29542168, -316.9751004 ], [ -316.9751004, 377.35131 ]]
832  ell2_test = [[ 2.52000000e+03, -3.00000000e+00],[ -3.00000000e+00, 1.78571429e-02]]
833  ell1 = Bunch.build_ellipse_2d(420., 0.5, 6., 200., Common.pdg_pid_to_mass[13], False)
834  ell2 = Bunch.build_ellipse_2d(420., 0.5, 6., 200., Common.pdg_pid_to_mass[13], True)
835  for i in range(2):
836  for j in range(2):
837  if abs(ell1[i,j]-ell1_test[i][j]) > Common.float_tolerance: return 'fail'
838  for i in range(2):
839  for j in range(2):
840  if abs(ell2[i,j]-ell2_test[i][j]) > Common.float_tolerance: return 'fail'
841  return 'pass'
842 
844  mu_bunch = Bunch()
845  test_ellipse = numpy.matrix([[ 1055.52708433, -0., 0., -632.87811819],
846  [ -0., 760.21504469, 632.87811819, 0. ],
847  [ 0., 632.87811819, 1055.52708433, -0. ],
848  [ -632.87811819, 0., -0., 760.21504469]])
849  ellipse = Bunch.build_penn_ellipse(6., Common.pdg_pid_to_mass[13], 333., 0., 200., 0., 4.*Common.units['T'], +1.)
850  for i,a in enumerate(ellipse.flatten().tolist()):
851  x = ellipse.flatten().tolist()[0][i]
852  y = test_ellipse.flatten().tolist()[0][i]
853  if abs(x - y) > __float_tol*1e3:
854  print 'FAILED',i,x,y,x-y,'\noutput ellipse\n',ellipse,'\ntest ellipse\n',test_ellipse
855  return 'fail'
856 
857  test_ellipse = numpy.matrix([[ 0.0306599, -4.08799, 0, -5.08241],
858  [ -4.08799, 1421.63, 5.08241, 0],
859  [ 0, 5.08241, 0.0306599, -4.08799],
860  [ -5.08241, 0, -4.08799, 1421.63]])
861  ellipse = Bunch.build_penn_ellipse(2., Common.pdg_pid_to_mass[11], 3., 4., 100., 5., 6., -1.)
862  for i,a in enumerate(ellipse.flatten().tolist()):
863  x = ellipse.flatten().tolist()[0][i]
864  y = test_ellipse.flatten().tolist()[0][i]
865  if abs(x - y) > __float_tol*1e3:
866  print 'FAILED',i,x,y,x-y,'\noutput ellipse\n',ellipse,'\ntest ellipse\n',test_ellipse
867  return 'fail'
868  return 'pass'
869 
871  if bunch.covariances_set() == is_set and bunch.means_set() == is_set: return 'pass'
872  return 'fail'
873 
874 def bunch_test(bunch):
875  test_results = []
876  tests = [bunch_delete_test, bunch_str_test, bunch_repr_test, bunch_copy_test, bunch_deepcopy_test, bunch_len_test,
877  bunch_getitem_test, bunch_setitem_test, bunch_conditional_remove_test,
878  bunch_hits_test, bunch_get_hits_test, bunch_mean_test, bunch_moment_test, bunch_standard_deviation_test, bunch_axis_list_to_covariance_list_test, bunch_convert_string_to_axis_list_test,
879  bunch_set_geometric_momentum_test, bunch_get_axes_test, bunch_get_test, bunch_get_hit_variable_test, bunch_list_get_hit_variable_test,
880  bunch_builtin_io_test, bunch_user_io_test, bunch_get_dispersion_rsquared_test, bunch_canonical_angular_momentum_test, bunch_period_transformation_test,
881  bunch_bunch_weight_test, bunch_clear_local_weights_test, bunch_clear_global_weights_test, bunch_clear_weights_test, bunch_cut_test,
882  bunch_transmission_cut_test, bunch_histogram_var_bins_test, bunch_histogram_test]
883  run_test_group(test_results, tests, [(bunch,)]*len(tests))
884 
885  try:
886  Common.has_numpy()
887  tests = [bunch_covariance_matrix_test, bunch_translate_test, bunch_abelian_transformation_test,bunch_transform_to_test, bunch_get_amplitude_test]
888  args = [(bunch,)]*len(tests)
889  for i in range(len(tests)):
890  run_test(test_results, tests[i], args[i])
891  except ImportError:
892  test_results.append('Warning - could not find NumPy. Skipping NumPy dependent tests')
893  try:
894  Common.has_root()
895  tests = [bunch_root_histogram_test, bunch_root_scatter_graph_test]
896  args = [(bunch,)]*len(tests)
897  for i in range(len(tests)):
898  run_test(test_results, tests[i], args[i])
899  except ImportError:
900  test_results.append('Warning - could not find ROOT. Skipping ROOT dependent tests')
901  try:
902  Common.has_matplot()
903  tests = [bunch_matplot_histogram_test, bunch_matplot_scatter_graph_test]
904  args = [(bunch,)]*len(tests)
905  for i in range(len(tests)):
906  run_test(test_results, tests[i], args[i])
907  except ImportError:
908  test_results.append('Warning - could not find matplotlib. Skipping matplotlib dependent tests')
909  return parse_tests(test_results)
910 
911 def cut_bc_test():
912  """
913  Test the private member function _cut and accompanying Bunchcore._cut_double
914  """
915  bunch = Bunch()
916  test = bunch.cut({"x":0.}, operator.ge) == None # regression for bug where
917  # hit initialisation wasn't
918  # done properly
919  hit_1 = Hit.new_from_dict({'x':-1., 'local_weight':1., 'event_number':1})
920  hit_2 = Hit.new_from_dict({'x':+1., 'local_weight':1., 'event_number':2})
921  bunch = Bunch.new_from_hits([hit_1, hit_2])
922  bunch.cut({"x":0.}, operator.ge)
923  test &= abs(bunch[0]['local_weight']-1.) < 1e-9
924  test &= abs(bunch[1]['local_weight']) < 1.e-9
925  test &= abs(bunch[0]['global_weight']-1.) < 1e-9
926  test &= abs(bunch[1]['global_weight']-1.) < 1.e-9
927 
928  bunch.clear_weights()
929  bunch.cut({"x":0.}, operator.ge, global_cut = True)
930  test &= abs(bunch[0]['local_weight']-1.) < 1e-9
931  test &= abs(bunch[1]['local_weight']-1.) < 1e-9
932  test &= abs(bunch[0]['global_weight']-1.) < 1e-9
933  test &= abs(bunch[1]['global_weight']) < 1.e-9
934 
935  bunch.clear_weights()
936  bunch.cut({"event_number":2}, operator.eq, global_cut = True)
937  test &= abs(bunch[0]['local_weight']-1.) < 1e-9
938  test &= abs(bunch[1]['local_weight']-1.) < 1e-9
939  test &= abs(bunch[0]['global_weight']-1.) < 1e-9
940  test &= abs(bunch[1]['global_weight']) < 1.e-9
941 
942  bunch.clear_weights()
943  bunch.cut({"event_number":2}, operator.eq)
944  test &= abs(bunch[0]['local_weight']-1.) < 1e-9
945  test &= abs(bunch[1]['local_weight']) < 1.e-9
946  test &= abs(bunch[0]['global_weight']-1.) < 1e-9
947  test &= abs(bunch[1]['global_weight']-1.) < 1e-9
948 
949  bunch.cut({"x":0}, operator.ge) # can we pass integers instead of floats
950  try:
951  bunch.cut({"x":0}, bunch.cut) # incorrect function type (not a comparator)
952  test = False
953  except ValueError:
954  pass
955 
956  try:
957  bunch.cut({"x":0}, 1) # incorrect type (not a function)
958  test = False
959  except ValueError:
960  pass
961 
962  try:
963  bunch.cut({"x":"fish"}, operator.eq) # incorrect type
964  test = False
965  except TypeError:
966  pass
967 
968  if test:
969  return 'pass'
970  else:
971  return 'fail'
972 
973 def test_bunch():
974  test_results = []
975  run_test(test_results, bunchcore_test, ())
976  run_test(test_results, cut_bc_test, ())
977  run_test(test_results, bunch_muon1_csv_read_test, ())
978  Hit.clear_global_weights()
979  hit_list1 = [Hit.new_from_dict({'t':1.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':0, 'charge':-1.},'pz'),
980  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':227.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':1, 'charge':-1.},'pz'),
981  Hit.new_from_dict({'t':0.,'x':1.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':2, 'charge':-1.},'pz'),
982  Hit.new_from_dict({'t':0.,'x':0.,'y':1.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':3, 'charge':-1.},'pz'),
983  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':1.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':4, 'charge':-1.},'pz'),
984  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':1.,'energy':226.,'mass':Common.pdg_pid_to_mass[11], 'pid':11, 'eventNumber':5, 'charge':-1.},'pz'),
985  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':1.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'eventNumber':6, 'charge':-1.},'pz')]
986  hit_list2 = [Hit.new_from_dict({'t':2.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':0, 'station':1, 'charge':-1.},'pz'),
987  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':228.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':1, 'station':1, 'charge':-1.},'pz'),
988  Hit.new_from_dict({'t':0.,'x':2.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':2, 'station':1, 'charge':-1.},'pz'),
989  Hit.new_from_dict({'t':0.,'x':0.,'y':2.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':3, 'station':1, 'charge':-1.},'pz'),
990  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':2.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':4, 'station':1, 'charge':-1.},'pz'),
991  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':2.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':100., 'eventNumber':5, 'station':1, 'charge':-1.},'pz')]
992  hit_list3 = [Hit.new_from_dict({'t':3.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':0, 'station':2, 'charge':-1.},'pz'),
993  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':0.,'energy':229.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':1, 'station':2, 'charge':-1.},'pz'),
994  Hit.new_from_dict({'t':0.,'x':3.,'y':0.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':2, 'station':2, 'charge':-1.},'pz'),
995  Hit.new_from_dict({'t':0.,'x':0.,'y':3.,'px':0.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':3, 'station':2, 'charge':-1.},'pz'),
996  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':3.,'py':0.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':4, 'station':2, 'charge':-1.},'pz'),
997  Hit.new_from_dict({'t':0.,'x':0.,'y':0.,'px':0.,'py':3.,'energy':226.,'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200., 'eventNumber':5, 'station':2, 'charge':-1.},'pz')]
998  hit_list4 = []
999  for i in range(-100,100):
1000  hit_list4.append(Hit.new_from_dict({'t':i,'x':i*i,'y':2./(i*i+0.1),'px':1./(i+0.1),'py':10./(i+0.5),'energy':226.+abs(i),'mass':Common.pdg_pid_to_mass[13], 'pid':13, 'z':200.+i, 'eventNumber':i+100, 'station':2, 'charge':-1.},'pz'))
1001  bunch = Bunch.new_from_hits(hit_list1)
1002  bunch2 = bunch
1003  bunch3 = bunch.deepcopy()
1004  bunch3.set_covariance_matrix()
1005  bunch4 = bunch3.deepcopy()
1006  bunch_dict = {0:bunch, 1:bunch2, 2:Bunch.new_from_hits(hit_list3), 3:Bunch.new_from_hits(hit_list2), 4:bunch3, 5:bunch4, 6:Bunch.new_from_hits(hit_list4)}
1007  bunch_list = [bunch, bunch2, Bunch.new_from_hits(hit_list3), Bunch.new_from_hits(hit_list2), bunch3, bunch4, Bunch.new_from_hits(hit_list4)]
1008  set_dict = {0:False, 1:False, 2:False, 3:False, 4:True, 5:True, 6:False}
1009  (passes, fails, warns) = (0,0,0)
1010 
1011  for key,a_bunch in bunch_dict.iteritems():
1012  print '\nTesting bunch',key
1013  (my_passes, my_fails, my_warns) = bunch_test(a_bunch)
1014  passes += my_passes
1015  fails += my_fails
1016  warns += my_warns
1017  print '\n\n'
1018 
1019  for key,a_bunch in bunch_dict.iteritems():
1020  run_test(test_results, bunch_set_covariance_matrix_test, (a_bunch, set_dict[key]))
1021  tests = [ bunch_equality_test, bunch_equality_test, bunch_equality_test,
1022  bunch_equality_test, bunch_equality_test, bunch_equality_test,
1023  bunch_list_get_test, bunch_dict_builtin_io_test,
1024  bunch_build_ellipse_2d_test, bunch_build_ellipse_penn_test ]
1025  args = [(bunch, bunch2, True), (bunch, bunch3, False), (bunch3, bunch4, True),
1026  (bunch, bunch2, True), (bunch, bunch3, False), (bunch3, bunch4, True),
1027  (bunch_dict,), (bunch_dict,),
1028  (), ()]
1029  run_test_group(test_results, tests, args)
1030 
1031  try:
1032  Common.has_numpy()
1033  run_test(test_results,bunch_new_hit_shell_test, () )
1034  except ImportError:
1035  test_results.append('Warning - could not find numpy. Skipping numpy dependent tests')
1036  try:
1037  Common.has_root()
1038  run_test(test_results,bunch_root_graph_test, (bunch_dict,) )
1039  run_test(test_results,bunch_root_graph_test, (bunch_list,) )
1040  except ImportError:
1041  test_results.append('Warning - could not find ROOT. Skipping ROOT dependent tests')
1042  try:
1043  Common.has_matplot()
1044  run_test(test_results, bunch_matplot_graph_test, (bunch_dict,) )
1045  run_test(test_results, bunch_matplot_graph_test, (bunch_list,) )
1046  except ImportError:
1047  test_results.append('Warning - could not find matplot. Skipping matplot dependent tests')
1048  try:
1049  Common.has_maus()
1050  run_test(test_results, bunch_maus_root_io_test, (bunch_dict,) )
1051  except ImportError:
1052  test_results.append('Warning - could not find matplot. Skipping matplot dependent tests')
1053  (passesEq, failsEq, warnsEq) = parse_tests(test_results)
1054  passes += passesEq
1055  fails += failsEq
1056  warns += warnsEq
1057  Bunch.clear_global_weights()
1058  print '\n==============\n|| BUNCH ||\n=============='
1059  print 'Passed ',passes,' tests\nFailed ',fails,' tests\n',warns,' warnings\n\n\n'
1060  return (passes,fails,warns)
1061 
1062 if __name__ == "__main__":
1063  test_bunch()
1064 
1065