Belle II Software  release-05-02-19
constraints_generator.py
1 '''
2 Created on 28 May 2018
3 
4 @author: kleinwrt
5 '''
6 import basf2 as b2
7 
8 import os
9 import sys
10 import pickle
11 
12 # Needed for pickling of constraints
13 from alignment.constraints import * # noqa
14 
15 
16 class ConstraintsGenerator(b2.Module):
17  """
18  basf2 python module to generate a file with a single set of constraints
19 
20  (Needed especially for time-dependent constraints)
21  """
22 
23  def __init__(self, generator):
24  """
25  Constructor
26 
27  Parameters
28  ----------
29  generator : alignment.Constraints
30  Set of constraints to be generated (defines also filename)
31 
32  """
33 
34  super(ConstraintsGenerator, self).__init__()
35 
36  self.generator = generator
37 
38  self.consts = dict()
39  super(ConstraintsGenerator, self).set_name(super(ConstraintsGenerator, self).name() + ":" + generator.filename)
40 
41  def event(self):
42  """ Event: attempt to add constraints
43  At each event, attempt to generate the constraints and add them
44  to the dict (same cheksum - using only labels - means it gets overriden!)
45  Thus is only works if the time-dependence is set (otherwise you always get the same labels and
46  only the last constraint coefficients will be stored)
47  """
48  consts = self.generator.generate()
49  for const in consts:
50  self.consts[const.get_checksum()] = const.data
51 
52  def terminate(self):
53  """ Terminate: write-down collect constraint data to file
54  """
55  if len(self.consts):
56  with open(self.generator.filename, 'w') as f:
57  for checksum, data in self.consts.items():
58  f.write('Constraint 0.\n')
59  for label, coef in data:
60  f.write('{} {}\n'.format(label, coef))
61  print("Finished file: ", self.generator.filename)
62  print("Done: ", self.generator.filename)
63 
64 
65 def save_config(constraint_sets, timedep_config=None, global_tags=None, init_event=None):
66  """
67  Save constraints configuration to a file (using pickle)
68 
69  Parameters
70  ----------
71  constraint_sets : list (alignment.Constraints)
72  List of sets of constraints
73  timedep_config : list(tuple(list(int), list(tuple(int, int, int))))
74  Time-depence configuration.
75  Each list item is 2-tuple with list of parameter numbers (use alignment.parameters to get them) and
76  the (event, run, exp) numbers at which values of these parameters can change.
77  global_tags : list (str)
78  List of global tag names and/or (absolute) file paths to local databases
79  init_event : tuple( int, int, int)
80  Event (event, run, exp) at which to initialize time-INdependent constraints
81 
82  """
83  a = {'constraint_sets': constraint_sets, 'timedep_config': timedep_config, 'global_tags': global_tags, 'init_event': init_event}
84 
85  file_Name = "constraint_config.pickled"
86  # open the file for writing
87  with open(file_Name, 'wb') as fileObject:
88 
89  # this writes the object a to the
90  # file named 'testfile'
91  pickle.dump(a, fileObject)
92  print('Saving config to: ', file_Name)
93 
94  return file_Name
95 
96 
97 def read_config(filename='constraint_config.pickled'):
98  """ Read the pickled constraint configuration from a file
99 
100  filename : str
101  File with pickled constraints' configuration
102  """
103  print('Reading config from: ', filename)
104  with open(filename, 'rb') as fileObject:
105  # load the object from the file into var b
106  b = pickle.load(fileObject)
107  return (b['constraint_sets'], b['timedep_config'], b['global_tags'], b['init_event'])
108 
109 
110 def gen_constraints(constraint_sets, timedep_config=None, global_tags=None, init_event=None):
111  """
112  Generate "event files" from timedep config and run over them to collect constraint data
113  and write it to files.
114 
115  This is a bit tricky. I did not found a way to run basf2 over just a specified list of events
116  other than generating the files with metadata for event and then running over the files.
117  This uses unning basf2 multiple times in one script - seems to work despite the warnings - but only
118  for the generation of "event files"
119 
120  Parameters
121  ----------
122  constraint_sets : list (alignment.Constraints)
123  List of sets of constraints
124  timedep_config : list(tuple(list(int), list(tuple(int, int, int))))
125  Time-depence configuration.
126  Each list item is 2-tuple with list of parameter numbers (use alignment.parameters to get them) and
127  the (event, run, exp) numbers at which values of these parameters can change.
128  global_tags : list (str)
129  List of global tag names and/or (absolute) file paths to local databases
130  init_event : tuple( int, int, int)
131  Event (event, run, exp) at which to initialize time-INdependent constraints
132 
133  """
134  if timedep_config is None:
135  timedep_config = []
136 
137  if global_tags is None:
138  global_tags = [tag for tag in b2.conditions.default_globaltags]
139 
140  events = []
141  for (labels, events_) in timedep_config:
142  events += [event for event in events_]
143 
144  if not len(timedep_config):
145  if init_event is None:
146  init_event = (0, 0, 0)
147  events = [init_event]
148 
149  events = [(exp, run, ev) for (ev, run, exp) in events]
150  events = sorted(list(set(events)))
151  events = [(ev_, run_, exp_) for (exp_, run_, ev_) in events]
152 
153  fileName = 'TimedepConfigEvent_exp{}run{}ev{}.root'
154  files = []
155 
156  print('Global tags:')
157  print(global_tags)
158  print('Global tags reversed (this will be used for b2.conditions.override_globaltags(...)):')
159  print([tag for tag in reversed(global_tags)])
160 
161  for tag in [tag for tag in reversed(global_tags)]:
162  if os.path.exists(tag):
163  b2.conditions.append_testing_payloads(os.path.abspath(tag))
164  else:
165  b2.conditions.append_globaltag(tag)
166 
167  for index, event in enumerate(events):
168  # conditions.reset()
169 
170  ev, run, exp = event
171  path = b2.create_path()
172  path.add_module("EventInfoSetter",
173  skipNEvents=ev,
174  evtNumList=[ev + 1],
175  runList=[run],
176  expList=[exp])
177  path.add_module('Progress')
178  this_filename = fileName.format(exp, run, ev)
179  path.add_module('RootOutput', outputFileName=this_filename, ignoreCommandLineOverride=True)
180  files.append(this_filename)
181  b2.process(path)
182  print(b2.statistics)
183 
184  print(files)
185 
186  # conditions.override_globaltags(global_tags)
187 
188  path = b2.create_path()
189  path.add_module("RootInput", inputFileNames=files, ignoreCommandLineOverride=True)
190  path.add_module('HistoManager')
191  path.add_module('Progress')
192  path.add_module('Gearbox')
193  path.add_module('Geometry')
194 
195  collector = path.add_module('MillepedeCollector',
196  timedepConfig=timedep_config)
197 
198  constraint_files = []
199  for constraint_set in constraint_sets:
200  constraint_set.configure_collector(collector)
201  constraint_files.append(constraint_set.filename)
202  path.add_module(ConstraintsGenerator(constraint_set))
203 
204  b2.process(path)
205  print(b2.statistics)
206 
207  return [os.path.abspath(file) for file in constraint_files]
208 
209 
210 def gen_constraints_from_config(filename='constraint_config.pickled'):
211  """
212  Generate constraints from pickled configuration
213  """
214  constraint_sets, timedep_config, global_tags, init_event = read_config(filename)
215  gen_constraints(constraint_sets, timedep_config, global_tags, init_event)
216 
217 
218 if __name__ == '__main__':
219  if len(sys.argv) != 2:
220  print('Usage: basf2 constraints_generator.py config_filename')
221  exit()
222 
223 
224  filename = sys.argv[1]
alignment.constraints_generator.ConstraintsGenerator.generator
generator
The generator: alignment.Constraints.
Definition: constraints_generator.py:36
alignment.constraints_generator.read_config
def read_config(filename='constraint_config.pickled')
Definition: constraints_generator.py:97
alignment.constraints_generator.ConstraintsGenerator.terminate
def terminate(self)
Definition: constraints_generator.py:52
alignment.constraints_generator.gen_constraints_from_config
def gen_constraints_from_config(filename='constraint_config.pickled')
Definition: constraints_generator.py:210
alignment.constraints_generator.ConstraintsGenerator.consts
consts
Generated constraints (key is constraint checksum)
Definition: constraints_generator.py:38
alignment.constraints
Definition: constraints.py:1
alignment.constraints_generator.ConstraintsGenerator
Definition: constraints_generator.py:16
alignment.constraints_generator.ConstraintsGenerator.event
def event(self)
Definition: constraints_generator.py:41
alignment.constraints_generator.ConstraintsGenerator.__init__
def __init__(self, generator)
Definition: constraints_generator.py:23
alignment.constraints_generator.gen_constraints
def gen_constraints(constraint_sets, timedep_config=None, global_tags=None, init_event=None)
Definition: constraints_generator.py:110
alignment.constraints_generator.save_config
def save_config(constraint_sets, timedep_config=None, global_tags=None, init_event=None)
Definition: constraints_generator.py:65