-
Notifications
You must be signed in to change notification settings - Fork 1.5k
[tutorials][ML] Add RDataLoader tutorial for XGBoost #22847
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| ## \file | ||
| ## \ingroup tutorial_ml | ||
| ## \notebook -nodraw | ||
| ## This tutorial demonstrates XGBoost training using RDataLoader to load data | ||
| ## directly from ROOT files without intermediate preparation steps. | ||
| ## This is an alternative implementation to tutorial tmva101_Training.py, | ||
| ## which stores intermediate files after train/test splitting | ||
| ## and uses numpy arrays to load data. | ||
|
Comment on lines
+4
to
+8
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I would not mention |
||
| ## | ||
| ## \macro_code | ||
| ## \macro_output | ||
| ## | ||
| ## \date July 2026 | ||
| ## \author Silia Taider | ||
|
|
||
| import ROOT | ||
|
|
||
| variables = ["Muon_pt_1", "Muon_pt_2", "Electron_pt_1", "Electron_pt_2"] | ||
|
|
||
|
|
||
| def filter_events(df): | ||
| """Reduce initial dataset to only events which shall be used for training""" | ||
| return df.Filter("nElectron>=2 && nMuon>=2", "At least two electrons and two muons") | ||
|
|
||
|
|
||
| def define_variables(df): | ||
| """Define the variables which shall be used for training""" | ||
| return ( | ||
| df | ||
| .Define("Muon_pt_1", "Muon_pt[0]") | ||
| .Define("Muon_pt_2", "Muon_pt[1]") | ||
| .Define("Electron_pt_1", "Electron_pt[0]") | ||
| .Define("Electron_pt_2", "Electron_pt[1]") | ||
| ) | ||
|
|
||
|
|
||
| def prepare_rdf(filename, label_value): | ||
| """Load, filter, define variables, and add label column""" | ||
| filepath = "root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/" + filename | ||
| df = ROOT.RDataFrame("Events", filepath) | ||
| df = filter_events(df) | ||
| df = define_variables(df) | ||
| df = df.Define("label", f"{label_value}.0") | ||
| return df | ||
|
|
||
|
|
||
| def load_data(): | ||
| """Load signal and background data""" | ||
| rdf_sig = prepare_rdf("SMHiggsToZZTo4L.root", 1) | ||
| rdf_bkg = prepare_rdf("ZZTo2e2mu.root", 0) | ||
|
|
||
| # Compute class-balancing weights | ||
| num_sig = rdf_sig.Count().GetValue() | ||
| num_bkg = rdf_bkg.Count().GetValue() | ||
| num_all = num_sig + num_bkg | ||
|
|
||
| rdf_sig = rdf_sig.Define("weight", f"{num_all}.0/{num_sig}.0") | ||
| rdf_bkg = rdf_bkg.Define("weight", f"{num_all}.0/{num_bkg}.0") | ||
|
|
||
| loader = ROOT.Experimental.ML.RDataLoader( | ||
| [rdf_sig, rdf_bkg], | ||
| columns=variables + ["label", "weight"], | ||
| target="label", | ||
| weights="weight", | ||
| batch_size=num_all, # Load all data in one batch | ||
| drop_remainder=False, | ||
| set_seed=42, | ||
| ) | ||
|
|
||
| # Split into training and testing sets | ||
| train, test = loader.train_test_split(test_size=0.5) | ||
|
|
||
| X_train, y_train, w_train = next(iter(train.as_numpy())) | ||
| X_test, y_test, w_test = next(iter(test.as_numpy())) | ||
|
|
||
| # Flatten target and weights from (n,1) to (n,) as expected by XGBoost | ||
| return (X_train, y_train.ravel(), w_train.ravel(), X_test, y_test.ravel(), w_test.ravel()) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| from sklearn.metrics import roc_auc_score | ||
| from xgboost import XGBClassifier | ||
|
|
||
| X_train, y_train, w_train, X_test, y_test, w_test = load_data() | ||
|
|
||
| print(f"Training events: {X_train.shape[0]}") | ||
| print(f"Testing events: {X_test.shape[0]}") | ||
|
|
||
| bdt = XGBClassifier(max_depth=3, n_estimators=500) | ||
| bdt.fit(X_train, y_train, sample_weight=w_train) | ||
|
|
||
| # Evaluate on test set | ||
| y_proba = bdt.predict_proba(X_test)[:, 1] | ||
| auc = roc_auc_score(y_test, y_proba) | ||
|
|
||
| print(f"Training done. ROC AUC: {auc:.4f}") | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The pattern you show generalized very well to other classification algos, so I would make that clear before people ask you to re-write the same tutorial for catboost or sckit-learn algorithms