Commit 3c012dbc authored by Dennis Noll's avatar Dennis Noll
Browse files

[utils] coffea: TreeExporter no lookup (now happens later)

parent 68847e94
......@@ -650,12 +650,6 @@ class TreeExporter(ArrayExporter):
output = "sync.root"
dtype = None
@classmethod
def look_up(cls, name):
for was, gets in cls.lut.items():
name = name.replace(was, gets)
return name
@classmethod
def arrays_to_tree(cls, arrays, target, vars, **kwargs):
import uproot
......@@ -680,11 +674,10 @@ class TreeExporter(ArrayExporter):
v = vals[:, j]
tree[var_name].append(v)
for cat_flag in arrays.keys():
tree[cat_flag].append(np.ones(var_array.shape[0]) * (cat == cat_flag))
tree[f"is_{cat_flag}"].append(np.ones(var_array.shape[0]) * (cat == cat_flag))
outtree = {}
for v, k in tree.items():
var_name = cls.look_up(v)
outtree[var_name] = np.concatenate(k, axis=-1)
outtree[v] = np.concatenate(k, axis=-1)
file["tree"] = uproot.newtree({n: v.dtype for n, v in outtree.items()})
file["tree"].extend(outtree)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment