Commit 4b198b13 authored by Dennis Noll's avatar Dennis Noll
Browse files

[utils] coffea: tree exporter, changes struct

parent 86eb2f82
......@@ -692,18 +692,14 @@ class TreeExporter(ArrayExporter):
if var not in cls.tensors.fget(cls).keys():
continue
var_array = _var_array.value
parts = [-1] if len(var_array.shape) == 2 else range(var_array.shape[1])
for i in parts:
if i == -1:
name = ""
vals = var_array
else:
name = f"{var}{i}_"
vals = var_array[:, i]
for j in range(vals.shape[-1]):
var_name = f"{name}{vars[var][2][j]}"
v = vals[:, j]
tree[var_name].append(v)
if vars[var][1] == 0:
name = f"{var}_" if "part" in vars[var][-1].get("groups", {}) else ""
for i, feat in enumerate(vars[var][2]):
tree[f"{name}{feat}"].append(var_array[:, i])
else:
for i in range(vars[var][1]):
for j, feat in enumerate(vars[var][2]):
tree[f"{var}{i}_{feat}"].append(var_array[:, i, j])
for cat_flag in arrays.keys():
tree[f"is_{cat_flag}"].append(np.ones(var_array.shape[0]) * (cat == cat_flag))
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