Merge pull request #71 from FAST-HEP/BK_issue_67-interp_jagged_arrays

Fix interpretation of user-defined JaggedArrays

CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Removed
 
+## [0.13.4] - 2019-09-21
+### Changed
+- Fixed interpretation of user-defined variables for uproot, issue #67, PR #71 [@benkrikler](https://github.com/benkrikler)
+
 ## [0.13.3] - 2019-09-12
 ### Changed
 - Add changes to support uproot 3.9.1 and greater, issue #68 [@benkrikler](https://github.com/benkrikler)

fast_carpenter/tree_wrapper.py

@@ -7,6 +7,7 @@
 """
 import uproot
 from uproot.interp.jagged import asjagged
+from uproot.interp.numerical import asdtype
 import copy
 import awkward
 
@@ -19,7 +20,7 @@ def wrapped_interpret(branch, *args, **kwargs):
 
     if isinstance(branch, WrappedTree.FakeBranch):
         if isinstance(branch._values, awkward.JaggedArray):
-            return asjagged(branch._values.content)
+            return asjagged(asdtype(branch._values.content.dtype.fields))
         else:
             return branch._values
 

fast_carpenter/version.py

@@ -12,5 +12,5 @@ def split_version(version):
     return tuple(result)
 
 
-__version__ = '0.13.3'
+__version__ = '0.13.4'
 version_info = split_version(__version__) # noqa

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.13.3
+current_version = 0.13.4
 commit = True
 tag = False
 

tests/summary/test_binned_dataframe.py

@@ -1,4 +1,5 @@
 import pandas as pd
+import copy
 import numpy as np
 import pytest
 import fast_carpenter.summary.binned_dataframe as bdf
@@ -150,6 +151,26 @@ def test_BinnedDataframe_numexpr(binned_df_3, tmpdir):
     assert len(binned_df_3._weights) == 1
 
 
+def test_BinnedDataframe_user_var_run(config_3, tmpdir, full_wrapped_tree):
+    config_4 = copy.deepcopy(config_3)
+    config_4["binning"][0]["in"] = "Electron_Pt"
+    binned_df_4 = bdf.BinnedDataframe("binned_df_4", out_dir="somewhere", **config_4)
+
+    chunk = FakeBEEvent(full_wrapped_tree, "mc")
+    px, py = chunk.tree.arrays(["Electron_Px", "Electron_Py"], outputtype=tuple)
+    pt = np.hypot(px, py)
+    chunk.tree.new_variable("Electron_Pt", pt)
+    collector = binned_df_4.collector()
+
+    binned_df_4.event(chunk)
+    dataset_readers_list = (("test_dataset", (binned_df_4,)),)
+    results = collector._prepare_output(dataset_readers_list)
+
+    bin_centers = pd.IntervalIndex(results.index.get_level_values('electron_pT')).mid
+    mean = np.sum((bin_centers[1:-1] * results['n'][1:-1]) / results['n'][1:-1].sum())
+    assert mean == pytest.approx(44.32584)
+
+
 def test_BinnedDataframe_numexpr_run_mc(binned_df_3, tmpdir, infile):
     chunk = FakeBEEvent(infile, "mc")
     collector = binned_df_3.collector()