xboa
_maus_json_hit_factory.py
Go to the documentation of this file.
1 
2 # This file is a part of xboa
3 #
4 # xboa is free software: you can redistribute it and/or modify
5 # it under the terms of the GNU General Public License as published by
6 # the Free Software Foundation, either version 3 of the License, or
7 # (at your option) any later version.
8 #
9 # xboa is distributed in the hope that it will be useful,
10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 # GNU General Public License for more details.
13 #
14 # You should have received a copy of the GNU General Public License
15 # along with xboa in the doc folder. If not, see
16 # <http://www.gnu.org/licenses/>.
17 
18 import json
19 import xboa.common as common
20 import xboa.hit
21 from xboa.hit.factory import HitFactoryBase
22 
23 class MausJsonHitFactory(HitFactoryBase):
24  """
25  MausJsonHitFactory strips hits of a specified type from a json_document and
26  returns to the user
27  """
28  def __init__(self, file_handle, format):
29  """
30  Initialise the hit factory
31  - file_handle: a file handle containing a set of json documents, one per
32  line
33  - format: format of data to be extracted from the json document
34  - spill_number: spill parameter that will be assigned to the hits
35  """
36  if format in self.file_types():
37  format = format[10:] # string maus_json_
38  else:
39  if format in self.deprecated_file_types():
40  print "Warning - using deprecated file type", format
41  format = format[5:] # strip maus_
42  else:
43  raise KeyError("Did not recognise format "+str(format))
44  self.fin = file_handle
45  self.hits = []
46  self.format = format
47  self.spill = None
48 
49 
50  def new_spill(self):
51  """
52  Read a new spill off the file handle
53  """
54  # if we have no more hits, look for the next Spill event and attempt to
55  # load MC
56  found_spill = False
57  if len(self.hits) == 0:
58  while not found_spill:
59  line = self.fin.readline()
60  if line == "":
61  raise xboa.hit.BadEventError("Reached the end of the file")
62  new_line = json.loads(line)
63  if "maus_event_type" in new_line and \
64  new_line["maus_event_type"] == "Spill":
65  found_spill = True
66  self.spill = new_line
67  try:
68  self._read_maus_spill()
69  except KeyError:
70  pass # badly formed spill or no data
71 
72  def make_hit(self):
73  """
74  Return the next hit from the spill file_handle
75  """
76  # try to produce a new hit
77  try:
78  while len(self.hits) == 0:
79  self.new_spill()
80  hit = self.hits.pop(0)
81  return hit
82  except IndexError:
83  raise xboa.hit.BadEventError("Run out of events")
84 
85  def _read_maus_spill(self):
86  """
87  Read a MAUS spill, converting to MAUS objects
88  """
89  if self.format == 'virtual_hit':
90  self._read_virtual_hits()
91  elif self.format == 'primary':
92  self._read_primaries()
93 
94  def _read_virtual_hits(self):
95  """
96  Read a virtual hits from the json_document
97  """
98  for ev, mc_event in enumerate(self.spill["mc_events"]):
99  for virtual_hit in mc_event["virtual_hits"]:
100  hit = self.hit_from_maus_object("virtual_hit", virtual_hit, ev)
101  hit["spill"] = self.spill["spill_number"]
102  self.hits.append(hit)
104  def _read_primaries(self):
105  """
106  Read primaries from the json_document
107  """
108  for ev, mc_event in enumerate(self.spill["mc_events"]):
109  hit = self.hit_from_maus_object("primary", mc_event["primary"], ev)
110  hit["spill"] = self.spill["spill_number"]
111  self.hits.append(hit)
112 
113  @classmethod
114  def hit_from_maus_object(cls, format, maus_dict, event_number):
115  """
116  Convert a MAUS object into a Hit according to specified formats
117  """
118  xboa_dict = {}
119  three_vec_conversions = cls._maus_three_vec_conversions[format]
120  conversion_dict = cls._maus_variable_conversions[format]
121  for maus_name, xboa_suffix in three_vec_conversions.iteritems():
122  for maus_xyz, value in maus_dict[maus_name].iteritems():
123  xboa_dict[xboa_suffix+maus_xyz] = value
124  for maus_key, xboa_key in conversion_dict.iteritems():
125  xboa_dict[xboa_key] = maus_dict[maus_key]
126  xboa_dict['event_number'] = event_number
127  if 'mass' not in xboa_dict.keys():
128  try:
129  xboa_dict['mass'] = common.pdg_pid_to_mass[abs(xboa_dict['pid'])]
130  except KeyError:
131  xboa_dict['mass'] = 0.
132  cls.bad_pid(xboa_dict['pid'])
133  if 'charge' not in xboa_dict.keys():
134  try:
135  xboa_dict['charge'] = common.pdg_pid_to_charge[xboa_dict['pid']]
136  except KeyError:
137  xboa_dict['charge'] = 0.
138  cls.bad_pid(xboa_dict['pid'])
139  hit = xboa.hit.Hit.new_from_dict(xboa_dict, cls._file_mass_shell[format])
140  return hit
141 
142 
143  @classmethod
144  def file_types(cls):
145  """Return a list of file types that the factory can read"""
146  return ["maus_json_virtual_hit", "maus_json_primary"]
147 
148  @classmethod
149  def deprecated_file_types(cls):
150  """Return a list of deprecated file types (no longer used)"""
151  return ["maus_virtual_hit", "maus_primary"]
152 
153  _maus_three_vec_conversions = { # maus three vectors are sub-dicts of virtual_hit
154  "virtual_hit":{"position":"", "momentum":"p", "b_field":"b", "e_field":"e"},
155  "primary":{"position":"", "momentum":"p"}
156  }
157 
158  _maus_variable_conversions = {
159  "virtual_hit":{"station_id":"station", "particle_id":"pid", "track_id":"particle_number", "time":"t", "mass":"mass", "charge":"charge", "proper_time":"proper_time", "path_length":"path_length"},
160  "primary":{"particle_id":"pid", "time":"t", "energy":"energy"} # we also force "mass" from "pid"
161  }
162 
163  _file_mass_shell = {
164  "virtual_hit":"energy",
165  "primary":"p"
166  }
167 
MausJsonHitFactory strips hits of a specified type from a json_document and returns to the user...
The hit factory module defines a number of factory classes used for generating hit objects...
Definition: __init__.py:1
common module defines common utility data and functions that are used elsewhere
Definition: __init__.py:1
def _read_virtual_hits
Read a virtual hits from the json_document.
def _read_primaries
Read primaries from the json_document.
Implemented within this module:
Definition: __init__.py:1
def _read_maus_spill
Read a MAUS spill, converting to MAUS objects.
def file_types
Return a list of file types that the factory can read.
def deprecated_file_types
Return a list of deprecated file types (no longer used)
def make_hit
Return the next hit from the spill file_handle.
def hit_from_maus_object
Convert a MAUS object into a Hit according to specified formats.