test: add convert script and update benchmark to use a standardized t…

…est file

bin/convert

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+"""
+Convert ROOT file to Parquet.
+
+Usage:
+  convert <root_file_path> [--output <output_path>] [--tree=<tree_name>] [-n=<n_events>] [--eps]
+  convert -h | --help
+
+Options:
+  -h --help              Show this help message and exit.
+  --output <output_path> Specify the output path for the Parquet file.
+  --tree=<tree_name>     Specify the name of the TTree in the ROOT file.
+  -n <n_events>          Only convert the first <n_events>.
+  --eps                  Generate an EPS branch randomly.
+"""
+
+import pandas as pd
+import uproot
+import numpy as np
+from docopt import docopt
+
+
+def convert(root_file_path, output_parquet_path=None, tree_name=None, n_events=None, eps=False):
+    with uproot.open(root_file_path) as root_file:
+        tree_name = root_file.keys()[0] if tree_name is None else tree_name
+
+        tree = root_file[tree_name]
+        numpy_arrays = tree.arrays(library='np')
+        dataframe = pd.DataFrame(numpy_arrays)
+        n_tot = len(dataframe)
+        if n_events is not None and n_events <= n_tot:
+            dataframe = dataframe.head(n_events)
+        else:
+            n_events = n_tot
+
+        if eps:
+            rng = np.random.default_rng()
+            p_gamma = rng.uniform(0.0, 0.4, size=n_events)
+            BigPhi = rng.choice([0.0, 45 * np.pi / 180, 90 * np.pi / 180, 135 * np.pi / 180], size=n_events)
+            eps_x = p_gamma * np.cos(BigPhi)
+            eps_y = p_gamma * np.sin(BigPhi)
+            dataframe['EPS'] = [np.array([ex, ey, 0.0], dtype=np.float32) for ex, ey in zip(eps_x, eps_y)]
+        else:
+            dataframe['EPS'] = [
+                np.array([ex, ey, 0.0], dtype=np.float32) for ex, ey in zip(dataframe['Px_Beam'], dataframe['Py_Beam'])
+            ]
+        dataframe['Px_Beam'] = np.zeros_like(dataframe['Px_Beam'])
+        dataframe['Py_Beam'] = np.zeros_like(dataframe['Py_Beam'])
+
+    if output_parquet_path is None:
+        output_parquet_path = f"{root_file_path.replace('.root', '.parquet')}"
+
+    dataframe.to_parquet(output_parquet_path)
+    print(f'Converted {n_events}/{n_tot} events from ROOT to Parquet')
+
+
+if __name__ == '__main__':
+    arguments = docopt(__doc__)
+    root_file_path = arguments['<root_file_path>']
+    output_parquet_path = arguments['--output']
+    tree_name = arguments['--tree']
+    n_events = int(arguments['-n'])
+    eps = arguments['--eps']
+
+    convert(root_file_path, output_parquet_path, tree_name, n_events, eps)

crates/rustitude/benches/kmatrix.rs

@@ -5,7 +5,7 @@ use rustitude::gluex::utils::{Frame, Reflectivity, Wave};
 use rustitude::prelude::*;
 
 pub fn criterion_kmatrix(c: &mut Criterion) {
-    let dataset = Dataset::from_parquet("benches/data_pol.parquet");
+    let dataset = Dataset::from_parquet("benches/test_data.parquet");
     let f0p: AmpOp = amplitude!("f0+", KMatrixF0::new(2));
     let f0n: AmpOp = amplitude!("f0-", KMatrixF0::new(2));
     let f2: AmpOp = amplitude!("f2", KMatrixF2::new(2));