vidyasagarv.com
Python bindings to oxyroot. Makes reading .root files blazing fast 🚀
First attempt at a reader with few supported types
+167
-22
+4
-2
Cargo.toml
+4
-2
Cargo.toml
+3
-1
pyproject.toml
+3
-1
pyproject.toml
+7
python/oxyroot/__init__.py
+7
python/oxyroot/__init__.py
-6
python/py_oxyroot/__init__.py
-6
python/py_oxyroot/__init__.py
+35
python/tests/read_ntuples.py
+35
python/tests/read_ntuples.py
···
···
1
+
import oxyroot
2
+
3
+
if __name__ == '__main__':
4
+
import uproot
5
+
import numpy as np
6
+
import time
7
+
8
+
file_name = "ntuples.root"
9
+
tree_name = 'mu_mc'
10
+
11
+
up_start_time = time.time()
12
+
up_tree = uproot.open(file_name)[tree_name]
13
+
for branch in up_tree.keys():
14
+
# print(branch, up_tree[branch].typename)
15
+
if up_tree[branch].typename != "std::string":
16
+
up_values = up_tree[branch].array(library="np")
17
+
print(f"Uproot read {branch} into a {type(up_values)} and it has a mean of {np.nanmean(up_values):.2f}")
18
+
up_end_time = time.time()
19
+
20
+
print("\n")
21
+
22
+
oxy_start_time = time.time()
23
+
oxy_branches = oxyroot.read_root(file_name, tree_name=tree_name)
24
+
for branch in oxy_branches:
25
+
oxyroot.read_root(file_name, tree_name=tree_name, branch=branch)
26
+
oxy_values = oxyroot.read_root(file_name, tree_name=tree_name, branch=branch)
27
+
if type(oxy_values) is np.ndarray:
28
+
print(f"Oxyroot read {branch} into a {type(oxy_values)} and it has a mean of {np.nanmean(oxy_values):.2f}")
29
+
else:
30
+
print(f"Oxyroot read {branch} into a {type(oxy_values)} and it has a length of {len(oxy_values)}")
31
+
oxy_end_time = time.time()
32
+
33
+
print("\n Total time")
34
+
print(f"Uproot took: {up_end_time - up_start_time:.3}s")
35
+
print(f"Oxyroot took: {oxy_end_time - oxy_start_time:.3}s")
-6
python/tests/test_all.py
-6
python/tests/test_all.py
+22
python/tests/test_read_from_uproot.py
+22
python/tests/test_read_from_uproot.py
···
···
1
+
import pytest
2
+
import oxyroot
3
+
import uproot
4
+
import numpy as np
5
+
import os
6
+
7
+
def test_read_from_uproot():
8
+
# Create a dummy ROOT file for testing
9
+
10
+
input = np.array([4.1, 5.2, 6.3])
11
+
file_name = "test.root"
12
+
13
+
with uproot.recreate(file_name) as f:
14
+
f.mktree("tree1", {"branch1": np.float64})
15
+
f["tree1"].extend({"branch1": input})
16
+
17
+
18
+
output = oxyroot.read_root(file_name, tree_name="tree1", branch="branch1")
19
+
assert(type(output) is np.ndarray)
20
+
assert(np.array_equal(input, output))
21
+
22
+
os.remove(file_name)
+96
-7
src/lib.rs
+96
-7
src/lib.rs
···
1
-
use pyo3::prelude::*;
2
3
-
/// Formats the sum of two numbers as string.
4
#[pyfunction]
5
-
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
6
-
Ok((a + b).to_string())
7
}
8
9
-
/// A Python module implemented in Rust.
10
#[pymodule]
11
-
fn py_oxyroot(m: &Bound<'_, PyModule>) -> PyResult<()> {
12
-
m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
13
Ok(())
14
}
···
1
+
use ::oxyroot::{Named, RootFile};
2
+
use numpy::ToPyArray;
3
+
use pyo3::{exceptions::PyValueError, prelude::*, types::PyModule, IntoPyObjectExt};
4
+
5
+
#[pyfunction]
6
+
fn version() -> PyResult<String> {
7
+
Ok(env!("CARGO_PKG_VERSION").to_string())
8
+
}
9
10
+
/// Read a ROOT file and return the list of trees, the branches of a tree, or the values of a branch.
11
#[pyfunction]
12
+
#[pyo3(signature = (path, tree_name = None, branch = None))]
13
+
fn read_root(
14
+
path: String,
15
+
tree_name: Option<String>,
16
+
branch: Option<String>,
17
+
) -> PyResult<Py<PyAny>> {
18
+
let mut file = RootFile::open(&path).unwrap();
19
+
let keys: Vec<String> = file
20
+
.keys()
21
+
.into_iter()
22
+
.map(|k| k.name().to_string())
23
+
.collect();
24
+
25
+
Python::attach(|py| -> PyResult<Py<PyAny>> {
26
+
match tree_name {
27
+
Some(name) => {
28
+
if let Ok(tree) = file.get_tree(&name) {
29
+
let branches_available: Vec<String> =
30
+
tree.branches().map(|b| b.name().to_string()).collect();
31
+
32
+
match branch {
33
+
Some(bs) => {
34
+
if let Some(branch) = tree.branch(&bs) {
35
+
match branch.item_type_name().as_str() {
36
+
"f32" => {
37
+
let data =
38
+
branch.as_iter::<f32>().unwrap().collect::<Vec<_>>();
39
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
40
+
}
41
+
"double" => {
42
+
let data =
43
+
branch.as_iter::<f64>().unwrap().collect::<Vec<_>>();
44
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
45
+
}
46
+
"int32_t" => {
47
+
let data =
48
+
branch.as_iter::<i32>().unwrap().collect::<Vec<_>>();
49
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
50
+
}
51
+
"int64_t" => {
52
+
let data =
53
+
branch.as_iter::<i64>().unwrap().collect::<Vec<_>>();
54
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
55
+
}
56
+
"uint32_t" => {
57
+
let data =
58
+
branch.as_iter::<u32>().unwrap().collect::<Vec<_>>();
59
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
60
+
}
61
+
"uint64_t" => {
62
+
let data =
63
+
branch.as_iter::<u64>().unwrap().collect::<Vec<_>>();
64
+
Ok(data.to_pyarray(py).into_py_any(py).unwrap())
65
+
}
66
+
"string" => {
67
+
let data =
68
+
branch.as_iter::<String>().unwrap().collect::<Vec<_>>();
69
+
Ok(data.into_py_any(py).unwrap())
70
+
}
71
+
other => Err(PyValueError::new_err(format!(
72
+
"Unsupported branch type: {}",
73
+
other
74
+
))),
75
+
}
76
+
} else {
77
+
Err(PyValueError::new_err(format!(
78
+
"Branch '{}' not found. Available branches are: {:?}",
79
+
bs, branches_available
80
+
)))
81
+
}
82
+
}
83
+
None => Ok(branches_available.into_py_any(py).unwrap()),
84
+
}
85
+
} else {
86
+
Err(PyValueError::new_err(format!(
87
+
"Tree '{}' not found. Available trees are: {:?}",
88
+
name, keys
89
+
)))
90
+
}
91
+
}
92
+
None => Ok(keys.into_py_any(py).unwrap()),
93
+
}
94
+
})
95
}
96
97
+
/// A Python module to read root files implemented in Rust.
98
#[pymodule]
99
+
fn oxyroot(m: &Bound<'_, PyModule>) -> PyResult<()> {
100
+
m.add_function(wrap_pyfunction!(version, m)?)?;
101
+
m.add_function(wrap_pyfunction!(read_root, m)?)?;
102
Ok(())
103
}