Skip to content

Commit

Permalink
Merge pull request #3 from EnginEren/features/byPeter
Browse files Browse the repository at this point in the history
removed number of steps condition to allow read in of arbitrary numbe…
  • Loading branch information
p-mckeown authored Aug 8, 2019
2 parents 4a9bd04 + d127965 commit 7876d8c
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions generation/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import numpy as np
from h5py import File as HDF5File

LAYER_SPECS = [(3, 96), (12, 12), (12, 6)]
LAYER_SPECS = [(12, 12), (12, 12), (12, 12)]

LAYER_DIV = np.cumsum(map(np.prod, LAYER_SPECS)).tolist()
LAYER_DIV = zip([0] + LAYER_DIV, LAYER_DIV)
Expand All @@ -34,7 +34,7 @@ def write_out_file(infile, outfile, tree=None):

assert len(cells) == sum(map(np.prod, LAYER_SPECS)) + OVERFLOW_BINS

for df in uproot.pandas.iterate(infile, "fancy_tree;1", branches = cells, entrysteps = len(cells)):
for df in uproot.pandas.iterate(infile, "fancy_tree;1", branches = cells):
X = df

for df in uproot.pandas.iterate(infile, "fancy_tree;1", branches = b'TotalEnergy'):
Expand All @@ -44,11 +44,12 @@ def write_out_file(infile, outfile, tree=None):
E = E.values.ravel()
#X = tree.pd.DataFrame(tree2array(T, branches=cells)).values
#E = tree.pd.DataFrame(tree2array(T, branches=['TotalEnergy'])).values.ravel()

print(X.shape)

with HDF5File(outfile, 'w') as h5:
for layer, (sh, (l, u)) in enumerate(zip(LAYER_SPECS, LAYER_DIV)):
h5['layer_{}'.format(layer)] = X[:, l:u].reshape((-1, ) + sh)
print(u)

h5['overflow'] = X[:, -OVERFLOW_BINS:]
h5['energy'] = E.reshape(-1, 1)
Expand Down

0 comments on commit 7876d8c

Please sign in to comment.