class RootFile(DataFile):
extension = ".root"
def __init__(self, filename, rw_if_possible=False):
self.filename = filename
self.num_frames = 0
self.file = TFile(filename, "READ")
if not self.file.IsOpen():
raise IOError(-1, "Cannot open ROOT file for reading", filename)
if rw_if_possible:
self.file.Close()
self.file = TFile(filename, "UPDATE")
if not self.file.IsWritable():
raise IOError(-1, "Cannot open ROOT file in RW mode", filename)
self.tree = self.file.Get("data")
self.ch = [TGraph(), TGraph(), TGraph(), TGraph()]
for i in xrange(1, 5):
self.tree.SetBranchAddress("ch{0}".format(i), self.ch[i-1])
self.num_frames = self.tree.GetEntries()
self.next_frame_idx = 0
def __del__(self):
if self.file:
self.file.Close()
def __enter__(self):
return self
def __exit__(self, type, value, tb):
if self.file:
self.file.Close()
def next_frame(self, next_idx=None):
if next_idx is not None:
self.next_frame_idx = next_idx
frame_idx = self.next_frame_idx
self.next_frame_idx += 1
if frame_idx >= self.num_frames:
return None
t = None
x = [None] * 4
self.tree.GetEntry(frame_idx)
for ch_num in xrange(4):
n_points = self.ch[ch_num].GetN()
if not n_points:
continue
if t is None:
t = np.frombuffer(self.ch[ch_num].GetX(),
dtype=float,
count=n_points)
x[ch_num] = np.frombuffer(self.ch[ch_num].GetY(),
dtype=float,
count=n_points)
return frame_idx, t, x
def getRootFile(self):
return self.file
