# Warning!. This script need Python 2.4 because it uses generator expresions.

from math import sqrt
from tables import *
from numarray import dot

codes = {
    "electron": 0,
    "photon": 1,
    "neutrino": 2
    }

class Particle(IsDescription):
    fourMomentum = FloatCol( dflt=0., shape=4 )
    idCode       = IntCol( dflt=0, shape=1 )
    productionVtx= FloatCol( dflt=0., shape=3 )
    decayVtx     = FloatCol( dflt=0., shape=3 )

fileh = openFile("particles.h5", mode = "w")
table = fileh.createTable(fileh.root, 'particles', Particle,
                          "A table for keeping particles", Filters(1))
# Fill the table with synthetic data
particle = table.row
for i in xrange(100):
    particle['fourMomentum'] = range(i,i+4)
    particle['idCode'] = i % 3
    particle.append()
table.flush()

# Get the row IDs for the electron particles
electronsRID = table.getWhereList(table.cols.idCode == codes["electron"])
print "Row ID for electrons-->", electronsRID

# Get the sum of energy for all neutrinos
idc = table.cols.idCode
nenergy = sum(sqrt(dot(n['fourMomentum'], n['fourMomentum']))
                   for n in table.where(idc == codes["neutrino"]))
print "Energy for all neutrinos-->", nenergy

fileh.close()