diff --git a/bindings/pyroot/pythonizations/test/CMakeLists.txt b/bindings/pyroot/pythonizations/test/CMakeLists.txt index e40a2175e84c4..d08cc2c6bac6d 100644 --- a/bindings/pyroot/pythonizations/test/CMakeLists.txt +++ b/bindings/pyroot/pythonizations/test/CMakeLists.txt @@ -128,7 +128,13 @@ ROOT_ADD_PYUNITTEST(pyroot_memory memory.py) # ML dataloader tests if (dataframe) - ROOT_ADD_PYUNITTEST(ml_dataloader ml_dataloader.py PYTHON_DEPS numpy tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader_single_dataframe ml_dataloader_single_dataframe.py GENERIC PYTHON_DEPS jax numpy pytest tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader_multiple_dataframes ml_dataloader_multiple_dataframes.py GENERIC PYTHON_DEPS numpy pytest tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader_undersampling ml_dataloader_undersampling.py GENERIC PYTHON_DEPS numpy pytest tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader_oversampling ml_dataloader_oversampling.py GENERIC PYTHON_DEPS numpy pytest tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader_save ml_dataloader_save.py GENERIC PYTHON_DEPS numpy pytest) + ROOT_ADD_PYUNITTEST(ml_dataloader_implicit_mt ml_dataloader_implicit_mt.py GENERIC PYTHON_DEPS numpy pytest) + ROOT_ADD_PYUNITTEST(ml_dataloader_from_files ml_dataloader_from_files.py GENERIC PYTHON_DEPS numpy pytest) endif() ROOT_ADD_PYUNITTEST(regression_18441 regression_18441.py) diff --git a/bindings/pyroot/pythonizations/test/conftest.py b/bindings/pyroot/pythonizations/test/conftest.py index f5b0289f4d122..576c025a138cd 100644 --- a/bindings/pyroot/pythonizations/test/conftest.py +++ b/bindings/pyroot/pythonizations/test/conftest.py @@ -1,3 +1,5 @@ +from itertools import count + import pytest import ROOT from ROOT._pythonization._th1 import _th1_derived_classes_to_pythonize @@ -5,6 +7,7 @@ from ROOT._pythonization._th3 import _th3_derived_classes_to_pythonize from ROOT._pythonization._uhi.indexing import _temporarily_disable_add_directory +# UHI tests th1_classes = { ("test_hist", "Test Histogram", 10, 1, 4): { "classes": [getattr(ROOT, klass) for klass in _th1_derived_classes_to_pythonize], @@ -49,3 +52,72 @@ def _iterate_bins(hist, flow=False): range(0 if flow else 1, hist.GetNbinsZ() + (2 if flow else 1)) if hist.GetDimension() > 2 else [0], ] yield from ((i, j, k)[: hist.GetDimension()] for i in ranges[0] for j in ranges[1] for k in ranges[2]) + + +# RDataLoader tests +TREE_NAME = "mytree" + +_file_ids = count() + + +@pytest.fixture(autouse=True) +def in_tmp_path(tmp_path, monkeypatch): + """Run each test in its own temporary directory managed by pytest.""" + monkeypatch.chdir(tmp_path) + + +@pytest.fixture(params=[False, True], ids=["lazy", "eager"]) +def load_eager(request): + """Run a test twice with the data loader's 2 loading modes.""" + return request.param + + +def from_files(*dfs): + """Write each RDataFrame to its own file and return an RDataFrame reading these files.""" + file_names = [f"input_{next(_file_ids)}.root" for _ in dfs] + for df, file_name in zip(dfs, file_names): + df.Snapshot(TREE_NAME, file_name) + return ROOT.RDataFrame(TREE_NAME, file_names if len(file_names) > 1 else file_names[0]) + + +def squares(num_of_entries, b1="(int) rdfentry_"): + """RDataFrame with the column b1, defined by the given expression, and b2 = b1 * b1.""" + return ROOT.RDataFrame(num_of_entries).Define("b1", b1).Define("b2", "(double) b1 * b1") + + +def weighted(num_of_entries, b1="(int) rdfentry_"): + """RDataFrame with the column b1, defined by the given expression, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return ( + ROOT + .RDataFrame(num_of_entries) + .Define("b1", b1) + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + .Define("b4", "(double) b3 * 10") + ) + + +def vectors(num_of_entries, b1="(int) rdfentry_"): + """RDataFrame with the column b1, defined by the given expression, and the vector columns v1 = {b1, b1 * 10} and + v2 = {b1 * 100, b1 * 1000}.""" + return ( + ROOT + .RDataFrame(num_of_entries) + .Define("b1", b1) + .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") + .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") + ) + + +def big_data(num_of_entries, b1="(int) rdfentry_"): + """RDataFrame with the column b1, defined by the given expression, and columns computed from the entry number: + b2 = entry * 2, b3 = entry + 10192, b4 = -entry and b5 = -entry - 10192.""" + return ( + ROOT + .RDataFrame(num_of_entries) + .Define("b1", b1) + .Define("b2", "(double) rdfentry_ * 2") + .Define("b3", "(int) rdfentry_ + 10192") + .Define("b4", "(int) -rdfentry_") + .Define("b5", "(double) -rdfentry_ - 10192") + ) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py deleted file mode 100644 index b19298bcb3990..0000000000000 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ /dev/null @@ -1,8341 +0,0 @@ -import math -import os -import unittest -from random import randrange, uniform - -import numpy as np -import ROOT - - -class DataLoaderMultipleFiles(unittest.TestCase): - file_name1 = "first_half.root" - file_name2 = "second_half.root" - file_name3 = "vector_columns.root" - tree_name = "mytree" - - # default constants - n_train_batch = 2 - n_val_batch = 1 - val_remainder = 1 - - # Helpers - def define_rdf(self, num_of_entries=10): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def create_file(self, num_of_entries=10): - self.define_rdf(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_5_entries_file(self): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 10") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name2) - ) - - def create_vector_file(self, num_of_entries=10): - ( - ROOT.RDataFrame(10) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name3) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - entries_before = df.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - entries_after = df.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframe is correctly reset - self.assertTrue(np.array_equal(entries_before, entries_after)) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=True, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test04_dropping_remainder(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x = [] - collected_y = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test05_more_than_one_file(self): - self.create_file() - self.create_5_entries_file() - - try: - df = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] - results_x_val = [9.0, 10.0, 11.0, 12.0, 13.0, 14.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0] - results_y_val = [81.0, 100.0, 121.0, 144.0, 169.0, 196.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test06_multiple_target_columns(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test07_multiple_input_columns(self): - file_name = "multiple_input_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 0.0, 1.0, 10.0, 2.0, 20.0, 3.0, 30.0, 4.0, 40.0, 5.0, 50.0] - results_x_val = [6.0, 60.0, 7.0, 70.0, 8.0, 80.0, 9.0, 90.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test08_filtered(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dff = df.Filter("b1 % 2 == 0", "name") - - filter_entries_before = dff.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0] - results_x_val = [6.0, 8.0] - results_y_train = [0.0, 4.0, 16.0] - results_y_val = [36.0, 64.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - filter_entries_after = dff.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframe is correctly reset - self.assertTrue(np.array_equal(filter_entries_before, filter_entries_after)) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test09_filtered_last_chunk(self): - file_name = "filtered_last_chunk.root" - tree_name = "myTree" - - ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name - ) - - try: - df = ROOT.RDataFrame(tree_name, file_name) - - dff = df.Filter("b1 % 2 == 0", "name") - - gen_train = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_y_train = [] - - train_iter = iter(gen_train.as_numpy()) - - for _ in range(3): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test10_two_epochs_shuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 0) - self.assertEqual(gen_validation.last_batch_no_of_rows, 1) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test12_PyTorch(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch(device="cpu")) - iter_val = iter(gen_validation.as_torch()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test13_TensorFlow(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0, 0.0, 0.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0, 0.0, 0.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test14_big_data(self): - file_name = "big_data.root" - tree_name = "myTree" - - entries_in_rdf = randrange(10000, 30000) - batch_size = randrange(100, 501) - batches_in_memory = randrange(5, 20) - - error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" - - def define_rdf(num_of_entries): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, batches_in_memory, num_of_entries): - define_rdf(num_of_entries) - - try: - df = ROOT.RDataFrame(tree_name, file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=size_of_batch, - batches_in_memory=batches_in_memory, - target=["b3", "b5"], - weights="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collect_x = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}") - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}") - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}") - - collect_x.extend(list(x[:, 0])) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message) - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message) - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message) - - collect_x.extend(list(x[:, 0])) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), - f"collected length: {len(set(collect_x))}\ - generated length {len(set(i for i in range(num_of_entries)))}", - ) - - except: - self.teardown_file(file_name) - raise - - test(batch_size, batches_in_memory, entries_in_rdf) - - def test15_two_runs_set_seed(self): - self.create_file() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - - def test16_vector_padding(self): - self.create_vector_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name3) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0, - 200.0, - 2000.0, - 3.0, - 30.0, - 0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0, - 400.0, - 4000.0, - 5.0, - 50.0, - 0, - 500.0, - 5000.0, - ] - results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [ - 6.0, - 60.0, - 0.0, - 600.0, - 6000.0, - 7.0, - 70.0, - 0.0, - 700.0, - 7000.0, - 8.0, - 80.0, - 0.0, - 800.0, - 8000.0, - 9.0, - 90.0, - 0.0, - 900.0, - 9000.0, - ] - results_y_val = [6.0, 7.0, 8.0, 9.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name3) - raise - - def test17_JAX(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_jax(device="cpu")) - iter_val = iter(gen_validation.as_jax()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - finally: - self.teardown_file(file_name) - - def test18_one_batch_in_memory(self): - # batches_in_memory=1 used to deadlock before delivering the first batch - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=2, - batches_in_memory=1, - target="b2", - shuffle=False, - drop_remainder=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - self.assertEqual(sum(1 for _ in gen_train.as_numpy()), 3) - self.assertEqual(sum(1 for _ in gen_validation.as_numpy()), 2) - - finally: - self.teardown_file(self.file_name1) - - def test19_defined_columns(self): - # columns defined after reading the file are loaded like the stored ones - self.create_file() - - try: - df = ( - ROOT - .RDataFrame(self.tree_name, self.file_name1) - .Define("b3", "b1 * 3") - .Define("v1", "ROOT::RVecD{b2, b2 + 1}") - ) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - columns=["b1", "v1", "b3"], - max_vec_sizes={"v1": 2}, - target="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - batches = list(gen_train.as_numpy()) + list(gen_validation.as_numpy()) - x = np.concatenate([x for x, _ in batches]) - y = np.concatenate([y for _, y in batches]) - - entries = np.arange(10, dtype=np.float32) - self.assertTrue(np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1]))) - self.assertTrue(np.array_equal(y, (entries * 3).reshape(-1, 1))) - - finally: - self.teardown_file(self.file_name1) - - -class DataLoaderEagerLoading(unittest.TestCase): - file_name1 = "first_half.root" - file_name2 = "second_half.root" - file_name3 = "vector_columns.root" - tree_name = "mytree" - - # default constants - n_train_batch = 2 - n_val_batch = 1 - val_remainder = 1 - - # Helpers - def define_rdf(self, num_of_entries=10): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def create_file(self, num_of_entries=10): - self.define_rdf(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_5_entries_file(self): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 10") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name2) - ) - - def create_vector_file(self, num_of_entries=10): - ( - ROOT.RDataFrame(10) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name3) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test04_dropping_remainder(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=True, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x = [] - collected_y = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test05_more_than_one_file(self): - self.create_file() - self.create_5_entries_file() - - try: - df = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] - results_x_val = [9.0, 10.0, 11.0, 12.0, 13.0, 14.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0] - results_y_val = [81.0, 100.0, 121.0, 144.0, 169.0, 196.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test06_multiple_target_columns(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test07_multiple_input_columns(self): - file_name = "multiple_input_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 0.0, 1.0, 10.0, 2.0, 20.0, 3.0, 30.0, 4.0, 40.0, 5.0, 50.0] - results_x_val = [6.0, 60.0, 7.0, 70.0, 8.0, 80.0, 9.0, 90.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test08_filtered(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dff = df.Filter("b1 % 2 == 0", "name") - - dl = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0] - results_x_val = [6.0, 8.0] - results_y_train = [0.0, 4.0, 16.0] - results_y_val = [36.0, 64.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test09_filtered_last_chunk(self): - file_name = "filtered_last_chunk.root" - tree_name = "myTree" - - ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name - ) - - try: - df = ROOT.RDataFrame(tree_name, file_name) - - dff = df.Filter("b1 % 2 == 0", "name") - - gen_train = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_y_train = [] - - train_iter = iter(gen_train.as_numpy()) - - for _ in range(3): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test10_two_epochs_shuffled(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 0) - self.assertEqual(gen_validation.last_batch_no_of_rows, 1) - - self.teardown_file(self.file_name1) - - except: - self.teardown_file(self.file_name1) - raise - - def test12_PyTorch(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test13_TensorFlow(self): - file_name = "multiple_target_columns.root" - - ROOT.RDataFrame(10).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Define( - "b3", "(double) rdfentry_ * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name) - - try: - df = ROOT.RDataFrame("myTree", file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0, 0.0, 0.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0, 0.0, 0.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise - - def test14_big_data(self): - file_name = "big_data.root" - tree_name = "myTree" - - entries_in_rdf = randrange(10000, 30000) - batch_size = randrange(100, 501) - batches_in_memory = randrange(5, 20) - - error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" - - def define_rdf(num_of_entries): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, batches_in_memory, num_of_entries): - define_rdf(num_of_entries) - - try: - df = ROOT.RDataFrame(tree_name, file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=size_of_batch, - batches_in_memory=batches_in_memory, - target=["b3", "b5"], - weights="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collect_x = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}") - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}") - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}") - - collect_x.extend(list(x[:, 0])) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message) - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message) - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message) - - collect_x.extend(list(x[:, 0])) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), - f"collected length: {len(set(collect_x))}\ - generated length {len(set(i for i in range(num_of_entries)))}", - ) - - except: - self.teardown_file(file_name) - raise - - test(batch_size, batches_in_memory, entries_in_rdf) - - def test15_two_runs_set_seed(self): - self.create_file() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - - def test16_vector_padding(self): - self.create_vector_file() - - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name3) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0, - 200.0, - 2000.0, - 3.0, - 30.0, - 0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0, - 400.0, - 4000.0, - 5.0, - 50.0, - 0, - 500.0, - 5000.0, - ] - results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [ - 6.0, - 60.0, - 0.0, - 600.0, - 6000.0, - 7.0, - 70.0, - 0.0, - 700.0, - 7000.0, - 8.0, - 80.0, - 0.0, - 800.0, - 8000.0, - 9.0, - 90.0, - 0.0, - 900.0, - 9000.0, - ] - results_y_val = [6.0, 7.0, 8.0, 9.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name3) - raise - - def test17_defined_columns(self): - # columns defined after reading the file are loaded like the stored ones - self.create_file() - - try: - df = ( - ROOT - .RDataFrame(self.tree_name, self.file_name1) - .Define("b3", "b1 * 3") - .Define("v1", "ROOT::RVecD{b2, b2 + 1}") - ) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=3, - batches_in_memory=2, - columns=["b1", "v1", "b3"], - max_vec_sizes={"v1": 2}, - target="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - batches = list(gen_train.as_numpy()) + list(gen_validation.as_numpy()) - x = np.concatenate([x for x, _ in batches]) - y = np.concatenate([y for _, y in batches]) - - entries = np.arange(10, dtype=np.float32) - self.assertTrue(np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1]))) - self.assertTrue(np.array_equal(y, (entries * 3).reshape(-1, 1))) - - finally: - self.teardown_file(self.file_name1) - - -class DataLoaderEagerLoadingMultipleDataframes(unittest.TestCase): - file_name1 = "first_half.root" - file_name2 = "second_half.root" - file_name3 = "second_file.root" - file_name4 = "vector_columns_1.root" - file_name5 = "vector_columns_2.root" - tree_name = "mytree" - - # default constants - n_train_batch = 2 - n_val_batch = 1 - val_remainder = 1 - - # Helpers - def define_rdf1(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def define_rdf2(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(double) b1*b1") - - return df - - def create_file1(self, num_of_entries=5): - self.define_rdf1(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_file2(self, num_of_entries=5): - self.define_rdf2(num_of_entries).Snapshot(self.tree_name, self.file_name2) - - def create_5_entries_file(self): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 10") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name3) - ) - - def create_vector_file1(self, num_of_entries=5): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name4) - ) - - def create_vector_file2(self, num_of_entries=5): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 5") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name5) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] - results_x_val = [3.0, 4.0, 8.0, 9.0] - results_y_train = [0.0, 1.0, 4.0, 25.0, 36.0, 49.0] - results_y_val = [9.0, 16.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(df1_entries_before, df1_entries_after)) - self.assertTrue(np.array_equal(df2_entries_before, df2_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test04_dropping_remainder(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=True, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x = [] - collected_y = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test05_more_than_one_file(self): - self.create_file1() - self.create_file2() - self.create_5_entries_file() - - try: - df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 11.0, 12.0] - results_x_val = [6.0, 7.0, 8.0, 9.0, 13.0, 14.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 100.0, 121.0, 144.0] - results_y_val = [36.0, 49.0, 64.0, 81.0, 169.0, 196.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise - - def test06_multiple_target_columns(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] - results_x_val = [3.0, 4.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] - results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] - results_z_val = [30.0, 40.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test07_multiple_input_columns(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name1) - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 0.0, 1.0, 10.0, 2.0, 20.0, 5.0, 50.0, 6.0, 60.0, 7.0, 70.0] - results_x_val = [3.0, 30.0, 4.0, 40.0, 8.0, 80.0, 9.0, 90.0] - results_y_train = [0.0, 1.0, 4.0, 25.0, 36.0, 49.0] - results_y_val = [9.0, 16.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test08_filtered(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dff1 = df1.Filter("b1 % 2 == 0", "name") - dff2 = df2.Filter("b1 % 2 != 0", "name") - - dff1_entries_before = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_before = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 5.0] - results_x_val = [4.0, 9.0] - results_y_train = [0.0, 4.0, 25.0] - results_y_val = [16.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(dff1_entries_before, dff1_entries_after)) - self.assertTrue(np.array_equal(dff2_entries_before, dff2_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test09_filtered_last_chunk(self): - file_name1 = "filtered_last_chunk_1.root" - file_name2 = "filtered_last_chunk_2.root" - tree_name = "myTree" - - ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name1 - ) - - ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_ + 10").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name2 - ) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dff1 = df1.Filter("b1 % 2 == 0", "name") - dff2 = df2.Filter("b1 % 2 == 0", "name") - - gen_train = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_y_train = [] - - train_iter = iter(gen_train.as_numpy()) - - for _ in range(3): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test10_two_epochs_shuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 0) - self.assertEqual(gen_validation.last_batch_no_of_rows, 1) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test12_PyTorch(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] - results_x_val = [3.0, 4.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] - results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] - results_z_val = [30.0, 40.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch(device="cpu")) - iter_val = iter(gen_validation.as_torch()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test13_TensorFlow(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] - results_x_val = [3.0, 4.0, 8.0, 9.0, 0.0, 0.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] - results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] - results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] - results_z_val = [30.0, 40.0, 80.0, 90.0, 0.0, 0.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test14_big_data(self): - file_name1 = "big_data_1.root" - file_name2 = "big_data_2.root" - tree_name = "myTree" - - entries_in_rdf = randrange(10000, 30000) - batch_size = randrange(100, 501) - batches_in_memory = randrange(5, 20) - - error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" - - def define_rdf(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, batches_in_memory, num_of_entries): - define_rdf(num_of_entries, file_name1) - define_rdf(num_of_entries, file_name2) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - batches_in_memory=batches_in_memory, - target=["b3", "b5"], - weights="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collect_x = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}") - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}") - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}") - - collect_x.extend(list(x[:, 0])) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message) - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message) - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message) - - collect_x.extend(list(x[:, 0])) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), - f"collected length: {len(set(collect_x))}\ - generated length {len(set(i for i in range(num_of_entries)))}", - ) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - test(batch_size, batches_in_memory, entries_in_rdf) - - def test15_two_runs_set_seed(self): - self.create_file1() - self.create_file2() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test16_vector_padding(self): - self.create_vector_file1() - self.create_vector_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name4) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name5) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0, - 200.0, - 2000.0, - 5.0, - 50.0, - 0, - 500.0, - 5000.0, - 6.0, - 60.0, - 0.0, - 600.0, - 6000.0, - 7.0, - 70.0, - 0.0, - 700.0, - 7000.0, - ] - results_y_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] - results_x_val = [ - 3.0, - 30.0, - 0.0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0.0, - 400.0, - 4000.0, - 8.0, - 80.0, - 0.0, - 800.0, - 8000.0, - 9.0, - 90.0, - 0.0, - 900.0, - 9000.0, - ] - results_y_val = [3.0, 4.0, 8.0, 9.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise - - def test17_dataframes_without_cluster_metadata(self): - # the dataframes are not read from a file, so their entries are counted with Count() - dl = ROOT.Experimental.ML.RDataLoader( - [self.define_rdf1(), self.define_rdf2()], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - gen_train, gen_validation = dl.train_test_split(0.4) - - x_train = [v for x, _ in gen_train.as_numpy() for v in x.flatten().tolist()] - x_val = [v for x, _ in gen_validation.as_numpy() for v in x.flatten().tolist()] - self.assertEqual(x_train, [0.0, 1.0, 2.0, 5.0, 6.0, 7.0]) - self.assertEqual(x_val, [3.0, 4.0, 8.0, 9.0]) - - -class DataLoaderLazyLoadingMultipleDataframes(unittest.TestCase): - file_name1 = "first_half.root" - file_name2 = "second_half.root" - file_name3 = "second_file.root" - file_name4 = "vector_columns_1.root" - file_name5 = "vector_columns_2.root" - tree_name = "mytree" - - # default constants - n_train_batch = 2 - n_val_batch = 1 - val_remainder = 1 - - # Helpers - def define_rdf1(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def define_rdf2(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(double) b1*b1") - - return df - - def create_file1(self, num_of_entries=5): - self.define_rdf1(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_file2(self, num_of_entries=5): - self.define_rdf2(num_of_entries).Snapshot(self.tree_name, self.file_name2) - - def create_5_entries_file(self): - ( - ROOT - .RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 10") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name3) - ) - - def create_vector_file1(self, num_of_entries=5): - ( - ROOT - .RDataFrame(5) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name4) - ) - - def create_vector_file2(self, num_of_entries=5): - ( - ROOT - .RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 5") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name5) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(df1_entries_before, df1_entries_after)) - self.assertTrue(np.array_equal(df2_entries_before, df2_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test04_dropping_remainder(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x = [] - collected_y = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test05_more_than_one_file(self): - self.create_file1() - self.create_file2() - self.create_5_entries_file() - - try: - df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] - results_x_val = [9.0, 10.0, 11.0, 12.0, 13.0, 14.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0] - results_y_val = [81.0, 100.0, 121.0, 144.0, 169.0, 196.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for x, y in gen_validation.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise - - def test06_multiple_target_columns(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test07_multiple_input_columns(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name1) - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 1.0, - 10.0, - 2.0, - 20.0, - 3.0, - 30.0, - 4.0, - 40.0, - 5.0, - 50.0, - ] - results_x_val = [6.0, 60.0, 7.0, 70.0, 8.0, 80.0, 9.0, 90.0] - results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] - results_y_val = [36.0, 49.0, 64.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test08_filtered(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dff1 = df1.Filter("b1 % 2 == 0", "name") - dff2 = df2.Filter("b1 % 2 != 0", "name") - - dff1_entries_before = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_before = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=4, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - # dff1 keeps the even entries 0, 2, 4 and dff2 the odd entries 5, 7, 9. Unshuffled, the first - # 4 of these 6 entries (0, 2, 4, 5) go to training and the last 2 (7, 9) to validation - results_x_train = [0.0, 2.0, 4.0, 5.0] - results_x_val = [7.0, 9.0] - results_y_train = [0.0, 4.0, 16.0, 25.0] - results_y_val = [49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (4, 1)) - self.assertTrue(y.shape == (4, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(dff1_entries_before, dff1_entries_after)) - self.assertTrue(np.array_equal(dff2_entries_before, dff2_entries_after)) - - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test09_filtered_last_chunk(self): - file_name1 = "filtered_last_chunk_1.root" - file_name2 = "filtered_last_chunk_2.root" - tree_name = "myTree" - - ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name1 - ) - - ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_ + 10").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name2 - ) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dff1 = df1.Filter("b1 % 2 == 0", "name") - dff2 = df2.Filter("b1 % 2 == 0", "name") - - gen_train = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - results_x_train = [0.0, 2.0, 4.0, 10.0, 12.0, 14.0, 6.0, 8.0, 16.0, 18.0] - results_y_train = [0.0, 4.0, 16.0, 100.0, 144.0, 196.0, 36.0, 64.0, 256.0, 324.0] - - collected_x_train = [] - collected_y_train = [] - - train_iter = iter(gen_train.as_numpy()) - - for _ in range(3): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test10_two_epochs_shuffled(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 0) - self.assertEqual(gen_validation.last_batch_no_of_rows, 1) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test12_PyTorch(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test13_TensorFlow(self): - file_name1 = "multiple_target_columns_1.root" - file_name2 = "multiple_target_columns_2.root" - - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) rdfentry_ + 5").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - - try: - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [6.0, 7.0, 8.0, 9.0, 0.0, 0.0] - results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] - results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] - results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] - results_z_val = [60.0, 70.0, 80.0, 90.0, 0.0, 0.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test14_big_data(self): - file_name1 = "big_data_1.root" - file_name2 = "big_data_2.root" - tree_name = "myTree" - - entries_in_rdf = randrange(10000, 30000) - batch_size = randrange(100, 501) - batches_in_memory = randrange(5, 20) - - error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" - - def define_rdf(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, batches_in_memory, num_of_entries): - define_rdf(num_of_entries, file_name1) - define_rdf(num_of_entries, file_name2) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - batches_in_memory=batches_in_memory, - target=["b3", "b5"], - weights="b2", - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collect_x = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}") - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}") - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}") - - collect_x.extend(list(x[:, 0])) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - - self.assertTrue(np.all(x[:, 0] * (-1) == x[:, 1]), error_message) - self.assertTrue(np.all(x[:, 0] + 10192 == y[:, 0]), error_message) - # self.assertTrue(np.all(x[:,0]*(-1)-10192==y[:,1]), error_message) - self.assertTrue(np.all(x[:, 0] * 2 == z[:, 0]), error_message) - - collect_x.extend(list(x[:, 0])) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collect_x.extend(list(x[:, 0])) - - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), - f"collected length: {len(set(collect_x))}\ - generated length {len(set(i for i in range(num_of_entries)))}", - ) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - test(batch_size, batches_in_memory, entries_in_rdf) - - def test15_two_runs_set_seed(self): - self.create_file1() - self.create_file2() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_val) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 6) - self.assertEqual(len(flat_x_val), 4) - self.assertEqual(len(flat_y_train), 6) - self.assertEqual(len(flat_y_val), 4) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test16_vector_padding(self): - self.create_vector_file1() - self.create_vector_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name4) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name5) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0, - 200.0, - 2000.0, - 3.0, - 30.0, - 0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0.0, - 400.0, - 4000.0, - 5.0, - 50.0, - 0.0, - 500.0, - 5000.0, - ] - results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_x_val = [ - 6.0, - 60.0, - 0.0, - 600.0, - 6000.0, - 7.0, - 70.0, - 0.0, - 700.0, - 7000.0, - 8.0, - 80.0, - 0.0, - 800.0, - 8000.0, - 9.0, - 90.0, - 0.0, - 900.0, - 9000.0, - ] - results_y_val = [6.0, 7.0, 8.0, 9.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(val_iter) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise - - def test17_shuffled_split_varies_with_seed(self): - self.create_file1() - self.create_file2() - - try: - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl1 = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - ) - - dl2 = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=43, - ) - - _, gen_val1 = dl1.train_test_split(0.4) - _, gen_val2 = dl2.train_test_split(0.4) - - val_1_collected = sorted([v for x, y in gen_val1.as_numpy() for v in x.flatten().tolist()]) - val_2_collected = sorted([v for x, y in gen_val2.as_numpy() for v in x.flatten().tolist()]) - - self.assertNotEqual(val_1_collected, val_2_collected) - - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - -class DataLoaderRandomUndersampling(unittest.TestCase): - file_name1 = "major.root" - file_name2 = "minor.root" - file_name3 = "second_file.root" - file_name4 = "vector_columns_major.root" - file_name5 = "vector_columns_minor.root" - tree_name = "mytree" - - # default constants - n_train_batch = 4 - n_val_batch = 3 - train_remainder = 1 - - # Helpers - def define_rdf_even(self, num_of_entries=20): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def define_rdf_odd(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(double) b1*b1") - - return df - - def create_file_major(self, num_of_entries=20): - self.define_rdf_even(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_file_minor(self, num_of_entries=5): - self.define_rdf_odd(num_of_entries).Snapshot(self.tree_name, self.file_name2) - - def create_5_entries_file(self): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) 2 * (rdfentry_ + 20)") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name3) - ) - - def create_vector_file_major(self, num_of_entries=20): - ( - ROOT.RDataFrame(20) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name4) - ) - - def create_vector_file_minor(self, num_of_entries=5): - ( - ROOT.RDataFrame(5) - .Define("b1", "(int) rdfentry_ + 20") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name5) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] - results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] - results_y_val = [576.0, 676.0, 784.0, 900.0, 49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - # check if there are no duplicate entries (replacement=False) - self.assertEqual(len(set(flat_x_train)), len(flat_x_train)) - self.assertEqual(len(set(flat_x_val)), len(flat_x_val)) - self.assertEqual(len(set(flat_y_train)), len(flat_y_train)) - self.assertEqual(len(set(flat_y_val)), len(flat_y_val)) - - # check if correct sampling_ratio (0.5 = minor/major) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test01_each_element_is_generated_unshuffled_replacement(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.2, - replacement=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 2.0, - 4.0, - 6.0, - 8.0, - 10.0, - 12.0, - 14.0, - 16.0, - 18.0, - 20.0, - 22.0, - 0.0, - 2.0, - 4.0, - 1.0, - 3.0, - 5.0, - ] - results_x_val = [24.0, 26.0, 28.0, 30.0, 32.0, 34.0, 36.0, 38.0, 24.0, 26.0, 7.0, 9.0] - results_y_train = [ - 0.0, - 4.0, - 16.0, - 36.0, - 64.0, - 100.0, - 144.0, - 196.0, - 256.0, - 324.0, - 400.0, - 484.0, - 0.0, - 4.0, - 16.0, - 1.0, - 9.0, - 25.0, - ] - results_y_val = [576.0, 676.0, 784.0, 900.0, 1024.0, 1156.0, 1296.0, 1444.0, 576.0, 676.0, 49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(6): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(9): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - # check if there are duplicate entries (replacement=True) - self.assertLess(len(set(flat_x_train)), len(flat_x_train)) - self.assertLess(len(set(flat_x_val)), len(flat_x_val)) - self.assertLess(len(set(flat_y_train)), len(flat_y_train)) - self.assertLess(len(set(flat_y_val)), len(flat_y_val)) - - # check if correct sampling_ratio (0.2 = minor/major) - self.assertEqual(num_major_train, 15) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 10) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test04_dropping_remainder(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - collected_x = [] - collected_y = [] - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x.append(x) - collected_y.append(y) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 7) - self.assertEqual(len(collected_y), 7) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test05_more_than_one_file(self): - self.create_file_major() - self.create_file_minor() - self.create_5_entries_file() - - try: - df_major = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name3]) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] - results_x_val = [30.0, 32.0, 34.0, 36.0, 7.0, 9.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] - results_y_val = [900.0, 1024.0, 1156.0, 1296.0, 49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise - - def test06_multiple_target_columns(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(20).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] - results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] - results_y_train = [ - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - 1.0, - 100.0, - 9.0, - 300.0, - 25.0, - 500.0, - ] - results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] - results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0] - results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test07_multiple_input_columns(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(20).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name1) - - ROOT.RDataFrame(5).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name2) - - try: - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 2.0, - 20.0, - 4.0, - 40.0, - 6.0, - 60.0, - 8.0, - 80.0, - 10.0, - 100.0, - 1.0, - 10.0, - 3.0, - 30.0, - 5.0, - 50.0, - ] - results_x_val = [24.0, 240.0, 26.0, 260.0, 28.0, 280.0, 30.0, 300.0, 7.0, 70.0, 9.0, 90.0] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] - results_y_val = [576.0, 676.0, 784.0, 900.0, 49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 2)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test08_filtered(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - df_major_filter = df_major.Filter("b1 % 8 != 0", "name") - - major_filter_entries_before = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major_filter, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [2.0, 4.0, 6.0, 10.0, 12.0, 14.0, 1.0, 3.0, 5.0] - results_x_val = [26.0, 28.0, 30.0, 34.0, 7.0, 9.0] - results_y_train = [4.0, 16.0, 36.0, 100.0, 144.0, 196.0, 1.0, 9.0, 25.0] - results_y_val = [676.0, 784.0, 900.0, 1156.0, 49.0, 81.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = np.sum((np.asarray(flat_x_train) % 2 == 0) & (np.asarray(flat_x_train) % 8 != 0)) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = np.sum((np.asarray(flat_x_val) % 2 == 0) & (np.asarray(flat_x_val) % 8 != 0)) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - major_filter_entries_after = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(major_filter_entries_before, major_filter_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test10_two_epochs_shuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 1) - self.assertEqual(gen_validation.last_batch_no_of_rows, 0) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test12_PyTorch(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(20).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] - results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] - results_y_train = [ - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - 1.0, - 100.0, - 9.0, - 300.0, - 25.0, - 500.0, - ] - results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] - results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0] - results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch(device="cpu")) - iter_val = iter(gen_validation.as_torch()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test13_TensorFlow(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(20).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(5).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0, 0.0] - results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] - results_y_train = [ - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - 1.0, - 100.0, - 9.0, - 300.0, - 25.0, - 500.0, - 0.0, - 0.0, - ] - results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] - results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0, 0.0] - results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 7) # num_major_train = 6 + 1 padded 0 from the last batch - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test14_big_data_replacement_false(self): - file_name1 = "big_data_major.root" - file_name2 = "big_data_minor.root" - tree_name = "myTree" - - entries_in_rdf_major = randrange(10000, 30000) - entries_in_rdf_minor = randrange(8000, 9999) - batch_size = randrange(100, 501) - min_allowed_sampling_ratio = entries_in_rdf_minor / entries_in_rdf_major - sampling_ratio = math.ceil(uniform(min_allowed_sampling_ratio, 2) * 100) / 100 - - error_message = f"\n Batch size: {batch_size}\ - Number of major entries: {entries_in_rdf_major} \ - Number of minor entries: {entries_in_rdf_minor}" - - def define_rdf_major(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def define_rdf_minor(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_ + 1").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): - define_rdf_major(num_of_entries_major, file_name1) - define_rdf_minor(num_of_entries_minor, file_name2) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - target=["b3", "b5"], - weights="b1", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=sampling_ratio, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collected_z_train = [] - collected_z_val = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_train.append(z.tolist()) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collected_z_train.append(z.tolist()) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_val.append(z.tolist()) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collected_z_val.append(z.tolist()) - - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - num_major_train = sum(np.array(flat_z_train) % 2 == 0) - num_minor_train = sum(np.array(flat_z_train) % 2 != 0) - num_major_val = sum(np.array(flat_z_val) % 2 == 0) - num_minor_val = sum(np.array(flat_z_val) % 2 != 0) - - # check if there are no duplicate entries (replacement=False) - self.assertEqual(len(set(flat_z_train)), len(flat_z_train)) - self.assertEqual(len(set(flat_z_val)), len(flat_z_val)) - - # check if the sampling stategy is correct - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test_raises(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): - define_rdf_major(num_of_entries_major, file_name1) - define_rdf_minor(num_of_entries_minor, file_name2) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - with self.assertRaisesRegex( - Exception, r"The sampling_ratio is too low: not enough entries in the majority class to sample from." - ): - ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - target=["b3", "b5"], - weights="b1", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=sampling_ratio, - replacement=False, - ).train_test_split(0.3) - - # test the functionality with a proper sampling ratio - test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) - - bad_sampling_ratio = round(max(min_allowed_sampling_ratio - 0.01, 0.01), 2) - # test that an error is raised when the sampling ratio is too low - test_raises(batch_size, entries_in_rdf_major, entries_in_rdf_minor, bad_sampling_ratio) - - def test14_big_data_replacement_true(self): - file_name1 = "big_data_major.root" - file_name2 = "big_data_minor.root" - tree_name = "myTree" - - entries_in_rdf_major = randrange(10000, 30000) - entries_in_rdf_minor = randrange(8000, 9999) - batch_size = randrange(100, 501) - - # max samling strategy to guarantee duplicate sampled entires - max_sampling_ratio = entries_in_rdf_minor / entries_in_rdf_major - sampling_ratio = round(uniform(0.01, max_sampling_ratio - 0.01), 2) - - error_message = f"\n Batch size: {batch_size}\ - Number of major entries: {entries_in_rdf_major} \ - Number of minor entries: {entries_in_rdf_minor}" - - def define_rdf_major(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def define_rdf_minor(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_ + 1").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): - define_rdf_major(num_of_entries_major, file_name1) - define_rdf_minor(num_of_entries_minor, file_name2) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - target=["b3", "b5"], - weights="b1", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=sampling_ratio, - replacement=True, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collected_z_train = [] - collected_z_val = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_train.append(z.tolist()) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collected_z_train.append(z.tolist()) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_val.append(z.tolist()) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collected_z_val.append(z.tolist()) - - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - num_major_train = sum(np.array(flat_z_train) % 2 == 0) - num_minor_train = sum(np.array(flat_z_train) % 2 != 0) - num_major_val = sum(np.array(flat_z_val) % 2 == 0) - num_minor_val = sum(np.array(flat_z_val) % 2 != 0) - - # check if there are duplicate entries (replacement=True) - self.assertLess(len(set(flat_z_train)), len(flat_z_train)) - self.assertLess(len(set(flat_z_val)), len(flat_z_val)) - - # check if the sampling stategy is correct - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) - - def test15_two_runs_set_seed(self): - self.create_file_major() - self.create_file_minor() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} - flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} - flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} - flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} - - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test16_vector_padding(self): - self.create_vector_file_major() - self.create_vector_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name4) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name5) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="undersampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0.0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0.0, - 200.0, - 2000.0, - 3.0, - 30.0, - 0.0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0.0, - 400.0, - 4000.0, - 5.0, - 50.0, - 0.0, - 500.0, - 5000.0, - 20.0, - 200.0, - 0.0, - 2000.0, - 20000.0, - 21.0, - 210.0, - 0.0, - 2100.0, - 21000.0, - 22.0, - 220.0, - 0.0, - 2200.0, - 22000.0, - ] - results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 20.0, 21.0, 22.0] - results_x_val = [ - 12.0, - 120.0, - 0.0, - 1200.0, - 12000.0, - 13.0, - 130.0, - 0.0, - 1300.0, - 13000.0, - 14.0, - 140.0, - 0.0, - 1400.0, - 14000.0, - 15.0, - 150.0, - 0.0, - 1500.0, - 15000.0, - 23.0, - 230.0, - 0.0, - 2300.0, - 23000.0, - 24.0, - 240.0, - 0.0, - 2400.0, - 24000.0, - ] - results_y_val = [12.0, 13.0, 14.0, 15.0, 23.0, 24.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 5)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_y_train) < 20) - num_minor_train = sum(np.array(flat_y_train) >= 20) - num_major_val = sum(np.array(flat_y_val) < 20) - num_minor_val = sum(np.array(flat_y_val) >= 20) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise - - -class DataLoaderRandomOversampling(unittest.TestCase): - file_name1 = "major.root" - file_name2 = "minor.root" - file_name3 = "second_file.root" - file_name4 = "vector_columns_major.root" - file_name5 = "vector_columns_minor.root" - tree_name = "mytree" - - # default constants - n_train_batch = 4 - n_val_batch = 3 - train_remainder = 1 - - # Helpers - def define_rdf_even(self, num_of_entries=20): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(double) b1*b1") - - return df - - def define_rdf_odd(self, num_of_entries=5): - df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(double) b1*b1") - - return df - - def create_file_major(self, num_of_entries=10): - self.define_rdf_even(num_of_entries).Snapshot(self.tree_name, self.file_name1) - - def create_file_minor(self, num_of_entries=3): - self.define_rdf_odd(num_of_entries).Snapshot(self.tree_name, self.file_name2) - - def create_extra_entry_file(self): - ( - ROOT.RDataFrame(1) - .Define("b1", "(int) 2 * (rdfentry_ + 3) + 1") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, self.file_name3) - ) - - def create_vector_file_major(self, num_of_entries=10): - ( - ROOT.RDataFrame(10) - .Define("b1", "(int) rdfentry_") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name4) - ) - - def create_vector_file_minor(self, num_of_entries=3): - ( - ROOT.RDataFrame(3) - .Define("b1", "(int) rdfentry_ + 10") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - .Define("v2", "ROOT::VecOps::RVec{ b1 * 100, b1 * 1000}") - .Snapshot(self.tree_name, self.file_name5) - ) - - def teardown_file(self, file): - os.remove(file) - - def test01_each_element_is_generated_unshuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] - results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] - results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - # check if there are no duplicate entries (oversampling) - self.assertLess(len(set(flat_x_train)), len(flat_x_train)) - self.assertLess(len(set(flat_x_val)), len(flat_x_val)) - self.assertLess(len(set(flat_y_train)), len(flat_y_train)) - self.assertLess(len(set(flat_y_val)), len(flat_y_val)) - - # check if correct sampling_ratio (0.5 = minor/major) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test02_each_element_is_generated_shuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - # check if there are no duplicate entries (oversampling) - self.assertLess(len(set(flat_x_train)), len(flat_x_train)) - self.assertLess(len(set(flat_x_val)), len(flat_x_val)) - self.assertLess(len(set(flat_y_train)), len(flat_y_train)) - self.assertLess(len(set(flat_y_val)), len(flat_y_val)) - - # check if correct sampling_ratio (0.5 = minor/major) - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test04_dropping_remainder(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - collected_x = [] - collected_y = [] - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x.append(x) - collected_y.append(y) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x.append(x) - collected_y.append(y) - - self.assertEqual(len(collected_x), 7) - self.assertEqual(len(collected_y), 7) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test05_more_than_one_file(self): - self.create_file_major() - self.create_file_minor() - self.create_extra_entry_file() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, [self.file_name2, self.file_name3]) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 5.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] - results_x_val = [7.0, 7.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [1.0, 9.0, 25.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] - results_y_val = [49.0, 49.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise - - def test06_multiple_target_columns(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(10).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(3).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] - results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [ - 1.0, - 100.0, - 9.0, - 300.0, - 1.0, - 100.0, - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - ] - results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] - results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0] - results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test07_multiple_input_columns(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(10).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name1) - - ROOT.RDataFrame(3).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Snapshot("myTree", file_name2) - - try: - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 1.0, - 10.0, - 3.0, - 30.0, - 1.0, - 10.0, - 0.0, - 0.0, - 2.0, - 20.0, - 4.0, - 40.0, - 6.0, - 60.0, - 8.0, - 80.0, - 10.0, - 100.0, - ] - results_x_val = [5.0, 50.0, 5.0, 50.0, 12.0, 120.0, 14.0, 140.0, 16.0, 160.0, 18.0, 180.0] - results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] - results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 2)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test08_filtered(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor_duplicate = ROOT.RDataFrame(self.tree_name, [self.file_name2, self.file_name2, self.file_name2]) - - df_minor_filter = df_minor_duplicate.Filter("rdfentry_ < 3", "name") - - major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_filter_entries_before = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor_filter], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] - results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] - results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_filter_entries_after = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - - # check if the dataframes are correctly reset - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_filter_entries_before, minor_filter_entries_after)) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test10_two_epochs_shuffled(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] - - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - both_epochs_collected_x_val.append(collected_x_val) - both_epochs_collected_y_val.append(collected_y_val) - - self.assertEqual(both_epochs_collected_x_val[0], both_epochs_collected_x_val[1]) - self.assertEqual(both_epochs_collected_y_val[0], both_epochs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test11_number_of_training_and_validation_batches_remainder(self): - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - number_of_training_batches = 0 - number_of_validation_batches = 0 - - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 - - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 - - self.assertEqual(gen_train.num_batches, number_of_training_batches) - self.assertEqual(gen_validation.num_batches, number_of_validation_batches) - self.assertEqual(gen_train.last_batch_no_of_rows, 1) - self.assertEqual(gen_validation.last_batch_no_of_rows, 0) - - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise - - def test12_PyTorch(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(10).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(3).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] - results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [ - 1.0, - 100.0, - 9.0, - 300.0, - 1.0, - 100.0, - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - ] - results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] - results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0] - results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - collected_z_val.append(z.tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - collected_z_train.append(z.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test13_TensorFlow(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" - - ROOT.RDataFrame(10).Define("b1", "(int) 2 * rdfentry_").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name1) - ROOT.RDataFrame(3).Define("b1", "(int) 2 * rdfentry_ + 1").Define("b2", "(int) b1 * b1").Define( - "b3", "(double) b1 * 10" - ).Define("b4", "(double) b3 * 10").Snapshot("myTree", file_name2) - try: - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], - batch_size=2, - target=["b2", "b4"], - weights="b3", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 0.0] - results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] - results_y_train = [ - 1.0, - 100.0, - 9.0, - 300.0, - 1.0, - 100.0, - 0.0, - 0.0, - 4.0, - 200.0, - 16.0, - 400.0, - 36.0, - 600.0, - 64.0, - 800.0, - 100.0, - 1000.0, - 0.0, - 0.0, - ] - results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] - results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 0.0] - results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] - - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) - - for _ in range(self.n_train_batch): - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - for _ in range(self.n_val_batch): - x, y, z = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_val.append(x.numpy().tolist()) - collected_y_val.append(y.numpy().tolist()) - collected_z_val.append(z.numpy().tolist()) - - x, y, z = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) - collected_x_train.append(x.numpy().tolist()) - collected_y_train.append(y.numpy().tolist()) - collected_z_train.append(z.numpy().tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - self.assertEqual(results_z_train, flat_z_train) - self.assertEqual(results_z_val, flat_z_val) - - num_major_train = sum(np.array(flat_x_train) % 2 == 0) - num_minor_train = sum(np.array(flat_x_train) % 2 != 0) - num_major_val = sum(np.array(flat_x_val) % 2 == 0) - num_minor_val = sum(np.array(flat_x_val) % 2 != 0) - - self.assertEqual(num_major_train, 7) # num_major_train = 6 + 1 padded 0 from the last batch - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - def test14_big_data_replacement_false(self): - file_name1 = "big_data_major.root" - file_name2 = "big_data_minor.root" - tree_name = "myTree" - - entries_in_rdf_major = randrange(10000, 30000) - entries_in_rdf_minor = randrange(8000, 9999) - batch_size = randrange(100, 501) - sampling_ratio = round(uniform(0.1, 2), 2) - - error_message = f"\n Batch size: {batch_size}\ - Number of major entries: {entries_in_rdf_major} \ - Number of minor entries: {entries_in_rdf_minor}" - - def define_rdf_major(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def define_rdf_minor(num_of_entries, file_name): - ROOT.RDataFrame(num_of_entries).Define("b1", "(int) 2 * rdfentry_ + 1").Define( - "b2", "(double) rdfentry_ * 2" - ).Define("b3", "(int) rdfentry_ + 10192").Define("b4", "(int) -rdfentry_").Define( - "b5", "(double) -rdfentry_ - 10192" - ).Snapshot(tree_name, file_name) - - def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): - define_rdf_major(num_of_entries_major, file_name1) - define_rdf_minor(num_of_entries_minor, file_name2) - - try: - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=size_of_batch, - target=["b3", "b5"], - weights="b1", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=sampling_ratio, - ) - - gen_train, gen_validation = dl.train_test_split(0.3) - - collected_z_train = [] - collected_z_val = [] - - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows - - n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches - n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for i in range(n_train_batches): - x, y, z = next(iter_train) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_train.append(z.tolist()) - - if train_remainder: - x, y, z = next(iter_train) - self.assertTrue(x.shape == (train_remainder, 2), error_message) - self.assertTrue(y.shape == (train_remainder, 2), error_message) - self.assertTrue(z.shape == (train_remainder, 1), error_message) - collected_z_train.append(z.tolist()) - - for _ in range(n_val_batches): - x, y, z = next(iter_val) - - self.assertTrue(x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}") - self.assertTrue(y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}") - self.assertTrue(z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}") - collected_z_val.append(z.tolist()) - - if val_remainder: - x, y, z = next(iter_val) - self.assertTrue(x.shape == (val_remainder, 2), error_message) - self.assertTrue(y.shape == (val_remainder, 2), error_message) - self.assertTrue(z.shape == (val_remainder, 1), error_message) - collected_z_val.append(z.tolist()) - - flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] - flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] - - num_major_train = sum(np.array(flat_z_train) % 2 == 0) - num_minor_train = sum(np.array(flat_z_train) % 2 != 0) - num_major_val = sum(np.array(flat_z_val) % 2 == 0) - num_minor_val = sum(np.array(flat_z_val) % 2 != 0) - - # check if there are no duplicate entries (replacement=False) - self.assertLessEqual(len(set(flat_z_train)), len(flat_z_train)) - self.assertLessEqual(len(set(flat_z_val)), len(flat_z_val)) - - # check if the sampling stategy is correct - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise - - test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) - - def test15_two_runs_set_seed(self): - self.create_file_major() - self.create_file_minor() - - try: - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - for _ in range(2): - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) - - for _ in range(self.n_train_batch): - x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - for _ in range(self.n_val_batch): - x, y = next(iter_val) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - x, y = next(iter_train) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(len(flat_x_train), 9) - self.assertEqual(len(flat_x_val), 6) - self.assertEqual(len(flat_y_train), 9) - self.assertEqual(len(flat_y_val), 6) - - both_runs_collected_x_val.append(collected_x_val) - both_runs_collected_y_val.append(collected_y_val) - self.assertEqual(both_runs_collected_x_val[0], both_runs_collected_x_val[1]) - self.assertEqual(both_runs_collected_y_val[0], both_runs_collected_y_val[1]) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - def test16_vector_padding(self): - self.create_vector_file_major() - self.create_vector_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name4) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name5) - max_vec_sizes = {"v1": 3, "v2": 2} - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, - drop_remainder=False, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=0.5, - replacement=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - results_x_train = [ - 10.0, - 100.0, - 0.0, - 1000.0, - 10000.0, - 11.0, - 110.0, - 0.0, - 1100.0, - 11000.0, - 10.0, - 100.0, - 0.0, - 1000.0, - 10000.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 10.0, - 0.0, - 100.0, - 1000.0, - 2.0, - 20.0, - 0.0, - 200.0, - 2000.0, - 3.0, - 30.0, - 0.0, - 300.0, - 3000.0, - 4.0, - 40.0, - 0.0, - 400.0, - 4000.0, - 5.0, - 50.0, - 0.0, - 500.0, - 5000.0, - ] - results_x_val = [ - 12.0, - 120.0, - 0.0, - 1200.0, - 12000.0, - 12.0, - 120.0, - 0.0, - 1200.0, - 12000.0, - 6.0, - 60.0, - 0.0, - 600.0, - 6000.0, - 7.0, - 70.0, - 0.0, - 700.0, - 7000.0, - 8.0, - 80.0, - 0.0, - 800.0, - 8000.0, - 9.0, - 90.0, - 0.0, - 900.0, - 9000.0, - ] - results_y_train = [10.0, 11.0, 10.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0] - results_y_val = [12.0, 12.0, 6.0, 7.0, 8.0, 9.0] - - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) - - for _ in range(self.n_val_batch): - x, y = next(val_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) - - for _ in range(self.n_train_batch): - x, y = next(train_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - x, y = next(train_iter) - self.assertTrue(x.shape == (self.train_remainder, 5)) - self.assertTrue(y.shape == (self.train_remainder, 1)) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) - - flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] - flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] - flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] - flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] - - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_x_val, flat_x_val) - self.assertEqual(results_y_train, flat_y_train) - self.assertEqual(results_y_val, flat_y_val) - - num_major_train = sum(np.array(flat_y_train) < 10) - num_minor_train = sum(np.array(flat_y_train) >= 10) - num_major_val = sum(np.array(flat_y_val) < 10) - num_minor_val = sum(np.array(flat_y_val) >= 10) - - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) - - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise - - def test17_set_seed_resamples_different_entries(self): - # with a fixed seed, the resampled minority entries used to all be the same one - self.create_file_major() - self.create_file_minor() - - try: - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df_major, df_minor], - batch_size=2, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=42, - load_eager=True, - sampling_type="oversampling", - sampling_ratio=1.0, - ) - - x = np.concatenate([x for x, _ in dl.as_numpy()]).flatten() - minor = x[x % 2 == 1] - - self.assertEqual(len(minor), 10) - self.assertGreater(len(np.unique(minor)), 1) - - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - - -class DataLoaderSave(unittest.TestCase): - file_name = "dataloader_save_input.root" - file_name2 = "dataloader_save_input2.root" - tree_name = "mytree" - - out_name = "dataloader_save_output.root" - out_name2 = "dataloader_save_output2.root" - - # Helpers - def create_file(self, file_name=None, num_of_entries=20, offset=0): - ( - ROOT - .RDataFrame(num_of_entries) - .Define("b1", f"(int) rdfentry_ + {offset}") - .Define("b2", "b1 + 1") - .Snapshot(self.tree_name, file_name or self.file_name) - ) - - def tearDown(self): - for file in (self.file_name, self.file_name2, self.out_name, self.out_name2): - if os.path.exists(file): - os.remove(file) - - def read_back(self, file_name, column): - """Read back one column of a saved file: a 1D array, or a 2D array for a vector column.""" - values = ROOT.RDataFrame(self.tree_name, file_name).AsNumpy([column])[column] - # a vector column comes back as one RVec per entry, all of the same padded size - return np.stack([np.asarray(v) for v in values]) if values.dtype == object else values - - def saved_columns(self, file_name): - return {str(c) for c in ROOT.RDataFrame(self.tree_name, file_name).GetColumnNames()} - - def test01_unshuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name).Define("w", "(double) 0.5") - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=5, - batches_in_memory=2, - columns=["b1", "b2", "w"], - target="b2", - weights="w", - shuffle=False, - drop_remainder=False, - ) - - dl.save(self.tree_name, self.out_name) - - # check all features, target and weights - self.assertEqual(self.saved_columns(self.out_name), {"b1", "b2", "w"}) - np.testing.assert_array_equal(self.read_back(self.out_name, "b1"), np.arange(20)) - np.testing.assert_array_equal(self.read_back(self.out_name, "b2"), np.arange(1, 21)) - np.testing.assert_array_equal(self.read_back(self.out_name, "w"), np.full(20, 0.5)) - - def test02_shuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=4, - batches_in_memory=2, - columns=["b1"], - shuffle=True, - set_seed=42, - drop_remainder=False, - ) - - gen_train, gen_validation = dl.train_test_split(0.4) - - gen_train.save(self.tree_name, self.out_name) - gen_validation.save(self.tree_name, self.out_name2) - - train_b1 = self.read_back(self.out_name, "b1") - val_b1 = self.read_back(self.out_name2, "b1") - - # the two splits cover the dataset exactly once - all_b1 = np.concatenate([train_b1, val_b1]) - np.testing.assert_array_equal(np.sort(all_b1), np.arange(20)) - - # and the entries are not written in their original order - self.assertTrue((train_b1 != np.sort(train_b1)).any()) - - def test03_vector_columns(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name).Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=5, - batches_in_memory=2, - columns=["b1", "v1"], - max_vec_sizes={"v1": 3}, - vec_padding=42.0, - shuffle=False, - drop_remainder=False, - ) - - dl.save(self.tree_name, self.out_name) - - # v1 is saved as a vector column, not exploded into v1_0... - self.assertEqual(self.saved_columns(self.out_name), {"b1", "v1"}) - - v1 = self.read_back(self.out_name, "v1") - self.assertEqual(v1.shape, (20, 3)) - # v1 contains 2 elements padded up to max_vec_sizes["v1"] (3) with vec_padding (42) - np.testing.assert_array_equal(v1[3], [3.0, 30.0, 42.0]) - - def test04_ttree_and_rntuple_agree(self): - self.create_file() - - def save_as(file_name, output_format): - df = ROOT.RDataFrame(self.tree_name, self.file_name).Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=5, - batches_in_memory=2, - columns=["b1", "v1"], - max_vec_sizes={"v1": 3}, - shuffle=False, - drop_remainder=False, - ) - - dl.save(self.tree_name, file_name, output_format=output_format) - - save_as(self.out_name, "ttree") - save_as(self.out_name2, "rntuple") - - for column in ["b1", "v1"]: - np.testing.assert_array_equal( - self.read_back(self.out_name, column), - self.read_back(self.out_name2, column), - ) - - def test05_filtered(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name).Filter("b1 % 2 == 0", "even") - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=2, - batches_in_memory=2, - columns=["b1"], - shuffle=False, - drop_remainder=False, - ) - - dl.save(self.tree_name, self.out_name) - - saved_b1 = self.read_back(self.out_name, "b1") - np.testing.assert_array_equal(np.sort(saved_b1), np.arange(0, 20, 2)) - - def test06_multiple_dataframes(self): - self.create_file(self.file_name, num_of_entries=10) - self.create_file(self.file_name2, num_of_entries=10, offset=100) - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - - dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=5, - batches_in_memory=2, - columns=["b1"], - shuffle=False, - drop_remainder=False, - ) - - dl.save(self.tree_name, self.out_name) - - saved_b1 = self.read_back(self.out_name, "b1") - expected = np.concatenate([np.arange(10), np.arange(100, 110)]) - np.testing.assert_array_equal(np.sort(saved_b1), expected) - - def test07_eager_loading(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name) - - dl = ROOT.Experimental.ML.RDataLoader( - df, - batch_size=5, - columns=["b1", "b2"], - shuffle=False, - drop_remainder=False, - load_eager=True, - ) - - dl.save(self.tree_name, self.out_name) - - np.testing.assert_array_equal(self.read_back(self.out_name, "b1"), np.arange(20)) - np.testing.assert_array_equal(self.read_back(self.out_name, "b2"), np.arange(1, 21)) - - def test08_round_trip(self): - self.create_file() - - df = ( - ROOT - .RDataFrame(self.tree_name, self.file_name) - .Define("w", "(double) 0.5") - .Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10 }") - ) - - loader_args = dict( - batch_size=5, - batches_in_memory=2, - columns=["b1", "v1", "b2", "w"], - max_vec_sizes={"v1": 3}, - vec_padding=42.0, - target="b2", - weights="w", - shuffle=False, - drop_remainder=False, - ) - - dl = ROOT.Experimental.ML.RDataLoader(df, **loader_args) - dl.save(self.tree_name, self.out_name) - - # the saved file is fed straight back into a new RDataLoader with the same columns/target - reloaded = ROOT.Experimental.ML.RDataLoader(ROOT.RDataFrame(self.tree_name, self.out_name), **loader_args) - - x, y, w = (np.concatenate(arrs) for arrs in zip(*reloaded.as_numpy())) - np.testing.assert_array_equal(x[:, 0], np.arange(20)) # b1 - np.testing.assert_array_equal(x[:, 1], np.arange(20)) # v1[0] == b1 - np.testing.assert_array_equal(x[:, 2], np.arange(20) * 10) # v1[1] == b1 * 10 - np.testing.assert_array_equal(x[:, 3], np.full(20, 42.0)) # v1[2], padding - np.testing.assert_array_equal(y.flatten(), np.arange(1, 21)) - np.testing.assert_array_equal(w.flatten(), np.full(20, 0.5)) - - -class DataLoaderImplicitMT(unittest.TestCase): - file_name = "dataloader_imt.root" - tree_name = "mytree" - - def tearDown(self): - ROOT.DisableImplicitMT() - os.remove(self.file_name) - - @unittest.skipUnless("imt" in ROOT.gROOT.GetConfigFeatures(), "ROOT was built without implicit multi-threading") - def test01_implicit_mt_not_supported(self): - ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Snapshot(self.tree_name, self.file_name) - ROOT.EnableImplicitMT(2) - df = ROOT.RDataFrame(self.tree_name, self.file_name) - - with self.assertRaisesRegex(RuntimeError, "Cannot use RDataLoader with implicit multi-threading enabled"): - ROOT.Experimental.ML.RDataLoader(df, batch_size=2).as_numpy() - - -class DataLoaderFromFiles(unittest.TestCase): - file_name1 = "dataloader_from_files1.root" - file_name2 = "dataloader_from_files2.root" - tree_name = "mytree" - - # Helpers - def create_file(self, file_name, num_of_entries=10, offset=0, output_format="ttree"): - opts = ROOT.RDF.RSnapshotOptions() - if output_format == "rntuple": - opts.fOutputFormat = ROOT.RDF.ESnapshotOutputFormat.kRNTuple - ( - ROOT - .RDataFrame(num_of_entries) - .Define("b1", f"(int) rdfentry_ + {offset}") - .Define("b2", "(double) b1 * b1") - .Snapshot(self.tree_name, file_name, ["b1", "b2"], opts) - ) - - def tearDown(self): - for file in (self.file_name1, self.file_name2): - if os.path.exists(file): - os.remove(file) - - def load_all(self, inputs, **kwargs): - """Concatenate the features and the target of all the batches of an unshuffled loader.""" - dl = ROOT.Experimental.ML.RDataLoader( - inputs, batch_size=3, batches_in_memory=2, target="b2", shuffle=False, drop_remainder=False, **kwargs - ) - batches = list(dl.as_numpy()) - x = np.concatenate([x for x, _ in batches]).flatten() - y = np.concatenate([y for _, y in batches]).flatten() - return x, y - - def test01_single_path(self): - self.create_file(self.file_name1) - - for load_eager in (False, True): - with self.subTest(load_eager=load_eager): - x, y = self.load_all((self.tree_name, self.file_name1), load_eager=load_eager) - np.testing.assert_array_equal(x, np.arange(10)) - np.testing.assert_array_equal(y, np.arange(10) ** 2) - - def test02_pathlike(self): - from pathlib import Path - - self.create_file(self.file_name1) - - x, y = self.load_all((self.tree_name, Path(self.file_name1))) - np.testing.assert_array_equal(x, np.arange(10)) - np.testing.assert_array_equal(y, np.arange(10) ** 2) - - def test03_list_of_paths(self): - from pathlib import Path - - self.create_file(self.file_name1, num_of_entries=9) - self.create_file(self.file_name2, num_of_entries=6, offset=9) - - # the files are read as a single dataset, in the given order - x, y = self.load_all((self.tree_name, [self.file_name1, Path(self.file_name2)])) - np.testing.assert_array_equal(x, np.arange(15)) - np.testing.assert_array_equal(y, np.arange(15) ** 2) - - def test04_glob(self): - self.create_file(self.file_name1) - self.create_file(self.file_name2, num_of_entries=5, offset=10) - - x, _ = self.load_all((self.tree_name, "dataloader_from_files*.root")) - np.testing.assert_array_equal(np.sort(x), np.arange(15)) - - def test05_rntuple(self): - self.create_file(self.file_name1, output_format="rntuple") - - x, y = self.load_all((self.tree_name, self.file_name1)) - np.testing.assert_array_equal(x, np.arange(10)) - np.testing.assert_array_equal(y, np.arange(10) ** 2) - - def test06_train_test_split(self): - self.create_file(self.file_name1) - - dl = ROOT.Experimental.ML.RDataLoader( - (self.tree_name, self.file_name1), - batch_size=2, - columns=["b1"], - shuffle=False, - drop_remainder=False, - ) - gen_train, gen_validation = dl.train_test_split(0.4) - - train_x = np.concatenate(list(gen_train.as_numpy())).flatten() - val_x = np.concatenate(list(gen_validation.as_numpy())).flatten() - np.testing.assert_array_equal(train_x, np.arange(6)) - np.testing.assert_array_equal(val_x, np.arange(6, 10)) - - def test07_invalid_inputs(self): - self.create_file(self.file_name1) - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - - # file paths need the name of the dataset to read - with self.assertRaisesRegex(TypeError, "dataset_name, files"): - self.load_all(self.file_name1) - - with self.assertRaisesRegex(ValueError, "at least one valid input is required"): - self.load_all([]) - with self.assertRaisesRegex(ValueError, "at least one valid input is required"): - self.load_all(df for df in []) - - # all the arguments but the inputs are keyword-only - with self.assertRaises(TypeError): - ROOT.Experimental.ML.RDataLoader(df, 3) - - def test08_several_datasets(self): - self.create_file(self.file_name1, num_of_entries=9) - self.create_file(self.file_name2, num_of_entries=6, offset=9) - - # one RDataFrame per tuple, which can be mixed with RDataFrames - df = ROOT.RDataFrame(self.tree_name, self.file_name2) - for inputs in ( - [(self.tree_name, self.file_name1), (self.tree_name, self.file_name2)], - [(self.tree_name, self.file_name1), df], - ): - with self.subTest(inputs=inputs): - x, y = self.load_all(inputs) - np.testing.assert_array_equal(x, np.arange(15)) - np.testing.assert_array_equal(y, np.arange(15) ** 2) - - -if __name__ == "__main__": - unittest.main() diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_from_files.py b/bindings/pyroot/pythonizations/test/ml_dataloader_from_files.py new file mode 100644 index 0000000000000..737595761038c --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_from_files.py @@ -0,0 +1,125 @@ +import numpy as np +import pytest +import ROOT + + +class TestDataLoaderFromFiles: + file_name1 = "dataloader_from_files1.root" + file_name2 = "dataloader_from_files2.root" + tree_name = "mytree" + + # Helpers + def create_file(self, file_name, num_of_entries=10, offset=0, output_format="ttree"): + opts = ROOT.RDF.RSnapshotOptions() + if output_format == "rntuple": + opts.fOutputFormat = ROOT.RDF.ESnapshotOutputFormat.kRNTuple + ( + ROOT + .RDataFrame(num_of_entries) + .Define("b1", f"(int) rdfentry_ + {offset}") + .Define("b2", "(double) b1 * b1") + .Snapshot(self.tree_name, file_name, ["b1", "b2"], opts) + ) + + def load_all(self, inputs, **kwargs): + """Concatenate the features and the target of all the batches of an unshuffled loader.""" + dl = ROOT.Experimental.ML.RDataLoader( + inputs, batch_size=3, batches_in_memory=2, target="b2", shuffle=False, drop_remainder=False, **kwargs + ) + batches = list(dl.as_numpy()) + x = np.concatenate([x for x, _ in batches]).flatten() + y = np.concatenate([y for _, y in batches]).flatten() + return x, y + + def test01_single_path(self, load_eager): + self.create_file(self.file_name1) + + x, y = self.load_all((self.tree_name, self.file_name1), load_eager=load_eager) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test02_pathlike(self): + from pathlib import Path + + self.create_file(self.file_name1) + + x, y = self.load_all((self.tree_name, Path(self.file_name1))) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test03_list_of_paths(self): + from pathlib import Path + + self.create_file(self.file_name1, num_of_entries=9) + self.create_file(self.file_name2, num_of_entries=6, offset=9) + + # the files are read as a single dataset, in the given order + x, y = self.load_all((self.tree_name, [self.file_name1, Path(self.file_name2)])) + np.testing.assert_array_equal(x, np.arange(15)) + np.testing.assert_array_equal(y, np.arange(15) ** 2) + + def test04_glob(self): + self.create_file(self.file_name1) + self.create_file(self.file_name2, num_of_entries=5, offset=10) + + x, _ = self.load_all((self.tree_name, "dataloader_from_files*.root")) + np.testing.assert_array_equal(np.sort(x), np.arange(15)) + + def test05_rntuple(self): + self.create_file(self.file_name1, output_format="rntuple") + + x, y = self.load_all((self.tree_name, self.file_name1)) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test06_train_test_split(self): + self.create_file(self.file_name1) + + dl = ROOT.Experimental.ML.RDataLoader( + (self.tree_name, self.file_name1), + batch_size=2, + columns=["b1"], + shuffle=False, + drop_remainder=False, + ) + gen_train, gen_validation = dl.train_test_split(0.4) + + train_x = np.concatenate(list(gen_train.as_numpy())).flatten() + val_x = np.concatenate(list(gen_validation.as_numpy())).flatten() + np.testing.assert_array_equal(train_x, np.arange(6)) + np.testing.assert_array_equal(val_x, np.arange(6, 10)) + + def test07_invalid_inputs(self): + self.create_file(self.file_name1) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) + + # file paths need the name of the dataset to read + with pytest.raises(TypeError, match="dataset_name, files"): + self.load_all(self.file_name1) + + with pytest.raises(ValueError, match="at least one valid input is required"): + self.load_all([]) + with pytest.raises(ValueError, match="at least one valid input is required"): + self.load_all(df for df in []) + + # all the arguments but the inputs are keyword-only + with pytest.raises(TypeError): + ROOT.Experimental.ML.RDataLoader(df, 3) + + def test08_several_datasets(self): + self.create_file(self.file_name1, num_of_entries=9) + self.create_file(self.file_name2, num_of_entries=6, offset=9) + + # one RDataFrame per tuple, which can be mixed with RDataFrames + df = ROOT.RDataFrame(self.tree_name, self.file_name2) + for inputs in ( + [(self.tree_name, self.file_name1), (self.tree_name, self.file_name2)], + [(self.tree_name, self.file_name1), df], + ): + x, y = self.load_all(inputs) + np.testing.assert_array_equal(x, np.arange(15)) + np.testing.assert_array_equal(y, np.arange(15) ** 2) + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_implicit_mt.py b/bindings/pyroot/pythonizations/test/ml_dataloader_implicit_mt.py new file mode 100644 index 0000000000000..c0a2cf5425a43 --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_implicit_mt.py @@ -0,0 +1,25 @@ +import pytest +import ROOT + + +class TestDataLoaderImplicitMT: + file_name = "dataloader_imt.root" + tree_name = "mytree" + + def teardown_method(self): + ROOT.DisableImplicitMT() + + @pytest.mark.skipif( + "imt" not in ROOT.gROOT.GetConfigFeatures(), reason="ROOT was built without implicit multi-threading" + ) + def test01_implicit_mt_not_supported(self): + ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Snapshot(self.tree_name, self.file_name) + ROOT.EnableImplicitMT(2) + df = ROOT.RDataFrame(self.tree_name, self.file_name) + + with pytest.raises(RuntimeError, match="Cannot use RDataLoader with implicit multi-threading enabled"): + ROOT.Experimental.ML.RDataLoader(df, batch_size=2).as_numpy() + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_multiple_dataframes.py b/bindings/pyroot/pythonizations/test/ml_dataloader_multiple_dataframes.py new file mode 100644 index 0000000000000..ed22e9dd12f0e --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_multiple_dataframes.py @@ -0,0 +1,1123 @@ +from random import randrange + +import numpy as np +import pytest +import ROOT +from conftest import big_data, from_files, squares, vectors, weighted + + +class TestDataLoaderMultipleDataframes: + # default constants + n_train_batch = 2 + n_val_batch = 1 + val_remainder = 1 + + @pytest.fixture + def df1(self): + """5 entries: b1 = 0, ..., 4 and b2 = b1 * b1.""" + return from_files(squares(5)) + + @pytest.fixture + def df2(self): + """5 entries: b1 = 5, ..., 9 and b2 = b1 * b1.""" + return from_files(squares(5, "(int) rdfentry_ + 5")) + + @pytest.fixture + def weighted_df1(self): + """5 entries: b1 = 0, ..., 4, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(5)) + + @pytest.fixture + def weighted_df2(self): + """5 entries: b1 = 5, ..., 9, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(5, "(int) rdfentry_ + 5")) + + def test01_each_element_is_generated_unshuffled(self, df1, df2, load_eager): + df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + results_x_val = [3.0, 4.0, 8.0, 9.0] + results_y_train = [0.0, 1.0, 4.0, 25.0, 36.0, 49.0] + results_y_val = [9.0, 16.0, 64.0, 81.0] + else: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] + results_y_val = [36.0, 49.0, 64.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(df1_entries_before, df1_entries_after) + assert np.array_equal(df2_entries_before, df2_entries_after) + + def test02_each_element_is_generated_shuffled(self, df1, df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + def test04_dropping_remainder(self, df1, df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=True, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x = [] + collected_y = [] + + for x, y in gen_train.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x.append(x) + collected_y.append(y) + + for x, y in gen_validation.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x.append(x) + collected_y.append(y) + + assert len(collected_x) == 3 + assert len(collected_y) == 3 + + def test05_more_than_one_file(self, load_eager): + df1 = from_files(squares(5), squares(5, "(int) rdfentry_ + 5")) + df2 = from_files(squares(5, "(int) rdfentry_ + 10")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 11.0, 12.0] + results_x_val = [6.0, 7.0, 8.0, 9.0, 13.0, 14.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 100.0, 121.0, 144.0] + results_y_val = [36.0, 49.0, 64.0, 81.0, 169.0, 196.0] + else: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + results_x_val = [9.0, 10.0, 11.0, 12.0, 13.0, 14.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0] + results_y_val = [81.0, 100.0, 121.0, 144.0, 169.0, 196.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + for x, y in gen_train.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for x, y in gen_validation.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test06_multiple_target_columns(self, weighted_df1, weighted_df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df1, weighted_df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + results_x_val = [3.0, 4.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] + results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] + results_z_val = [30.0, 40.0, 80.0, 90.0] + else: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 2) + assert z.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test07_multiple_input_columns(self, load_eager): + df1 = from_files( + ROOT + .RDataFrame(5) + .Define("b1", "(int) rdfentry_") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + df2 = from_files( + ROOT + .RDataFrame(5) + .Define("b1", "(int) rdfentry_ + 5") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 0.0, 1.0, 10.0, 2.0, 20.0, 5.0, 50.0, 6.0, 60.0, 7.0, 70.0] + results_x_val = [3.0, 30.0, 4.0, 40.0, 8.0, 80.0, 9.0, 90.0] + results_y_train = [0.0, 1.0, 4.0, 25.0, 36.0, 49.0] + results_y_val = [9.0, 16.0, 64.0, 81.0] + else: + results_x_train = [ + 0.0, + 0.0, + 1.0, + 10.0, + 2.0, + 20.0, + 3.0, + 30.0, + 4.0, + 40.0, + 5.0, + 50.0, + ] + results_x_val = [6.0, 60.0, 7.0, 70.0, 8.0, 80.0, 9.0, 90.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] + results_y_val = [36.0, 49.0, 64.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 2) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 2) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 2) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test08_filtered(self, df1, df2, load_eager): + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 != 0", "name") + + dff1_entries_before = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] + dff2_entries_before = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=3 if load_eager else 4, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 2.0, 5.0] + results_x_val = [4.0, 9.0] + results_y_train = [0.0, 4.0, 25.0] + results_y_val = [16.0, 81.0] + else: + # dff1 keeps the even entries 0, 2, 4 and dff2 the odd entries 5, 7, 9. Unshuffled, the first + # 4 of these 6 entries (0, 2, 4, 5) go to training and the last 2 (7, 9) to validation + results_x_train = [0.0, 2.0, 4.0, 5.0] + results_x_val = [7.0, 9.0] + results_y_train = [0.0, 4.0, 16.0, 25.0] + results_y_val = [49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + x, y = next(train_iter) + assert x.shape == (len(results_x_train), 1) + assert y.shape == (len(results_y_train), 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] + dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(dff1_entries_before, dff1_entries_after) + assert np.array_equal(dff2_entries_before, dff2_entries_after) + + def test09_filtered_last_chunk(self, load_eager): + df1 = from_files(ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Define("b2", "(UShort_t) b1 * b1")) + df2 = from_files(ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_ + 10").Define("b2", "(UShort_t) b1 * b1")) + + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 == 0", "name") + + gen_train = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + if load_eager: + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] + else: + results_x_train = [0.0, 2.0, 4.0, 10.0, 12.0, 14.0, 6.0, 8.0, 16.0, 18.0] + results_y_train = [0.0, 4.0, 16.0, 100.0, 144.0, 196.0, 36.0, 64.0, 256.0, 324.0] + + collected_x_train = [] + collected_y_train = [] + + train_iter = iter(gen_train.as_numpy()) + + for _ in range(3): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (1, 1) + assert y.shape == (1, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train + + def test10_two_epochs_shuffled(self, df1, df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] + + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + both_epochs_collected_x_val.append(collected_x_val) + both_epochs_collected_y_val.append(collected_y_val) + + assert both_epochs_collected_x_val[0] == both_epochs_collected_x_val[1] + assert both_epochs_collected_y_val[0] == both_epochs_collected_y_val[1] + + def test11_number_of_training_and_validation_batches_remainder(self, df1, df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + number_of_training_batches = 0 + number_of_validation_batches = 0 + + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 + + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 + + assert gen_train.num_batches == number_of_training_batches + assert gen_validation.num_batches == number_of_validation_batches + assert gen_train.last_batch_no_of_rows == 0 + assert gen_validation.last_batch_no_of_rows == 1 + + def test12_PyTorch(self, weighted_df1, weighted_df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df1, weighted_df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + results_x_val = [3.0, 4.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] + results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] + results_z_val = [30.0, 40.0, 80.0, 90.0] + else: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_torch()) + iter_val = iter(gen_validation.as_torch(device="cpu")) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 2) + assert z.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test13_TensorFlow(self, weighted_df1, weighted_df2, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df1, weighted_df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + results_x_val = [3.0, 4.0, 8.0, 9.0, 0.0, 0.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 25.0, 500.0, 36.0, 600.0, 49.0, 700.0] + results_y_val = [9.0, 300.0, 16.0, 400.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] + results_z_train = [0.0, 10.0, 20.0, 50.0, 60.0, 70.0] + results_z_val = [30.0, 40.0, 80.0, 90.0, 0.0, 0.0] + else: + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0, 0.0, 0.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0, 0.0, 0.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test14_big_data(self, load_eager): + entries_in_rdf = randrange(10000, 30000) + batch_size = randrange(100, 501) + batches_in_memory = randrange(5, 20) + + error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" + + def test(size_of_batch, batches_in_memory, num_of_entries): + df1 = from_files(big_data(num_of_entries)) + df2 = from_files(big_data(num_of_entries)) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=size_of_batch, + batches_in_memory=batches_in_memory, + target=["b3", "b5"], + weights="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.3) + + collect_x = [] + + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows + + n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches + n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for i in range(n_train_batches): + x, y, z = next(iter_train) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + + assert np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}" + assert np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}" + # assert np.all(x[:,0]*(-1)-10192==y[:,1]), error_message + assert np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}" + + collect_x.extend(list(x[:, 0])) + + if train_remainder: + x, y, z = next(iter_train) + assert x.shape == (train_remainder, 2), error_message + assert y.shape == (train_remainder, 2), error_message + assert z.shape == (train_remainder, 1), error_message + collect_x.extend(list(x[:, 0])) + + for _ in range(n_val_batches): + x, y, z = next(iter_val) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + + assert np.all(x[:, 0] * (-1) == x[:, 1]), error_message + assert np.all(x[:, 0] + 10192 == y[:, 0]), error_message + # assert np.all(x[:,0]*(-1)-10192==y[:,1]), error_message + assert np.all(x[:, 0] * 2 == z[:, 0]), error_message + + collect_x.extend(list(x[:, 0])) + + if val_remainder: + x, y, z = next(iter_val) + assert x.shape == (val_remainder, 2), error_message + assert y.shape == (val_remainder, 2), error_message + assert z.shape == (val_remainder, 1), error_message + collect_x.extend(list(x[:, 0])) + + assert set(collect_x) == set(i for i in range(num_of_entries)), ( + f"collected length: {len(set(collect_x))}\ + generated length {len(set(i for i in range(num_of_entries)))}" + ) + + test(batch_size, batches_in_memory, entries_in_rdf) + + def test15_two_runs_set_seed(self, df1, df2, load_eager): + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] + + for _ in range(2): + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=42, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + both_runs_collected_x_val.append(collected_x_val) + both_runs_collected_y_val.append(collected_y_val) + assert both_runs_collected_x_val[0] == both_runs_collected_x_val[1] + assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] + + def test16_vector_padding(self, load_eager): + df1 = from_files(vectors(5)) + df2 = from_files(vectors(5, "(int) rdfentry_ + 5")) + max_vec_sizes = {"v1": 3, "v2": 2} + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b1", + max_vec_sizes=max_vec_sizes, + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + if load_eager: + results_x_train = [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 10.0, + 0, + 100.0, + 1000.0, + 2.0, + 20.0, + 0, + 200.0, + 2000.0, + 5.0, + 50.0, + 0, + 500.0, + 5000.0, + 6.0, + 60.0, + 0.0, + 600.0, + 6000.0, + 7.0, + 70.0, + 0.0, + 700.0, + 7000.0, + ] + results_y_train = [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + results_x_val = [ + 3.0, + 30.0, + 0.0, + 300.0, + 3000.0, + 4.0, + 40.0, + 0.0, + 400.0, + 4000.0, + 8.0, + 80.0, + 0.0, + 800.0, + 8000.0, + 9.0, + 90.0, + 0.0, + 900.0, + 9000.0, + ] + results_y_val = [3.0, 4.0, 8.0, 9.0] + else: + results_x_train = [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 10.0, + 0, + 100.0, + 1000.0, + 2.0, + 20.0, + 0, + 200.0, + 2000.0, + 3.0, + 30.0, + 0, + 300.0, + 3000.0, + 4.0, + 40.0, + 0.0, + 400.0, + 4000.0, + 5.0, + 50.0, + 0.0, + 500.0, + 5000.0, + ] + results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [ + 6.0, + 60.0, + 0.0, + 600.0, + 6000.0, + 7.0, + 70.0, + 0.0, + 700.0, + 7000.0, + 8.0, + 80.0, + 0.0, + 800.0, + 8000.0, + 9.0, + 90.0, + 0.0, + 900.0, + 9000.0, + ] + results_y_val = [6.0, 7.0, 8.0, 9.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 5) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 5) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 5) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test17_shuffled_split_varies_with_seed(self, df1, df2): + dl1 = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=42, + ) + + dl2 = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=43, + ) + + _, gen_val1 = dl1.train_test_split(0.4) + _, gen_val2 = dl2.train_test_split(0.4) + + val_1_collected = sorted([v for x, y in gen_val1.as_numpy() for v in x.flatten().tolist()]) + val_2_collected = sorted([v for x, y in gen_val2.as_numpy() for v in x.flatten().tolist()]) + + assert val_1_collected != val_2_collected + + def test17_dataframes_without_cluster_metadata(self): + # the dataframes are not read from a file, so their entries are counted with Count() + dl = ROOT.Experimental.ML.RDataLoader( + [squares(5), squares(5, "(int) rdfentry_ + 5")], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) + gen_train, gen_validation = dl.train_test_split(0.4) + + x_train = [v for x, _ in gen_train.as_numpy() for v in x.flatten().tolist()] + x_val = [v for x, _ in gen_validation.as_numpy() for v in x.flatten().tolist()] + assert x_train == [0.0, 1.0, 2.0, 5.0, 6.0, 7.0] + assert x_val == [3.0, 4.0, 8.0, 9.0] + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_oversampling.py b/bindings/pyroot/pythonizations/test/ml_dataloader_oversampling.py new file mode 100644 index 0000000000000..f6cc20af39f8b --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_oversampling.py @@ -0,0 +1,1164 @@ +from random import randrange, uniform + +import numpy as np +import pytest +import ROOT +from conftest import big_data, from_files, squares, vectors, weighted + + +class TestDataLoaderRandomOversampling: + # default constants + n_train_batch = 4 + n_val_batch = 3 + train_remainder = 1 + + @pytest.fixture + def df_major(self): + """The majority class, 10 entries: b1 = 0, 2, ..., 18 and b2 = b1 * b1.""" + return from_files(squares(10, "(int) 2 * rdfentry_")) + + @pytest.fixture + def df_minor(self): + """The minority class, 3 entries: b1 = 1, 3, ..., 5 and b2 = b1 * b1.""" + return from_files(squares(3, "(int) 2 * rdfentry_ + 1")) + + @pytest.fixture + def weighted_df_major(self): + """The majority class, 10 entries: b1 = 0, 2, ..., 18, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(10, "(int) 2 * rdfentry_")) + + @pytest.fixture + def weighted_df_minor(self): + """The minority class, 3 entries: b1 = 1, 3, ..., 5, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(3, "(int) 2 * rdfentry_ + 1")) + + def test01_each_element_is_generated_unshuffled(self, df_major, df_minor): + major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] + results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + # check if there are no duplicate entries (oversampling) + assert len(set(flat_x_train)) < len(flat_x_train) + assert len(set(flat_x_val)) < len(flat_x_val) + assert len(set(flat_y_train)) < len(flat_y_train) + assert len(set(flat_y_val)) < len(flat_y_val) + + # check if correct sampling_ratio (0.5 = minor/major) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) + + def test02_each_element_is_generated_shuffled(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + # check if there are no duplicate entries (oversampling) + assert len(set(flat_x_train)) < len(flat_x_train) + assert len(set(flat_x_val)) < len(flat_x_val) + assert len(set(flat_y_train)) < len(flat_y_train) + assert len(set(flat_y_val)) < len(flat_y_val) + + # check if correct sampling_ratio (0.5 = minor/major) + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + def test04_dropping_remainder(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + collected_x = [] + collected_y = [] + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x.append(x) + collected_y.append(y) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x.append(x) + collected_y.append(y) + + assert len(collected_x) == 7 + assert len(collected_y) == 7 + + def test05_more_than_one_file(self, df_major): + df_minor = from_files(squares(3, "(int) 2 * rdfentry_ + 1"), squares(1, "(int) 2 * (rdfentry_ + 3) + 1")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 5.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + results_x_val = [7.0, 7.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [1.0, 9.0, 25.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] + results_y_val = [49.0, 49.0, 144.0, 196.0, 256.0, 324.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test06_multiple_target_columns(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_minor, weighted_df_major], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [ + 1.0, + 100.0, + 9.0, + 300.0, + 1.0, + 100.0, + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + ] + results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] + results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0] + results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 2) + assert z.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test07_multiple_input_columns(self): + df_major = from_files( + ROOT + .RDataFrame(10) + .Define("b1", "(int) 2 * rdfentry_") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + df_minor = from_files( + ROOT + .RDataFrame(3) + .Define("b1", "(int) 2 * rdfentry_ + 1") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 1.0, + 10.0, + 3.0, + 30.0, + 1.0, + 10.0, + 0.0, + 0.0, + 2.0, + 20.0, + 4.0, + 40.0, + 6.0, + 60.0, + 8.0, + 80.0, + 10.0, + 100.0, + ] + results_x_val = [5.0, 50.0, 5.0, 50.0, 12.0, 120.0, 14.0, 140.0, 16.0, 160.0, 18.0, 180.0] + results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] + results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 2) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 2) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 2) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test08_filtered(self, df_major): + minor = squares(3, "(int) 2 * rdfentry_ + 1") + df_minor_duplicate = from_files(minor, minor, minor) + + df_minor_filter = df_minor_duplicate.Filter("rdfentry_ < 3", "name") + + major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_filter_entries_before = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor_filter], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [1.0, 9.0, 1.0, 0.0, 4.0, 16.0, 36.0, 64.0, 100.0] + results_y_val = [25.0, 25.0, 144.0, 196.0, 256.0, 324.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_filter_entries_after = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_filter_entries_before, minor_filter_entries_after) + + def test10_two_epochs_shuffled(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] + + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + both_epochs_collected_x_val.append(collected_x_val) + both_epochs_collected_y_val.append(collected_y_val) + + assert both_epochs_collected_x_val[0] == both_epochs_collected_x_val[1] + assert both_epochs_collected_y_val[0] == both_epochs_collected_y_val[1] + + def test11_number_of_training_and_validation_batches_remainder(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + number_of_training_batches = 0 + number_of_validation_batches = 0 + + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 + + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 + + assert gen_train.num_batches == number_of_training_batches + assert gen_validation.num_batches == number_of_validation_batches + assert gen_train.last_batch_no_of_rows == 1 + assert gen_validation.last_batch_no_of_rows == 0 + + def test12_PyTorch(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_major, weighted_df_minor], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [ + 1.0, + 100.0, + 9.0, + 300.0, + 1.0, + 100.0, + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + ] + results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] + results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0] + results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_torch()) + iter_val = iter(gen_validation.as_torch(device="cpu")) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 2) + assert z.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test13_TensorFlow(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_major, weighted_df_minor], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [1.0, 3.0, 1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 0.0] + results_x_val = [5.0, 5.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [ + 1.0, + 100.0, + 9.0, + 300.0, + 1.0, + 100.0, + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + 0.0, + 0.0, + ] + results_y_val = [25.0, 500.0, 25.0, 500.0, 144.0, 1200.0, 196.0, 1400.0, 256.0, 1600.0, 324.0, 1800.0] + results_z_train = [10.0, 30.0, 10.0, 0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 0.0] + results_z_val = [50.0, 50.0, 120.0, 140.0, 160.0, 180.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 7 # num_major_train = 6 + 1 padded 0 from the last batch + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test14_big_data_replacement_false(self): + entries_in_rdf_major = randrange(10000, 30000) + entries_in_rdf_minor = randrange(8000, 9999) + batch_size = randrange(100, 501) + sampling_ratio = round(uniform(0.1, 2), 2) + + error_message = f"\n Batch size: {batch_size}\ + Number of major entries: {entries_in_rdf_major} \ + Number of minor entries: {entries_in_rdf_minor}" + + def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): + df1 = from_files(big_data(num_of_entries_major, "(int) 2 * rdfentry_")) + df2 = from_files(big_data(num_of_entries_minor, "(int) 2 * rdfentry_ + 1")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=size_of_batch, + target=["b3", "b5"], + weights="b1", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=sampling_ratio, + ) + + gen_train, gen_validation = dl.train_test_split(0.3) + + collected_z_train = [] + collected_z_val = [] + + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows + + n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches + n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for i in range(n_train_batches): + x, y, z = next(iter_train) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_train.append(z.tolist()) + + if train_remainder: + x, y, z = next(iter_train) + assert x.shape == (train_remainder, 2), error_message + assert y.shape == (train_remainder, 2), error_message + assert z.shape == (train_remainder, 1), error_message + collected_z_train.append(z.tolist()) + + for _ in range(n_val_batches): + x, y, z = next(iter_val) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_val.append(z.tolist()) + + if val_remainder: + x, y, z = next(iter_val) + assert x.shape == (val_remainder, 2), error_message + assert y.shape == (val_remainder, 2), error_message + assert z.shape == (val_remainder, 1), error_message + collected_z_val.append(z.tolist()) + + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + num_major_train = sum(np.array(flat_z_train) % 2 == 0) + num_minor_train = sum(np.array(flat_z_train) % 2 != 0) + num_major_val = sum(np.array(flat_z_val) % 2 == 0) + num_minor_val = sum(np.array(flat_z_val) % 2 != 0) + + # check if there are no duplicate entries (replacement=False) + assert len(set(flat_z_train)) <= len(flat_z_train) + assert len(set(flat_z_val)) <= len(flat_z_val) + + # check if the sampling stategy is correct + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio + + test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) + + def test15_two_runs_set_seed(self, df_major, df_minor): + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] + + for _ in range(2): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + both_runs_collected_x_val.append(collected_x_val) + both_runs_collected_y_val.append(collected_y_val) + assert both_runs_collected_x_val[0] == both_runs_collected_x_val[1] + assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] + + def test16_vector_padding(self): + df_major = from_files(vectors(10)) + df_minor = from_files(vectors(3, "(int) rdfentry_ + 10")) + max_vec_sizes = {"v1": 3, "v2": 2} + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b1", + max_vec_sizes=max_vec_sizes, + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 10.0, + 100.0, + 0.0, + 1000.0, + 10000.0, + 11.0, + 110.0, + 0.0, + 1100.0, + 11000.0, + 10.0, + 100.0, + 0.0, + 1000.0, + 10000.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 10.0, + 0.0, + 100.0, + 1000.0, + 2.0, + 20.0, + 0.0, + 200.0, + 2000.0, + 3.0, + 30.0, + 0.0, + 300.0, + 3000.0, + 4.0, + 40.0, + 0.0, + 400.0, + 4000.0, + 5.0, + 50.0, + 0.0, + 500.0, + 5000.0, + ] + results_x_val = [ + 12.0, + 120.0, + 0.0, + 1200.0, + 12000.0, + 12.0, + 120.0, + 0.0, + 1200.0, + 12000.0, + 6.0, + 60.0, + 0.0, + 600.0, + 6000.0, + 7.0, + 70.0, + 0.0, + 700.0, + 7000.0, + 8.0, + 80.0, + 0.0, + 800.0, + 8000.0, + 9.0, + 90.0, + 0.0, + 900.0, + 9000.0, + ] + results_y_train = [10.0, 11.0, 10.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_y_val = [12.0, 12.0, 6.0, 7.0, 8.0, 9.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 5) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 5) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 5) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_y_train) < 10) + num_minor_train = sum(np.array(flat_y_train) >= 10) + num_major_val = sum(np.array(flat_y_val) < 10) + num_minor_val = sum(np.array(flat_y_val) >= 10) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test17_set_seed_resamples_different_entries(self, df_major, df_minor): + # with a fixed seed, the resampled minority entries used to all be the same one + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=42, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=1.0, + ) + + x = np.concatenate([x for x, _ in dl.as_numpy()]).flatten() + minor = x[x % 2 == 1] + + assert len(minor) == 10 + assert len(np.unique(minor)) > 1 + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_save.py b/bindings/pyroot/pythonizations/test/ml_dataloader_save.py new file mode 100644 index 0000000000000..6176d8f8ab574 --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_save.py @@ -0,0 +1,210 @@ +import numpy as np +import pytest +import ROOT +from conftest import from_files + + +class TestDataLoaderSave: + tree_name = "mytree" + + out_name = "dataloader_save_output.root" + out_name2 = "dataloader_save_output2.root" + + @pytest.fixture + def df(self): + """20 entries: b1 = 0, ..., 19 and b2 = b1 + 1.""" + return from_files(self.define_rdf()) + + # Helpers + def define_rdf(self, num_of_entries=20, offset=0): + return ROOT.RDataFrame(num_of_entries).Define("b1", f"(int) rdfentry_ + {offset}").Define("b2", "b1 + 1") + + def read_back(self, file_name, column): + """Read back one column of a saved file: a 1D array, or a 2D array for a vector column.""" + values = ROOT.RDataFrame(self.tree_name, file_name).AsNumpy([column])[column] + # a vector column comes back as one RVec per entry, all of the same padded size + return np.stack([np.asarray(v) for v in values]) if values.dtype == object else values + + def saved_columns(self, file_name): + return {str(c) for c in ROOT.RDataFrame(self.tree_name, file_name).GetColumnNames()} + + def test01_unshuffled(self, df): + df = df.Define("w", "(double) 0.5") + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=5, + batches_in_memory=2, + columns=["b1", "b2", "w"], + target="b2", + weights="w", + shuffle=False, + drop_remainder=False, + ) + + dl.save(self.tree_name, self.out_name) + + # check all features, target and weights + assert self.saved_columns(self.out_name) == {"b1", "b2", "w"} + np.testing.assert_array_equal(self.read_back(self.out_name, "b1"), np.arange(20)) + np.testing.assert_array_equal(self.read_back(self.out_name, "b2"), np.arange(1, 21)) + np.testing.assert_array_equal(self.read_back(self.out_name, "w"), np.full(20, 0.5)) + + def test02_shuffled(self, df): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=4, + batches_in_memory=2, + columns=["b1"], + shuffle=True, + set_seed=42, + drop_remainder=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + gen_train.save(self.tree_name, self.out_name) + gen_validation.save(self.tree_name, self.out_name2) + + train_b1 = self.read_back(self.out_name, "b1") + val_b1 = self.read_back(self.out_name2, "b1") + + # the two splits cover the dataset exactly once + all_b1 = np.concatenate([train_b1, val_b1]) + np.testing.assert_array_equal(np.sort(all_b1), np.arange(20)) + + # and the entries are not written in their original order + assert (train_b1 != np.sort(train_b1)).any() + + def test03_vector_columns(self, df): + df = df.Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=5, + batches_in_memory=2, + columns=["b1", "v1"], + max_vec_sizes={"v1": 3}, + vec_padding=42.0, + shuffle=False, + drop_remainder=False, + ) + + dl.save(self.tree_name, self.out_name) + + # v1 is saved as a vector column, not exploded into v1_0... + assert self.saved_columns(self.out_name) == {"b1", "v1"} + + v1 = self.read_back(self.out_name, "v1") + assert v1.shape == (20, 3) + # v1 contains 2 elements padded up to max_vec_sizes["v1"] (3) with vec_padding (42) + np.testing.assert_array_equal(v1[3], [3.0, 30.0, 42.0]) + + def test04_ttree_and_rntuple_agree(self): + def save_as(file_name, output_format): + df = from_files(self.define_rdf()).Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=5, + batches_in_memory=2, + columns=["b1", "v1"], + max_vec_sizes={"v1": 3}, + shuffle=False, + drop_remainder=False, + ) + + dl.save(self.tree_name, file_name, output_format=output_format) + + save_as(self.out_name, "ttree") + save_as(self.out_name2, "rntuple") + + for column in ["b1", "v1"]: + np.testing.assert_array_equal( + self.read_back(self.out_name, column), + self.read_back(self.out_name2, column), + ) + + def test05_filtered(self, df): + df = df.Filter("b1 % 2 == 0", "even") + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=2, + batches_in_memory=2, + columns=["b1"], + shuffle=False, + drop_remainder=False, + ) + + dl.save(self.tree_name, self.out_name) + + saved_b1 = self.read_back(self.out_name, "b1") + np.testing.assert_array_equal(np.sort(saved_b1), np.arange(0, 20, 2)) + + def test06_multiple_dataframes(self): + df1 = from_files(self.define_rdf(num_of_entries=10)) + df2 = from_files(self.define_rdf(num_of_entries=10, offset=100)) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=5, + batches_in_memory=2, + columns=["b1"], + shuffle=False, + drop_remainder=False, + ) + + dl.save(self.tree_name, self.out_name) + + saved_b1 = self.read_back(self.out_name, "b1") + expected = np.concatenate([np.arange(10), np.arange(100, 110)]) + np.testing.assert_array_equal(np.sort(saved_b1), expected) + + def test07_eager_loading(self, df): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=5, + columns=["b1", "b2"], + shuffle=False, + drop_remainder=False, + load_eager=True, + ) + + dl.save(self.tree_name, self.out_name) + + np.testing.assert_array_equal(self.read_back(self.out_name, "b1"), np.arange(20)) + np.testing.assert_array_equal(self.read_back(self.out_name, "b2"), np.arange(1, 21)) + + def test08_round_trip(self, df): + df = df.Define("w", "(double) 0.5").Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10 }") + + loader_args = dict( + batch_size=5, + batches_in_memory=2, + columns=["b1", "v1", "b2", "w"], + max_vec_sizes={"v1": 3}, + vec_padding=42.0, + target="b2", + weights="w", + shuffle=False, + drop_remainder=False, + ) + + dl = ROOT.Experimental.ML.RDataLoader(df, **loader_args) + dl.save(self.tree_name, self.out_name) + + # the saved file is fed straight back into a new RDataLoader with the same columns/target + reloaded = ROOT.Experimental.ML.RDataLoader(ROOT.RDataFrame(self.tree_name, self.out_name), **loader_args) + + x, y, w = (np.concatenate(arrs) for arrs in zip(*reloaded.as_numpy())) + np.testing.assert_array_equal(x[:, 0], np.arange(20)) # b1 + np.testing.assert_array_equal(x[:, 1], np.arange(20)) # v1[0] == b1 + np.testing.assert_array_equal(x[:, 2], np.arange(20) * 10) # v1[1] == b1 * 10 + np.testing.assert_array_equal(x[:, 3], np.full(20, 42.0)) # v1[2], padding + np.testing.assert_array_equal(y.flatten(), np.arange(1, 21)) + np.testing.assert_array_equal(w.flatten(), np.full(20, 0.5)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_single_dataframe.py b/bindings/pyroot/pythonizations/test/ml_dataloader_single_dataframe.py new file mode 100644 index 0000000000000..6ace3563922cf --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_single_dataframe.py @@ -0,0 +1,1060 @@ +from random import randrange + +import numpy as np +import pytest +import ROOT +from conftest import big_data, from_files, squares, vectors + + +class TestDataLoaderSingleDataframe: + # default constants + n_train_batch = 2 + n_val_batch = 1 + val_remainder = 1 + + @pytest.fixture + def df(self): + """10 entries: b1 = 0, ..., 9 and b2 = b1 * b1.""" + return from_files(squares(10)) + + @pytest.fixture + def weighted_df(self): + """10 entries: the Short_t b1 = 0, ..., 9, the UShort_t b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files( + ROOT + .RDataFrame(10) + .Define("b1", "(Short_t) rdfentry_") + .Define("b2", "(UShort_t) b1 * b1") + .Define("b3", "(double) rdfentry_ * 10") + .Define("b4", "(double) b3 * 10") + ) + + def test01_each_element_is_generated_unshuffled(self, df, load_eager): + entries_before = df.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] + results_y_val = [36.0, 49.0, 64.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + entries_after = df.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframe is correctly reset + assert np.array_equal(entries_before, entries_after) + + def test02_each_element_is_generated_shuffled(self, df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=True, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + def test04_dropping_remainder(self, df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=True, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x = [] + collected_y = [] + + for x, y in gen_train.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x.append(x) + collected_y.append(y) + + for x, y in gen_validation.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x.append(x) + collected_y.append(y) + + assert len(collected_x) == 3 + assert len(collected_y) == 3 + + def test05_more_than_one_file(self, load_eager): + df = from_files(squares(10), squares(5, "(int) rdfentry_ + 10")) + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + results_x_val = [9.0, 10.0, 11.0, 12.0, 13.0, 14.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0] + results_y_val = [81.0, 100.0, 121.0, 144.0, 169.0, 196.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + for x, y in gen_train.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for x, y in gen_validation.as_numpy(): + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test06_multiple_target_columns(self, weighted_df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + weighted_df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 2) + assert z.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test07_multiple_input_columns(self, load_eager): + df = from_files( + ROOT + .RDataFrame(10) + .Define("b1", "(Short_t) rdfentry_") + .Define("b2", "(UShort_t) b1 * b1") + .Define("b3", "(double) rdfentry_ * 10") + ) + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 0.0, 1.0, 10.0, 2.0, 20.0, 3.0, 30.0, 4.0, 40.0, 5.0, 50.0] + results_x_val = [6.0, 60.0, 7.0, 70.0, 8.0, 80.0, 9.0, 90.0] + results_y_train = [0.0, 1.0, 4.0, 9.0, 16.0, 25.0] + results_y_val = [36.0, 49.0, 64.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 2) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 2) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 2) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test08_filtered(self, df, load_eager): + dff = df.Filter("b1 % 2 == 0", "name") + + filter_entries_before = dff.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + dff, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0] + results_x_val = [6.0, 8.0] + results_y_train = [0.0, 4.0, 16.0] + results_y_val = [36.0, 64.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + filter_entries_after = dff.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframe is correctly reset + assert np.array_equal(filter_entries_before, filter_entries_after) + + def test09_filtered_last_chunk(self, load_eager): + df = from_files(ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1")) + + dff = df.Filter("b1 % 2 == 0", "name") + + gen_train = ROOT.Experimental.ML.RDataLoader( + dff, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0] + results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] + + collected_x_train = [] + collected_y_train = [] + + train_iter = iter(gen_train.as_numpy()) + + for _ in range(3): + x, y = next(train_iter) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (1, 1) + assert y.shape == (1, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train + + def test10_two_epochs_shuffled(self, df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] + + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + both_epochs_collected_x_val.append(collected_x_val) + both_epochs_collected_y_val.append(collected_y_val) + + assert both_epochs_collected_x_val[0] == both_epochs_collected_x_val[1] + assert both_epochs_collected_y_val[0] == both_epochs_collected_y_val[1] + + def test11_number_of_training_and_validation_batches_remainder(self, df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + number_of_training_batches = 0 + number_of_validation_batches = 0 + + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 + + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 + + assert gen_train.num_batches == number_of_training_batches + assert gen_validation.num_batches == number_of_validation_batches + assert gen_train.last_batch_no_of_rows == 0 + assert gen_validation.last_batch_no_of_rows == 1 + + def test12_PyTorch(self, weighted_df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + weighted_df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_torch(device="cpu")) + iter_val = iter(gen_validation.as_torch()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 2) + assert z.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test13_TensorFlow(self, weighted_df, load_eager): + dl = ROOT.Experimental.ML.RDataLoader( + weighted_df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0, 0.0, 0.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0, 0.0, 0.0, 0.0, 0.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0, 0.0, 0.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test14_big_data(self, load_eager): + entries_in_rdf = randrange(10000, 30000) + batch_size = randrange(100, 501) + batches_in_memory = randrange(5, 20) + + error_message = f"\n batch_size={batch_size} batches_in_memory={batches_in_memory} entries={entries_in_rdf}" + + def test(size_of_batch, batches_in_memory, num_of_entries): + df = from_files(big_data(num_of_entries)) + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=size_of_batch, + batches_in_memory=batches_in_memory, + target=["b3", "b5"], + weights="b2", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.3) + + collect_x = [] + + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows + + n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches + n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for i in range(n_train_batches): + x, y, z = next(iter_train) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + + assert np.all(x[:, 0] * (-1) == x[:, 1]), error_message + f" row: {i}" + assert np.all(x[:, 0] + 10192 == y[:, 0]), error_message + f" row: {i}" + # assert np.all(x[:,0]*(-1)-10192==y[:,1]), error_message + assert np.all(x[:, 0] * 2 == z[:, 0]), error_message + f" row: {i}" + + collect_x.extend(list(x[:, 0])) + + if train_remainder: + x, y, z = next(iter_train) + assert x.shape == (train_remainder, 2), error_message + assert y.shape == (train_remainder, 2), error_message + assert z.shape == (train_remainder, 1), error_message + collect_x.extend(list(x[:, 0])) + + for _ in range(n_val_batches): + x, y, z = next(iter_val) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + + assert np.all(x[:, 0] * (-1) == x[:, 1]), error_message + assert np.all(x[:, 0] + 10192 == y[:, 0]), error_message + # assert np.all(x[:,0]*(-1)-10192==y[:,1]), error_message + assert np.all(x[:, 0] * 2 == z[:, 0]), error_message + + collect_x.extend(list(x[:, 0])) + + if val_remainder: + x, y, z = next(iter_val) + assert x.shape == (val_remainder, 2), error_message + assert y.shape == (val_remainder, 2), error_message + assert z.shape == (val_remainder, 1), error_message + collect_x.extend(list(x[:, 0])) + + assert set(collect_x) == set(i for i in range(num_of_entries)), ( + f"collected length: {len(set(collect_x))}\ + generated length {len(set(i for i in range(num_of_entries)))}" + ) + + test(batch_size, batches_in_memory, entries_in_rdf) + + def test15_two_runs_set_seed(self, df, load_eager): + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] + + for _ in range(2): + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=42, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 + + both_runs_collected_x_val.append(collected_x_val) + both_runs_collected_y_val.append(collected_y_val) + assert both_runs_collected_x_val[0] == both_runs_collected_x_val[1] + assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] + + def test16_vector_padding(self, load_eager): + df = from_files(vectors(10)) + max_vec_sizes = {"v1": 3, "v2": 2} + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b1", + max_vec_sizes=max_vec_sizes, + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 10.0, + 0, + 100.0, + 1000.0, + 2.0, + 20.0, + 0, + 200.0, + 2000.0, + 3.0, + 30.0, + 0, + 300.0, + 3000.0, + 4.0, + 40.0, + 0, + 400.0, + 4000.0, + 5.0, + 50.0, + 0, + 500.0, + 5000.0, + ] + results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [ + 6.0, + 60.0, + 0.0, + 600.0, + 6000.0, + 7.0, + 70.0, + 0.0, + 700.0, + 7000.0, + 8.0, + 80.0, + 0.0, + 800.0, + 8000.0, + 9.0, + 90.0, + 0.0, + 900.0, + 9000.0, + ] + results_y_val = [6.0, 7.0, 8.0, 9.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (3, 5) + assert y.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (3, 5) + assert y.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(val_iter) + assert x.shape == (self.val_remainder, 5) + assert y.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test17_JAX(self, weighted_df): + dl = ROOT.Experimental.ML.RDataLoader( + weighted_df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] + results_x_val = [6.0, 7.0, 8.0, 9.0] + results_y_train = [0.0, 0.0, 1.0, 100.0, 4.0, 200.0, 9.0, 300.0, 16.0, 400.0, 25.0, 500.0] + results_y_val = [36.0, 600.0, 49.0, 700.0, 64.0, 800.0, 81.0, 900.0] + results_z_train = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0] + results_z_val = [60.0, 70.0, 80.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_jax(device="cpu")) + iter_val = iter(gen_validation.as_jax()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (3, 1) + assert y.shape == (3, 2) + assert z.shape == (3, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_val) + assert x.shape == (self.val_remainder, 1) + assert y.shape == (self.val_remainder, 2) + assert z.shape == (self.val_remainder, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + def test18_one_batch_in_memory(self, df): + # batches_in_memory=1 used to deadlock before delivering the first batch + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=2, + batches_in_memory=1, + target="b2", + shuffle=False, + drop_remainder=True, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + assert sum(1 for _ in gen_train.as_numpy()) == 3 + assert sum(1 for _ in gen_validation.as_numpy()) == 2 + + def test19_defined_columns(self, df, load_eager): + # columns defined after reading the file are loaded like the stored ones + + df = df.Define("b3", "b1 * 3").Define("v1", "ROOT::RVecD{b2, b2 + 1}") + + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + columns=["b1", "v1", "b3"], + max_vec_sizes={"v1": 2}, + target="b3", + shuffle=False, + drop_remainder=False, + load_eager=load_eager, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + batches = list(gen_train.as_numpy()) + list(gen_validation.as_numpy()) + x = np.concatenate([x for x, _ in batches]) + y = np.concatenate([y for _, y in batches]) + + entries = np.arange(10, dtype=np.float32) + assert np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1])) + assert np.array_equal(y, (entries * 3).reshape(-1, 1)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__])) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader_undersampling.py b/bindings/pyroot/pythonizations/test/ml_dataloader_undersampling.py new file mode 100644 index 0000000000000..0abf841c4da86 --- /dev/null +++ b/bindings/pyroot/pythonizations/test/ml_dataloader_undersampling.py @@ -0,0 +1,1373 @@ +import math +from random import randrange, uniform + +import numpy as np +import pytest +import ROOT +from conftest import big_data, from_files, squares, vectors, weighted + + +class TestDataLoaderRandomUndersampling: + # default constants + n_train_batch = 4 + n_val_batch = 3 + train_remainder = 1 + + @pytest.fixture + def df_major(self): + """The majority class, 20 entries: b1 = 0, 2, ..., 38 and b2 = b1 * b1.""" + return from_files(squares(20, "(int) 2 * rdfentry_")) + + @pytest.fixture + def df_minor(self): + """The minority class, 5 entries: b1 = 1, 3, ..., 9 and b2 = b1 * b1.""" + return from_files(squares(5, "(int) 2 * rdfentry_ + 1")) + + @pytest.fixture + def weighted_df_major(self): + """The majority class, 20 entries: b1 = 0, 2, ..., 38, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(20, "(int) 2 * rdfentry_")) + + @pytest.fixture + def weighted_df_minor(self): + """The minority class, 5 entries: b1 = 1, 3, ..., 9, b2 = b1 * b1, b3 = b1 * 10 and b4 = b3 * 10.""" + return from_files(weighted(5, "(int) 2 * rdfentry_ + 1")) + + def test01_each_element_is_generated_unshuffled(self, df_major, df_minor): + major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] + results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] + results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] + results_y_val = [576.0, 676.0, 784.0, 900.0, 49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + # check if there are no duplicate entries (replacement=False) + assert len(set(flat_x_train)) == len(flat_x_train) + assert len(set(flat_x_val)) == len(flat_x_val) + assert len(set(flat_y_train)) == len(flat_y_train) + assert len(set(flat_y_val)) == len(flat_y_val) + + # check if correct sampling_ratio (0.5 = minor/major) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) + + def test01_each_element_is_generated_unshuffled_replacement(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.2, + replacement=True, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 0.0, + 2.0, + 4.0, + 6.0, + 8.0, + 10.0, + 12.0, + 14.0, + 16.0, + 18.0, + 20.0, + 22.0, + 0.0, + 2.0, + 4.0, + 1.0, + 3.0, + 5.0, + ] + results_x_val = [24.0, 26.0, 28.0, 30.0, 32.0, 34.0, 36.0, 38.0, 24.0, 26.0, 7.0, 9.0] + results_y_train = [ + 0.0, + 4.0, + 16.0, + 36.0, + 64.0, + 100.0, + 144.0, + 196.0, + 256.0, + 324.0, + 400.0, + 484.0, + 0.0, + 4.0, + 16.0, + 1.0, + 9.0, + 25.0, + ] + results_y_val = [576.0, 676.0, 784.0, 900.0, 1024.0, 1156.0, 1296.0, 1444.0, 576.0, 676.0, 49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(6): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(9): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + # check if there are duplicate entries (replacement=True) + assert len(set(flat_x_train)) < len(flat_x_train) + assert len(set(flat_x_val)) < len(flat_x_val) + assert len(set(flat_y_train)) < len(flat_y_train) + assert len(set(flat_y_val)) < len(flat_y_val) + + # check if correct sampling_ratio (0.2 = minor/major) + assert num_major_train == 15 + assert num_minor_train == 3 + assert num_major_val == 10 + assert num_minor_val == 2 + + def test02_each_element_is_generated_shuffled(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + def test04_dropping_remainder(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + collected_x = [] + collected_y = [] + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x.append(x) + collected_y.append(y) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x.append(x) + collected_y.append(y) + + assert len(collected_x) == 7 + assert len(collected_y) == 7 + + def test05_more_than_one_file(self, df_minor): + df_major = from_files(squares(20, "(int) 2 * rdfentry_"), squares(5, "(int) 2 * (rdfentry_ + 20)")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] + results_x_val = [30.0, 32.0, 34.0, 36.0, 7.0, 9.0] + results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] + results_y_val = [900.0, 1024.0, 1156.0, 1296.0, 49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test06_multiple_target_columns(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_minor, weighted_df_major], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] + results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] + results_y_train = [ + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + 1.0, + 100.0, + 9.0, + 300.0, + 25.0, + 500.0, + ] + results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] + results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0] + results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 2) + assert z.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test07_multiple_input_columns(self): + df_major = from_files( + ROOT + .RDataFrame(20) + .Define("b1", "(int) 2 * rdfentry_") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + df_minor = from_files( + ROOT + .RDataFrame(5) + .Define("b1", "(int) 2 * rdfentry_ + 1") + .Define("b2", "(int) b1 * b1") + .Define("b3", "(double) b1 * 10") + ) + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 0.0, + 0.0, + 2.0, + 20.0, + 4.0, + 40.0, + 6.0, + 60.0, + 8.0, + 80.0, + 10.0, + 100.0, + 1.0, + 10.0, + 3.0, + 30.0, + 5.0, + 50.0, + ] + results_x_val = [24.0, 240.0, 26.0, 260.0, 28.0, 280.0, 30.0, 300.0, 7.0, 70.0, 9.0, 90.0] + results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 1.0, 9.0, 25.0] + results_y_val = [576.0, 676.0, 784.0, 900.0, 49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 2) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 2) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 2) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + def test08_filtered(self, df_major, df_minor): + df_major_filter = df_major.Filter("b1 % 8 != 0", "name") + + major_filter_entries_before = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major_filter, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [2.0, 4.0, 6.0, 10.0, 12.0, 14.0, 1.0, 3.0, 5.0] + results_x_val = [26.0, 28.0, 30.0, 34.0, 7.0, 9.0] + results_y_train = [4.0, 16.0, 36.0, 100.0, 144.0, 196.0, 1.0, 9.0, 25.0] + results_y_val = [676.0, 784.0, 900.0, 1156.0, 49.0, 81.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = np.sum((np.asarray(flat_x_train) % 2 == 0) & (np.asarray(flat_x_train) % 8 != 0)) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = np.sum((np.asarray(flat_x_val) % 2 == 0) & (np.asarray(flat_x_val) % 8 != 0)) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + major_filter_entries_after = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + + # check if the dataframes are correctly reset + assert np.array_equal(major_filter_entries_before, major_filter_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) + + def test10_two_epochs_shuffled(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] + + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + both_epochs_collected_x_val.append(collected_x_val) + both_epochs_collected_y_val.append(collected_y_val) + + assert both_epochs_collected_x_val[0] == both_epochs_collected_x_val[1] + assert both_epochs_collected_y_val[0] == both_epochs_collected_y_val[1] + + def test11_number_of_training_and_validation_batches_remainder(self, df_major, df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + number_of_training_batches = 0 + number_of_validation_batches = 0 + + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 + + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 + + assert gen_train.num_batches == number_of_training_batches + assert gen_validation.num_batches == number_of_validation_batches + assert gen_train.last_batch_no_of_rows == 1 + assert gen_validation.last_batch_no_of_rows == 0 + + def test12_PyTorch(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_major, weighted_df_minor], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0] + results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] + results_y_train = [ + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + 1.0, + 100.0, + 9.0, + 300.0, + 25.0, + 500.0, + ] + results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] + results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0] + results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_torch(device="cpu")) + iter_val = iter(gen_validation.as_torch()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + collected_z_val.append(z.tolist()) + + x, y, z = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 2) + assert z.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + collected_z_train.append(z.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test13_TensorFlow(self, weighted_df_major, weighted_df_minor): + dl = ROOT.Experimental.ML.RDataLoader( + [weighted_df_major, weighted_df_minor], + batch_size=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 1.0, 3.0, 5.0, 0.0] + results_x_val = [24.0, 26.0, 28.0, 30.0, 7.0, 9.0] + results_y_train = [ + 0.0, + 0.0, + 4.0, + 200.0, + 16.0, + 400.0, + 36.0, + 600.0, + 64.0, + 800.0, + 100.0, + 1000.0, + 1.0, + 100.0, + 9.0, + 300.0, + 25.0, + 500.0, + 0.0, + 0.0, + ] + results_y_val = [576.0, 2400.0, 676.0, 2600.0, 784.0, 2800.0, 900.0, 3000.0, 49.0, 700.0, 81.0, 900.0] + results_z_train = [0.0, 20.0, 40.0, 60.0, 80.0, 100.0, 10.0, 30.0, 50.0, 0.0] + results_z_val = [240.0, 260.0, 280.0, 300.0, 70.0, 90.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] + + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) + + for _ in range(self.n_train_batch): + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + for _ in range(self.n_val_batch): + x, y, z = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_val.append(x.numpy().tolist()) + collected_y_val.append(y.numpy().tolist()) + collected_z_val.append(z.numpy().tolist()) + + x, y, z = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 2) + assert z.shape == (2, 1) + collected_x_train.append(x.numpy().tolist()) + collected_y_train.append(y.numpy().tolist()) + collected_z_train.append(z.numpy().tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + assert results_z_train == flat_z_train + assert results_z_val == flat_z_val + + num_major_train = sum(np.array(flat_x_train) % 2 == 0) + num_minor_train = sum(np.array(flat_x_train) % 2 != 0) + num_major_val = sum(np.array(flat_x_val) % 2 == 0) + num_minor_val = sum(np.array(flat_x_val) % 2 != 0) + + assert num_major_train == 7 # num_major_train = 6 + 1 padded 0 from the last batch + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + def test14_big_data_replacement_false(self): + entries_in_rdf_major = randrange(10000, 30000) + entries_in_rdf_minor = randrange(8000, 9999) + batch_size = randrange(100, 501) + min_allowed_sampling_ratio = entries_in_rdf_minor / entries_in_rdf_major + sampling_ratio = math.ceil(uniform(min_allowed_sampling_ratio, 2) * 100) / 100 + + error_message = f"\n Batch size: {batch_size}\ + Number of major entries: {entries_in_rdf_major} \ + Number of minor entries: {entries_in_rdf_minor}" + + def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): + df1 = from_files(big_data(num_of_entries_major, "(int) 2 * rdfentry_")) + df2 = from_files(big_data(num_of_entries_minor, "(int) 2 * rdfentry_ + 1")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=size_of_batch, + target=["b3", "b5"], + weights="b1", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=sampling_ratio, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.3) + + collected_z_train = [] + collected_z_val = [] + + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows + + n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches + n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for i in range(n_train_batches): + x, y, z = next(iter_train) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_train.append(z.tolist()) + + if train_remainder: + x, y, z = next(iter_train) + assert x.shape == (train_remainder, 2), error_message + assert y.shape == (train_remainder, 2), error_message + assert z.shape == (train_remainder, 1), error_message + collected_z_train.append(z.tolist()) + + for _ in range(n_val_batches): + x, y, z = next(iter_val) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_val.append(z.tolist()) + + if val_remainder: + x, y, z = next(iter_val) + assert x.shape == (val_remainder, 2), error_message + assert y.shape == (val_remainder, 2), error_message + assert z.shape == (val_remainder, 1), error_message + collected_z_val.append(z.tolist()) + + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + num_major_train = sum(np.array(flat_z_train) % 2 == 0) + num_minor_train = sum(np.array(flat_z_train) % 2 != 0) + num_major_val = sum(np.array(flat_z_val) % 2 == 0) + num_minor_val = sum(np.array(flat_z_val) % 2 != 0) + + # check if there are no duplicate entries (replacement=False) + assert len(set(flat_z_train)) == len(flat_z_train) + assert len(set(flat_z_val)) == len(flat_z_val) + + # check if the sampling stategy is correct + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio + + def test_raises(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): + df1 = from_files(big_data(num_of_entries_major, "(int) 2 * rdfentry_")) + df2 = from_files(big_data(num_of_entries_minor, "(int) 2 * rdfentry_ + 1")) + + with pytest.raises( + Exception, + match=r"The sampling_ratio is too low: not enough entries in the majority class to sample from.", + ): + ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=size_of_batch, + target=["b3", "b5"], + weights="b1", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=sampling_ratio, + replacement=False, + ).train_test_split(0.3) + + # test the functionality with a proper sampling ratio + test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) + + bad_sampling_ratio = round(max(min_allowed_sampling_ratio - 0.01, 0.01), 2) + # test that an error is raised when the sampling ratio is too low + test_raises(batch_size, entries_in_rdf_major, entries_in_rdf_minor, bad_sampling_ratio) + + def test14_big_data_replacement_true(self): + entries_in_rdf_major = randrange(10000, 30000) + entries_in_rdf_minor = randrange(8000, 9999) + batch_size = randrange(100, 501) + + # max samling strategy to guarantee duplicate sampled entires + max_sampling_ratio = entries_in_rdf_minor / entries_in_rdf_major + sampling_ratio = round(uniform(0.01, max_sampling_ratio - 0.01), 2) + + error_message = f"\n Batch size: {batch_size}\ + Number of major entries: {entries_in_rdf_major} \ + Number of minor entries: {entries_in_rdf_minor}" + + def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_ratio): + df1 = from_files(big_data(num_of_entries_major, "(int) 2 * rdfentry_")) + df2 = from_files(big_data(num_of_entries_minor, "(int) 2 * rdfentry_ + 1")) + + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=size_of_batch, + target=["b3", "b5"], + weights="b1", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=sampling_ratio, + replacement=True, + ) + + gen_train, gen_validation = dl.train_test_split(0.3) + + collected_z_train = [] + collected_z_val = [] + + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows + + n_train_batches = gen_train.num_batches - 1 if train_remainder else gen_train.num_batches + n_val_batches = gen_validation.num_batches - 1 if val_remainder else gen_validation.num_batches + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for i in range(n_train_batches): + x, y, z = next(iter_train) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_train.append(z.tolist()) + + if train_remainder: + x, y, z = next(iter_train) + assert x.shape == (train_remainder, 2), error_message + assert y.shape == (train_remainder, 2), error_message + assert z.shape == (train_remainder, 1), error_message + collected_z_train.append(z.tolist()) + + for _ in range(n_val_batches): + x, y, z = next(iter_val) + + assert x.shape == (size_of_batch, 2), error_message + f" row: {i} x shape: {x.shape}" + assert y.shape == (size_of_batch, 2), error_message + f" row: {i} y shape: {y.shape}" + assert z.shape == (size_of_batch, 1), error_message + f" row: {i} z shape: {z.shape}" + collected_z_val.append(z.tolist()) + + if val_remainder: + x, y, z = next(iter_val) + assert x.shape == (val_remainder, 2), error_message + assert y.shape == (val_remainder, 2), error_message + assert z.shape == (val_remainder, 1), error_message + collected_z_val.append(z.tolist()) + + flat_z_train = [z for zl in collected_z_train for zs in zl for z in zs] + flat_z_val = [z for zl in collected_z_val for zs in zl for z in zs] + + num_major_train = sum(np.array(flat_z_train) % 2 == 0) + num_minor_train = sum(np.array(flat_z_train) % 2 != 0) + num_major_val = sum(np.array(flat_z_val) % 2 == 0) + num_minor_val = sum(np.array(flat_z_val) % 2 != 0) + + # check if there are duplicate entries (replacement=True) + assert len(set(flat_z_train)) < len(flat_z_train) + assert len(set(flat_z_val)) < len(flat_z_val) + + # check if the sampling stategy is correct + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio + + test(batch_size, entries_in_rdf_major, entries_in_rdf_minor, sampling_ratio) + + def test15_two_runs_set_seed(self, df_major, df_minor): + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] + + for _ in range(2): + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) + + for _ in range(self.n_train_batch): + x, y = next(iter_train) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + for _ in range(self.n_val_batch): + x, y = next(iter_val) + assert x.shape == (2, 1) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + x, y = next(iter_train) + assert x.shape == (self.train_remainder, 1) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = {x for xl in collected_x_train for xs in xl for x in xs} + flat_x_val = {x for xl in collected_x_val for xs in xl for x in xs} + flat_y_train = {y for yl in collected_y_train for ys in yl for y in ys} + flat_y_val = {y for yl in collected_y_val for ys in yl for y in ys} + + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 + + both_runs_collected_x_val.append(collected_x_val) + both_runs_collected_y_val.append(collected_y_val) + assert both_runs_collected_x_val[0] == both_runs_collected_x_val[1] + assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] + + def test16_vector_padding(self): + df_major = from_files(vectors(20)) + df_minor = from_files(vectors(5, "(int) rdfentry_ + 20")) + max_vec_sizes = {"v1": 3, "v2": 2} + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b1", + max_vec_sizes=max_vec_sizes, + shuffle=False, + drop_remainder=False, + load_eager=True, + sampling_type="undersampling", + sampling_ratio=0.5, + replacement=False, + ) + + gen_train, gen_validation = dl.train_test_split(0.4) + + results_x_train = [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 10.0, + 0.0, + 100.0, + 1000.0, + 2.0, + 20.0, + 0.0, + 200.0, + 2000.0, + 3.0, + 30.0, + 0.0, + 300.0, + 3000.0, + 4.0, + 40.0, + 0.0, + 400.0, + 4000.0, + 5.0, + 50.0, + 0.0, + 500.0, + 5000.0, + 20.0, + 200.0, + 0.0, + 2000.0, + 20000.0, + 21.0, + 210.0, + 0.0, + 2100.0, + 21000.0, + 22.0, + 220.0, + 0.0, + 2200.0, + 22000.0, + ] + results_y_train = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 20.0, 21.0, 22.0] + results_x_val = [ + 12.0, + 120.0, + 0.0, + 1200.0, + 12000.0, + 13.0, + 130.0, + 0.0, + 1300.0, + 13000.0, + 14.0, + 140.0, + 0.0, + 1400.0, + 14000.0, + 15.0, + 150.0, + 0.0, + 1500.0, + 15000.0, + 23.0, + 230.0, + 0.0, + 2300.0, + 23000.0, + 24.0, + 240.0, + 0.0, + 2400.0, + 24000.0, + ] + results_y_val = [12.0, 13.0, 14.0, 15.0, 23.0, 24.0] + + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) + + for _ in range(self.n_val_batch): + x, y = next(val_iter) + assert x.shape == (2, 5) + assert y.shape == (2, 1) + collected_x_val.append(x.tolist()) + collected_y_val.append(y.tolist()) + + for _ in range(self.n_train_batch): + x, y = next(train_iter) + assert x.shape == (2, 5) + assert y.shape == (2, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + x, y = next(train_iter) + assert x.shape == (self.train_remainder, 5) + assert y.shape == (self.train_remainder, 1) + collected_x_train.append(x.tolist()) + collected_y_train.append(y.tolist()) + + flat_x_train = [x for xl in collected_x_train for xs in xl for x in xs] + flat_x_val = [x for xl in collected_x_val for xs in xl for x in xs] + flat_y_train = [y for yl in collected_y_train for ys in yl for y in ys] + flat_y_val = [y for yl in collected_y_val for ys in yl for y in ys] + + assert results_x_train == flat_x_train + assert results_x_val == flat_x_val + assert results_y_train == flat_y_train + assert results_y_val == flat_y_val + + num_major_train = sum(np.array(flat_y_train) < 20) + num_minor_train = sum(np.array(flat_y_train) >= 20) + num_major_val = sum(np.array(flat_y_val) < 20) + num_minor_val = sum(np.array(flat_y_val) >= 20) + + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 + + +if __name__ == "__main__": + raise SystemExit(pytest.main(args=[__file__]))