xboa
_common.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 """
19 common contains some global physics and other parameters, plus some useful functions
20 - icool_pid_to_pdg = dict from icool particle identification indices to pdg indices
21 - pdg_pid_to_icool = inverse of above
22 - constants = string to number dict of physical constants
23 - units = string to number dict of some units
24 Functions documented below
25 """
26 
27 import operator
28 import copy
29 import bisect
30 import sys
31 import string
32 import math
33 import os
34 import pydoc
35 import bisect
36 import atexit
37 import time
38 import signal
39 import ctypes
40 
41 xboa_version = '0.16.2'
42 float_tolerance = 1.e-9
43 kill_subprocesses_at_exit = True
44 
45 __has_maus = False
46 try:
47  import libMausCpp
48  __has_maus = True
49 except:
50  print 'Maus library not detected'
51 
52 __has_root = True
53 try:
54  import ROOT
55  ROOT.gStyle.SetPalette(1)
56 except ImportError:
57  print 'PyRoot not detected - ROOT graphing functions will not be available'
58  __has_root = False
59 
60 __has_scipy = True
61 try:
62  import scipy
63 except ImportError:
64  print 'SciPy not detected - a few analysis functions will not be available'
65  __has_scipy = False
66 
67 
68 __has_numpy = True
69 try:
70  import numpy
71  from numpy import linalg
72  from numpy import matrix
73 except ImportError:
74  print 'NumPy not detected - many analysis functions will not be available'
75  __has_numpy = False
76 
77 __has_matplot = True
78 try:
79  import matplotlib
80  import pylab
81  from matplotlib import pyplot
82 except ImportError:
83  print 'MatPlotLib not detected - matplotlib graphing functions will not be available'
84  __has_matplot = False
85 
86 __has_multiprocessing = True
87 try:
88  import multiprocessing
89 except ImportError:
90  print 'multiprocessing library not detected - subprocess library will not be available'
91  __has_multiprocessing = False
92 
93 __has_json = True
94 try:
95  import json
96 except ImportError:
97  print 'json library not detected - maus import will not be available'
98  __has_json = False
99 
100 pdg_pid_to_muon1 = {-13:'mu+'}
101 muon1_pid_to_pdg = {}
102 for k,v in pdg_pid_to_muon1.iteritems():
103  muon1_pid_to_pdg[v] = k
104 
105 icool_pid_to_pdg = {0:0, 1:-11, 2:-13, 3:211, 4:321, 5:2212, -1:11, -2:13, -3:-211, -4:-321, -5:-2212};
106 pdg_pid_to_icool = {};
107 for k,v in icool_pid_to_pdg.iteritems():
108  pdg_pid_to_icool[v] = k;
109 
110 #* MARS PARTICLE ID (JJ):
111 #* 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
112 #* p n pi+ pi- K+ K- mu+ mu- g e- e+ pbar pi0 d t He3 He4 num nuam nue nuae
113 #* 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40
114 #* K0L K0S K0 AK0 LAM ALAM SIG+ SIG0 SIG- nbar KSI0 KSI- OM- sib- si0b sib+ AKSI0 ksib+ omb+
115 
116 mars_pid_to_pdg = {1:2212, 2:2112, 3:211, 4:-211, 5:321, 6:-321, 7:-13, 8:13, 9:22, 10:11, 11:-11, 12:-2212, 13:111, 14:1000010020, 15:1000010030, 16:1000020030, 17:1000020040, 18:14, 19:-14, 20:12, 21:-12, 22:130, 23:310, 24:311, 25:-311,26:3122,27:-3122,28:0,29:0,30:0,31:-2112,32:0,33:0,34:0,35:0,36:0,38:0,39:0,40:0};
117 pdg_pid_to_mars = {}
118 for k,v in mars_pid_to_pdg.iteritems():
119  pdg_pid_to_mars[v] = k;
120 
121 pdg_pid_to_mass = {0:0, 11:0.510998910, 12:0., 13:105.6583668, 14:0., 22:0., 111:134.9766, 211:139.57018, 321:493.667, 2112:939.56536, 2212:938.271996, 1000010020:1876.1239, 1000010030:2809.432, 1000020030:2809.41346, 1000020040: 3728.4001, 130:497.614, 310:497.614, 311:497.614, 3122:1115.683}
122 pdg_pid_to_name = {0:'none', 11:'e-', 12:'electron neutrino', 13:'mu-', 14:'muon neutrino', 22:'photon', 111:'pi0', 211:'pi+', 321:'K+', 2112:'neutron', 2212:'proton', 1000010020:'deuterium', 1000010030:'tritium', 1000020030:'3He', 1000020040:'4He', 130:'K0L', 310:'K0S', 311:'K0', 3122:'lambda',
123 -11:'e+', -12:'electron antineutrino', -13:'mu+', -14:'muon antineutrino', -211:'pi-', -321:'K-', -2112:'antineutron', -2212:'antiproton', -3122:'antilambda'}
124 
125 pdg_pid_to_charge = {0:0, 11:-1, 12:0, 13:-1, 14:0, 22:0, 111:0, 211:+1, 321:+1, 2112:0, 2212:+1, 1000010020:0, 1000010030:0, 1000020030:0, 1000020040:0, 130:0, 310:0, 311:0, 3122:0, -11:+1, -12:0, -13:+1, -14:0, -211:-1, -321:-1, -2112:0, -2212:-1, -3122:0}
127 
128 constants = {'c_light':299.792458, 'pi':3.14159265,'echarge':1}
129 units = {'':1.,
130  'mum':1.e-3, 'mm':1., 'cm':10., 'm':1.e3, 'km':1.e6,
131  'ns':1., 'mus':1.e3, 'ms':1e6, 's':1.e9,
132  'eV':1e-6, 'keV':1e-3, 'MeV':1., 'GeV':1.e3, 'TeV':1e6,
133  'eV/c':1e-6, 'keV/c':1e-3, 'MeV/c':1., 'GeV/c':1.e3, 'TeV/c':1e6,
134  'eV/c2':1e-6, 'keV/c2':1e-3, 'MeV/c2':1., 'GeV/c2':1.e3, 'TeV/c2':1e6,
135  'Gauss':1.e-7, 'mT':1.e-6, 'T':1.e-3, 'kT':1.,
136  'V':1.e-6, 'kV':1.e-3, 'MV':1., 'GV':1.e3,
137  'kHz':1.e-6, 'MHz':1.e-3, 'GHz':1.,
138  'GV/m':1.,'GV/mm':1.e3,
139  'kW':6.24150974e6, 'MW':6.24150974e9, 'GW':6.24150974e12,
140  'degrees':2.*constants['pi']/360., 'radians':1., 'deg':2.*constants['pi']/360., 'rad':1., 'degree':2.*constants['pi']/360., 'radian':1.,
141  'echarge':1., 'Coulomb':6.24150974*10.**18.
142 }
143 
144 ### root globals
145 #line_color_int=1, line_style_int=1, line_width_int=2, fill_color_int=None, stats_bool=False, hist_title_string=''
146 class rg:
147  """
148  Container to hold some details of root global style information
149  """
150  canvas_border_mode = 0
151  canvas_highlight_color = 2
152  canvas_fill_color = 10
153  hist_fill_color = canvas_fill_color
154  line_width = 2
155  line_color = 1
156  line_style = 1
157  fill_color = 0
158  graph_fill_color = 10
159  stats = False
160  histo_margin = 0.0 #margin around edge of a histogram inside axes
161  graph_margin = 0.1 #margin around edge of a graph inside axes
162  fit_color = 4
163  fit_style = 2
165 ## privates
166 _canvas_persistent = []
167 _hist_persistent = []
168 _graph_persistent = []
169 _legend_persistent = []
170 _function_persistent = []
172 ## pyplot globals
173 _figure_index = 1
174 
175 def substitute(file_name_in, file_name_out, switch_dict):
176  """
177  Read in file_name_in and write to file_name_out, replacing key with value in switch_dict. Must be a built in function somewhere to do same...
178 
179  - file_name_in = string name of the input file
180  - file_name_out = string name of the output file
181  - switch_dict = dict of values to be swapped to the values they will be swapped for
182 
183  e.g. common.substitute('file.in', 'file.out', {'energy':'momentum'})
184  """
185  fin = open(file_name_in, 'r')
186  fout = open(file_name_out, 'w')
187  for line in fin:
188  for key, value in switch_dict.iteritems():
189  line = line.replace(str(key), str(value))
190  fout.write(line)
191  fin.close()
192  fout.close()
193 
194 wrapped_y_function = None
195 def __y_function_wrapper(x_list_of_lists):
196  global wrapped_y_function
197  out = []
198  for value in x_list_of_lists: out.append(wrapped_y_function(value))
199  return out
200 
201 def nd_newton_raphson1(y_function, y_tolerances_list, x_start_values_list, x_deltas_list, max_iteration=10, x_upper_limits=None, x_lower_limits = None, verbose=True):
202  """
203  Root finding in an arbitrary dimensional system. Returns x-value for y(x) = 0; caveat is dimension of y must equal dimension of x.
204  If you use this, you might find more and better root finding functions in SciPy module
205 
206  - y_function is a reference to the function to be minimised i.e. y(x); it takes a list of x-values; and returns a list of y-values
207  - y_tolerances_list is a list of tolerances; the iteration will stop when abs(value) < value_tolerance
208  _ x_start_values_list is a list of the values I will try to start with
209  - x_deltas_list is a list of the initial estimates of the error on x
210  - max_iteration is the maximum number of iterations allowed
211  - x_upper_limits is for future development
212  - x_lower_limits is for future development
213 
214  e.g. nd_newton_raphson(some_function, [0.1, 0.1], [0,0], [1,1]) will find root to y(x) < (0.1,0.1); starting at x=(0,0); initial error estimated to be [1,1].
215  some_function would be called like some_function([x_0, x_1]) and should return a list like [y_0,y_1]
216  """
217  global wrapped_y_function
218  wrapped_y_function = y_function
219  return nd_newton_raphson2( __y_function_wrapper, y_tolerances_list, x_start_values_list, x_deltas_list, max_iteration, x_upper_limits, x_lower_limits, verbose)
221 def nd_newton_raphson2(y_function, y_tolerances_list, x_start_values_list, x_deltas_list, max_iteration=10, x_upper_limits=None, x_lower_limits = None, verbose=True):
222  """
223  Alternative version of nd_newton_raphson1. Here y_function takes a list of lists of x_values, of length dimension+1 and returns a list of lists of y_values
224  Optimisation for when y_function can be made faster by running several jobs at once...
225 
226  - y_function is a reference to the function to be minimised i.e. y(x); it takes a list of lists of x-values; and returns a list of lists y-values
227  - y_tolerances_list is a list of tolerances; the iteration will stop when abs(value) < value_tolerance
228  - x_start_values_list is a list of the values I will try to start with
229  - x_deltas_list is a list of the initial estimates of the error on x
230  - max_iteration is the maximum number of iterations allowed
231  - x_upper_limits is for future development
232  - x_lower_limits is for future development
233 
234  e.g. nd_newton_raphson(some_function, [0.1, 0.1], [0,0], [1,1]) will find root to y(x) < (0.1,0.1); starting at x=(0,0); initial error estimated to be [1,1].
235  some_function would be called like some_function([[x_00, x_01] ,[x_10, x_11], [x_20, x_21]) and should return a list like [[y_00,y_01],[y_10,y_11],[y_20,y_21]]
236  """
237  has_numpy()
238  done = False
239  x_list = []
240  for key in range( len(x_start_values_list)+1 ):
241  x_list.append( copy.deepcopy(x_start_values_list) )
242  delta_x = numpy.matrix( str(x_deltas_list) )
243  count = copy.deepcopy(max_iteration)
244  limits = type(x_upper_limits) == type([]) and type(x_lower_limits) == type([])
245  if limits:
246  delta_limit = []
247  for i in range(len(x_start_values_list)): delta_limit.append(x_upper_limits[i]-x_lower_limits[i])
248  while not done and count > 0:
249  count -= 1
250  jacobian = numpy.matrix( numpy.zeros( (len(x_start_values_list), len(x_start_values_list)) ) )
251  for i in range( len(x_start_values_list) ):
252  if limits: #make the function periodic with the limits
253  x_list[0][i] -= delta_limit[i]*math.floor(x_list[0][i]/delta_limit[i]-x_lower_limits[i]);
254  for i in range( len(x_start_values_list) ):
255  x_list[i+1] = copy.deepcopy(x_list[0])
256  x_list[i+1][i] += delta_x[0,i]
257  y_x = y_function(x_list)
258  try:
259  y_x0 = numpy.matrix( str(y_x[0]) )
260  except:
261  raise RuntimeError('Newton-Raphson failed to evaluate function with input '+str(x_list)+' and output '+str(y_x))
262  # check that I move closer to the root... should stabilise the function I hope
263  y_magnitude_squared = numpy.vdot(y_x0, y_x0)
264  for i in range( len(x_start_values_list) ):
265  if abs(delta_x[0,i]) < float_tolerance:
266  print 'Warning - returned without convergence; delta_x',str(delta_x),'is below float tolerance'
267  return x_list[0]
268  for j in range( len(y_tolerances_list) ):
269  jacobian[j,i] = (y_x[i+1][j] - y_x[0][j])/delta_x[0,i]
270  x_list[i+1][i] -= delta_x[0,i]
271  try:
272  delta_x = -(linalg.inv(jacobian)*y_x0.transpose()).transpose()
273  except:
274  print 'Newton-Raphson failed with (singular?) jacobian\n',jacobian,'\nbailing out'
275  print 'x:',x_list[0],' y(x):',y_x[0],' dx:',delta_x
276  return x_list[0]
277  done = True
278  if verbose: print 'x:',x_list[0],' y(x):',y_x[0],' dx:',delta_x
279  for i in range( len(x_start_values_list) ):
280  if abs(y_x[0][i]) > y_tolerances_list[i]:
281  done = False
282  x_list[0][i] += delta_x[0,i]
283  y_x0 = y_function([x_list[0]])
284  return x_list[0]
285 
286 def has_multiprocessing():
287  """Raise an exception if multiprocessing libraries have not been imported properly"""
288  if not __has_multiprocessing:
289  raise ImportError("Attempt to use multiprocessing when library has not been imported - multiprocessing requires matplot >= 2.6")
290  return True
291 
292 def has_maus():
293  """Raise an exception if MAUS tracking library has not been imported properly"""
294  if(not __has_maus):
295  raise ImportError("Attempt to use maus when library has not been imported - check your maus installation")
296  return True
297 
298 def has_root():
299  """Raise an exception if ROOT graphics libraries have not been imported properly"""
300  if(not __has_root):
301  raise ImportError("Attempt to use root when library has not been imported - check your root installation")
302  return True
303 
304 def has_numpy():
305  """Raise an exception if NumPy numerical algebra libraries have not been imported properly"""
306  if(not __has_numpy):
307  raise ImportError("Attempt to use numpy when library has not been imported - check your numpy installation")
308  return True
309 
310 def has_scipy():
311  """Raise an exception if SciPy math/analysis libraries have not been imported properly"""
312  if(not __has_scipy):
313  raise ImportError("Attempt to use scipy when library has not been imported - check your scipy installation")
314  return True
315 
316 
317 def has_matplot():
318  """Raise an exception if NumPy numerical algebra libraries have not been imported properly"""
319  if(not __has_matplot):
320  raise ImportError("Attempt to use matplotlib when library has not been imported - check your matplotlib installation")
321  return True
322 
323 def has_json():
324  """Raise an exception if json data libraries have not been imported properly"""
325  if(not __has_json):
326  raise ImportError("Attempt to use json when library has not been imported - check your json installation (python >= 2.6)")
327  return True
328 
330 def min_max(x_float_list, weight_list=[], margin = 0.1, xmin=None, xmax=None):
331  """
332  Return minimum and maximum of a list (i) discarding values with ~0 weight and (ii) adding a margin. For making histograms.
333 
334  - x_float_list = return minimum and maximum of this list of floats
335  - weight_list = ignore items in x_float_list if weight is 0. Ignored if weight_list is not same length as x_float_list
336  - margin = add a margin given by (x_max-x_min)*margin
337  - xmin = if set, will override the xmin value
338  - xmax = if set, will override the xmax value
339 
340  e.g. common.min_max([0.1,0.2,0.3,0.4], [0,1,1,1], 0.2) will return [0.16,0.44]
341  """
342  new_floats = x_float_list
343  if len(weight_list) == len(x_float_list):
344  new_floats = []
345  for ind in range(len(weight_list)):
346  if weight_list[ind] > 1e-6:
347  new_floats.append(x_float_list[ind])
348  if len(new_floats) == 0: x = [0.,0.]
349  else: x = [min(new_floats), max(new_floats)]
350  delta = (x[1]-x[0])*margin
351  x[0] -= delta
352  x[1] += delta
353  if(x[1] - x[0] < 1e-9):
354  x[0] -= 1.
355  x[1] += 1.
356  if xmin!=None: x[0] = xmin
357  if xmax!=None: x[1] = xmax
358  return x
359 
360 def multisort(list_of_lists):
361  """Sort a list of lists by the first list"""
362  #first convert from vertical lists to horizontal lists
363  horizontal_list = []
364  for i in range( len(list_of_lists[0]) ):
365  horizontal_list.append([])
366  for j in range( len(list_of_lists) ):
367  horizontal_list[-1].append(list_of_lists[j][i])
368  #now sort by first element
369  getitem = operator.itemgetter(0)
370  horizontal_list = sorted(horizontal_list, key=getitem)
371  #now convert back to list_of_lists
372  for i in range( len(list_of_lists) ):
373  for j in range( len(list_of_lists[0]) ):
374  list_of_lists[i][j] = horizontal_list[j][i]
375  return list_of_lists
376 
377 def n_bins(n_points, nx_bins=None, ny_bins=None, nz_bins=None, n_dimensions=1):
378  """
379  Dynamically decide a number of bins depending on the number of points in the histogram
380 
381  - n_points = number of data points in the histogram
382  - nx_bins = set to an integer to override the automatic selection for number of x bins
383  - ny_bins = set to an integer to override the automatic selection for number of y bins
384  - n_dimensions = set to number of dimensions in the histogram
385 
386  Return value is a tuple (nx_bins, ny_bins, nz_bins), setting 0 to values that are out of the dimension range
387  """
388  out = [nx_bins, ny_bins, nz_bins]
389  num_events = float(n_points)
390  if nx_bins==None and n_dimensions==1: out[0] = int(num_events/10.)+1
391  if nx_bins==None and n_dimensions==2: out[0] = int(num_events**0.7/10.)+1
392  if ny_bins==None and n_dimensions==2: out[1] = int(num_events**0.7/10.)+1
393  if nx_bins==None and n_dimensions==3: out[0] = int(num_events**0.5/10.)+1
394  if ny_bins==None and n_dimensions==3: out[1] = int(num_events**0.5/10.)+1
395  if nz_bins==None and n_dimensions==3: out[2] = int(num_events**0.5/10.)+1
396  for i in range(3):
397  if out[i] == None: out[i]= 0 #set out to 0 for values that are out of the dimension range
398  return tuple(out)
399 
400 def histogram(x_values, x_bins, y_values=None, y_bins=None, weights=None):
401  """
402  Get a 1d or 2d list of bin weights from a set of data, weights and bin edges
403 
404  - x_values = list of x values to be binned
405  - x_bins = list of x bin edges
406  - y_values = list of y values to be binned. Set to None to make a 1d binning
407  - y_bins = list of y bin edges of same length as x_values
408  - weights = list of statistical weights of same length as x_values. Set to None to make all weights default to 1.
409 
410  Return value is a tuple of (bin_weights, x_bins, y_bins)
411  """
412  has_numpy()
413  is_1d = False
414  if weights == None:
415  weights = [1.]*len(x_values)
416  if y_values == None:
417  is_1d = True
418  y_values = [0.]*len(x_values)
419  y_bins = [-1.,1.]
420  contents = numpy.zeros((len(x_bins)-1, len(y_bins)-1), )
421  for i,x in enumerate(x_values):
422  p = bisect.bisect_right(x_bins,x_values[i])-1
423  q = bisect.bisect_right(y_bins,y_values[i])-1
424  if p>=0 and p<len(x_bins)-1 and q>=0 and q<len(y_bins)-1:
425  contents[p,q] += weights[i]
426  out = (contents, x_bins, y_bins)
427  if is_1d: return (contents, x_bins, [])
428  else: return (contents, x_bins, y_bins)
429 
430 def get_bin_edges(list_of_variables, number_of_bins, xmin=None, xmax=None):
431  """
432  Get a sorted list of equally spaced bin edges from a list of floats
433 
434  - list_of_variables = list of floats to be binned
435  - number_of_bins = number of bins to make; note that there will be number_of_bins+1 edges
436  - xmin = lower edge of all the bins (set to None to auto-detect)
437  - xmax = upper edge of all the bins (set to None to auto-detect)
438  """
439  my_bins = []
440  mm = min_max(list_of_variables, margin=0.0)
441  if xmin!=None: mm[0]=xmin
442  if xmax!=None: mm[1]=xmax
443  delta = (mm[1] - mm[0])/float(number_of_bins)
444  for i in range(number_of_bins+1):
445  my_bins.append(mm[0]+delta*i)
446  return my_bins
447 
448 def make_root_canvas(name_string, title_string=None, bg_color=rg.canvas_fill_color, highlight_color=rg.canvas_highlight_color,
449  border_mode=rg.canvas_border_mode, frame_fill_color=rg.hist_fill_color):
450  """
451  Make a root canvas with name canvas_name_string-<index> where <index> is a unique integer starting from 0
452 
453  - name_string = Name of the canvas. Due to a ROOT bug, xboa adds a unique identifier to the name_string to ensure the canvas is drawn.
454  - title_string = Title of the canvas (as displayed in canvas window title bar); if set to None will use the name_string.
455  - bg_color = Fill color of the canvas.
456  - highlight_color = When a canvas is selected, ROOT draws a border with a particular color to indicate that the canvas is selected.
457  - border_mode = When a canvas is selected, ROOT draws a border if border mode is not set to 0.
458  - frame_fill_color = Fill color of frames drawn on the canvas (e.g. histograms).
459  """
460  has_root()
461  canvas_name_string = name_string+"-"+str(len(_canvas_persistent))
462  if title_string == None: title_string = name_string
463  canvas = ROOT.TCanvas(canvas_name_string, title_string)
464  canvas.SetHighLightColor(highlight_color)
465  canvas.SetFillColor(bg_color)
466  canvas.SetBorderMode(border_mode)
467  canvas.SetFrameFillColor(frame_fill_color)
468  _canvas_persistent.append(canvas)
469  return canvas
470 
471 def make_root_histogram(name_string, x_float_list, x_axis_string, n_x_bins, y_float_list=[], y_axis_string='', n_y_bins=0, weight_list=[], xmin=None, xmax=None, ymin=None, ymax=None,
472  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=''):
473  """
474  Make a root histogram with data taken from float lists and axes named after the axis strings.
475 
476  - name_string = name given to the histogram
477  - x_float_list = list of x-data
478  - x_axis_string = string used to label the x-axis
479  - n_x_bins = number of bins in x direction
480  - y_float_list = list of y-data. If number of items in y list not equal to number in x list, will build 1d histogram
481  - y_axis_string = string used to label the y-axis
482  - n_y_bins = number of y bins
483  - weight_list = if present, each item will be filled with weight taken from this list
484  - xmin = float that overrides auto-detection of minimum x-axis value
485  - xmax = float that overrides auto-detection of maximum x-axis value
486  - ymin = float that overrides auto-detection of minimum y-axis value
487  - ymax = float that overrides auto-detection of maximum y-axis value
488  - line_color = int that sets the line colour of the histogram
489  - line_style = int that sets the line style of the histogram
490  - line_width = int that sets the line width of the histogram
491  - fill_color = int that sets the fill color of the histogram
492  - stats = set to True to plot a stats box on the histogram
493  - hist_title_string = specify the string that will appear as a title on the canvas
494 
495  Return value is the histogram
496  """
497  has_root()
498  name_string += " "+str(len(_hist_persistent))
499  if len(y_float_list) == len(x_float_list):
500  x_min_max = min_max(x_float_list, weight_list, margin=rg.histo_margin, xmin=xmin, xmax=xmax)
501  y_min_max = min_max(y_float_list, weight_list, margin=rg.histo_margin, xmin=ymin, xmax=ymax)
502  hist = ROOT.TH2D(name_string, hist_title_string+';'+x_axis_string+';'+y_axis_string, n_x_bins, x_min_max[0], x_min_max[1], n_y_bins, y_min_max[0], y_min_max[1])
503  if(len(weight_list) == len(x_float_list)):
504  for i in range( len(x_float_list) ):
505  hist.Fill(x_float_list[i], y_float_list[i], weight_list[i])
506  else:
507  for i in range( len(x_float_list) ):
508  hist.Fill(x_float_list[i], y_float_list[i])
509  else:
510  x_min_max = min_max(x_float_list, weight_list, margin=rg.histo_margin, xmin=xmin, xmax=xmax)
511  hist = ROOT.TH1D(name_string, hist_title_string+';'+x_axis_string, n_x_bins, x_min_max[0], x_min_max[1])
512  if(len(weight_list) == len(x_float_list)):
513  for i in range( len(x_float_list) ):
514  hist.Fill(x_float_list[i], weight_list[i])
515  else:
516  for i in range( len(x_float_list) ):
517  hist.Fill(x_float_list[i])
518  _hist_persistent.append(hist)
519  hist.SetLineColor(line_color)
520  hist.SetLineStyle(line_style)
521  if fill_color!=None: hist.SetFillColor(fill_color)
522  hist.SetStats(stats)
523  return hist
524 
525 def make_root_legend(canvas, root_item_list):
526  """
527  Build a legend for the canvas
528  """
529  has_root()
530  if len(root_item_list) == 0:
531  raise KeyError("No items for ROOT legend")
532  canvas.cd()
533  leg = ROOT.TLegend()
534  for index, leg in enumerate(root_item_list):
535  leg_min = 0.89-0.08*len(root_item_list)
536  if leg_min < 0.1:
537  leg_min = 0.1
538  leg = ROOT.TLegend(0.0, leg_min, 0.4, 0.90) # pylint: disable=E1101
539  leg.SetEntrySeparation(0.6)
540  for i, hist in enumerate(root_item_list):
541  leg.AddEntry(hist, root_item_list[i].GetName())
542  leg.SetFillColor(10)
543  leg.SetBorderSize(0)
544  leg.Draw()
545  canvas.Update()
546  _legend_persistent.append(leg)
547  return leg
549 
550 #def make_root_graph_2d(name_string, x_float_list, x_axis_string, y_float_list, y_axis_string, z_float_list, z_axis_string, sort=True, xmin=None, xmax=None, ymin=None, ymax=None, zmin=None, zmax=None,
551 # line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.graph_fill_color, hist_title_string=''):
552 # """
553 # Make a root graph with data taken from float lists and axes named after the axis strings. Return value is a tuple of (hist, graph)
554 # name_string = name given to the histogram
555 # x_float_list = list of x-data
556 # x_axis_string = string used to label the x-axis
557 # y_float_list = list of y-data
558 # y_axis_string = string used to label the y-axis
559 # z_float_list = list of z-data
560 # z_axis_string = string used to label the z-axis
561 # sort = boolean - set to true to automatically sort input data
562 # xmin = float that overrides auto-detection of minimum x-axis value
563 # xmax = float that overrides auto-detection of maximum x-axis value
564 # ymin = float that overrides auto-detection of minimum y-axis value
565 # ymax = float that overrides auto-detection of maximum y-axis value
566 # zmin = float that overrides auto-detection of minimum z-axis value
567 # zmax = float that overrides auto-detection of maximum z-axis value
568 # line_color = int that sets the line colour of the graph
569 # line_style = int that sets the line style of the graph
570 # line_width = int that sets the line width of the graph
571 # fill_color = graphs dont usually get a fill, but sometimes the fill colour turns up in e.g. legend drawing
572 # hist_title_string = specify the string that will appear as a title
573 # Return value is a tuple of (histogram, graph, graph_2) where graph is the TGraph2D and graph_2 is a graph of x vs y
574 # """
575 # has_root()
576 # if(len(x_float_list) == 0 or len(x_float_list) != len(y_float_list) or len(x_float_list) != len(z_float_list)):
577 # raise IndexError('Attempt to draw graph with no x-points, or different number of x to y to z points.')
578 # multilist = [x_float_list, y_float_list, z_float_list]
579 # if sort: multisort(multilist)
580 # x_min_max = min_max(multilist[0], margin=rg.graph_margin, xmin=xmin, xmax=xmax)
581 # y_min_max = min_max(multilist[1], margin=rg.graph_margin, xmin=ymin, xmax=ymax)
582 # z_min_max = min_max(multilist[1], margin=rg.graph_margin, xmin=zmin, xmax=zmax)
583 # hist = make_root_histogram(name_string, [], x_axis_string, 1000, [], y_axis_string, 1000, [], x_min_max[0], x_min_max[1], y_min_max[0], y_min_max[1],
584 # line_color=rg.line_color, line_style=rg.line_style, line_width=0, fill_color=None, stats=False, hist_title_string=hist_title_string)
585 # hist.SetMinimum(z_min_max[0])
586 # hist.SetMaximum(z_min_max[1])
587 # graph = ROOT.TGraph2D(len(x_float_list))
588 # graph.SetTitle(name_string)
589 # graph.SetName(name_string)
590 # graph.SetHistogram(hist)
591 # graph_2 = ROOT.TGraph(len(x_float_list))
592 # for i in range(len(x_float_list)):
593 # graph.SetPoint(i, x_float_list[i], y_float_list[i], z_float_list[i])
594 # graph_2.SetPoint(i, x_float_list[i], y_float_list[i])
595 # _graph_persistent.append(graph)
596 # _graph_persistent.append(graph_2)
597 # return (hist, graph, graph_2)
598 
599 
600 def make_matplot_histogram(x_float_list, x_axis_string, n_x_bins, y_float_list=[], y_axis_string='', n_y_bins=0, weight_list=[]):
601  """
602  Make a matplot graph with data taken from float lists and axes naemd after the axis strings. Return value is a tuple of (hist, graph)
603  matplot can format using tex expressions - use '$some math expression$' to include math text in your labels
604 
605  - x_float_list = list of x-data
606  - x_axis_string = string used to label the x-axis
607  - y_float_list = list of y-data
608  - y_axis_string = string used to label the y-axis
609 
610  After building the graph, use matplotlib.pyplot.show() to show something on the screen
611  """
612  has_matplot()
613  has_numpy()
614  global _figure_index
615  pyplot.figure(_figure_index)
616  _figure_index += 1
617  if(len(x_float_list) == 0):
618  raise IndexError('Attempt to draw histogram with no x-points')
619  if not len(y_float_list) == len(x_float_list):
620 # x_min_max = min_max(x_float_list, weight_list, margin=histo_margin)
621 # my_bins = range(int(n_x_bins))
622 # for i in range(len(my_bins)): my_bins[i] = x_min_max[0]+float(i)*(x_min_max[1]-x_min_max[0])/float(len(my_bins))
623  if(weight_list == []):
624  (n, my_bins) = numpy.histogram(a=x_float_list, bins=n_x_bins)
625  else:
626  (n, my_bins) = numpy.histogram(a=x_float_list, bins=n_x_bins, weights=weight_list)
627  new_bins = []
628  new_n = []
629  index = 0
630  while(abs(n[index])<1.e-9 and index < len(n)): index += 1
631  for i in range(index, len(n)-1):
632  new_bins.append((my_bins[i] + my_bins[i+1])/2.)
633  new_n.append(n[i])
634  hist = pylab.plot(new_bins, new_n)
635  pylab.xlabel(x_axis_string)
636  else:
637  if len(weight_list) == len(x_float_list): hist = pyplot.hexbin(x_float_list, y_float_list, weight_list, gridsize=(n_x_bins,n_y_bins))
638  else: hist = pyplot.hexbin(x_float_list, y_float_list, gridsize=(n_x_bins,n_y_bins))
639  pylab.xlabel(x_axis_string)
640  pylab.ylabel(y_axis_string)
641  return hist
642 
643 def make_root_graph(name_string, x_float_list, x_axis_string, y_float_list, y_axis_string, sort=True, xmin=None, xmax=None, ymin=None, ymax=None,
644  line_color=rg.line_color, line_style=rg.line_style, line_width=rg.line_width, fill_color=rg.graph_fill_color, hist_title_string=''):
645  """
646  Make a root graph with data taken from float lists and axes named after the axis strings. Return value is a tuple of (hist, graph)
647 
648  - name_string = name given to the histogram
649  - x_float_list = list of x-data
650  - x_axis_string = string used to label the x-axis
651  - y_float_list = list of y-data
652  - y_axis_string = string used to label the y-axis
653  - sort = boolean - set to true to automatically sort input data
654  - xmin = float that overrides auto-detection of minimum x-axis value
655  - xmax = float that overrides auto-detection of maximum x-axis value
656  - ymin = float that overrides auto-detection of minimum y-axis value
657  - ymax = float that overrides auto-detection of maximum y-axis value
658  - line_color = int that sets the line colour of the graph
659  - line_style = int that sets the line style of the graph
660  - line_width = int that sets the line width of the graph
661  - fill_color = graphs dont usually get a fill, but sometimes the fill colour turns up in e.g. legend drawing
662  - hist_title_string = specify the string that will appear as a title
663 
664  Return value is a tuple of (histogram, graph)
665  """
666  has_root()
667  x_float_list = copy.deepcopy(x_float_list)
668  y_float_list = copy.deepcopy(y_float_list)
669  if(len(x_float_list) == 0 or len(x_float_list) != len(y_float_list)):
670  raise IndexError('Attempt to draw graph with no x-points, or different number of x to y points')
671 # name_string += " "+str(len(_hist_persistent))
672  multilist = [x_float_list, y_float_list]
673  if sort: multisort(multilist)
674  x_min_max = min_max(multilist[0], margin=rg.graph_margin, xmin=xmin, xmax=xmax)
675  y_min_max = min_max(multilist[1], margin=rg.graph_margin, xmin=ymin, xmax=ymax)
676  hist = make_root_histogram(name_string, [], x_axis_string, 1000, [], y_axis_string, 1000, [], x_min_max[0], x_min_max[1], y_min_max[0], y_min_max[1],
677  line_color=rg.line_color, line_style=rg.line_style, line_width=0, fill_color=None, stats=False, hist_title_string=hist_title_string)
678  graph = ROOT.TGraph(len(x_float_list))
679  graph.SetTitle(name_string)
680  graph.SetName(name_string)
681  graph.SetLineColor(line_color)
682  graph.SetLineStyle(line_style)
683  graph.SetLineWidth(line_width)
684  graph.SetFillColor(fill_color)
685  for i in range(len(x_float_list)):
686  graph.SetPoint(i, x_float_list[i], y_float_list[i])
687  _graph_persistent.append(graph)
688  return (hist, graph)
689 
690 def make_matplot_graph(x_float_list, x_axis_string, y_float_list, y_axis_string, sort=True):
691  """
692  Make a matplot graph with data taken from float lists and axes naemd after the axis strings. Return value is a tuple of (hist, graph)
693  matplot can format using tex expressions - use '$some math expression$' to include math text in your labels
694 
695  - x_float_list = list of x-data
696  - x_axis_string = string used to label the x-axis
697  - y_float_list = list of y-data
698  - y_axis_string = string used to label the y-axis
699  - sort = boolean - set to true to automatically sort input data
700 
701  After building the graph, use matplotlib.pyplot.show() to show something on the screen
702  """
703  has_matplot()
704  global _figure_index
705  pyplot.figure(_figure_index)
706  _figure_index += 1
707  if(len(x_float_list) == 0 or len(x_float_list) != len(y_float_list)):
708  raise IndexError('Attempt to draw graph with no x-points, or different number of x to y points')
709  multilist = [x_float_list, y_float_list]
710  multisort(multilist)
711  x_min_max = min_max(multilist[0], margin=rg.graph_margin)
712  y_min_max = min_max(multilist[1], margin=rg.graph_margin)
713  myplot = pylab.plot(x_float_list, y_float_list)
714  pylab.xlabel(x_axis_string)
715  pylab.ylabel(y_axis_string)
716  pylab.xlim (x_min_max[0], x_min_max[1])
717  pylab.ylim (y_min_max[0], y_min_max[1])
718  matplotlib.pyplot.draw()
719 
720 def make_root_multigraph(name_string, x_float_list_of_lists, x_axis_string, y_float_list_of_lists, y_axis_string):
721  """
722  Print several different graphs on the same canvas. Some default colour scheme is applied, but it may not be the best...
723 
724  - name_string = name that will be given to the axes (histogram)
725  - x_float_list_of_lists = list of lists. Each list will be used as the x-axis for a graph
726  - x_axis_string = string that will be used to label the x_axis
727  - y_float_list_of_lists = list of lists. Each list will be used as the y-axis for a graph
728  - y_axis_string = string that will be used to label the y_axis
729 
730  E.g. common.make_root_multigraph('example', [[1.,2.,3.,4.], [1.,4.,9.,16.]], 'x', [[1.,2.,3.,4.],[1.,2.,3.,4.]], 'f(x)') will make a graph of f = x and f = x^0.5
731  """
732  has_root()
733  name_string += " "+str(len(_hist_persistent))
734  total_x_list = []
735  total_y_list = []
736  for a_list in x_float_list_of_lists: total_x_list += a_list
737  for a_list in y_float_list_of_lists: total_y_list += a_list
738  x_min_max = min_max(total_x_list, margin=rg.graph_margin)
739  y_min_max = min_max(total_y_list, margin=rg.graph_margin)
740  hist = ROOT.TH2D(name_string, ';'+x_axis_string+';'+y_axis_string, 1000, x_min_max[0], x_min_max[1], 1000, y_min_max[0], y_min_max[1])
741  graphs = []
742  for index in range( len(x_float_list_of_lists) ):
743  graphs.append(ROOT.TGraph( len(x_float_list_of_lists[index]) ))
744  ROOT.gStyle.SetPalette(int(len(x_float_list_of_lists)*1.25))
745  graphs[-1].SetLineColor(ROOT.gStyle.GetColorPalette(index + int(len(x_float_list_of_lists)*0.25) )) #funny algebra to darken the colors slightly
746  for i in range(len(x_float_list_of_lists[index])):
747  graphs[-1].SetPoint(i, x_float_list_of_lists[index][i], y_float_list_of_lists[index][i])
748  _graph_persistent.append(graphs[-1])
749  _hist_persistent.append(hist)
750  hist.SetStats(False)
751  return (hist, graphs)
752 
753 def make_matplot_multigraph(x_float_list_of_lists, x_axis_string, y_float_list_of_lists, y_axis_string):
754  """
755  Print several different graphs on the same axes. Some default colour scheme is applied, but it may not be the best...
756 
757  - x_float_list_of_lists = list of lists. Each list will be used as the x-axis for a graph
758  - x_axis_string = string that will be used to label the x_axis
759  - y_float_list_of_lists = list of lists. Each list will be used as the y-axis for a graph
760  - y_axis_string = string that will be used to label the y_axis
761 
762  E.g. common.make_matplot_multigraph('example', [[1.,2.,3.,4.], [1.,4.,9.,16.]], 'x', [[1.,2.,3.,4.],[1.,2.,3.,4.]], 'f(x)') will make a graph of f = x and f = x^0.5
763  """
764  has_matplot()
765  global _figure_index
766  pyplot.figure(_figure_index)
767  _figure_index += 1
768  total_x_list = []
769  total_y_list = []
770  for a_list in x_float_list_of_lists: total_x_list += a_list
771  for a_list in y_float_list_of_lists: total_y_list += a_list
772  x_min_max = min_max(total_x_list, margin=rg.graph_margin)
773  y_min_max = min_max(total_y_list, margin=rg.graph_margin)
774  for index in range( len(x_float_list_of_lists) ):
775  multilist = [x_float_list_of_lists[index], y_float_list_of_lists[index]]
776  multisort(multilist)
777  myplot = pylab.plot(multilist[0], multilist[1])
778  pylab.xlabel(x_axis_string)
779  pylab.ylabel(y_axis_string)
780  pylab.xlim (x_min_max[0], x_min_max[1])
781  pylab.ylim (y_min_max[0], y_min_max[1])
782  matplotlib.pyplot.draw()
783 
784 def make_matplot_scatter(x_float_list, x_axis_string, y_float_list, y_axis_string):
785  """
786  Make a matplot scatter graph with data taken from float lists and axes naemd after the axis strings.
787  matplot can format using tex expressions - use '$some math expression$' to include math text in your labels
789  - x_float_list = list of x-data
790  - x_axis_string = string used to label the x-axis
791  - y_float_list = list of y-data
792  - y_axis_string = string used to label the y-axis
793 
794  After building the graph, use matplotlib.pyplot.show() to show something on the screen
795  """
796  has_matplot()
797  global _figure_index
798  pyplot.figure(_figure_index)
799  _figure_index += 1
800  if(len(x_float_list) == 0):
801  raise IndexError('Attempt to draw histogram with no x-points')
802  hist = pyplot.scatter(x_float_list, y_float_list, s=1)
803  pylab.xlabel(x_axis_string)
804  pylab.ylabel(x_axis_string)
805  return hist
806 
807 def wait_for_root():
808  """
809  Force python to halt processing until ROOT windows are closed
810  """
811  print 'Close ROOT windows to continue'
812  root_done = False
813  while not root_done:
814  root_done = True
815  for canvas in _canvas_persistent:
816  if not canvas == None:
817  print canvas
818  root_done = False
819  time.sleep(0.01)
820 
822  """
823  Close root plots (and free memory)
824  """
825  global _canvas_persistent, _hist_persistent, _graph_persistent
826  for canv in _canvas_persistent: del(canv)
827  for hist in _hist_persistent: del(hist)
828  for graph in _graph_persistent: del(graph)
829  _canvas_persistent = []
830  _hist_persistent = []
831  _canvas_persistent = []
832 
833 def wait_for_matplot():
834  """Show any plots made using matplotlib on the screen"""
835  has_matplot()
836  matplotlib.pyplot.show()
839  """Show matplotlib plots and return to the script"""
840  has_matplot()
841  try:
842  subprocess( matplotlib.pyplot.show, () )
843  except:
844  raise ImportError('matplot_show_and_continue is only available is (i) ROOT is not installed or (ii) you have python > 2.6')
845 
846 __mp_subprocesses = []
847 __fk_subprocesses = []
848 def subprocess(function, args):
849  """
850  Make a function call in a subprocess; return the subprocess pid
851  - This uses the multiprocessing libary in the first instance - if multiprocessing is not available
852  the function tries to use os.fork(); however, os.fork is not compatible with ROOT. If
853  - multiprocessing is not available and ROOT is installed, this routine will throw an ImportError
854  """
855  try:
857  __mp_subprocesses.append( multiprocessing.Process(target=function, args=args) )
858  __mp_subprocesses[-1].start()
859  return __mp_subprocesses[-1]
860  except:
861  try:
862  has_root()
863  raise ImportError('Error - attempt to make a subprocess using fork when ROOT library is active')
864  except ImportError:
865  pid = os.fork()
866  if pid == 0:
867  try: function(*args)
868  except: os._exit(1)
869  os._exit(0)
870  else:
871  __fk_subprocesses.append(pid)
872  return pid
873 
875  """
876  Kill all subprocesses generated by common.subprocess call; automatically called at exit unless kill_subprocesses_at_exit is set to False
877  Note that this makes a call to the os, which may take a bit of time to respond.
878  """
879  for ps in __mp_subprocesses:
880  ps.terminate()
881  time.sleep(0.01) #dont ask me why this is necessary...
882  for pid in __fk_subprocesses:
883  os.kill(pid, signal.SIGKILL)
884 
885 def kolmogorov_smirnov_test(list_1, list_2):
886  """
887  Convenience wrapper for ROOT kolmogorov-smirnov test.
888  - list_1 = list of floats that are sampled from some parent probability distribution
889  - list_2 = list of floats that are sampled from other some parent distribution
890  Returns double between 0. and 1. giving probability that list_1 and list_2 have the same parent distribution.
891  """
892  list_1 = sorted(list_1)
893  list_2 = sorted(list_2)
894  c_array_1 = ctypes.c_double*len(list_1)
895  c_array_2 = ctypes.c_double*len(list_2)
896  ca1 = c_array_1()
897  ca2 = c_array_2()
898  for i in range(len(list_1)): ca1[i]=ctypes.c_double(list_1[i])
899  for i in range(len(list_2)): ca2[i]=ctypes.c_double(list_2[i])
900  ks = ROOT.TMath.KolmogorovTest(len(list_1), ca1, len(list_2), ca2, '')
901  return ks
902 
903 def __atexit():
904  """Calls some functions automatically at exit"""
905  if(kill_subprocesses_at_exit):
907 atexit.register(__atexit)
908 
909 
910 def make_grid(n_dimensions, n_per_dimension):
911  """
912  Make a rectangular n_dimensional grid of points evenly spaced between [-1,1] in each
913  dimension
914 
915  - n_dimensions = number of dimensions
916  - n_per_dimension = number of points in each dimension; total number of points will be
917  n_per_dimension^n_dimensions
918 
919  Return value is a list of numpy.matrices with shape (n_dimensions, 1)
920  """
921  has_numpy()
922  pos_vector = [numpy.matrix( [-1.]*n_dimensions )]
923  for i in range(n_dimensions):
924  pos_vector_copy = copy.deepcopy(pos_vector)
925  for j in range(1, n_per_dimension):
926  pos = (2.*j)/float(n_per_dimension-1)-1.
927  for vec in pos_vector_copy:
928  vec = copy.deepcopy(vec)
929  vec[0,i] = pos
930  pos_vector.append(vec)
931  return pos_vector
932 
933 def make_shell(n_per_dimension, ellipse):
934  """
935  Make a shell of points that sit on a hyper-ellipsoid defined by ellipse matrix
936 
937  - n_per_dimension = number of points in each dimension; total number of points will be
938  n_per_dimension^n_dimensions if n_per_dimension is even or
939  n_per_dimension^n_dimensions-1 if n_per_dimension is odd
940  - ellipse = matrix that defines the ellipse on which vector sits; should be a
941  numpy.matrix with shape (vec_length, vec_length)
942 
943  Points are defined on a (hyper-)cuboidal grid and then compressed so that lengths are
944  all 1. Doesn't necessarily mean points are evenly spaced. Return value is a list of
945  numpy.matrices with shape (n_dimensions, 1)
946  """
947  has_numpy()
948  n_dimensions = numpy.shape(ellipse)[0]
949  grid = make_grid(n_dimensions, n_per_dimension)
950  shell = []
951  ellipse_inv = numpy.linalg.inv(ellipse)
952  for vector in grid:
953  if numpy.vdot(vector,vector) > float_tolerance:
954  shell.append(normalise_vector(vector, ellipse_inv))
955  return shell
956 
957 def normalise_vector(vector, matrix_inverse):
958  """
959  Normalise vector so that vector.T() * matrix.inverse * vector = 1
960 
961  - vector = the vector to normalise; should be a numpy.matrix with shape (vec_length,1)
962  - matrix_inverse = inverse of matrix that defines the ellipse on which vector sits; should be a
963  numpy.matrix with shape (vec_length, vec_length)
964 
965  Return value is a list of numpy.matrices with shape (vec_length,1)
966  """
967  has_numpy()
968  scale = 1./(vector * matrix_inverse * vector.transpose())[0,0]**0.5
969  if scale != scale: raise(ValueError("Bad input - vector or matrix magnitude 0")) #isnan and isinf not available in 2.5
970  vector = vector*scale
971  return vector
972 
973 def __function_with_queue(args):
974  """
975  Wrapper function to put multiprocessing output into a queue
976 
977  - args tuple of (function_call, queue, function_arg) where
978  - function_call is a reference to the function that we want to wrap\n
979  - queue is the queue that will hold return values from the function\n
980  - function_arg is the argument to function_call\n
981  - index is an index indicating function_arg's position in the inputs
983  tuple of (index, Output) is placed into queue; if function_call throws an
984  exception, the exception is placed on the queue instead
985  """
986  (function_call, queue, function_arg, index) = args
987  try:
988  queue.put((index, function_call(*function_arg)))
989  except:
990  queue.put((index, sys.exc_info()[1]))
991 
992 def process_list(function_call, list_of_args, max_n_processes):
993  """
994  Run multiprocessing on a list of arguments
995 
996  - function_call multiprocess this function call
997  - list_of_args list of tuples of arguments for function_call
998  - max_n_processes maximum number of concurrent processes to use
999 
1000  Returns list of return values, one for each function call. List is always
1001  sorted into same order as input.
1002 
1003  e.g. process_list(time.sleep, [(3, ), (6, ), (2, )], 2) will multiprocess
1004  the time.sleep function with inputs 3, 6 and 2 across 2 cores and return
1005  list like [None, None, None].
1006  """
1007  manager = multiprocessing.Manager()
1008  queue = manager.Queue() # queue stores return values
1009  pool = multiprocessing.Pool(max_n_processes) # pool runs multiprocess
1010  # each process needs:
1011  # reference to queue to put returns values in
1012  # index (to sort return values)
1013  # plus function call and argument
1014  new_list_of_args = [(function_call, queue, x, i) for i,x in enumerate(list_of_args)]
1015  # run the processes
1016  pool.map(__function_with_queue, new_list_of_args)
1017  # sort output into a list
1018  out_list = []
1019  while not queue.empty():
1020  out_list.append(queue.get())
1021  out_list.sort()
1022  for i,item in enumerate(out_list):
1023  out_list[i] = item[1]
1024  # cleanup
1025  pool.close()
1026  return out_list
1027 
1028 def _sum_weight(weights, in_cut):
1029  n_points = len(in_cut)
1030  return sum([weights[i] for i in range(n_points) if in_cut[i]])
1031 
1032 def _get_fit_ellipse_covariance_matrix(points, weights, in_cut):
1033  n_points = len(points)
1034  sum_weight = _sum_weight(weights, in_cut)
1035  dimension = len(points[0])
1036  mean = numpy.zeros(dimension)
1037  cov = numpy.zeros([dimension, dimension])
1038  for i in range(n_points):
1039  if in_cut[i]:
1040  a_point = points[i]
1041  for j in range(dimension):
1042  mean[j] += a_point[j]/sum_weight*weights[i]
1043  for k in range(dimension):
1044  cov[j, k] += a_point[j]*a_point[k]/sum_weight*weights[i]
1045  for i in range(dimension):
1046  for j in range(dimension):
1047  cov[i, j] -= mean[i]*mean[j]
1048  return mean, cov
1049 
1050 def _update_fit_ellipse_cut(points, cov_inv, mean, eps_cut):
1051  n_points = len(points)
1052  n_dim = 2
1053  in_cut = [True for i in range(n_points)]
1054  sqrt_det_inv = numpy.linalg.det(cov_inv)**(1./float(n_dim))
1055  for i in range(n_points):
1056  vec = numpy.array(points[i])-mean
1057  vec_t = numpy.transpose(vec)
1058  eps = numpy.dot(vec_t, numpy.dot(cov_inv, vec))/sqrt_det_inv
1059  in_cut[i] = eps < eps_cut
1060  return in_cut
1061 
1062 def fit_ellipse(points, eps_cut, weights=None, max_number_of_iterations = 10, verbose = True):
1063  """
1064  Fit an ellipse of arbitrary dimension n to a set of points
1065 
1066  - points iterable of points; each point should be an iterable of floats of
1067  length n
1068  - eps_cut float cut value; only particles with
1069  eps = x^T*V^{-1}*x*|V|**{1/n} < eps_{cut}
1070  are considered. The cut value is a normalised chi-squared. It
1071  should be positive.
1072  - weights list of floats, one for each point; use statistically weighted
1073  particles for the fit. Set to None to ignore weights.
1074  - max_number_of_iterations integer number of iterations; the ellipse fitting
1075  procedure is repeated several times to attempt to improve the
1076  fit of the ellipse. This sets a maximum number of repetitions.
1077  - verbose set to True to provide some verbose output during fitting
1078 
1079  Returns ellipse centre vector <x> and defining matrix V. The ensemble of
1080  points on the ellipse is given by (x-<x>)^T*V^{-1}*(x-<x>) where x is a
1081  particle vector and V is the defining matrix.
1082  """
1083  n_points = len(points)
1084  old_det = 1.
1085  mean, cov = numpy.ones([2]), numpy.ones([2, 2])
1086  in_cut = [True for i in range(n_points)]
1087  if weights == None:
1088  weights = [1.,]*n_points
1089  iterations = 0
1090  try:
1091  while abs(numpy.linalg.det(cov)-old_det) > float_tolerance and \
1092  iterations < max_number_of_iterations:
1093  iterations += 1
1094  my_sum = 0
1095  for i in in_cut:
1096  if i == True:
1097  my_sum += 1
1098  old_det = numpy.linalg.det(cov)
1099  mean, cov = _get_fit_ellipse_covariance_matrix(points, weights, in_cut)
1100  mat_inv = numpy.linalg.inv(cov)
1101  in_cut = _update_fit_ellipse_cut(points, mat_inv, mean, eps_cut)
1102  except Exception:
1103  if verbose:
1104  print "xboa.common.fit_ellipse(...) failed to converge; final iteration as follows."
1105  sys.excepthook(*sys.exc_info())
1106  if verbose:
1107  print 'Weight in cut:\n', _sum_weight(weights, in_cut)
1108  print "Means:\n", mean
1109  print "Ellipse:\n", cov
1110  print "Number of iterations:\n", iterations
1111  return mean, cov
1112 
1113 def make_root_ellipse_function(mean, cov, contours=None, xmin=-1e3, xmax=1e3, ymin=-1e3, ymax=1e3):
1114  """
1115  Make a ROOT TFunction for a given beam ellipse
1116 
1117  - mean sequence type of length>=2 (must allow reference by integer [index])
1118  - cov ellipse >= 2x2 symmetric numpy.array (must allow referencing by integer
1119  indices like [a, b])
1120  - contours iterable; draw contours at values of x^T V^{-1} x given by
1121  elements. Uses ROOT default if equal to None.
1123  Returns a ROOT TFunction. Note ROOT TFunction does not like drawing
1124  small ellipses if the ellipse matrix is highly correlated (i.e. determinant
1125  near zero). A workaround is to make a TGraph with appropriate points (using
1126  e.g. formula given by this function). May be possible to fix by increasing
1127  the number of points in plot (func.SetNpy).
1128  """
1129  has_root()
1130  mat_inv = numpy.linalg.inv(cov)
1131  fit_func_str = "[2]*(x-[0])**2+[4]*(y-[1])**2+2*[3]*(x-[0])*(y-[1])"
1132  fit_func = ROOT.TF2("ellipse", fit_func_str)
1133  if contours != None:
1134  fit_func.SetContour(len(contours))
1135  for i, contour in enumerate(contours):
1136  fit_func.SetContourLevel(i, contour)
1137  if xmax == None or xmin == None or ymax == None or ymin == None:
1138  xmin = -1e3+mean[0]
1139  xmax = 1e3+mean[0]
1140  ymin = -1e3+mean[1]
1141  ymax = 1e3+mean[1]
1142 
1143  fit_func.SetRange(xmin, ymin, xmax, ymax)
1144  fit_func.SetParameter(0, mean[0])
1145  fit_func.SetParameter(1, mean[1])
1146  fit_func.SetParameter(2, mat_inv[0, 0])
1147  fit_func.SetParameter(3, mat_inv[0, 1])
1148  fit_func.SetParameter(4, mat_inv[1, 1])
1149  fit_func.SetLineColor(rg.fit_color)
1150  fit_func.SetLineStyle(rg.fit_style)
1151  _function_persistent.append(fit_func)
1152  return fit_func
1153 
1154 
1155 def common_overview_doc(verbose = False):
1156  """Creates some summary documentation for the common module. If verbose is True then will also print any functions or data not included in summary"""
1157  common_doc = '\cCommon module contains a number of useful interfaces, defaults, data and ancillary functions that support the rest of the XBOA package.\n'
1158 
1159  name_list = ['math', 'root', 'matplot', 'data', 'defaults', 'other_stuff']
1160  function_list = {
1161  'math' : ['min_max', 'multisort', 'nd_newton_raphson1', 'nd_newton_raphson2'],
1162  'root' : ['has_root','make_root_graph', 'make_root_histogram', 'make_root_multigraph', 'clear_root', 'wait_for_root', 'make_root_canvas'],
1163  'matplot' : ['has_matplot', 'make_matplot_graph', 'make_matplot_histogram', 'make_matplot_multigraph', 'make_matplot_scatter', 'wait_for_matplot', 'matplot_show_and_continue'],
1164  'other_stuff' : ['get_bin_edges', 'histogram', 'substitute', 'has_numpy', 'build_installation', 'make_grid', 'make_shell', 'normalise_vector', 'kolmogorov_smirnov_test', 'kill_all_subprocesses', 'has_multiprocessing', 'has_json'],
1165  'data' : ['constants', 'pdg_pid_to_icool', 'pdg_pid_to_mars', 'pdg_pid_to_mass', 'pdg_pid_to_name', 'pdg_pid_to_charge', 'icool_pid_to_pdg', 'mars_pid_to_pdg', 'units'],
1166  'defaults' : ['canvas_fill_color', 'canvas_highlight_color', 'default_margin', 'float_tolerance', 'graph_margin', 'histo_margin', 'python_version', 'xboa_version', 'kill_subprocesses_at_exit']
1167  }
1168 
1169  function_doc = {
1170  'math' : 'Maths functions:',
1171  'root' : 'Interfaces to ROOT plotting library:',
1172  'matplot' : 'Interfaces to matplotlib plotting library:',
1173  'other_stuff' : 'Some other useful functions:',
1174  'data' : 'Physics data:',
1175  'defaults' : 'Defaults for e.g. root canvases, etc:',
1176  }
1177 
1178  if verbose:
1179  dir_common = dir(sys.modules[__name__])
1180  print 'The following functions and data are in common but not in common_overview_doc:'
1181  for func in dir_common:
1182  found = False
1183  for func_sublist in function_list.values():
1184  if func in func_sublist: found = True
1185  if not found: print func,
1186  print '\n'
1187 
1188  print 'The following functions and data are in common_overview_doc but not in Common:'
1189  for func_sublist in function_list.values():
1190  for func in func_sublist:
1191  if func not in dir_common:
1192  print func,
1193  print
1194 
1195  doc = common_doc
1196  for key in name_list:
1197  doc = doc + function_doc[key]+'\n'
1198  for item in function_list[key]:
1199  doc = doc+' '+item+'\n'
1200  return doc
1201 
1202 __doc__ = common_overview_doc()
1203 
1204 
def has_json
Raise an exception if json data libraries have not been imported properly.
Definition: _common.py:336
def has_root
Raise an exception if ROOT graphics libraries have not been imported properly.
Definition: _common.py:307
def __function_with_queue
Wrapper function to put multiprocessing output into a queue.
Definition: _common.py:1023
def multisort
Sort a list of lists by the first list.
Definition: _common.py:375
def make_root_histogram
Make a root histogram with data taken from float lists and axes named after the axis strings...
Definition: _common.py:515
def normalise_vector
Normalise vector so that vector.T() * matrix.inverse * vector = 1.
Definition: _common.py:1003
def make_root_multigraph
Print several different graphs on the same canvas.
Definition: _common.py:755
def make_matplot_multigraph
Print several different graphs on the same axes.
Definition: _common.py:788
def substitute
Read in file_name_in and write to file_name_out, replacing key with value in switch_dict.
Definition: _common.py:187
def kill_all_subprocesses
Kill all subprocesses generated by common.subprocess call; automatically called at exit unless kill_s...
Definition: _common.py:910
def make_matplot_scatter
Make a matplot scatter graph with data taken from float lists and axes naemd after the axis strings...
Definition: _common.py:821
def __y_function_wrapper
Definition: _common.py:198
def wait_for_matplot
Show any plots made using matplotlib on the screen.
Definition: _common.py:863
def has_maus
Raise an exception if MAUS tracking library has not been imported properly.
Definition: _common.py:300
def make_shell
Make a shell of points that sit on a hyper-ellipsoid defined by ellipse matrix.
Definition: _common.py:982
def min_max
Return minimum and maximum of a list (i) discarding values with ~0 weight and (ii) adding a margin...
Definition: _common.py:354
def _update_fit_ellipse_cut
Definition: _common.py:1089
def kolmogorov_smirnov_test
Convenience wrapper for ROOT kolmogorov-smirnov test.
Definition: _common.py:924
def make_grid
Make a rectangular n_dimensional grid of points evenly spaced between [-1,1] in each dimension...
Definition: _common.py:955
def histogram
Get a 1d or 2d list of bin weights from a set of data, weights and bin edges.
Definition: _common.py:427
def nd_newton_raphson1
Root finding in an arbitrary dimensional system.
Definition: _common.py:220
def has_multiprocessing
Raise an exception if multiprocessing libraries have not been imported properly.
Definition: _common.py:293
def fit_ellipse
Fit an ellipse of arbitrary dimension n to a set of points.
Definition: _common.py:1122
def has_numpy
Raise an exception if NumPy numerical algebra libraries have not been imported properly.
Definition: _common.py:314
def common_overview_doc
Creates some summary documentation for the common module.
Definition: _common.py:1198
def nd_newton_raphson2
Alternative version of nd_newton_raphson1.
Definition: _common.py:241
def has_matplot
Raise an exception if NumPy numerical algebra libraries have not been imported properly.
Definition: _common.py:329
def has_scipy
Raise an exception if SciPy math/analysis libraries have not been imported properly.
Definition: _common.py:321
def make_root_legend
Build a legend for the canvas.
Definition: _common.py:548
def wait_for_root
Force python to halt processing until ROOT windows are closed.
Definition: _common.py:837
def make_root_graph
Make a root graph with data taken from float lists and axes named after the axis strings.
Definition: _common.py:687
def n_bins
Dynamically decide a number of bins depending on the number of points in the histogram.
Definition: _common.py:402
def make_root_ellipse_function
Make a ROOT TFunction for a given beam ellipse.
Definition: _common.py:1169
def _get_fit_ellipse_covariance_matrix
Definition: _common.py:1071
def clear_root
Close root plots (and free memory)
Definition: _common.py:852
def make_root_canvas
Make a root canvas with name canvas_name_string-<index> where <index> is a unique integer starting fr...
Definition: _common.py:477
def matplot_show_and_continue
Show matplotlib plots and return to the script.
Definition: _common.py:869
def process_list
Run multiprocessing on a list of arguments.
Definition: _common.py:1045
def __atexit
Calls some functions automatically at exit.
Definition: _common.py:938
def get_bin_edges
Get a sorted list of equally spaced bin edges from a list of floats.
Definition: _common.py:455
def make_matplot_histogram
Make a matplot graph with data taken from float lists and axes naemd after the axis strings...
Definition: _common.py:632
def make_matplot_graph
Make a matplot graph with data taken from float lists and axes naemd after the axis strings...
Definition: _common.py:725
root globals line_color_int=1, line_style_int=1, line_width_int=2, fill_color_int=None, stats_bool=False, hist_title_string=''
Definition: _common.py:151
def subprocess
Make a function call in a subprocess; return the subprocess pid.
Definition: _common.py:885