gambit is hosted by Hepforge, IPPP Durham
GAMBIT  v1.5.0-2191-ga4742ac
a Global And Modular Bsm Inference Tool
colouring.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 #
3 # GAMBIT: Global and Modular BSM Inference Tool
4 #*********************************************
5 # \file
6 #
7 # The GAMBIT Colouring-in Book
8 #
9 # A script for priming pippi runs and patching
10 # together their outputs. Designed to allow
11 # one to highlight specific regions of preferred
12 # parameter space, and then overplot them
13 # (e.g. relic density mechanisms).
14 #
15 #*********************************************
16 #
17 # Authors (add name and date if you modify):
18 #
19 # \author Pat Scott
20 # (p.scott@imperial.ac.uk)
21 # \date 2017 Feb
22 #
23 #*********************************************
24 
25 import sys
26 import yaml
27 import os
28 import copy
29 import re
30 import datetime
31 import stat
32 from multiprocessing import cpu_count
33 
34 # Global constants
35 verbose = False
36 default_contour = 95.4
37 default_contour_style = "Solid"
38 default_contour_width = 2
39 default_colour = "blue"
40 default_n_contours = 150
41 n_threads = cpu_count()
42 
43 # Checks for the presence of one or more keys in a YAML node
44 def check_node(node, name, keys):
45  is_ok = False
46  if isinstance(keys, basestring):
47  error = " Region \""+name+"\" must have key \""+str(keys)+"\"."
48  is_ok = (keys in node)
49  else:
50  error = " Region \""+name+"\" must have at least one of the keys "+str(keys)+"."
51  for x in keys:
52  if (x in node):
53  is_ok = True
54  break
55  if not is_ok: sys.exit(error)
56 
57 
58 def usage():
59  sys.exit('\n Usage: colouring [prime/combine] file1 file2\n where file1 is a YAML file that defines the regions you want to find.\n file2 is the pip file that defines the pippi run to modify.\n')
60 
61 
62 def main(arguments):
63  #Starting routine
64 
65  commandLineOptions = { 'prime':prime, 'combine':combine }
66 
67  if (len(arguments) is not 4):
68  #Print usage info and quit if pippi is invoked without an argument
69  usage()
70 
71  else:
72 
73  try:
74  #Check if colouring has been invoked with one of the two known specific commands
75  command = commandLineOptions[arguments[1]]
76  command(arguments[2], arguments[3])
77  except KeyError:
78  print
79  print 'Unrecognised colouring mode: '+arguments[1]
80  usage()
81 
82  sys.exit()
83 
84 
85 # Prepare separate pippi runs for each region
86 def prime(region_file, pip_file):
87 
88  print
89 
90  # Load up the regions from the yaml file, and the template for pippi runs from the pip file.
91  yaml_loader = yaml.full_load if hasattr(yaml, 'full_load') else yaml.load
92  regions = yaml_loader(file(region_file, 'r'))
93  pip = file(pip_file, 'r').read()
94 
95  # Get the drawing order
96  if regions["drawing_order"]:
97  drawing_order = regions["drawing_order"]
98  else:
99  sys.exit('ERROR: No drawing_order specified in '+region_file)
100 
101  # Retrieve the best fit in the base scan.
102  #First work out where the best fit file is. Start with the dir.
103  parse_dir = re.search('parse_dir\s*=\s*("(.*)"|\'(.*)\')', pip)
104  parse_dir = parse_dir.group(2 if parse_dir.group(2) else 3)
105  # Now infer the filename from the chain name.
106  chain_name = re.search('main_chain\s*=\s*("(.*)"|\'(.*)\')', pip)
107  chain_name = chain_name.group(2 if chain_name.group(2) else 3)
108  chain_name = re.sub(r'.*/|\..?.?.?$', '', chain_name)
109  best_filename = parse_dir+"/"+chain_name+".best"
110  # Now open the file and read in the best fit.
111  best_file = file(best_filename, 'r')
112  for x in best_file.readlines():
113  if not x.startswith('#'):
114  parts = x.split(':')
115  bestFit = float(parts[1])
116  break
117 
118 
119  # Iterate over the regions defined in the yaml file
120  for name in drawing_order:
121  if (verbose): print "Priming pippi run for extracting region: "+name
122  mechanism = regions[name]
123 
124  # Check that the mechanism is sufficiently specified
125  check_node(mechanism, name, ["extra_1D_plots", "extra_2D_plots"])
126  check_node(mechanism, name, "datastream")
127  check_node(mechanism, name, "cut")
128  check_node(mechanism, name, "label")
129 
130  # Make a new pip file for this region
131  newpip = copy.deepcopy(pip)
132 
133  # Replace the necessary parts of it, starting with the common section
134  contour = default_contour
135  if ("contour_levels" in mechanism): contour = mechanism["contour_levels"]
136  newpip = re.sub('contour_levels\s*=.*\n', 'contour_levels = '+str(contour)+'\n', newpip)
137  if ("extra_1D_plots" in mechanism): newpip = re.sub('oneD_plot_quantities\s*=\s*', 'oneD_plot_quantities = '+str(mechanism["extra_1D_plots"])+' ', newpip)
138  if ("extra_2D_plots" in mechanism): newpip = re.sub('twoD_plot_quantities\s*=\s*', 'twoD_plot_quantities = '+str(mechanism["extra_2D_plots"])+' ', newpip)
139  newpip = re.sub("bf_lnlike_for_profile_like\s*=\s*", "bf_lnlike_for_profile_like = "+str(bestFit)+" ", newpip)
140 
141  # Now the parse section
142  newpip = re.sub('parse_dir\s*=.*\n', 'parse_dir = \'parse_'+name+'\'\n', newpip)
143  newpip = re.sub('data_ranges\s*=\s*', 'data_ranges = '+mechanism["cut"]+' ', newpip)
144  if ("preamble" in mechanism): newpip = re.sub('preamble\s*=.*\n', 'preamble = \''+mechanism["preamble"]+'\'\n', newpip)
145  newpip = re.sub('assign_to_pippi_datastream\s*=\s*', 'assign_to_pippi_datastream = '
146  +mechanism["datastream"]+' \\\n ', newpip)
147  newpip = re.sub('quantity_labels\s*=\s*', 'quantity_labels = '
148  +mechanism["label"]+' \\\n ', newpip)
149 
150  # Now the script section
151  newpip = re.sub('script_dir\s*=.*\n', 'script_dir = \'scripts_'+name+'\'\n', newpip)
152  colour = default_colour
153  if ("colour" in mechanism): colour = mechanism["colour"]
154  newpip = re.sub('colour_scheme\s*=.*\n', 'colour_scheme = Blockshading_'+colour+'\n', newpip)
155 
156  # Now the plot section
157  newpip = re.sub('plot_dir\s*=.*\n', 'plot_dir = \'plots_'+name+'\'\n', newpip)
158 
159  # Output the new pip file
160  filename = os.path.splitext(pip_file)
161  filename = filename[0]+"_"+name+filename[1]
162  with open(filename,"w") as f:
163  f.write(newpip)
164  print " Generated new pip file "+filename
165 
166  print
167 
168 
169 def combine(region_file, pip_file):
170 
171  # Load up the regions from the yaml file, and the template for pippi runs from the pip file.
172  regions = yaml.load(file(region_file, 'r'))
173  pip = file(pip_file, 'r').read()
174 
175  # Get the drawing order
176  if regions["drawing_order"]:
177  drawing_order = regions["drawing_order"]
178  else:
179  sys.exit('ERROR: No drawing_order specified in '+region_file)
180 
181  # Find the filenames of all the 2D profile likelihood scripts
182  script_dir = re.search('script_dir\s*=\s*("(.*)"|\'(.*)\')', pip)
183  script_dir = script_dir.group(2 if script_dir.group(2) else 3)
184  like2D_scripts = [x for x in os.listdir(script_dir) if "like2D" in x]
185 
186  # Make a new scripts dir for the combined scripts
187  combo_dir = script_dir+'_combined'
188  if not os.path.isdir(combo_dir): os.mkdir(combo_dir)
189 
190  # Make a new plots dir for the combined plots
191  plot_dir = re.search('plot_dir\s*=\s*("(.*)"|\'(.*)\')', pip)
192  plot_dir = plot_dir.group(2 if plot_dir.group(2) else 3)
193  plot_combo_dir = plot_dir+'_combined'
194  if not os.path.isdir(plot_combo_dir): os.mkdir(plot_combo_dir)
195 
196  # Make a bulk script that calls all of the combined plotting scripts
197  bulkscript_name = 'make_combined_'+re.sub('\.pip', '', pip_file)+'_plots.bsh'
198  bulkscript = file(bulkscript_name, 'w')
199  bulkscript.write('cd '+combo_dir+'\n')
200 
201  # Iterate over the different scripts
202  for script_index, script_name in enumerate(like2D_scripts):
203 
204  # Initialise empty lists of the things that will vary with region
205  marker_commands = []
206  plot_commands = []
207  fill_colours = []
208  contour_commands = []
209  min_contours = []
210  region_names = []
211  regional_script = None
212 
213  # Iterate over all the regions mentioned in the drawing order
214  for name in drawing_order:
215 
216  # Open the script file, cycling if it doesn't exist
217  fullname = script_dir+'_'+name+'/'+script_name
218  if not os.path.isfile(fullname): continue
219  regional_script = file(fullname, 'r').read()
220 
221  # Save the name of the region
222  region_names.append(name)
223 
224  # Save the line style for the region
225  contour_style = regions[name]["contour_style"] if "contour_style" in regions[name] else default_contour_style
226 
227  # Extract the marker command
228  marker_command = re.findall('\s\s--draw-marker.*\n', regional_script)
229  if marker_command: marker_commands.append(marker_command[0])
230 
231  # Extract the fill colour
232  fill_colour = re.findall('--#...', regional_script)
233  fill_colours.append(fill_colour[1][2:])
234 
235  # Extract the plot command
236  plot_command = re.findall('\s\s--plot.*fill-transparency 1.*\n', regional_script)
237  plot_commands.append(plot_command[0])
238 
239  # Extract the contour command
240  contour_command_list = re.findall('\s\s--draw-contour.*\n', regional_script)
241  min_contour = 1.0
242  for i,x in enumerate(contour_command_list):
243  contour_command_list[i] = re.sub('width.*\n', 'width '+str(default_contour_width)+'\\\n', x)
244  contour_command_list[i] = re.sub('style\s*.*?\s', 'style '+contour_style+' ', contour_command_list[i])
245  min_contour = min(min_contour, float(re.search('--draw-contour\s(.*?)\s', contour_command_list[i]).group(1)))
246  contour_commands.append(contour_command_list)
247  min_contours.append(min_contour)
248 
249  # If this script is not present for any of the regions, go on to the next one.
250  if (regional_script is None): continue
251 
252  # Extract the axis styles from the last regional script
253  axis_styles = re.findall('\s\s--axis-style\s*[^y]*\n', regional_script)
254 
255  # Read in the original script
256  base_script = file(script_dir+'/'+script_name, 'r').readlines()
257 
258  # Extract preamble from original script
259  preamble = ''
260  for x in base_script[2:]:
261  preamble = preamble+x
262  if x.startswith(' --xyz-map'): break
263 
264  # Extract textual annotations from the original script
265  text = [re.sub('White','Black',x) for x in base_script if x.startswith(' --draw-text') or x.startswith(' --legend')]
266 
267  outfile = open(combo_dir+'/'+script_name,'w')
268  outfile.write('#!/usr/bin/env bash\n')
269  outfile.write('# This plot script created by the pippi scripter \'colouring\' on '+datetime.datetime.now().strftime('%c')+'\n')
270 
271  # Write the number of fill steps to be used
272  for region in region_names: outfile.write('n_contours_'+re.sub('-','_',region)+'='+str(default_n_contours)+'\n')
273 
274  # Write the fist half of the ctioga2 command
275  outfile.write(preamble)
276 
277  # Draw the filled regions, using many closely-packed contours
278  for region, plot_command in enumerate(plot_commands):
279  outfile.write(' --color \''+fill_colours[region]+'\' \\\n')
280  outfile.write(plot_command)
281  outfile.write(' $(echo $(for (( i=0; i<$n_contours_'+re.sub('-','_',region_names[region])+'; i++ )); do echo "--draw-contour '
282  '$(echo "'+str(min_contours[region])+' + '+str(1.0-min_contours[region])+'*$i/($n_contours_'
283  +re.sub('-','_',region_names[region])+'-1)" | bc -l) /style Solid /width '+str(default_contour_width)+' "; done)) \\\n')
284  outfile.write(' --draw-contour 1.000 /style Solid /width '+str(default_contour_width)+'\\\n')
285 
286  # Draw the actual outline contours on top of all the filled regions
287  for region, plot_command in enumerate(plot_commands):
288  outfile.write(plot_command)
289  for x in contour_commands[region]: outfile.write(x)
290 
291  # Draw any markers
292  for marker_command in marker_commands:
293  outfile.write(marker_command)
294 
295  # Write any textual annotations, and set the axes up in the manner of the region plots
296  for x in text: outfile.write(x)
297  for x in axis_styles: outfile.write(x)
298 
299  # Add an entry to the bulk plotting script
300  pdf_name = re.sub("\.bsh", ".pdf", script_name)
301  bulkscript.write('(echo "./'+script_name+'" && ./' + script_name +
302  ' && gs -sDEVICE=pdfwrite -dCompatibilityLevel=1.4 -dPDFSETTINGS=/screen -dNOPAUSE -dQUIET -dBATCH -sOutputFile='+
303  '../' + plot_combo_dir + '/' + pdf_name + ' ' + pdf_name +
304  ' && rm ' + pdf_name + ') &\n')
305  if (script_index%n_threads == n_threads - 1 ): bulkscript.write('wait\n')
306 
307  # Done! Close the new script file and make it executable for both user and group.
308  outfile.close
309  os.chmod(combo_dir+'/'+script_name, stat.S_IRWXU | stat.S_IRWXG)
310 
311  # All over red rover. Go write a paper.
312  bulkscript.write('wait\n')
313  bulkscript.write('echo "Generated all combined plots. Now go write a paper!"\n')
314  bulkscript.close
315  os.chmod(bulkscript_name, stat.S_IRWXU | stat.S_IRWXG)
316 
317 #Actual program launched on invocation
318 main(sys.argv)
std::string str
Shorthand for a standard string.
Definition: Analysis.hpp:35
def combine(region_file, pip_file)
Definition: colouring.py:169
def main(arguments)
Definition: colouring.py:62
def check_node(node, name, keys)
Definition: colouring.py:44
def usage()
Definition: colouring.py:58
def prime(region_file, pip_file)
Definition: colouring.py:86