scikit-hep / uproot3

ROOT I/O in pure Python and NumPy.
BSD 3-Clause "New" or "Revised" License
314 stars 67 forks source link

Import from pandas or numpy #545

Closed hayg25 closed 2 years ago

hayg25 commented 2 years ago

Dear All, is there a way to import a Dataframe or a numpy array in uproot, and create a TTree for example ?

Thx

jpivarski commented 2 years ago

Yes, but please use Uproot 4 for any new work:

% python
Python 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:59) 
[GCC 10.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pandas as pd
>>> df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [1.1, 2.2, 3.3, 4.4, 5.5]})
>>> df
   x    y
0  1  1.1
1  2  2.2
2  3  3.3
3  4  4.4
4  5  5.5
>>> import uproot   # this is Uproot 4, installed with "pip install uproot", not "uproot3"
>>> with uproot.recreate("new-file.root") as file:
...     file["tree"] = df
... 
>>> 
% python
Python 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:59) 
[GCC 10.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import ROOT
>>> file = ROOT.TFile("new-file.root")
>>> tree = file.Get("tree")
>>> tree.Scan()
************************************************
*    Row   * index.ind *       x.x *       y.y *
************************************************
*        0 *         0 *         1 *       1.1 *
*        1 *         1 *         2 *       2.2 *
*        2 *         2 *         3 *       3.3 *
*        3 *         3 *         4 *       4.4 *
*        4 *         4 *         5 *       5.5 *
************************************************
5

To write NumPy arrays into a TTree, make them a dict of NumPy arrays (e.g. just the {"x": np.array([1, 2, 3, 4, 5]), "y": np.array([1.1, 2.2, 3.3, 4.4, 5.5])} without the pd.DataFrame). The documentation is here.