# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# <br><br><br><br><br>
#
# # uproot: array-based ROOT I/O
#
# <br><br><br><br><br>

# +
# Uproot presents files and subdirectories with the same interface as Python dicts.

import uproot
file = uproot.open("http://scikit-hep.org/uproot/examples/nesteddirs.root")
file.keys()

# Exercise: how many objects named "tree" does this file contain?
# 
# (https://uproot.readthedocs.io/en/latest/root-io.html#uproot-rootio-rootdirectory)

# +
# TTrees are also presented like Python dicts.

events = uproot.open("data/Zmumu.root")["events"]
events.keys()

# +
# And there are functions like array (singular) and arrays (plural) to get one or more arrays.

print(events.array("E1"), end="\n\n\n\n")

print(events.arrays(["E1", "px1", "py1", "pz1"]))

# +
# Some tricks for arrays (plural): (1) specify an encoding to get strings, rather than raw bytes.
# 
# Note: Python 3 bytestrings start with `b`, strings do not.

events.arrays(["E1", "px1", "py1", "pz1"], namedecode="utf-8")

# +
# Some tricks for arrays (plural): (2) specify an output type to wrap them in a different container:

events.arrays(["E1", "px1", "py1", "pz1"], outputtype=tuple)

# +
# Some tricks for arrays (plural): (2) specify an output type to wrap them in a different container:
# (see also tree.pandas.df)

import pandas
events.arrays(["E1", "px1", "py1", "pz1"], outputtype=pandas.DataFrame)

# +
# Some tricks for arrays (plural): (3a) use filename wildcards (*, ?, [...]) to select by name-pattern.

events.arrays(["E*", "p[xyz]1"])

# +
# Some tricks for arrays (plural): (3b) use slashes inside the quotes to use full regular expressions.

events.arrays(["/E.*/", "/p[x-z]1/"])
# -

# <br>
#
# <p style="font-weight: bold; font-size: 1.875em; color: gray">Three ways to get data:</p>
#
# <table width="100%" style="font-size: 1.25em"><tr>
#     <td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Direct</p>
#         <p>Read the file and return an array.</p>
#         <ul>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#id11">TBranch.array</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#array">TTree.array</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#arrays">TTree.arrays</a></li>
#         </ul>
#     </td><td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Lazy</p>
#         <p>Get an object that reads on demand.</p>
#         <ul>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#id13">TBranch.lazyarray</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#lazyarray">TTree.lazyarray</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#lazyarrays">TTree.lazyarrays</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/opening-files.html#uproot-lazyarray-and-lazyarrays">uproot.lazyarray</a>*</li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/opening-files.html#uproot-lazyarray-and-lazyarrays">uproot.lazyarrays</a>*</li>
#         </ul>
#     </td><td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Iterative</p>
#         <p>Read arrays in batches of entries.</p>
#         <ul>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/ttree-handling.html#iterate">TTree.iterate</a></li>
#             <li style="margin-bottom: 0.3em"><a href="https://uproot.readthedocs.io/en/latest/opening-files.html#uproot-iterate">uproot.iterate</a>*</li>
#         </ul>
#     </td>
# </tr></table>
#
# <p>*Lazy arrays or iteration over sets of files.</p>

# +
# Direct:

events.array("E1")

# +
# Lazy:

array = events.lazyarray("E1", entrysteps=500)
# print([len(x) for x in array.chunks])

array

# +
# Iterative:

for chunk in events.iterate("E1", entrysteps=500):
    print(len(chunk[b"E1"]), chunk[b"E1"][:5])
# -

# <br><br>
#
# <p style="font-weight: bold; font-size: 1.875em; color: gray">Advantages and disadvantages of each:</p>
#
# <table width="100%" style="font-size: 1.25em"><tr>
#     <td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Direct</p>
#         <p>Simple; returns pure Numpy arrays if possible.</p>
#     </td><td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Lazy</p>
#         <p>Transparently work on data too large to fit into memory.</p>
#     </td><td width="33%" style="vertical-align: top">
#         <p style="font-weight: bold; font-size: 1.5em; margin-bottom: 0.5em">Iterative</p>
#         <p>Control the loading of data into and out of memory.</p>
#     </td>
# </tr></table>

# +
# Exercise: compute numpy.sqrt(E1**2 - px1**2 - py1**2 - pz1**2) in all three modes.
# 
# (https://docs.scipy.org/doc/numpy/reference/generated/numpy.concatenate.html)

import numpy

arrays = events.arrays(["E1", "p[xyz]1"], namedecode="utf-8")
E1, px1, py1, pz1 = arrays["E1"], arrays["px1"], arrays["py1"], arrays["pz1"]

# This is direct. What about lazy and iterative?
result = numpy.sqrt(E1**2 - px1**2 - py1**2 - pz1**2)
result

# +
# Controlling the chunk size:

print("Lazy or iteration steps as a fixed number of entries:")
for arrays in events.iterate(entrysteps=500):
    print(len(arrays[b"E1"]))

print("\nLazy or iteration steps as a fixed memory footprint:")
for arrays in events.iterate(entrysteps="100 kB"):
    print(len(arrays[b"E1"]))
# -

# <br><br><br><br><br>
#
# ## Application: feeding data into PyTorch
#
# <br><br><br><br><br>

# +
# Let's train a neural network! First, though, we need to set up a problem.

# 2-dimensional for simplicity (easier to visualize and we don't have much data).
E1, E2, px1, py1, pz1, px2, py2, pz2, Q1, Q2 = \
                                events.arrays(["E[12]", "p[xyz][12]", "Q[12]"], outputtype=tuple)
E = (E1 + E2)
p = numpy.sqrt((px1 + px2)**2 + (py1 + py2)**2 + (pz1 + pz2)**2)

# Need to predict opposite sign (deepskyblue) vs same sign (orange) using E and p.
# %matplotlib inline
import matplotlib.pyplot
matplotlib.pyplot.scatter(E[Q1 != Q2], p[Q1 != Q2], c="deepskyblue", edgecolor="black");
matplotlib.pyplot.scatter(E[Q1 == Q2], p[Q1 == Q2], c="orange", edgecolor="black");
matplotlib.pyplot.xlim(0, 200);
matplotlib.pyplot.ylim(0, 200);

# +
import torch                             # transform inputs to fit PyTorch's expected shape and type
X = torch.from_numpy(numpy.dstack([E, p])[0].astype(numpy.float32))
y = torch.from_numpy((Q1 != Q2).astype(numpy.float32).reshape(-1, 1))

neural_network = torch.nn.Sequential(    # the neural network topology:
    torch.nn.Linear(2, 5),               # input → hidden: 2 dimensions → 5 dimensions
    torch.nn.Sigmoid(),                  # non-linearity applied to each of the 5 components
    torch.nn.Linear(5, 1))               # hidden → output: 5 dimensions → 1 dimension

loss_fn = torch.nn.MSELoss(reduction="sum")
optimizer = torch.optim.Adam(neural_network.parameters(), lr=0.001)

for i in range(1000):                    # iterate 1000 times to minimize loss: y_pred - y
    y_pred = neural_network(X)           # neural_network is a function: X ↦ y
    loss = loss_fn(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
# -

grid_of_points = numpy.dstack(numpy.mgrid[0:220:20, 0:220:20].astype(numpy.float32))
Z = neural_network(torch.from_numpy(grid_of_points)).detach().numpy().reshape(grid_of_points.shape[:2])
matplotlib.pyplot.contourf(grid_of_points[:, :, 0], grid_of_points[:, :, 1], Z);
matplotlib.pyplot.scatter(E[Q1 != Q2], p[Q1 != Q2], c="deepskyblue", edgecolor="black");
matplotlib.pyplot.scatter(E[Q1 == Q2], p[Q1 == Q2], c="orange", edgecolor="black");
matplotlib.pyplot.xlim(0, 200);
matplotlib.pyplot.ylim(0, 200);

# +
# Now let's pretend the sample is so large, we can't load it into memory.
# What has changed? What's the same?

optimizer = torch.optim.Adam(neural_network.parameters(), lr=0.0001)  # learn slower: fewer data points...

for E1, E2, px1, py1, pz1, px2, py2, pz2, Q1, Q2 in \
        events.iterate(["E[12]", "p[xyz][12]", "Q[12]"], outputtype=tuple, entrysteps=500):
    E = (E1 + E2)
    p = numpy.sqrt((px1 + px2)**2 + (py1 + py2)**2 + (pz1 + pz2)**2)
    
    X = torch.from_numpy(numpy.dstack([E, p])[0].astype(numpy.float32))
    y = torch.from_numpy((Q1 != Q2).astype(numpy.float32).reshape(-1, 1))

    for i in range(1000):
        y_pred = neural_network(X)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# +
# Now do it from a sample of jagged arrays (variable number of values per event).

events2 = uproot.open("data/HZZ.root")["events"]
E_all, px_all, py_all, pz_all, q_all = \
                         events2.arrays(["Muon_E", "Muon_P[xyz]", "Muon_Charge"], outputtype=tuple)

E_all, px_all, py_all, pz_all, q_all

# +
# Getting two muons from events that have them:

print(E_all.counts)

E_all[E_all.counts >= 2, 0], E_all[E_all.counts >= 2, 1]

# +
# Or all pairs, without double-counting:

left, right = E_all.argchoose(2).i0, E_all.argchoose(2).i1

E_all[left].flatten(), E_all[right].flatten()

# +
# Exercise: set up and run the neural network with muons from the jagged sample.

X = FIXME
y = FIXME

neural_network = torch.nn.Sequential(torch.nn.Linear(2, 5), torch.nn.Sigmoid(), torch.nn.Linear(5, 1))
loss_fn = torch.nn.MSELoss(reduction="sum")
optimizer = torch.optim.Adam(neural_network.parameters(), lr=0.001)
for i in range(1000):
    y_pred = neural_network(X)
    loss = loss_fn(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
# -

grid_of_points = numpy.dstack(numpy.mgrid[0:520:20, 0:520:20].astype(numpy.float32))
Z = neural_network(torch.from_numpy(grid_of_points)).detach().numpy().reshape(grid_of_points.shape[:2])
matplotlib.pyplot.contourf(grid_of_points[:, :, 0], grid_of_points[:, :, 1], Z);
matplotlib.pyplot.scatter(E[Q1 != Q2], p[Q1 != Q2], c="deepskyblue", edgecolor="black");
matplotlib.pyplot.scatter(E[Q1 == Q2], p[Q1 == Q2], c="orange", edgecolor="black");
matplotlib.pyplot.xlim(0, 500);
matplotlib.pyplot.ylim(0, 500);

# <br><br><br><br><br>
#
# ## Caching
#
# <br><br><br><br><br>

# +
# Direct array-reading is simple: every time you ask for an array, it reads from the file (raw bytes may be cached).
# 
# You could avoid duplicate reading/decompressing/formatting by keeping a reference to previously read arrays.

# read
arrays = events.arrays(["E1", "p[xyz]1"], namedecode="utf-8")

# use
numpy.sqrt(arrays["px1"]**2 + arrays["py1"]**2)

# use again
numpy.sqrt(arrays["E1"]**2 - arrays["px1"]**2 - arrays["py1"]**2 - arrays["pz1"]**2)

# +
# But that would force you to re-arrange your analysis script to satisfy hardware constraints.
# 
# Instead, use uproot's caching mechanism: any dict-like object can be used as a cache.

mycache = {}

# read and use
E1, px1, py1, pz1 = events.arrays(["E1", "p[xyz]1"], outputtype=tuple, cache=mycache)
numpy.sqrt(px1**2 + py1**2)

# get from cache and use again
E1, px1, py1, pz1 = events.arrays(["E1", "p[xyz]1"], outputtype=tuple, cache=mycache)
numpy.sqrt(E1**2 - px1**2 - py1**2 - pz1**2)

# +
# The data are now in the mycache dict, which you can clear whenever you need to.

# mycache.clear()
mycache
# -

# <br><br>
#
# <p style="font-size: 1.25em">The point is that naive, exploratory code can become production-ready code <i>in small steps</i>.</p>
#
# <p style="font-size: 1.25em; font-weight: bold; color: gray">Development sequence:</p>
#
#    1. Read directly from the file in early exploration because it requires the least effort.
#    2. Insert `mycache = {}` and `cache=mycache` to avoid costly re-reading when you start looking at more data.
#    3. Use an `uproot.ArrayCache` to specify an upper limit when you start running out of memory.
#    4. Maybe use an external cache like [memcached](https://realpython.com/python-memcache-efficient-caching/) to quickly recover from crashing 90% through your script, to more easily debug that last 10%.
#    5. Maybe use a [diskcache](http://www.grantjenks.com/docs/diskcache/tutorial.html) to split data between a small, fast disk and a large, slow disk...
#
# Most cache libraries in the Python ecosystem use this dict of string → objects interface.
#
# <br><br>

# +
# Example of limited memory:

mycache = uproot.ArrayCache("100 kB")     # integer number of bytes or string with units

events.arrays(cache=mycache)

# arrays in cache    arrays in file
len(mycache),        len(events.keys())

# +
# With large lazy arrays, you SHOULD use a cache with an upper limit.
# (Otherwise, they'll hang onto data forever.)

data = uproot.lazyarray(
    # a bunch or files
    "data/sample-*-zlib.root",
    # TTree name in each file
    "sample",
    # branch(s) in each file for lazyarray(s)
    "Af8",
    # use this cache with an upper limit instead of reading exactly once
    cache=uproot.ArrayCache("5 kB"))

data

# +
# Compute something to do a pass over all chunks...
# data + 100

[chunk.ismaterialized for chunk in data.chunks]

# +
# With iteration, you SHOULD NOT use a cache because the iteration manages memory explicitly.
# You wouldn't want to keep the previous iteration in memory when you're done with it, right?

for chunk in uproot.iterate("data/sample-*-zlib.root", "sample", "Af8", entrysteps="0.5 kB"):
    print(len(chunk[b"Af8"]), chunk[b"Af8"][:5])
# -

# <br><br><br><br><br>
#
# ## Objects in arrays
#
# <br><br><br><br><br>

# +
# Most of the complexity of ROOT data is simplified by the fact that custom classes are "split" into branches.
# uproot sees the branches.

tree = uproot.open("http://scikit-hep.org/uproot/examples/Event.root")["T"]
tree.show()

# branch name              streamer type, if any      uproot's interpretation

# +
# In this view, class attributes are _not_ special types and can be read as arrays of numbers.

tree.array("fTemperature", entrystop=20)

# +
# Fixed-width matrices map onto multidimensional arrays,

tree.array("fMatrix[4][4]", entrystop=6)

# +
# branches with multiple leaves ("leaf-list") map onto Numpy record arrays,

uproot.open("http://scikit-hep.org/uproot/examples/leaflist.root")["tree"].array("leaflist")

# +
# and anything in variable-length lists become a JaggedArrays,

tree.array("fTracks.fMass2", entrystop=6)

# +
# even if they're fixed-width within jagged or whatever.

tree.array("fTracks.fTArray[3]", entrystop=6)

# +
# There are some types that ROOT does not split because they are too complex.
# For example, *histograms* inside a TTree:

tree.array("fH", entrystop=6)

# +
# Uproot can read objects like this because ROOT describes their layout in "streamers."
# Uproot reads the (most common types of) streamers and generates Python classes.
# Some of these Python classes have specialized, high-level methods.

for histogram in tree.array("fH", entrystop=3):
    print(histogram.title)
    print(histogram.values)
print("\n...\n")
for histogram in tree.array("fH", entrystart=-3):
    print(histogram.title)
    print(histogram.values)

# +
# One type, TLorentzVector, has MANY high-level methods.

events3 = uproot.open("data/HZZ-objects.root")["events"]
muons = events3.array("muonp4")
muons

# +
# A single TLorentzVector has the methods you'd expect (kinematics, delta_phi, delta_r),

muon1, muon2 = muons[0]

muon1, muon1.pt, muon1.delta_phi(muon2), (muon1 + muon2).mass

# +
# but they also apply vectorially (like Numpy functions):

muons1 = muons[muons.counts >= 2, 0]   # first muon in each event (that has two)
muons2 = muons[muons.counts >= 2, 1]   # second muon in each event (that has two)

muons1, muons1.pt, muons1.delta_phi(muons2), (muons1 + muons2).mass

# +
# Doubly nested lists like std::vector<std::vector<T>> aren't handled transparently
# because they're stored differently.
# 
# Note: it's an ObjectArray, an array containing list objects.

doubly_nested = uproot.open(
                   "http://scikit-hep.org/uproot/examples/vectorVectorDouble.root")["t"]["x"].array()
doubly_nested

# +
# The multidimensional slicing doesn't apply as it does for JaggedArrays.

try:
    doubly_nested[doubly_nested.counts > 0, 0]
except Exception as err:
    print(type(err), err)

# +
# If you need to work with this type, explicitly convert it into a JaggedArray of JaggedArrays.

import awkward
jagged = awkward.fromiter(doubly_nested)
jagged

# +
# Now you can do all the things.

print(f"jagged[jagged.counts > 0, 0]: {jagged[jagged.counts > 0, 0]}")
print(f"jagged.flatten():             {jagged.flatten()}")
print(f"jagged.flatten().flatten():   {jagged.flatten().flatten()}")
print(f"jagged.sum():                 {jagged.sum()}")
print(f"jagged.sum().sum():           {jagged.sum().sum()}")
# -

# <br><br><br><br><br>
#
# ## Reading and writing histograms
#
# <br><br><br><br><br>

# +
# As we've seen, histograms have some convenience methods.
# They're mostly for conversion to other formats, like Numpy.
# 
# Numpy "histograms" are a 2-tuple of counts and edges.

uproot.open("http://scikit-hep.org/uproot/examples/hepdata-example.root")["hpx"].numpy()

# +
# Similarly for 2-dimensional histograms.

uproot.open("http://scikit-hep.org/uproot/examples/hepdata-example.root")["hpxpy"].numpy()

# +
# It can also be useful to turn histograms into Pandas DataFrames (note the IntervalIndex).

uproot.open("http://scikit-hep.org/uproot/examples/Event.root")["htime"].pandas()

# +
# Or HEPData's YAML format. As Python objects, it's just a little work to make different formats.

print(uproot.open("http://scikit-hep.org/uproot/examples/Event.root")["htime"].hepdata())

# +
# At the moment, only two kinds of objects can be *written* to ROOT files: TObjString and histograms.
# 
# To write, open a file for writing (create/recreate/update) and assign to it like a dict:

file = uproot.recreate("tmp.root", compression=uproot.ZLIB(4))
file["name"] = "Some object, like a TObjString."

# +
import ROOT

pyroot_file = ROOT.TFile("tmp.root")
pyroot_file.Get("name")

# +
# During assignment, uproot recognizes some Pythonic types, such as Numpy histograms.

file["from_numpy"] = numpy.histogram(numpy.random.normal(0, 1, 10000))

# +
pyroot_file = ROOT.TFile("tmp.root")           # refresh the PyROOT file
pyroot_hist = pyroot_file.Get("from_numpy")

canvas = ROOT.TCanvas("canvas", "", 400, 300)
pyroot_hist.Draw("hist")
canvas.Draw()

# +
# 2-dimensional Numpy histograms.

file["from_numpy2d"] = numpy.histogram2d(numpy.random.normal(0, 1, 10000), numpy.random.normal(0, 1, 10000))

# +
pyroot_file = ROOT.TFile("tmp.root")           # refresh the PyROOT file
pyroot_hist = pyroot_file.Get("from_numpy2d")

pyroot_hist.Draw()
canvas.Draw()
# -

# <br><br><br><br><br>
#
# ## Writing array data
#
# <br><br><br><br><br>

# +
# Obviously, it would be nice to write array data to ROOT files as new TTrees.
# This is an IRIS-HEP summer fellowship project (<NAME>).
# 
# Meanwhile, there are many nice options for saving and sharing Numpy arrays.

# Numpy's NPZ files (really just ZIP)
numpy.savez_compressed("one_array.npz", events.array("E1"))
numpy.savez_compressed("many_arrays.npz", **events.arrays(namedecode="utf-8"))

# HDF5
import h5py
with h5py.File("tmp.hdf5", "w") as file:
    file["E1"] = events.array("E1")

# +
# If the arrays are jagged, lazy, or contain objects (i.e. not pure Numpy), use awkward-array.

# awkward's AWKD files (really just ZIP)
awkward.save("one_array.awkd", events2.array("Muon_Px"), mode="w")
awkward.save("many_arrays.awkd", events2.arrays("Muon_*", namedecode="utf-8"), mode="w")

# HDF5
with h5py.File("tmp.hdf5", "w") as file:
    awkward_file = awkward.hdf5(file)
    awkward_file["Muon_Px"] = events2.array("Muon_Px")

# +
# There are also industry-standard formats for jagged arrays.

awkward.toparquet("one_array.parquet", events2.array("Muon_Px"))
awkward.toparquet("many_arrays.parquet", events2.lazyarrays())

# +
# You can also do some crazy things like "saving the laziness" of lazy arrays to make lightweight skims:

# Read as lazy arrays with "persist virtual",
data = events.lazyarrays(["E*", "p[xyz]*"], persistvirtual=True)

# add a derived feature to the lazy array,
data["mass"] = numpy.sqrt((data["E1"] + data["E2"])**2 - (data["px1"] + data["px2"])**2 -
                          (data["py1"] + data["py2"])**2 - (data["pz1"] + data["pz2"])**2)

# and save it.
awkward.save("derived-feature.awkd", data, mode="w")

# +
# When you read it back, the original features from from the ROOT files and the derived from the AWKD:

data2 = awkward.load("derived-feature.awkd")

# Reads from derived-feature.awkd
print(data2["mass"])

# Reads from the original ROOT files (if you don't move them!)
print(data2["E1"])

# +
# Similarly, if you apply masks (cuts) to a lazy array, the cut is saved to AWKD but the original data are in ROOT:

selected = data[data["mass"] < 80]
print(selected)

awkward.save("selected-events.awkd", selected, mode="w")

print(awkward.load("selected-events.awkd"))
# -

# <br><br><br><br><br>
#
# <p style="font-size: 2em; text-align: center"><b>That's all!</b></p>
#
# <br><br><br><br><br>

StarcoderdataJupyterScriptsDedupFiltered