From 578d56c3535b88b4f38176dd99535ef47e554d44 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Fri, 23 Oct 2026 18:29:46 +0200 Subject: [PATCH 1/8] [ML] Add jax to the RDataLoader tests dependencies --- bindings/pyroot/pythonizations/test/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bindings/pyroot/pythonizations/test/CMakeLists.txt b/bindings/pyroot/pythonizations/test/CMakeLists.txt index e40a2175e84c4..8dd2a31c17db7 100644 --- a/bindings/pyroot/pythonizations/test/CMakeLists.txt +++ b/bindings/pyroot/pythonizations/test/CMakeLists.txt @@ -128,7 +128,7 @@ 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 ml_dataloader.py PYTHON_DEPS jax numpy tensorflow torch) endif() ROOT_ADD_PYUNITTEST(regression_18441 regression_18441.py) From b8c904ef17841c653c2f76c68e36d338aeef113e Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:17:01 +0200 Subject: [PATCH 2/8] [ML] Run the RDataLoader tests with pytest --- bindings/pyroot/pythonizations/test/CMakeLists.txt | 2 +- bindings/pyroot/pythonizations/test/ml_dataloader.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/CMakeLists.txt b/bindings/pyroot/pythonizations/test/CMakeLists.txt index 8dd2a31c17db7..a3aca94a92e6e 100644 --- a/bindings/pyroot/pythonizations/test/CMakeLists.txt +++ b/bindings/pyroot/pythonizations/test/CMakeLists.txt @@ -128,7 +128,7 @@ ROOT_ADD_PYUNITTEST(pyroot_memory memory.py) # ML dataloader tests if (dataframe) - ROOT_ADD_PYUNITTEST(ml_dataloader ml_dataloader.py PYTHON_DEPS jax numpy tensorflow torch) + ROOT_ADD_PYUNITTEST(ml_dataloader ml_dataloader.py GENERIC PYTHON_DEPS jax numpy pytest tensorflow torch) endif() ROOT_ADD_PYUNITTEST(regression_18441 regression_18441.py) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index b19298bcb3990..9a168d14e8061 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -4,6 +4,7 @@ from random import randrange, uniform import numpy as np +import pytest import ROOT @@ -8338,4 +8339,4 @@ def test08_several_datasets(self): if __name__ == "__main__": - unittest.main() + raise SystemExit(pytest.main(args=[__file__])) From 8fb8009744e58c3659a357ddaddc415f7e3191f0 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:19:58 +0200 Subject: [PATCH 3/8] [ML] Switch RDataLoader unittest-base test classes to pytest The test classes no longer derive from unittest.TestCase, their names get the "Test" prefix fot pytest to collect. Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 2266 ++++++++--------- 1 file changed, 1131 insertions(+), 1135 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 9a168d14e8061..3fee463f455d9 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -1,6 +1,5 @@ import math import os -import unittest from random import randrange, uniform import numpy as np @@ -8,7 +7,7 @@ import ROOT -class DataLoaderMultipleFiles(unittest.TestCase): +class TestDataLoaderMultipleFiles: file_name1 = "first_half.root" file_name2 = "second_half.root" file_name3 = "vector_columns.root" @@ -82,21 +81,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -105,15 +104,15 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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 - self.assertTrue(np.array_equal(entries_before, entries_after)) + assert np.array_equal(entries_before, entries_after) self.teardown_file(self.file_name1) @@ -148,21 +147,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -171,10 +170,10 @@ def test02_each_element_is_generated_shuffled(self): 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) + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 self.teardown_file(self.file_name1) @@ -203,19 +202,19 @@ def test04_dropping_remainder(self): collected_y = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) + assert len(collected_x) == 3 + assert len(collected_y) == 3 self.teardown_file(self.file_name1) @@ -252,14 +251,14 @@ def test05_more_than_one_file(self): collected_y_val = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -268,10 +267,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -321,26 +320,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -352,12 +351,12 @@ def test06_multiple_target_columns(self): 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) + 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 self.teardown_file(file_name) @@ -401,21 +400,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -424,10 +423,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name) @@ -470,14 +469,14 @@ def test08_filtered(self): val_iter = iter(gen_validation.as_numpy()) x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -486,15 +485,15 @@ def test08_filtered(self): 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) + 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 - self.assertTrue(np.array_equal(filter_entries_before, filter_entries_after)) + assert np.array_equal(filter_entries_before, filter_entries_after) self.teardown_file(self.file_name1) @@ -534,22 +533,22 @@ def test09_filtered_last_chunk(self): for _ in range(3): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) + 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] - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train self.teardown_file(file_name) @@ -588,21 +587,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -611,16 +610,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) @@ -650,10 +649,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) @@ -702,26 +701,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -733,12 +732,12 @@ def test12_PyTorch(self): 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) + 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 self.teardown_file(file_name) @@ -787,26 +786,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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()) @@ -818,12 +817,12 @@ def test13_TensorFlow(self): 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) + 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 self.teardown_file(file_name) @@ -880,49 +879,48 @@ def test(size_of_batch, batches_in_memory, num_of_entries): 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}") + 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}" - 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}") + 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) - 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) + 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) - 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}") + 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}" - 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) + 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) - 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) + 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])) - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), + 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)))}", + generated length {len(set(i for i in range(num_of_entries)))}" ) except: @@ -962,21 +960,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -985,15 +983,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) @@ -1083,21 +1081,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -1106,10 +1104,10 @@ def test16_vector_padding(self): 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) + 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 self.teardown_file(self.file_name3) @@ -1158,26 +1156,26 @@ def test17_JAX(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -1189,12 +1187,12 @@ def test17_JAX(self): 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) + 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 finally: self.teardown_file(file_name) @@ -1217,8 +1215,8 @@ def test18_one_batch_in_memory(self): 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) + assert sum(1 for _ in gen_train.as_numpy()) == 3 + assert sum(1 for _ in gen_validation.as_numpy()) == 2 finally: self.teardown_file(self.file_name1) @@ -1253,14 +1251,14 @@ def test19_defined_columns(self): 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))) + assert np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1])) + assert np.array_equal(y, (entries * 3).reshape(-1, 1)) finally: self.teardown_file(self.file_name1) -class DataLoaderEagerLoading(unittest.TestCase): +class TestDataLoaderEagerLoading: file_name1 = "first_half.root" file_name2 = "second_half.root" file_name3 = "vector_columns.root" @@ -1332,21 +1330,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -1355,10 +1353,10 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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 self.teardown_file(self.file_name1) @@ -1393,21 +1391,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -1416,10 +1414,10 @@ def test02_each_element_is_generated_shuffled(self): 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) + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 self.teardown_file(self.file_name1) @@ -1448,19 +1446,19 @@ def test04_dropping_remainder(self): collected_y = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) + assert len(collected_x) == 3 + assert len(collected_y) == 3 self.teardown_file(self.file_name1) @@ -1497,14 +1495,14 @@ def test05_more_than_one_file(self): collected_y_val = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -1513,10 +1511,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -1566,26 +1564,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -1597,12 +1595,12 @@ def test06_multiple_target_columns(self): 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) + 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 self.teardown_file(file_name) @@ -1646,21 +1644,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -1669,10 +1667,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name) @@ -1713,14 +1711,14 @@ def test08_filtered(self): val_iter = iter(gen_validation.as_numpy()) x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -1729,10 +1727,10 @@ def test08_filtered(self): 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) + 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 self.teardown_file(self.file_name1) @@ -1772,22 +1770,22 @@ def test09_filtered_last_chunk(self): for _ in range(3): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) + 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] - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train self.teardown_file(file_name) @@ -1826,21 +1824,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -1849,16 +1847,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) @@ -1888,10 +1886,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) @@ -1940,26 +1938,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -1971,12 +1969,12 @@ def test12_PyTorch(self): 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) + 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 self.teardown_file(file_name) @@ -2025,26 +2023,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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()) @@ -2056,12 +2054,12 @@ def test13_TensorFlow(self): 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) + 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 self.teardown_file(file_name) @@ -2119,49 +2117,48 @@ def test(size_of_batch, batches_in_memory, num_of_entries): 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}") + 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}" - 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}") + 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) - 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) + 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) - 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}") + 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}" - 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) + 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) - 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) + 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])) - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), + 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)))}", + generated length {len(set(i for i in range(num_of_entries)))}" ) except: @@ -2201,21 +2198,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -2224,15 +2221,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) @@ -2322,21 +2319,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -2345,10 +2342,10 @@ def test16_vector_padding(self): 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) + 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 self.teardown_file(self.file_name3) @@ -2387,14 +2384,14 @@ def test17_defined_columns(self): 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))) + assert np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1])) + assert np.array_equal(y, (entries * 3).reshape(-1, 1)) finally: self.teardown_file(self.file_name1) -class DataLoaderEagerLoadingMultipleDataframes(unittest.TestCase): +class TestDataLoaderEagerLoadingMultipleDataframes: file_name1 = "first_half.root" file_name2 = "second_half.root" file_name3 = "second_file.root" @@ -2490,21 +2487,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -2513,17 +2510,17 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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 - self.assertTrue(np.array_equal(df1_entries_before, df1_entries_after)) - self.assertTrue(np.array_equal(df2_entries_before, df2_entries_after)) + assert np.array_equal(df1_entries_before, df1_entries_after) + assert np.array_equal(df2_entries_before, df2_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -2562,21 +2559,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -2585,10 +2582,10 @@ def test02_each_element_is_generated_shuffled(self): 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) + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -2621,19 +2618,19 @@ def test04_dropping_remainder(self): collected_y = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) + assert len(collected_x) == 3 + assert len(collected_y) == 3 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -2674,14 +2671,14 @@ def test05_more_than_one_file(self): collected_y_val = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -2690,10 +2687,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -2750,26 +2747,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -2781,12 +2778,12 @@ def test06_multiple_target_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -2838,21 +2835,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -2861,10 +2858,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -2913,14 +2910,14 @@ def test08_filtered(self): val_iter = iter(gen_validation.as_numpy()) x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -2929,17 +2926,17 @@ def test08_filtered(self): 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) + 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 - self.assertTrue(np.array_equal(dff1_entries_before, dff1_entries_after)) - self.assertTrue(np.array_equal(dff2_entries_before, dff2_entries_after)) + assert np.array_equal(dff1_entries_before, dff1_entries_after) + assert np.array_equal(dff2_entries_before, dff2_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -2988,22 +2985,22 @@ def test09_filtered_last_chunk(self): for _ in range(3): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) + 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] - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -3046,21 +3043,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -3069,16 +3066,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3111,10 +3108,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3170,26 +3167,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -3201,12 +3198,12 @@ def test12_PyTorch(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -3262,26 +3259,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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()) @@ -3293,12 +3290,12 @@ def test13_TensorFlow(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -3361,49 +3358,48 @@ def test(size_of_batch, batches_in_memory, num_of_entries): 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}") + 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}" - 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}") + 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) - 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) + 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) - 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}") + 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}" - 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) + 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) - 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) + 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])) - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), + 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)))}", + generated length {len(set(i for i in range(num_of_entries)))}" ) except: @@ -3447,21 +3443,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -3470,15 +3466,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3571,21 +3567,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -3594,10 +3590,10 @@ def test16_vector_padding(self): 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) + 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 self.teardown_file(self.file_name4) self.teardown_file(self.file_name5) @@ -3621,11 +3617,11 @@ def test17_dataframes_without_cluster_metadata(self): 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]) + 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] -class DataLoaderLazyLoadingMultipleDataframes(unittest.TestCase): +class TestDataLoaderLazyLoadingMultipleDataframes: file_name1 = "first_half.root" file_name2 = "second_half.root" file_name3 = "second_file.root" @@ -3723,21 +3719,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -3746,17 +3742,17 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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 - self.assertTrue(np.array_equal(df1_entries_before, df1_entries_after)) - self.assertTrue(np.array_equal(df2_entries_before, df2_entries_after)) + assert np.array_equal(df1_entries_before, df1_entries_after) + assert np.array_equal(df2_entries_before, df2_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3794,21 +3790,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -3817,10 +3813,10 @@ def test02_each_element_is_generated_shuffled(self): 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) + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3852,19 +3848,19 @@ def test04_dropping_remainder(self): collected_y = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 3) - self.assertEqual(len(collected_y), 3) + assert len(collected_x) == 3 + assert len(collected_y) == 3 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3904,14 +3900,14 @@ def test05_more_than_one_file(self): collected_y_val = [] for x, y in gen_train.as_numpy(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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(): - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + assert x.shape == (3, 1) + assert y.shape == (3, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -3920,10 +3916,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -3979,26 +3975,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -4010,12 +4006,12 @@ def test06_multiple_target_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -4079,21 +4075,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 2)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 2)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -4102,10 +4098,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -4156,14 +4152,14 @@ def test08_filtered(self): val_iter = iter(gen_validation.as_numpy()) x, y = next(train_iter) - self.assertTrue(x.shape == (4, 1)) - self.assertTrue(y.shape == (4, 1)) + assert x.shape == (4, 1) + assert 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)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x_val.append(x.tolist()) collected_y_val.append(y.tolist()) @@ -4172,17 +4168,17 @@ def test08_filtered(self): 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) + 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 - self.assertTrue(np.array_equal(dff1_entries_before, dff1_entries_after)) - self.assertTrue(np.array_equal(dff2_entries_before, dff2_entries_after)) + assert np.array_equal(dff1_entries_before, dff1_entries_after) + assert np.array_equal(dff2_entries_before, dff2_entries_after) finally: self.teardown_file(self.file_name1) @@ -4226,22 +4222,22 @@ def test09_filtered_last_chunk(self): for _ in range(3): x, y = next(train_iter) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (1, 1)) - self.assertTrue(y.shape == (1, 1)) + 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] - self.assertEqual(results_x_train, flat_x_train) - self.assertEqual(results_y_train, flat_y_train) + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -4283,21 +4279,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -4306,16 +4302,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -4347,10 +4343,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -4405,26 +4401,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 2)) - self.assertTrue(z.shape == (self.val_remainder, 1)) + 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()) @@ -4436,12 +4432,12 @@ def test12_PyTorch(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -4496,26 +4492,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 2)) - self.assertTrue(z.shape == (3, 1)) + 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()) @@ -4527,12 +4523,12 @@ def test13_TensorFlow(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -4594,49 +4590,48 @@ def test(size_of_batch, batches_in_memory, num_of_entries): 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}") + 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}" - 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}") + 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) - 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) + 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) - 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}") + 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}" - 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) + 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) - 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) + 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])) - self.assertTrue( - set(collect_x) == set(i for i in range(num_of_entries)), + 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)))}", + generated length {len(set(i for i in range(num_of_entries)))}" ) except: @@ -4679,21 +4674,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 1)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 1)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -4702,15 +4697,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -4802,21 +4797,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (3, 5)) - self.assertTrue(y.shape == (3, 1)) + 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) - self.assertTrue(x.shape == (self.val_remainder, 5)) - self.assertTrue(y.shape == (self.val_remainder, 1)) + 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()) @@ -4825,10 +4820,10 @@ def test16_vector_padding(self): 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) + 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 self.teardown_file(self.file_name4) self.teardown_file(self.file_name5) @@ -4870,14 +4865,14 @@ def test17_shuffled_split_varies_with_seed(self): 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) + assert val_1_collected != val_2_collected finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) -class DataLoaderRandomUndersampling(unittest.TestCase): +class TestDataLoaderRandomUndersampling: file_name1 = "major.root" file_name2 = "minor.root" file_name3 = "second_file.root" @@ -4976,21 +4971,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -4999,10 +4994,10 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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) @@ -5010,23 +5005,23 @@ def test01_each_element_is_generated_unshuffled(self): 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)) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + 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 - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5111,15 +5106,15 @@ def test01_each_element_is_generated_unshuffled_replacement(self): for _ in range(6): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x_train.append(x.tolist()) collected_y_train.append(y.tolist()) @@ -5128,10 +5123,10 @@ def test01_each_element_is_generated_unshuffled_replacement(self): 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) + 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) @@ -5139,16 +5134,16 @@ def test01_each_element_is_generated_unshuffled_replacement(self): 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)) + 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) - self.assertEqual(num_major_train, 15) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 10) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 15 + assert num_minor_train == 3 + assert num_major_val == 10 + assert num_minor_val == 2 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5190,21 +5185,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -5213,10 +5208,10 @@ def test02_each_element_is_generated_shuffled(self): 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) + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5256,20 +5251,20 @@ def test04_dropping_remainder(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 7) - self.assertEqual(len(collected_y), 7) + assert len(collected_x) == 7 + assert len(collected_y) == 7 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5317,21 +5312,21 @@ def test05_more_than_one_file(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -5340,10 +5335,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5422,26 +5417,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) + 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()) @@ -5453,22 +5448,22 @@ def test06_multiple_target_columns(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -5542,21 +5537,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 2)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -5565,10 +5560,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -5620,21 +5615,21 @@ def test08_filtered(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -5643,27 +5638,27 @@ def test08_filtered(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + 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 - self.assertTrue(np.array_equal(major_filter_entries_before, major_filter_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) + assert np.array_equal(major_filter_entries_before, major_filter_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5709,21 +5704,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -5732,16 +5727,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5777,10 +5772,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -5857,26 +5852,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) + 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()) @@ -5888,22 +5883,22 @@ def test12_PyTorch(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -5982,26 +5977,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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()) @@ -6013,22 +6008,22 @@ def test13_TensorFlow(self): 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) + 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) - 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -6105,31 +6100,31 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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}") + 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) - 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) + 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) - 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}") + 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) - 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) + 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] @@ -6141,12 +6136,12 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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)) + 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 - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -6162,8 +6157,9 @@ def test_raises(size_of_batch, num_of_entries_major, num_of_entries_minor, sampl 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." + 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], @@ -6254,31 +6250,31 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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}") + 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) - 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) + 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) - 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}") + 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) - 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) + 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] @@ -6290,12 +6286,12 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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)) + 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 - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -6342,21 +6338,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -6365,15 +6361,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -6494,21 +6490,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 5)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -6517,20 +6513,20 @@ def test16_vector_padding(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(self.file_name4) self.teardown_file(self.file_name5) @@ -6541,7 +6537,7 @@ def test16_vector_padding(self): raise -class DataLoaderRandomOversampling(unittest.TestCase): +class TestDataLoaderRandomOversampling: file_name1 = "major.root" file_name2 = "minor.root" file_name3 = "second_file.root" @@ -6639,21 +6635,21 @@ def test01_each_element_is_generated_unshuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -6662,10 +6658,10 @@ def test01_each_element_is_generated_unshuffled(self): 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) + 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) @@ -6673,23 +6669,23 @@ def test01_each_element_is_generated_unshuffled(self): 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)) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + 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 - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_entries_before, minor_entries_after)) + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_entries_before, minor_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -6730,21 +6726,21 @@ def test02_each_element_is_generated_shuffled(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -6754,16 +6750,16 @@ def test02_each_element_is_generated_shuffled(self): 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)) + 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) - 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) + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -6802,20 +6798,20 @@ def test04_dropping_remainder(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + assert x.shape == (2, 1) + assert y.shape == (2, 1) collected_x.append(x) collected_y.append(y) - self.assertEqual(len(collected_x), 7) - self.assertEqual(len(collected_y), 7) + assert len(collected_x) == 7 + assert len(collected_y) == 7 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -6862,21 +6858,21 @@ def test05_more_than_one_file(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -6885,10 +6881,10 @@ def test05_more_than_one_file(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -6966,26 +6962,26 @@ def test06_multiple_target_columns(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) + 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()) @@ -6997,22 +6993,22 @@ def test06_multiple_target_columns(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -7085,21 +7081,21 @@ def test07_multiple_input_columns(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 2)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 2)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -7108,10 +7104,10 @@ def test07_multiple_input_columns(self): 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -7162,21 +7158,21 @@ def test08_filtered(self): for _ in range(self.n_train_batch): x, y = next(train_iter) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -7185,27 +7181,27 @@ def test08_filtered(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + 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 - self.assertTrue(np.array_equal(major_entries_before, major_entries_after)) - self.assertTrue(np.array_equal(minor_filter_entries_before, minor_filter_entries_after)) + assert np.array_equal(major_entries_before, major_entries_after) + assert np.array_equal(minor_filter_entries_before, minor_filter_entries_after) self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -7250,21 +7246,21 @@ def test10_two_epochs_shuffled(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -7273,16 +7269,16 @@ def test10_two_epochs_shuffled(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -7317,10 +7313,10 @@ def test11_number_of_training_and_validation_batches_remainder(self): 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) + 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 self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -7396,26 +7392,26 @@ def test12_PyTorch(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 2)) - self.assertTrue(z.shape == (self.train_remainder, 1)) + 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()) @@ -7427,22 +7423,22 @@ def test12_PyTorch(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -7520,26 +7516,26 @@ def test13_TensorFlow(self): 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)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 2)) - self.assertTrue(z.shape == (2, 1)) + 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()) @@ -7551,22 +7547,22 @@ def test13_TensorFlow(self): 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) + 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) - 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) + 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 self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -7641,31 +7637,31 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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}") + 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) - 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) + 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) - 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}") + 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) - 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) + 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] @@ -7677,12 +7673,12 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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)) + 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 - self.assertEqual(round((num_minor_train / num_major_train), 2), sampling_ratio) - self.assertEqual(round((num_minor_val / num_major_val), 2), sampling_ratio) + assert round((num_minor_train / num_major_train), 2) == sampling_ratio + assert round((num_minor_val / num_major_val), 2) == sampling_ratio self.teardown_file(file_name1) self.teardown_file(file_name2) @@ -7728,21 +7724,21 @@ def test15_two_runs_set_seed(self): for _ in range(self.n_train_batch): x, y = next(iter_train) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 1)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 1)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -7751,15 +7747,15 @@ def test15_two_runs_set_seed(self): 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) + 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) - 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]) + 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] finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) @@ -7880,21 +7876,21 @@ def test16_vector_padding(self): for _ in range(self.n_val_batch): x, y = next(val_iter) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (2, 5)) - self.assertTrue(y.shape == (2, 1)) + 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) - self.assertTrue(x.shape == (self.train_remainder, 5)) - self.assertTrue(y.shape == (self.train_remainder, 1)) + 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()) @@ -7903,20 +7899,20 @@ def test16_vector_padding(self): 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) + 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) - self.assertEqual(num_major_train, 6) - self.assertEqual(num_minor_train, 3) - self.assertEqual(num_major_val, 4) - self.assertEqual(num_minor_val, 2) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 self.teardown_file(self.file_name4) self.teardown_file(self.file_name5) @@ -7950,15 +7946,15 @@ def test17_set_seed_resamples_different_entries(self): 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) + assert len(minor) == 10 + assert len(np.unique(minor)) > 1 finally: self.teardown_file(self.file_name1) self.teardown_file(self.file_name2) -class DataLoaderSave(unittest.TestCase): +class TestDataLoaderSave: file_name = "dataloader_save_input.root" file_name2 = "dataloader_save_input2.root" tree_name = "mytree" @@ -7976,7 +7972,7 @@ def create_file(self, file_name=None, num_of_entries=20, offset=0): .Snapshot(self.tree_name, file_name or self.file_name) ) - def tearDown(self): + def teardown_method(self): for file in (self.file_name, self.file_name2, self.out_name, self.out_name2): if os.path.exists(file): os.remove(file) @@ -8009,7 +8005,7 @@ def test01_unshuffled(self): 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"}) + 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)) @@ -8042,7 +8038,7 @@ def test02_shuffled(self): 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()) + assert (train_b1 != np.sort(train_b1)).any() def test03_vector_columns(self): self.create_file() @@ -8063,10 +8059,10 @@ def test03_vector_columns(self): 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"}) + assert self.saved_columns(self.out_name) == {"b1", "v1"} v1 = self.read_back(self.out_name, "v1") - self.assertEqual(v1.shape, (20, 3)) + 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]) @@ -8194,25 +8190,27 @@ def test08_round_trip(self): np.testing.assert_array_equal(w.flatten(), np.full(20, 0.5)) -class DataLoaderImplicitMT(unittest.TestCase): +class TestDataLoaderImplicitMT: file_name = "dataloader_imt.root" tree_name = "mytree" - def tearDown(self): + def teardown_method(self): ROOT.DisableImplicitMT() os.remove(self.file_name) - @unittest.skipUnless("imt" in ROOT.gROOT.GetConfigFeatures(), "ROOT was built without implicit multi-threading") + @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 self.assertRaisesRegex(RuntimeError, "Cannot use RDataLoader with implicit multi-threading enabled"): + with pytest.raises(RuntimeError, match="Cannot use RDataLoader with implicit multi-threading enabled"): ROOT.Experimental.ML.RDataLoader(df, batch_size=2).as_numpy() -class DataLoaderFromFiles(unittest.TestCase): +class TestDataLoaderFromFiles: file_name1 = "dataloader_from_files1.root" file_name2 = "dataloader_from_files2.root" tree_name = "mytree" @@ -8230,7 +8228,7 @@ def create_file(self, file_name, num_of_entries=10, offset=0, output_format="ttr .Snapshot(self.tree_name, file_name, ["b1", "b2"], opts) ) - def tearDown(self): + def teardown_method(self): for file in (self.file_name1, self.file_name2): if os.path.exists(file): os.remove(file) @@ -8249,10 +8247,9 @@ 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) + 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 @@ -8310,16 +8307,16 @@ def test07_invalid_inputs(self): 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"): + with pytest.raises(TypeError, match="dataset_name, files"): self.load_all(self.file_name1) - with self.assertRaisesRegex(ValueError, "at least one valid input is required"): + with pytest.raises(ValueError, match="at least one valid input is required"): self.load_all([]) - with self.assertRaisesRegex(ValueError, "at least one valid input is required"): + 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 self.assertRaises(TypeError): + with pytest.raises(TypeError): ROOT.Experimental.ML.RDataLoader(df, 3) def test08_several_datasets(self): @@ -8332,10 +8329,9 @@ def test08_several_datasets(self): [(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) + 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__": From 4c4702a0c52616a87cc2c3ce34cb2b1a2d4b3e6a Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:24:52 +0200 Subject: [PATCH 4/8] [ML] Add fixtures for the inputs of the RDataLoader tests Each test now runs in its own temporary directory which pytest creates and cleans up. The next commits will use these. the fixtures are dupliucated in this commit because the tests are duplicated for lazy and eager loading modes, the test classes will be merged in a future commit and so will the fixtures. Assisted-by: Claude Opus 5.5 --- .../pyroot/pythonizations/test/conftest.py | 66 ++++++++++ .../pythonizations/test/ml_dataloader.py | 123 ++++++++++++++++++ 2 files changed, 189 insertions(+) diff --git a/bindings/pyroot/pythonizations/test/conftest.py b/bindings/pyroot/pythonizations/test/conftest.py index f5b0289f4d122..803cd7ff70e18 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,66 @@ 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) + + +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 index 3fee463f455d9..046f4b265d0af 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -5,6 +5,7 @@ import numpy as np import pytest import ROOT +from conftest import from_files, squares, weighted class TestDataLoaderMultipleFiles: @@ -18,6 +19,23 @@ class TestDataLoaderMultipleFiles: 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") + ) + # Helpers def define_rdf(self, num_of_entries=10): df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") @@ -1269,6 +1287,23 @@ class TestDataLoaderEagerLoading: 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") + ) + # Helpers def define_rdf(self, num_of_entries=10): df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") @@ -2404,6 +2439,26 @@ class TestDataLoaderEagerLoadingMultipleDataframes: 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")) + # Helpers def define_rdf1(self, num_of_entries=5): df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") @@ -3634,6 +3689,26 @@ class TestDataLoaderLazyLoadingMultipleDataframes: 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")) + # Helpers def define_rdf1(self, num_of_entries=5): df = ROOT.RDataFrame(num_of_entries).Define("b1", "(int) rdfentry_").Define("b2", "(double) b1*b1") @@ -4885,6 +4960,26 @@ class TestDataLoaderRandomUndersampling: 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")) + # 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") @@ -6550,6 +6645,26 @@ class TestDataLoaderRandomOversampling: 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")) + # 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") @@ -7962,7 +8077,15 @@ class TestDataLoaderSave: 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 create_file(self, file_name=None, num_of_entries=20, offset=0): ( ROOT From 3d50041f9c53de38536b4451add8c8d9517d0644 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:27:09 +0200 Subject: [PATCH 5/8] [ML] Remove the file cleanup from the RDataLoader tests This remooves the try/except/finally blocks in favour of using the fixtures defined previously. Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 10745 +++++++--------- 1 file changed, 5000 insertions(+), 5745 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 046f4b265d0af..586422f8a0d5a 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -1,5 +1,4 @@ import math -import os from random import randrange, uniform import numpy as np @@ -62,241 +61,208 @@ def create_vector_file(self, num_of_entries=10): .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, - ) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - gen_train, gen_validation = dl.train_test_split(0.4) + entries_before = df.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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] + 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()) - 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 + 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()) - entries_after = df.AsNumpy(["rdfentry_"])["rdfentry_"] + 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 the dataframe is correctly reset - assert np.array_equal(entries_before, entries_after) + 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 - self.teardown_file(self.file_name1) + entries_after = df.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - raise + # check if the dataframe is correctly reset + assert np.array_equal(entries_before, entries_after) 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, - ) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=True, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_val_batch): x, y = next(val_iter) - assert x.shape == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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 len(flat_x_train) == 6 - assert len(flat_x_val) == 4 - assert len(flat_y_train) == 6 - assert len(flat_y_val) == 4 + 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()) - self.teardown_file(self.file_name1) + 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} - except: - self.teardown_file(self.file_name1) - raise + 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): 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) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - collected_x = [] - collected_y = [] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=True, + ) - 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) + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + collected_x = [] + collected_y = [] - assert len(collected_x) == 3 - assert len(collected_y) == 3 + 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) - self.teardown_file(self.file_name1) + 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) - except: - self.teardown_file(self.file_name1) - raise + assert len(collected_x) == 3 + assert len(collected_y) == 3 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]) + 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] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): file_name = "multiple_target_columns.root" @@ -304,83 +270,76 @@ def test06_multiple_target_columns(self): 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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "multiple_input_columns.root" @@ -389,135 +348,121 @@ def test07_multiple_input_columns(self): "b3", "(double) rdfentry_ * 10" ).Snapshot("myTree", file_name) - try: - df = ROOT.RDataFrame("myTree", file_name) + 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, - ) + 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) + 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] + 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 = [] + 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()) + 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()) + 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 == (self.val_remainder, 2) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 2) + 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 + 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()) - self.teardown_file(file_name) + 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] - except: - self.teardown_file(file_name) - raise + 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): 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_"] + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - dl = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) + dff = df.Filter("b1 % 2 == 0", "name") - gen_train, gen_validation = dl.train_test_split(0.4) + filter_entries_before = dff.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + dl = ROOT.Experimental.ML.RDataLoader( + dff, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) - 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] + 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()) - 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 + 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()) - filter_entries_after = dff.AsNumpy(["rdfentry_"])["rdfentry_"] + 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 the dataframe is correctly reset - assert np.array_equal(filter_entries_before, filter_entries_after) + 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 - self.teardown_file(self.file_name1) + filter_entries_after = dff.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - raise + # check if the dataframe is correctly reset + assert np.array_equal(filter_entries_before, filter_entries_after) def test09_filtered_last_chunk(self): file_name = "filtered_last_chunk.root" @@ -527,156 +472,139 @@ def test09_filtered_last_chunk(self): tree_name, file_name ) - try: - df = ROOT.RDataFrame(tree_name, file_name) - - dff = df.Filter("b1 % 2 == 0", "name") + df = ROOT.RDataFrame(tree_name, file_name) - gen_train = ROOT.Experimental.ML.RDataLoader( - dff, - batch_size=3, - batches_in_memory=2, - target="b2", - shuffle=False, - drop_remainder=False, - ) + dff = df.Filter("b1 % 2 == 0", "name") - 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] + gen_train = ROOT.Experimental.ML.RDataLoader( + dff, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_y_train = [] + 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] - train_iter = iter(gen_train.as_numpy()) + collected_x_train = [] + collected_y_train = [] - 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()) + train_iter = iter(gen_train.as_numpy()) + for _ in range(3): x, y = next(train_iter) - assert x.shape == (1, 1) - assert y.shape == (1, 1) + assert x.shape == (3, 1) + assert y.shape == (3, 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 + 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()) - self.teardown_file(file_name) + 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] - except: - self.teardown_file(file_name) - raise + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train 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, - ) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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} + 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 + 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) + 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] - finally: - self.teardown_file(self.file_name1) + 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): 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) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - number_of_training_batches = 0 - number_of_validation_batches = 0 + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target="b2", + shuffle=False, + drop_remainder=False, + ) - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + number_of_training_batches = 0 + number_of_validation_batches = 0 - 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 + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 - self.teardown_file(self.file_name1) + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 - except: - self.teardown_file(self.file_name1) - raise + 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): file_name = "multiple_target_columns.root" @@ -685,83 +613,76 @@ def test12_PyTorch(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_torch(device="cpu")) - iter_val = iter(gen_validation.as_torch()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_torch(device="cpu")) + iter_val = iter(gen_validation.as_torch()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "multiple_target_columns.root" @@ -770,56 +691,47 @@ def test13_TensorFlow(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) - 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()) + 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) @@ -828,25 +740,27 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "big_data.root" @@ -868,270 +782,255 @@ def define_rdf(num_of_entries): def test(size_of_batch, batches_in_memory, num_of_entries): define_rdf(num_of_entries) - try: - df = ROOT.RDataFrame(tree_name, file_name) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.3) - collect_x = [] + collect_x = [] - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + 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 + 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()) + 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) + 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 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}" + 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])) + 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])) + 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) + 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 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 + 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])) + 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])) + 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))}\ + 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)))}" - ) - - 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) - 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] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, + target="b2", + shuffle=True, drop_remainder=False, + set_seed=42, ) 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()) + 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(val_iter) - assert x.shape == (3, 5) + 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()) - 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(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): + self.create_vector_file() + + 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) - assert x.shape == (self.val_remainder, 5) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 5) + 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name3) - raise + 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): file_name = "multiple_target_columns.root" @@ -1140,140 +1039,128 @@ def test17_JAX(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + batches_in_memory=2, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_jax(device="cpu")) - iter_val = iter(gen_validation.as_jax()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_jax(device="cpu")) + iter_val = iter(gen_validation.as_jax()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - finally: - self.teardown_file(file_name) + 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): # 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, - ) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=2, + batches_in_memory=1, + target="b2", + shuffle=False, + drop_remainder=True, + ) - assert sum(1 for _ in gen_train.as_numpy()) == 3 - assert sum(1 for _ in gen_validation.as_numpy()) == 2 + gen_train, gen_validation = dl.train_test_split(0.4) - finally: - self.teardown_file(self.file_name1) + 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): # 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, - ) + df = ( + ROOT + .RDataFrame(self.tree_name, self.file_name1) + .Define("b3", "b1 * 3") + .Define("v1", "ROOT::RVecD{b2, b2 + 1}") + ) - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - 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]) + gen_train, gen_validation = dl.train_test_split(0.4) - 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)) + 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]) - finally: - self.teardown_file(self.file_name1) + 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)) class TestDataLoaderEagerLoading: @@ -1330,234 +1217,201 @@ def create_vector_file(self, num_of_entries=10): .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) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - 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] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) + 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] - except: - self.teardown_file(self.file_name1) - raise + 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 test02_each_element_is_generated_shuffled(self): self.create_file() - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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()) + 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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 len(flat_x_train) == 6 - assert len(flat_x_val) == 4 - assert len(flat_y_train) == 6 - assert len(flat_y_val) == 4 + 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()) - self.teardown_file(self.file_name1) + 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} - except: - self.teardown_file(self.file_name1) - raise + 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): 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) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - collected_x = [] - collected_y = [] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=True, + load_eager=True, + ) - 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) + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + collected_x = [] + collected_y = [] - assert len(collected_x) == 3 - assert len(collected_y) == 3 + 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) - self.teardown_file(self.file_name1) + 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) - except: - self.teardown_file(self.file_name1) - raise + assert len(collected_x) == 3 + assert len(collected_y) == 3 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) + df = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - 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] + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): file_name = "multiple_target_columns.root" @@ -1565,83 +1419,76 @@ def test06_multiple_target_columns(self): 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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "multiple_input_columns.root" @@ -1650,287 +1497,256 @@ def test07_multiple_input_columns(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 2) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 2) + 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 + 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()) - self.teardown_file(file_name) + 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] - except: - self.teardown_file(file_name) - raise + 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): self.create_file() - try: - df = ROOT.RDataFrame(self.tree_name, self.file_name1) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - dff = df.Filter("b1 % 2 == 0", "name") + 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, - ) + 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) + 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] + 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 - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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 + ) + + 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()) - val_iter = iter(gen_validation.as_numpy()) + 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(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] + 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()) - 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 + 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.teardown_file(self.file_name1) + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train - except: - self.teardown_file(self.file_name1) - raise + def test10_two_epochs_shuffled(self): + self.create_file() - def test09_filtered_last_chunk(self): - file_name = "filtered_last_chunk.root" - tree_name = "myTree" + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1").Snapshot( - tree_name, file_name + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, ) - 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, - ) + 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] - results_y_train = [0.0, 4.0, 16.0, 36.0, 64.0, 100.0, 144.0, 196.0, 256.0, 324.0] + 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 = [] - train_iter = iter(gen_train.as_numpy()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - for _ in range(3): - x, y = next(train_iter) + 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()) - 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 - - 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) - 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()) - + for _ in range(self.n_val_batch): x, y = next(iter_val) - assert x.shape == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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} + 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 + 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) + 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] - finally: - self.teardown_file(self.file_name1) + 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): 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) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) - number_of_training_batches = 0 - number_of_validation_batches = 0 + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + number_of_training_batches = 0 + number_of_validation_batches = 0 - 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 + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 - self.teardown_file(self.file_name1) + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 - except: - self.teardown_file(self.file_name1) - raise + 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): file_name = "multiple_target_columns.root" @@ -1939,83 +1755,76 @@ def test12_PyTorch(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_torch()) + iter_val = iter(gen_validation.as_torch(device="cpu")) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "multiple_target_columns.root" @@ -2024,56 +1833,47 @@ def test13_TensorFlow(self): "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, - ) + df = ROOT.RDataFrame("myTree", file_name) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + df, + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) - 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()) + 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) @@ -2082,25 +1882,27 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name) - - except: - self.teardown_file(file_name) - raise + 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): file_name = "big_data.root" @@ -2122,308 +1924,289 @@ def define_rdf(num_of_entries): def test(size_of_batch, batches_in_memory, num_of_entries): define_rdf(num_of_entries) - try: - df = ROOT.RDataFrame(tree_name, file_name) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.3) - collect_x = [] + collect_x = [] - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + 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 + 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()) + 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) + 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 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}" + 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])) + 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])) + 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) + 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 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 + 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])) + 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])) + 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))}\ + 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)))}" - ) - - 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) - 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] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, + target="b2", + shuffle=True, drop_remainder=False, + set_seed=42, 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) - assert x.shape == (3, 5) - assert y.shape == (3, 1) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) for _ in range(self.n_train_batch): - x, y = next(train_iter) - assert x.shape == (3, 5) + 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()) - x, y = next(val_iter) - assert x.shape == (self.val_remainder, 5) + 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] + 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): + self.create_vector_file() + + 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) + 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()) - 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 + 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()) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name3) - raise + 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_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, - ) + df = ( + ROOT + .RDataFrame(self.tree_name, self.file_name1) + .Define("b3", "b1 * 3") + .Define("v1", "ROOT::RVecD{b2, b2 + 1}") + ) - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - 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]) + gen_train, gen_validation = dl.train_test_split(0.4) - 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)) + 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]) - finally: - self.teardown_file(self.file_name1) + 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)) class TestDataLoaderEagerLoadingMultipleDataframes: @@ -2502,260 +2285,219 @@ def create_vector_file2(self, num_of_entries=5): .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, - ) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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] + 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()) - 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 + 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()) - df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] + 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 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) + 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 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + # 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): 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, - ) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_val_batch): x, y = next(val_iter) - assert x.shape == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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 len(flat_x_train) == 6 - assert len(flat_x_val) == 4 - assert len(flat_y_train) == 6 - assert len(flat_y_val) == 4 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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} - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): 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) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - collected_x = [] - collected_y = [] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=True, + load_eager=True, + ) - 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) + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + collected_x = [] + collected_y = [] - assert len(collected_x) == 3 - assert len(collected_y) == 3 + 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) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + assert len(collected_x) == 3 + assert len(collected_y) == 3 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) + df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) - 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] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -2767,86 +2509,77 @@ def test06_multiple_target_columns(self): 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) + 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, - ) + 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) + 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] + 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 = [] + 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -2860,146 +2593,128 @@ def test07_multiple_input_columns(self): "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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 2) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 2) + 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 + 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()) - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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] - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): 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_"] + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - dl = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 != 0", "name") - gen_train, gen_validation = dl.train_test_split(0.4) + dff1_entries_before = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] + dff2_entries_before = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + dl = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + train_iter = iter(gen_train.as_numpy()) + val_iter = iter(gen_validation.as_numpy()) - 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] + 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()) - 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 + 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()) - dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] + 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 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) + 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 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] + dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + # 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): file_name1 = "filtered_last_chunk_1.root" @@ -3014,167 +2729,145 @@ def test09_filtered_last_chunk(self): 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") + df1 = ROOT.RDataFrame(tree_name, file_name1) + df2 = ROOT.RDataFrame(tree_name, file_name2) - gen_train = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - load_eager=True, - ) + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 == 0", "name") - 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] + gen_train = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - collected_x_train = [] - collected_y_train = [] + 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] - train_iter = iter(gen_train.as_numpy()) + collected_x_train = [] + collected_y_train = [] - 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()) + train_iter = iter(gen_train.as_numpy()) + for _ in range(3): x, y = next(train_iter) - assert x.shape == (1, 1) - assert y.shape == (1, 1) + assert x.shape == (3, 1) + assert y.shape == (3, 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] + 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()) - assert results_x_train == flat_x_train - assert results_y_train == flat_y_train + 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.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train 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, - ) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + for _ in range(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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} + 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 + 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) + 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] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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): 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) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - number_of_training_batches = 0 - number_of_validation_batches = 0 + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + number_of_training_batches = 0 + number_of_validation_batches = 0 - 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 + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -3187,86 +2880,77 @@ def test12_PyTorch(self): "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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_torch(device="cpu")) - iter_val = iter(gen_validation.as_torch()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_torch(device="cpu")) + iter_val = iter(gen_validation.as_torch()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -3279,57 +2963,48 @@ def test13_TensorFlow(self): "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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) - 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()) + 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) @@ -3338,27 +3013,27 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "big_data_1.root" @@ -3382,85 +3057,79 @@ 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) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.3) - collect_x = [] + collect_x = [] - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + 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 + 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()) + 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) + 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 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}" + 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])) + 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])) + 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) + 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 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 + 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])) + 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])) + 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))}\ + 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)))}" - ) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + ) test(batch_size, batches_in_memory, entries_in_rdf) @@ -3468,195 +3137,182 @@ 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) - 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] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, + target="b2", + shuffle=True, drop_remainder=False, + set_seed=42, 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()) + 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(val_iter) - assert x.shape == (3, 5) + 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()) - 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(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): + self.create_vector_file1() + self.create_vector_file2() + + 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) - assert x.shape == (self.val_remainder, 5) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 5) + 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) + 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] - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise + 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_dataframes_without_cluster_metadata(self): # the dataframes are not read from a file, so their entries are counted with Count() @@ -3755,256 +3411,215 @@ def create_vector_file2(self, num_of_entries=5): .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, - ) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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] + 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()) - 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 + 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()) - df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] - df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] + 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 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) + 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 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + df1_entries_after = df1.AsNumpy(["rdfentry_"])["rdfentry_"] + df2_entries_after = df2.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + # 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): 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, - ) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=True, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_val_batch): x, y = next(val_iter) - assert x.shape == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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} + 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()) - assert len(flat_x_train) == 6 - assert len(flat_x_val) == 4 - assert len(flat_y_train) == 6 - assert len(flat_y_val) == 4 + 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.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 + 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): 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) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - collected_x = [] - collected_y = [] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=True, + ) - 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) + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + collected_x = [] + collected_y = [] - assert len(collected_x) == 3 - assert len(collected_y) == 3 + 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) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + assert len(collected_x) == 3 + assert len(collected_y) == 3 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) + df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) - 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] + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -4016,85 +3631,76 @@ def test06_multiple_target_columns(self): 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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -4108,156 +3714,142 @@ def test07_multiple_input_columns(self): "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] + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - 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()) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + ) - 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()) + gen_train, gen_validation = dl.train_test_split(0.4) - 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()) + 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 == (self.val_remainder, 2) - assert y.shape == (self.val_remainder, 1) + assert x.shape == (3, 2) + 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 + 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()) - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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] - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): 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) - assert x.shape == (4, 1) - assert y.shape == (4, 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()) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - 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] + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 != 0", "name") - 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_before = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] + dff2_entries_before = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] - dff1_entries_after = dff1.AsNumpy(["rdfentry_"])["rdfentry_"] - dff2_entries_after = dff2.AsNumpy(["rdfentry_"])["rdfentry_"] + dl = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=4, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=False, + ) - # 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) + gen_train, gen_validation = dl.train_test_split(0.4) - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + # 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 == (4, 1) + assert y.shape == (4, 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): file_name1 = "filtered_last_chunk_1.root" @@ -4272,164 +3864,142 @@ def test09_filtered_last_chunk(self): 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") + df1 = ROOT.RDataFrame(tree_name, file_name1) + df2 = ROOT.RDataFrame(tree_name, file_name2) - gen_train = ROOT.Experimental.ML.RDataLoader( - [dff1, dff2], - batch_size=3, - target="b2", - shuffle=False, - drop_remainder=False, - ) + dff1 = df1.Filter("b1 % 2 == 0", "name") + dff2 = df2.Filter("b1 % 2 == 0", "name") - 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] + gen_train = ROOT.Experimental.ML.RDataLoader( + [dff1, dff2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + ) - collected_x_train = [] - collected_y_train = [] + 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] - train_iter = iter(gen_train.as_numpy()) + collected_x_train = [] + collected_y_train = [] - 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()) + train_iter = iter(gen_train.as_numpy()) + for _ in range(3): x, y = next(train_iter) - assert x.shape == (1, 1) - assert y.shape == (1, 1) + assert x.shape == (3, 1) + assert y.shape == (3, 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] + 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()) - assert results_x_train == flat_x_train - assert results_y_train == flat_y_train + 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.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + assert results_x_train == flat_x_train + assert results_y_train == flat_y_train 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) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.4) - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] + 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 = [] + 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 1) + 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} + 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 + 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) + 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] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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): 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) + df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) + df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - number_of_training_batches = 0 - number_of_validation_batches = 0 + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + ) - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + number_of_training_batches = 0 + number_of_validation_batches = 0 - 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 + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -4442,85 +4012,76 @@ def test12_PyTorch(self): "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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_torch()) + iter_val = iter(gen_validation.as_torch(device="cpu")) - 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()) + 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 == (self.val_remainder, 1) - assert y.shape == (self.val_remainder, 2) - assert z.shape == (self.val_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_1.root" @@ -4533,56 +4094,47 @@ def test13_TensorFlow(self): "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, - ) + df1 = ROOT.RDataFrame("myTree", file_name1) + df2 = ROOT.RDataFrame("myTree", file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + dl = ROOT.Experimental.ML.RDataLoader( + [df1, df2], + batch_size=3, + target=["b2", "b4"], + weights="b3", + shuffle=False, + drop_remainder=False, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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] - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] + collected_z_train = [] + collected_z_val = [] - 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()) + iter_train = iter(gen_train.as_tensorflow()) + iter_val = iter(gen_validation.as_tensorflow()) - 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()) + 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) @@ -4591,27 +4143,27 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "big_data_1.root" @@ -4635,84 +4187,78 @@ 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) + 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, - ) + 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) + gen_train, gen_validation = dl.train_test_split(0.3) - collect_x = [] + collect_x = [] - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + 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 + 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()) + 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) + 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 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}" + 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])) + 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])) + 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) + 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 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 + 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])) + 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])) + 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))}\ + 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)))}" - ) - - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + ) test(batch_size, batches_in_memory, entries_in_rdf) @@ -4720,234 +4266,216 @@ 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) - 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 = [] + both_runs_collected_y_val = [] - 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] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, - shuffle=False, + target="b2", + shuffle=True, drop_remainder=False, + set_seed=42, ) 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) - assert x.shape == (3, 5) - assert y.shape == (3, 1) - collected_x_val.append(x.tolist()) - collected_y_val.append(y.tolist()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) for _ in range(self.n_train_batch): - x, y = next(train_iter) - assert x.shape == (3, 5) + 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()) - x, y = next(val_iter) - assert x.shape == (self.val_remainder, 5) + 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] + 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 + assert len(flat_x_train) == 6 + assert len(flat_x_val) == 4 + assert len(flat_y_train) == 6 + assert len(flat_y_val) == 4 - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) + 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] - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise + def test16_vector_padding(self): + self.create_vector_file1() + self.create_vector_file2() - def test17_shuffled_split_varies_with_seed(self): - self.create_file1() - self.create_file2() + 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} - 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="b1", + max_vec_sizes=max_vec_sizes, + shuffle=False, + drop_remainder=False, + ) - dl1 = 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) - dl2 = ROOT.Experimental.ML.RDataLoader( - [df1, df2], - batch_size=3, - target="b2", - shuffle=True, - drop_remainder=False, - set_seed=43, - ) + 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()) - _, gen_val1 = dl1.train_test_split(0.4) - _, gen_val2 = dl2.train_test_split(0.4) + 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()) - 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()]) + 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()) - assert val_1_collected != val_2_collected + 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] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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): + self.create_file1() + self.create_file2() -class TestDataLoaderRandomUndersampling: + 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()]) + + assert val_1_collected != val_2_collected + + +class TestDataLoaderRandomUndersampling: file_name1 = "major.root" file_name2 = "minor.root" file_name3 = "second_file.root" @@ -5023,427 +4551,377 @@ def create_vector_file_minor(self, num_of_entries=5): .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_"] + 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, - ) + major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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] - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_train_batch): x, y = next(train_iter) - assert x.shape == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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 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) - - 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 + 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): 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) - 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 + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - 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) + 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, + ) - # 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) + gen_train, gen_validation = dl.train_test_split(0.4) - # 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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()) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): 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, - ) + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_train_batch): x, y = next(train_iter) - assert x.shape == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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 len(flat_x_train) == 9 - assert len(flat_x_val) == 6 - assert len(flat_y_train) == 9 - assert len(flat_y_val) == 6 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): 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) + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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, + ) - collected_x = [] - collected_y = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + 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.append(x) - collected_y.append(y) + collected_x = [] + collected_y = [] - assert len(collected_x) == 7 - assert len(collected_y) == 7 + 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) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + assert len(collected_x) == 7 + assert len(collected_y) == 7 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, - ) + df_major = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name3]) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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] - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -5455,118 +4933,109 @@ def test06_multiple_target_columns(self): 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] + df_major = ROOT.RDataFrame("myTree", file_name1) + df_minor = ROOT.RDataFrame("myTree", file_name2) - 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()) + 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, + ) - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 2) - assert z.shape == (self.train_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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()) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -5580,305 +5049,274 @@ def test07_multiple_input_columns(self): "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()) + df_major = ROOT.RDataFrame("myTree", file_name1) + df_minor = ROOT.RDataFrame("myTree", file_name2) - 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()) + 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, + ) - 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()) + 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 == (self.train_remainder, 2) - assert y.shape == (self.train_remainder, 1) + assert x.shape == (2, 2) + 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] + 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()) - 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 + 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()) - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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] - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): 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") + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - 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, - ) + df_major_filter = df_major.Filter("b1 % 8 != 0", "name") - gen_train, gen_validation = dl.train_test_split(0.4) + major_filter_entries_before = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_before = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + 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, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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] + 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()) - 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 + 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()) - 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) + 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 num_major_train == 6 - assert num_minor_train == 3 - assert num_major_val == 4 - assert num_minor_val == 2 + 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 - major_filter_entries_after = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] + 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) - # 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) + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + major_filter_entries_after = df_major_filter.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_entries_after = df_minor.AsNumpy(["rdfentry_"])["rdfentry_"] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + # 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): 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) + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - both_epochs_collected_x_val = [] - both_epochs_collected_y_val = [] + 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, + ) - for _ in range(2): - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + both_epochs_collected_x_val = [] + both_epochs_collected_y_val = [] - 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(2): + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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} + 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 + 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) + 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] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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): 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) + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - number_of_training_batches = 0 - number_of_validation_batches = 0 + 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, + ) - for _ in gen_train.as_numpy(): - number_of_training_batches += 1 + gen_train, gen_validation = dl.train_test_split(0.4) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + number_of_training_batches = 0 + number_of_validation_batches = 0 - 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 + for _ in gen_train.as_numpy(): + number_of_training_batches += 1 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + for _ in gen_validation.as_numpy(): + number_of_validation_batches += 1 - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -5890,118 +5328,109 @@ def test12_PyTorch(self): 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) + 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) - 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()) + 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, + ) - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 2) - assert z.shape == (self.train_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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()) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -6013,81 +5442,63 @@ def test13_TensorFlow(self): 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 = [] + df_minor = ROOT.RDataFrame("myTree", file_name1) + df_major = ROOT.RDataFrame("myTree", file_name2) - 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()) + 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, + ) - 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()) + 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) @@ -6096,37 +5507,46 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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()) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "big_data_major.root" @@ -6161,89 +5581,81 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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, - ) + df1 = ROOT.RDataFrame(tree_name, file_name1) + df2 = ROOT.RDataFrame(tree_name, file_name2) - gen_train, gen_validation = dl.train_test_split(0.3) + 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, + ) - collected_z_train = [] - collected_z_val = [] + gen_train, gen_validation = dl.train_test_split(0.3) - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + collected_z_train = [] + collected_z_val = [] - 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 + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + 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 - for i in range(n_train_batches): - x, y, z = next(iter_train) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + for i in range(n_train_batches): + x, y, z = next(iter_train) - 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()) + 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()) - for _ in range(n_val_batches): - x, y, z = next(iter_val) + 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()) - 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()) + for _ in range(n_val_batches): + x, y, z = next(iter_val) - 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()) + 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()) - 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] + 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()) - 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) + 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] - # 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) + 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 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 + # 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) - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + # 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): define_rdf_major(num_of_entries_major, file_name1) @@ -6311,89 +5723,81 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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, - ) + df1 = ROOT.RDataFrame(tree_name, file_name1) + df2 = ROOT.RDataFrame(tree_name, file_name2) - gen_train, gen_validation = dl.train_test_split(0.3) + 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, + ) - collected_z_train = [] - collected_z_val = [] + gen_train, gen_validation = dl.train_test_split(0.3) - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + collected_z_train = [] + collected_z_val = [] - 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 + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + 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 - for i in range(n_train_batches): - x, y, z = next(iter_train) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + for i in range(n_train_batches): + x, y, z = next(iter_train) - 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()) + 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()) - for _ in range(n_val_batches): - x, y, z = next(iter_val) + 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()) - 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()) + for _ in range(n_val_batches): + x, y, z = next(iter_val) - 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()) + 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()) - 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] + 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()) - 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) + 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] - # 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) + 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 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 + # 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) - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + # 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) @@ -6401,88 +5805,17 @@ 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()) + both_runs_collected_x_val = [] + both_runs_collected_y_val = [] - 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] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, + target="b2", shuffle=False, drop_remainder=False, load_eager=True, @@ -6493,152 +5826,210 @@ def test16_vector_padding(self): 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) for _ in range(self.n_train_batch): - x, y = next(train_iter) - assert x.shape == (2, 5) + 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()) - x, y = next(train_iter) - assert x.shape == (self.train_remainder, 5) + 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 + 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} - 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 len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 - assert num_major_train == 6 - assert num_minor_train == 3 - assert num_major_val == 4 - assert num_minor_val == 2 + 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] - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) + def test16_vector_padding(self): + self.create_vector_file_major() + self.create_vector_file_minor() - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise + 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, + ) -class TestDataLoaderRandomOversampling: - 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" + 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 + + +class TestDataLoaderRandomOversampling: + 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 @@ -6708,308 +6099,267 @@ def create_vector_file_minor(self, num_of_entries=3): .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) + 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_"] + 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) + 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, + ) - 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] + gen_train, gen_validation = dl.train_test_split(0.4) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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] - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_train_batch): x, y = next(train_iter) - assert x.shape == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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 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) - - 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 + 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): 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 = 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) + 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, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + 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_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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()) + for _ in range(self.n_train_batch): x, y = next(train_iter) - assert x.shape == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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] - - # 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 - - 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 + 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): 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 = 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()) + 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, + ) - collected_x = [] - collected_y = [] + gen_train, gen_validation = dl.train_test_split(0.4) - 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) + 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.append(x) - collected_y.append(y) + collected_x = [] + collected_y = [] - assert len(collected_x) == 7 - assert len(collected_y) == 7 + 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) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + assert len(collected_x) == 7 + assert len(collected_y) == 7 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]) + 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] + 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, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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] + 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()) - 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 + 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()) - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) + 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] - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - self.teardown_file(self.file_name3) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -7021,117 +6371,108 @@ def test06_multiple_target_columns(self): 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) + 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) - 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()) + 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, + ) - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 2) - assert z.shape == (self.train_remainder, 1) + 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()) - 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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()) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "multiple_target_columns_major.root" @@ -7145,425 +6486,272 @@ def test07_multiple_input_columns(self): "b3", "(double) b1 * 10" ).Snapshot("myTree", file_name2) - try: - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) + 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) - assert x.shape == (2, 2) - assert y.shape == (2, 1) - collected_x_train.append(x.tolist()) - collected_y_train.append(y.tolist()) + 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, + ) - 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()) + 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 == (self.train_remainder, 2) - assert y.shape == (self.train_remainder, 1) + assert x.shape == (2, 2) + 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] + 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()) - 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 + 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()) - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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] - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): 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") + 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]) - 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, - ) + df_minor_filter = df_minor_duplicate.Filter("rdfentry_ < 3", "name") - gen_train, gen_validation = dl.train_test_split(0.4) + major_entries_before = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_filter_entries_before = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - 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] + 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, + ) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] + gen_train, gen_validation = dl.train_test_split(0.4) - train_iter = iter(gen_train.as_numpy()) - val_iter = iter(gen_validation.as_numpy()) + 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] - 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()) + collected_x_train = [] + collected_x_val = [] + collected_y_train = [] + collected_y_val = [] - 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()) + 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 == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 1) + 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) - - 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) - - 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) - 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] - 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 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 gen_train.as_numpy(): - number_of_training_batches += 1 + 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()) - for _ in gen_validation.as_numpy(): - number_of_validation_batches += 1 + 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 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 + 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 - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + 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) - except: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) - raise + assert num_major_train == 6 + assert num_minor_train == 3 + assert num_major_val == 4 + assert num_minor_val == 2 - def test12_PyTorch(self): - file_name1 = "multiple_target_columns_major.root" - file_name2 = "multiple_target_columns_minor.root" + major_entries_after = df_major.AsNumpy(["rdfentry_"])["rdfentry_"] + minor_filter_entries_after = df_minor_filter.AsNumpy(["rdfentry_"])["rdfentry_"] - 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) + # 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) - 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, - ) + def test10_two_epochs_shuffled(self): + self.create_file_major() + self.create_file_minor() - gen_train, gen_validation = dl.train_test_split(0.4) + 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, + ) - 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] + 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 = [] - collected_z_train = [] - collected_z_val = [] - iter_train = iter(gen_train.as_torch()) - iter_val = iter(gen_validation.as_torch(device="cpu")) + 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) + x, y = next(iter_train) assert x.shape == (2, 1) - assert y.shape == (2, 2) - assert z.shape == (2, 1) + assert y.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) + x, y = next(iter_val) assert x.shape == (2, 1) - assert y.shape == (2, 2) - assert z.shape == (2, 1) + assert y.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) + x, y = next(iter_train) assert x.shape == (self.train_remainder, 1) - assert y.shape == (self.train_remainder, 2) - assert z.shape == (self.train_remainder, 1) + assert y.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 + assert len(flat_x_train) == 9 + assert len(flat_x_val) == 6 + assert len(flat_y_train) == 9 + assert len(flat_y_val) == 6 - 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) + both_epochs_collected_x_val.append(collected_x_val) + both_epochs_collected_y_val.append(collected_y_val) - assert num_major_train == 6 - assert num_minor_train == 3 - assert num_major_val == 4 - assert num_minor_val == 2 + 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] - self.teardown_file(file_name1) - self.teardown_file(file_name2) + def test11_number_of_training_and_validation_batches_remainder(self): + self.create_file_major() + self.create_file_minor() - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - def test13_TensorFlow(self): + 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): file_name1 = "multiple_target_columns_major.root" file_name2 = "multiple_target_columns_minor.root" @@ -7573,80 +6761,175 @@ def test13_TensorFlow(self): 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) + 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, - ) + 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) + 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] + 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()) - collected_x_train = [] - collected_x_val = [] - collected_y_train = [] - collected_y_val = [] - collected_z_train = [] - collected_z_val = [] + 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()) - iter_train = iter(gen_train.as_tensorflow()) - iter_val = iter(gen_validation.as_tensorflow()) + 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 - 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()) + def test13_TensorFlow(self): + file_name1 = "multiple_target_columns_major.root" + file_name2 = "multiple_target_columns_minor.root" - 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()) + 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) + 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) assert x.shape == (2, 1) assert y.shape == (2, 2) @@ -7655,37 +6938,46 @@ def test13_TensorFlow(self): 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 - - self.teardown_file(file_name1) - self.teardown_file(file_name2) + 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()) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + 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): file_name1 = "big_data_major.root" @@ -7719,88 +7011,80 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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, - ) + df1 = ROOT.RDataFrame(tree_name, file_name1) + df2 = ROOT.RDataFrame(tree_name, file_name2) - gen_train, gen_validation = dl.train_test_split(0.3) + 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, + ) - collected_z_train = [] - collected_z_val = [] + gen_train, gen_validation = dl.train_test_split(0.3) - train_remainder = gen_train.last_batch_no_of_rows - val_remainder = gen_validation.last_batch_no_of_rows + collected_z_train = [] + collected_z_val = [] - 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 + train_remainder = gen_train.last_batch_no_of_rows + val_remainder = gen_validation.last_batch_no_of_rows - iter_train = iter(gen_train.as_numpy()) - iter_val = iter(gen_validation.as_numpy()) + 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 - for i in range(n_train_batches): - x, y, z = next(iter_train) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) - 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()) + for i in range(n_train_batches): + x, y, z = next(iter_train) - 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()) + 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()) - for _ in range(n_val_batches): - x, y, z = next(iter_val) + 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()) - 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()) + for _ in range(n_val_batches): + x, y, z = next(iter_val) - 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()) + 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()) - 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] + 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()) - 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) + 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] - # 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) + 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 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 + # 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) - self.teardown_file(file_name1) - self.teardown_file(file_name2) - except: - self.teardown_file(file_name1) - self.teardown_file(file_name2) - raise + # 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) @@ -7808,203 +7092,50 @@ 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) - 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 = [] + both_runs_collected_y_val = [] - 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] - 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} + 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="b1", - max_vec_sizes=max_vec_sizes, + target="b2", 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()) + iter_train = iter(gen_train.as_numpy()) + iter_val = iter(gen_validation.as_numpy()) for _ in range(self.n_train_batch): - x, y = next(train_iter) - assert x.shape == (2, 5) + 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()) - x, y = next(train_iter) - assert x.shape == (self.train_remainder, 5) + 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()) @@ -8014,59 +7145,194 @@ def test16_vector_padding(self): 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 + 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): + self.create_vector_file_major() + self.create_vector_file_minor() + + 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) + 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()) - 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) + 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 num_major_train == 6 - assert num_minor_train == 3 - assert num_major_val == 4 - assert num_minor_val == 2 + 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 - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) + 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) - except: - self.teardown_file(self.file_name4) - self.teardown_file(self.file_name5) - raise + 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): # 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, - ) + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - x = np.concatenate([x for x, _ in dl.as_numpy()]).flatten() - minor = x[x % 2 == 1] + 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, + ) - assert len(minor) == 10 - assert len(np.unique(minor)) > 1 + x = np.concatenate([x for x, _ in dl.as_numpy()]).flatten() + minor = x[x % 2 == 1] - finally: - self.teardown_file(self.file_name1) - self.teardown_file(self.file_name2) + assert len(minor) == 10 + assert len(np.unique(minor)) > 1 class TestDataLoaderSave: @@ -8095,11 +7361,6 @@ def create_file(self, file_name=None, num_of_entries=20, offset=0): .Snapshot(self.tree_name, file_name or self.file_name) ) - def teardown_method(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] @@ -8319,7 +7580,6 @@ class TestDataLoaderImplicitMT: def teardown_method(self): ROOT.DisableImplicitMT() - os.remove(self.file_name) @pytest.mark.skipif( "imt" not in ROOT.gROOT.GetConfigFeatures(), reason="ROOT was built without implicit multi-threading" @@ -8351,11 +7611,6 @@ def create_file(self, file_name, num_of_entries=10, offset=0, output_format="ttr .Snapshot(self.tree_name, file_name, ["b1", "b2"], opts) ) - def teardown_method(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( From 1ae793fb691fd4560ca655276b032edfa3aa1eb3 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:37:07 +0200 Subject: [PATCH 6/8] [ML] Use the fixtures in the RDataLoader tests for creating the inputs Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 1377 +++-------------- 1 file changed, 215 insertions(+), 1162 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 586422f8a0d5a..61f8ec4fa2470 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -4,15 +4,10 @@ import numpy as np import pytest import ROOT -from conftest import from_files, squares, weighted +from conftest import big_data, from_files, squares, vectors, weighted class TestDataLoaderMultipleFiles: - 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 @@ -35,37 +30,7 @@ def weighted_df(self): .Define("b4", "(double) b3 * 10") ) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test01_each_element_is_generated_unshuffled(self, df): entries_before = df.AsNumpy(["rdfentry_"])["rdfentry_"] dl = ROOT.Experimental.ML.RDataLoader( @@ -127,11 +92,7 @@ def test01_each_element_is_generated_unshuffled(self): # check if the dataframe is correctly reset assert np.array_equal(entries_before, entries_after) - def test02_each_element_is_generated_shuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test02_each_element_is_generated_shuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -181,11 +142,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test04_dropping_remainder(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -216,10 +173,7 @@ def test04_dropping_remainder(self): assert len(collected_y) == 3 def test05_more_than_one_file(self): - self.create_file() - self.create_5_entries_file() - - df = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) + df = from_files(squares(10), squares(5, "(int) rdfentry_ + 10")) dl = ROOT.Experimental.ML.RDataLoader( df, @@ -264,16 +218,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df = ROOT.RDataFrame("myTree", file_name) - + def test06_multiple_target_columns(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, batches_in_memory=2, target=["b2", "b4"], @@ -342,13 +289,13 @@ def test06_multiple_target_columns(self): assert results_z_val == flat_z_val 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) - - df = ROOT.RDataFrame("myTree", file_name) + 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, @@ -404,11 +351,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test08_filtered(self, df): dff = df.Filter("b1 % 2 == 0", "name") filter_entries_before = dff.AsNumpy(["rdfentry_"])["rdfentry_"] @@ -465,14 +408,7 @@ def test08_filtered(self): assert np.array_equal(filter_entries_before, filter_entries_after) 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 - ) - - df = ROOT.RDataFrame(tree_name, file_name) + df = from_files(ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1")) dff = df.Filter("b1 % 2 == 0", "name") @@ -512,11 +448,7 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test10_two_epochs_shuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -576,11 +508,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test11_number_of_training_and_validation_batches_remainder(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -606,17 +534,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 0 assert gen_validation.last_batch_no_of_rows == 1 - 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) - - df = ROOT.RDataFrame("myTree", file_name) - + def test12_PyTorch(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, batches_in_memory=2, target=["b2", "b4"], @@ -684,17 +604,9 @@ def test12_PyTorch(self): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - 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) - - df = ROOT.RDataFrame("myTree", file_name) - + def test13_TensorFlow(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, batches_in_memory=2, target=["b2", "b4"], @@ -763,26 +675,14 @@ def test13_TensorFlow(self): assert results_z_val == flat_z_val 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) - - df = ROOT.RDataFrame(tree_name, file_name) + df = from_files(big_data(num_of_entries)) dl = ROOT.Experimental.ML.RDataLoader( df, @@ -856,13 +756,10 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self): - self.create_file() - + def test15_two_runs_set_seed(self, df): 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, @@ -920,9 +817,7 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name3) + df = from_files(vectors(10)) max_vec_sizes = {"v1": 3, "v2": 2} dl = ROOT.Experimental.ML.RDataLoader( @@ -1032,17 +927,9 @@ def test16_vector_padding(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - - df = ROOT.RDataFrame("myTree", file_name) - + def test17_JAX(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, batches_in_memory=2, target=["b2", "b4"], @@ -1110,11 +997,8 @@ def test17_JAX(self): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test18_one_batch_in_memory(self): + def test18_one_batch_in_memory(self, df): # batches_in_memory=1 used to deadlock before delivering the first batch - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) dl = ROOT.Experimental.ML.RDataLoader( df, @@ -1130,16 +1014,10 @@ def test18_one_batch_in_memory(self): 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): + def test19_defined_columns(self, df): # columns defined after reading the file are loaded like the stored ones - self.create_file() - df = ( - ROOT - .RDataFrame(self.tree_name, self.file_name1) - .Define("b3", "b1 * 3") - .Define("v1", "ROOT::RVecD{b2, b2 + 1}") - ) + df = df.Define("b3", "b1 * 3").Define("v1", "ROOT::RVecD{b2, b2 + 1}") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -1164,11 +1042,6 @@ def test19_defined_columns(self): class TestDataLoaderEagerLoading: - 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 @@ -1191,37 +1064,7 @@ def weighted_df(self): .Define("b4", "(double) b3 * 10") ) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test01_each_element_is_generated_unshuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -1276,11 +1119,7 @@ def test01_each_element_is_generated_unshuffled(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test02_each_element_is_generated_shuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test02_each_element_is_generated_shuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -1330,11 +1169,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test04_dropping_remainder(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -1365,10 +1200,7 @@ def test04_dropping_remainder(self): assert len(collected_y) == 3 def test05_more_than_one_file(self): - self.create_file() - self.create_5_entries_file() - - df = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) + df = from_files(squares(10), squares(5, "(int) rdfentry_ + 10")) dl = ROOT.Experimental.ML.RDataLoader( df, @@ -1413,16 +1245,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df = ROOT.RDataFrame("myTree", file_name) - + def test06_multiple_target_columns(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, target=["b2", "b4"], weights="b3", @@ -1491,13 +1316,13 @@ def test06_multiple_target_columns(self): assert results_z_val == flat_z_val 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) - - df = ROOT.RDataFrame("myTree", file_name) + 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, @@ -1553,11 +1378,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test08_filtered(self, df): dff = df.Filter("b1 % 2 == 0", "name") dl = ROOT.Experimental.ML.RDataLoader( @@ -1607,14 +1428,7 @@ def test08_filtered(self): assert results_y_val == flat_y_val 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 - ) - - df = ROOT.RDataFrame(tree_name, file_name) + df = from_files(ROOT.RDataFrame(20).Define("b1", "(Short_t) rdfentry_").Define("b2", "(UShort_t) b1 * b1")) dff = df.Filter("b1 % 2 == 0", "name") @@ -1654,11 +1468,7 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test10_two_epochs_shuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -1718,11 +1528,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name1) - + def test11_number_of_training_and_validation_batches_remainder(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -1748,17 +1554,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 0 assert gen_validation.last_batch_no_of_rows == 1 - 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) - - df = ROOT.RDataFrame("myTree", file_name) - + def test12_PyTorch(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, target=["b2", "b4"], weights="b3", @@ -1826,17 +1624,9 @@ def test12_PyTorch(self): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - 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) - - df = ROOT.RDataFrame("myTree", file_name) - + def test13_TensorFlow(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - df, + weighted_df, batch_size=3, target=["b2", "b4"], weights="b3", @@ -1905,26 +1695,14 @@ def test13_TensorFlow(self): assert results_z_val == flat_z_val 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) - - df = ROOT.RDataFrame(tree_name, file_name) + df = from_files(big_data(num_of_entries)) dl = ROOT.Experimental.ML.RDataLoader( df, @@ -1999,13 +1777,10 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self): - self.create_file() - + def test15_two_runs_set_seed(self, df): 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, @@ -2063,9 +1838,7 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file() - - df = ROOT.RDataFrame(self.tree_name, self.file_name3) + df = from_files(vectors(10)) max_vec_sizes = {"v1": 3, "v2": 2} dl = ROOT.Experimental.ML.RDataLoader( @@ -2175,16 +1948,10 @@ def test16_vector_padding(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test17_defined_columns(self): + def test17_defined_columns(self, df): # columns defined after reading the file are loaded like the stored ones - self.create_file() - df = ( - ROOT - .RDataFrame(self.tree_name, self.file_name1) - .Define("b3", "b1 * 3") - .Define("v1", "ROOT::RVecD{b2, b2 + 1}") - ) + df = df.Define("b3", "b1 * 3").Define("v1", "ROOT::RVecD{b2, b2 + 1}") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -2210,13 +1977,6 @@ def test17_defined_columns(self): class TestDataLoaderEagerLoadingMultipleDataframes: - 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 @@ -2242,56 +2002,7 @@ 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")) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test01_each_element_is_generated_unshuffled(self, df1, df2): df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] @@ -2356,13 +2067,7 @@ def test01_each_element_is_generated_unshuffled(self): 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): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test02_each_element_is_generated_shuffled(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -2412,13 +2117,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test04_dropping_remainder(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -2449,12 +2148,8 @@ def test04_dropping_remainder(self): assert len(collected_y) == 3 def test05_more_than_one_file(self): - self.create_file1() - self.create_file2() - self.create_5_entries_file() - - df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) + 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], @@ -2499,21 +2194,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test06_multiple_target_columns(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -2582,19 +2265,20 @@ def test06_multiple_target_columns(self): assert results_z_val == flat_z_val 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) + 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], @@ -2650,13 +2334,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test08_filtered(self, df1, df2): dff1 = df1.Filter("b1 % 2 == 0", "name") dff2 = df2.Filter("b1 % 2 != 0", "name") @@ -2717,20 +2395,8 @@ def test08_filtered(self): assert np.array_equal(dff2_entries_before, dff2_entries_after) 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 - ) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + 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") @@ -2771,13 +2437,7 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test10_two_epochs_shuffled(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -2837,13 +2497,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test11_number_of_training_and_validation_batches_remainder(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -2869,22 +2523,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 0 assert gen_validation.last_batch_no_of_rows == 1 - 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test12_PyTorch(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -2952,22 +2593,9 @@ def test12_PyTorch(self): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test13_TensorFlow(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -3036,29 +2664,15 @@ def test13_TensorFlow(self): assert results_z_val == flat_z_val 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) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + df1 = from_files(big_data(num_of_entries)) + df2 = from_files(big_data(num_of_entries)) dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], @@ -3133,16 +2747,10 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self): - self.create_file1() - self.create_file2() - + def test15_two_runs_set_seed(self, df1, df2): 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], @@ -3200,11 +2808,8 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file1() - self.create_vector_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name4) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name5) + 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( @@ -3317,7 +2922,7 @@ def test16_vector_padding(self): 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()], + [squares(5), squares(5, "(int) rdfentry_ + 5")], batch_size=3, target="b2", shuffle=False, @@ -3333,13 +2938,6 @@ def test17_dataframes_without_cluster_metadata(self): class TestDataLoaderLazyLoadingMultipleDataframes: - 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 @@ -3365,59 +2963,7 @@ 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")) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test01_each_element_is_generated_unshuffled(self, df1, df2): df1_entries_before = df1.AsNumpy(["rdfentry_"])["rdfentry_"] df2_entries_before = df2.AsNumpy(["rdfentry_"])["rdfentry_"] @@ -3481,13 +3027,7 @@ def test01_each_element_is_generated_unshuffled(self): 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): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test02_each_element_is_generated_shuffled(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -3536,13 +3076,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test04_dropping_remainder(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -3572,12 +3106,8 @@ def test04_dropping_remainder(self): assert len(collected_y) == 3 def test05_more_than_one_file(self): - self.create_file1() - self.create_file2() - self.create_5_entries_file() - - df1 = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name2]) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name3) + 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], @@ -3621,21 +3151,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test06_multiple_target_columns(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -3703,19 +3221,20 @@ def test06_multiple_target_columns(self): assert results_z_val == flat_z_val 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) + 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], @@ -3783,13 +3302,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test08_filtered(self, df1, df2): dff1 = df1.Filter("b1 % 2 == 0", "name") dff2 = df2.Filter("b1 % 2 != 0", "name") @@ -3852,20 +3365,8 @@ def test08_filtered(self): assert np.array_equal(dff2_entries_before, dff2_entries_after) 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 - ) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + 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") @@ -3905,13 +3406,7 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test10_two_epochs_shuffled(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -3970,13 +3465,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test11_number_of_training_and_validation_batches_remainder(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -4001,22 +3490,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 0 assert gen_validation.last_batch_no_of_rows == 1 - 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test12_PyTorch(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -4083,22 +3559,9 @@ def test12_PyTorch(self): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - 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) - - df1 = ROOT.RDataFrame("myTree", file_name1) - df2 = ROOT.RDataFrame("myTree", file_name2) - + def test13_TensorFlow(self, weighted_df1, weighted_df2): dl = ROOT.Experimental.ML.RDataLoader( - [df1, df2], + [weighted_df1, weighted_df2], batch_size=3, target=["b2", "b4"], weights="b3", @@ -4166,29 +3629,15 @@ def test13_TensorFlow(self): assert results_z_val == flat_z_val 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) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + df1 = from_files(big_data(num_of_entries)) + df2 = from_files(big_data(num_of_entries)) dl = ROOT.Experimental.ML.RDataLoader( [df1, df2], @@ -4262,16 +3711,10 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self): - self.create_file1() - self.create_file2() - + def test15_two_runs_set_seed(self, df1, df2): 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], @@ -4328,11 +3771,8 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file1() - self.create_vector_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name4) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name5) + 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( @@ -4441,13 +3881,7 @@ def test16_vector_padding(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test17_shuffled_split_varies_with_seed(self): - self.create_file1() - self.create_file2() - - df1 = ROOT.RDataFrame(self.tree_name, self.file_name1) - df2 = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test17_shuffled_split_varies_with_seed(self, df1, df2): dl1 = ROOT.Experimental.ML.RDataLoader( [df1, df2], batch_size=3, @@ -4476,13 +3910,6 @@ def test17_shuffled_split_varies_with_seed(self): class TestDataLoaderRandomUndersampling: - 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 @@ -4508,56 +3935,7 @@ 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")) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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_"] @@ -4642,13 +4020,7 @@ def test01_each_element_is_generated_unshuffled(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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, @@ -4755,13 +4127,7 @@ def test01_each_element_is_generated_unshuffled_replacement(self): assert num_major_val == 10 assert num_minor_val == 2 - def test02_each_element_is_generated_shuffled(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test02_each_element_is_generated_shuffled(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -4814,13 +4180,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 9 assert len(flat_y_val) == 6 - def test04_dropping_remainder(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test04_dropping_remainder(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -4858,13 +4218,8 @@ def test04_dropping_remainder(self): assert len(collected_x) == 7 assert len(collected_y) == 7 - def test05_more_than_one_file(self): - self.create_file_major() - self.create_file_minor() - self.create_5_entries_file() - - df_major = ROOT.RDataFrame(self.tree_name, [self.file_name1, self.file_name3]) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) + 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], @@ -4923,21 +4278,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - + def test06_multiple_target_columns(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_minor, weighted_df_major], batch_size=2, target=["b2", "b4"], weights="b3", @@ -5038,19 +4381,20 @@ def test06_multiple_target_columns(self): assert num_minor_val == 2 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) - - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) + 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], @@ -5128,13 +4472,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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_"] @@ -5214,13 +4552,7 @@ def test08_filtered(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test10_two_epochs_shuffled(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -5283,13 +4615,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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, @@ -5318,21 +4644,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 1 assert gen_validation.last_batch_no_of_rows == 0 - 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) - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - + def test12_PyTorch(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_major, weighted_df_minor], batch_size=2, target=["b2", "b4"], weights="b3", @@ -5432,21 +4746,9 @@ def test12_PyTorch(self): assert num_major_val == 4 assert num_minor_val == 2 - 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) - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - + def test13_TensorFlow(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_major, weighted_df_minor], batch_size=2, target=["b2", "b4"], weights="b3", @@ -5549,10 +4851,6 @@ def test13_TensorFlow(self): assert num_minor_val == 2 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) @@ -5563,26 +4861,9 @@ def test14_big_data_replacement_false(self): 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) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + 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], @@ -5658,11 +4939,8 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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): - 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) + 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, @@ -5689,10 +4967,6 @@ def test_raises(size_of_batch, num_of_entries_major, num_of_entries_minor, sampl 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) @@ -5705,26 +4979,9 @@ def test14_big_data_replacement_true(self): 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) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + 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], @@ -5801,16 +5058,10 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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() - + def test15_two_runs_set_seed(self, df_major, df_minor): 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], @@ -5870,11 +5121,8 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file_major() - self.create_vector_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name4) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name5) + 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( @@ -6024,13 +5272,6 @@ def test16_vector_padding(self): class TestDataLoaderRandomOversampling: - 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 @@ -6056,56 +5297,7 @@ 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")) - # 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 test01_each_element_is_generated_unshuffled(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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_"] @@ -6189,13 +5381,7 @@ def test01_each_element_is_generated_unshuffled(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test02_each_element_is_generated_shuffled(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -6254,13 +5440,7 @@ def test02_each_element_is_generated_shuffled(self): assert len(flat_y_train) == 9 assert len(flat_y_val) == 6 - def test04_dropping_remainder(self): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test04_dropping_remainder(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -6297,13 +5477,8 @@ def test04_dropping_remainder(self): assert len(collected_x) == 7 assert len(collected_y) == 7 - def test05_more_than_one_file(self): - self.create_file_major() - self.create_file_minor() - self.create_extra_entry_file() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, [self.file_name2, self.file_name3]) + 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], @@ -6361,21 +5536,9 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - 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) - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) - + def test06_multiple_target_columns(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_minor, weighted_df_major], batch_size=2, target=["b2", "b4"], weights="b3", @@ -6475,19 +5638,20 @@ def test06_multiple_target_columns(self): assert num_minor_val == 2 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) - - df_major = ROOT.RDataFrame("myTree", file_name1) - df_minor = ROOT.RDataFrame("myTree", file_name2) + 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], @@ -6564,12 +5728,9 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self): - self.create_file_major() - self.create_file_minor() - - 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]) + 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") @@ -6649,13 +5810,7 @@ def test08_filtered(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + def test10_two_epochs_shuffled(self, df_major, df_minor): dl = ROOT.Experimental.ML.RDataLoader( [df_major, df_minor], batch_size=2, @@ -6717,13 +5872,7 @@ def test10_two_epochs_shuffled(self): 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): - self.create_file_major() - self.create_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) - + 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, @@ -6751,21 +5900,9 @@ def test11_number_of_training_and_validation_batches_remainder(self): assert gen_train.last_batch_no_of_rows == 1 assert gen_validation.last_batch_no_of_rows == 0 - 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) - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - + def test12_PyTorch(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_major, weighted_df_minor], batch_size=2, target=["b2", "b4"], weights="b3", @@ -6864,21 +6001,9 @@ def test12_PyTorch(self): assert num_major_val == 4 assert num_minor_val == 2 - 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) - df_minor = ROOT.RDataFrame("myTree", file_name1) - df_major = ROOT.RDataFrame("myTree", file_name2) - + def test13_TensorFlow(self, weighted_df_major, weighted_df_minor): dl = ROOT.Experimental.ML.RDataLoader( - [df_minor, df_major], + [weighted_df_major, weighted_df_minor], batch_size=2, target=["b2", "b4"], weights="b3", @@ -6980,10 +6105,6 @@ def test13_TensorFlow(self): assert num_minor_val == 2 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) @@ -6993,26 +6114,9 @@ def test14_big_data_replacement_false(self): 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) - - df1 = ROOT.RDataFrame(tree_name, file_name1) - df2 = ROOT.RDataFrame(tree_name, file_name2) + 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], @@ -7088,16 +6192,10 @@ def test(size_of_batch, num_of_entries_major, num_of_entries_minor, sampling_rat 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() - + def test15_two_runs_set_seed(self, df_major, df_minor): 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], @@ -7156,11 +6254,8 @@ def test15_two_runs_set_seed(self): assert both_runs_collected_y_val[0] == both_runs_collected_y_val[1] def test16_vector_padding(self): - self.create_vector_file_major() - self.create_vector_file_minor() - - df_major = ROOT.RDataFrame(self.tree_name, self.file_name4) - df_minor = ROOT.RDataFrame(self.tree_name, self.file_name5) + 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( @@ -7308,13 +6403,8 @@ def test16_vector_padding(self): assert num_major_val == 4 assert num_minor_val == 2 - def test17_set_seed_resamples_different_entries(self): + 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 - self.create_file_major() - self.create_file_minor() - - 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], @@ -7336,8 +6426,6 @@ def test17_set_seed_resamples_different_entries(self): class TestDataLoaderSave: - file_name = "dataloader_save_input.root" - file_name2 = "dataloader_save_input2.root" tree_name = "mytree" out_name = "dataloader_save_output.root" @@ -7352,15 +6440,6 @@ def df(self): 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 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 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] @@ -7370,10 +6449,8 @@ def read_back(self, file_name, column): 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") + def test01_unshuffled(self, df): + df = df.Define("w", "(double) 0.5") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -7394,11 +6471,7 @@ def test01_unshuffled(self): 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) - + def test02_shuffled(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=4, @@ -7424,10 +6497,8 @@ def test02_shuffled(self): # and the entries are not written in their original order assert (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}") + def test03_vector_columns(self, df): + df = df.Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -7451,10 +6522,8 @@ def test03_vector_columns(self): 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}") + df = from_files(self.define_rdf()).Define("v1", "ROOT::VecOps::RVec{ b1, b1 * 10}") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -7477,10 +6546,8 @@ def save_as(file_name, output_format): 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") + def test05_filtered(self, df): + df = df.Filter("b1 % 2 == 0", "even") dl = ROOT.Experimental.ML.RDataLoader( df, @@ -7497,11 +6564,8 @@ def test05_filtered(self): 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) + 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], @@ -7518,11 +6582,7 @@ def test06_multiple_dataframes(self): 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) - + def test07_eager_loading(self, df): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=5, @@ -7537,15 +6597,8 @@ def test07_eager_loading(self): 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 }") - ) + 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, From 74521adb6fed54e07b79f32881047342d6cbc7bb Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Thu, 22 Oct 2026 18:49:10 +0200 Subject: [PATCH 7/8] [ML] Merge the RDataLoader eager and lazy test classes into one parametrized test each The eager and lazy versions of the tests were near copies of each other They are now merged into the sam etest class and use a parametrized fixture to switch between modes. The tests that only existed for one of the modes still only run for it only. Assisted-by: Claude Opus 5.5 --- .../pyroot/pythonizations/test/conftest.py | 6 + .../pythonizations/test/ml_dataloader.py | 2576 +++-------------- 2 files changed, 424 insertions(+), 2158 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/conftest.py b/bindings/pyroot/pythonizations/test/conftest.py index 803cd7ff70e18..576c025a138cd 100644 --- a/bindings/pyroot/pythonizations/test/conftest.py +++ b/bindings/pyroot/pythonizations/test/conftest.py @@ -66,6 +66,12 @@ def in_tmp_path(tmp_path, monkeypatch): 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] diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 61f8ec4fa2470..415a5404b31e6 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -7,7 +7,7 @@ from conftest import big_data, from_files, squares, vectors, weighted -class TestDataLoaderMultipleFiles: +class TestDataLoaderSingleDataframe: # default constants n_train_batch = 2 n_val_batch = 1 @@ -30,7 +30,7 @@ def weighted_df(self): .Define("b4", "(double) b3 * 10") ) - def test01_each_element_is_generated_unshuffled(self, df): + def test01_each_element_is_generated_unshuffled(self, df, load_eager): entries_before = df.AsNumpy(["rdfentry_"])["rdfentry_"] dl = ROOT.Experimental.ML.RDataLoader( @@ -40,6 +40,7 @@ def test01_each_element_is_generated_unshuffled(self, df): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -92,7 +93,7 @@ def test01_each_element_is_generated_unshuffled(self, df): # check if the dataframe is correctly reset assert np.array_equal(entries_before, entries_after) - def test02_each_element_is_generated_shuffled(self, df): + def test02_each_element_is_generated_shuffled(self, df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -100,6 +101,7 @@ def test02_each_element_is_generated_shuffled(self, df): target="b2", shuffle=True, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -142,7 +144,7 @@ def test02_each_element_is_generated_shuffled(self, df): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self, df): + def test04_dropping_remainder(self, df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -150,6 +152,7 @@ def test04_dropping_remainder(self, df): target="b2", shuffle=False, drop_remainder=True, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -172,7 +175,7 @@ def test04_dropping_remainder(self, df): assert len(collected_x) == 3 assert len(collected_y) == 3 - def test05_more_than_one_file(self): + def test05_more_than_one_file(self, load_eager): df = from_files(squares(10), squares(5, "(int) rdfentry_ + 10")) dl = ROOT.Experimental.ML.RDataLoader( @@ -182,6 +185,7 @@ def test05_more_than_one_file(self): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -218,7 +222,7 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test06_multiple_target_columns(self, weighted_df): + def test06_multiple_target_columns(self, weighted_df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( weighted_df, batch_size=3, @@ -227,6 +231,7 @@ def test06_multiple_target_columns(self, weighted_df): weights="b3", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -288,7 +293,7 @@ def test06_multiple_target_columns(self, weighted_df): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test07_multiple_input_columns(self): + def test07_multiple_input_columns(self, load_eager): df = from_files( ROOT .RDataFrame(10) @@ -304,6 +309,7 @@ def test07_multiple_input_columns(self): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -351,7 +357,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self, df): + def test08_filtered(self, df, load_eager): dff = df.Filter("b1 % 2 == 0", "name") filter_entries_before = dff.AsNumpy(["rdfentry_"])["rdfentry_"] @@ -363,6 +369,7 @@ def test08_filtered(self, df): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -407,7 +414,7 @@ def test08_filtered(self, df): # check if the dataframe is correctly reset assert np.array_equal(filter_entries_before, filter_entries_after) - def test09_filtered_last_chunk(self): + 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") @@ -419,6 +426,7 @@ def test09_filtered_last_chunk(self): 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] @@ -448,7 +456,7 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self, df): + def test10_two_epochs_shuffled(self, df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -456,6 +464,7 @@ def test10_two_epochs_shuffled(self, df): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -508,7 +517,7 @@ def test10_two_epochs_shuffled(self, df): 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): + def test11_number_of_training_and_validation_batches_remainder(self, df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( df, batch_size=3, @@ -516,6 +525,7 @@ def test11_number_of_training_and_validation_batches_remainder(self, df): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -534,7 +544,7 @@ def test11_number_of_training_and_validation_batches_remainder(self, df): assert gen_train.last_batch_no_of_rows == 0 assert gen_validation.last_batch_no_of_rows == 1 - def test12_PyTorch(self, weighted_df): + def test12_PyTorch(self, weighted_df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( weighted_df, batch_size=3, @@ -543,6 +553,7 @@ def test12_PyTorch(self, weighted_df): weights="b3", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -604,7 +615,7 @@ def test12_PyTorch(self, weighted_df): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test13_TensorFlow(self, weighted_df): + def test13_TensorFlow(self, weighted_df, load_eager): dl = ROOT.Experimental.ML.RDataLoader( weighted_df, batch_size=3, @@ -613,6 +624,7 @@ def test13_TensorFlow(self, weighted_df): weights="b3", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -674,7 +686,7 @@ def test13_TensorFlow(self, weighted_df): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test14_big_data(self): + def test14_big_data(self, load_eager): entries_in_rdf = randrange(10000, 30000) batch_size = randrange(100, 501) batches_in_memory = randrange(5, 20) @@ -692,6 +704,7 @@ def test(size_of_batch, batches_in_memory, num_of_entries): weights="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.3) @@ -756,7 +769,7 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self, df): + def test15_two_runs_set_seed(self, df, load_eager): both_runs_collected_x_val = [] both_runs_collected_y_val = [] @@ -769,6 +782,7 @@ def test15_two_runs_set_seed(self, df): shuffle=True, drop_remainder=False, set_seed=42, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -816,7 +830,7 @@ def test15_two_runs_set_seed(self, df): 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): + def test16_vector_padding(self, load_eager): df = from_files(vectors(10)) max_vec_sizes = {"v1": 3, "v2": 2} @@ -828,6 +842,7 @@ def test16_vector_padding(self): 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) @@ -851,2000 +866,21 @@ def test16_vector_padding(self): 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): - # 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, - ) - - 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)) - - -class TestDataLoaderEagerLoading: - # 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): - 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) - 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 - - def test02_each_element_is_generated_shuffled(self, df): - 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) - 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): - 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(): - 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): - df = from_files(squares(10), squares(5, "(int) rdfentry_ + 10")) - - 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(): - 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): - dl = ROOT.Experimental.ML.RDataLoader( - weighted_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) - 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): - 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, - 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) - 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): - 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) - 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 - - def test09_filtered_last_chunk(self): - 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, - 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) - 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): - 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) - 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): - 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 - - 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): - dl = ROOT.Experimental.ML.RDataLoader( - weighted_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) - 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): - dl = ROOT.Experimental.ML.RDataLoader( - weighted_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) - 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): - 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=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) - - 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): - both_runs_collected_x_val = [] - both_runs_collected_y_val = [] - - 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) - 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): - df = from_files(vectors(10)) - 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) - 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_defined_columns(self, df): - # 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=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) - assert np.array_equal(x, np.column_stack([entries, entries**2, entries**2 + 1])) - assert np.array_equal(y, (entries * 3).reshape(-1, 1)) - - -class TestDataLoaderEagerLoadingMultipleDataframes: - # 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): - 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) - 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): - 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) - 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): - 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(): - 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): - 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=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(): - 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): - dl = ROOT.Experimental.ML.RDataLoader( - [weighted_df1, weighted_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) - 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): - 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=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) - 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): - 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) - 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 - - 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): - 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=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) - 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): - 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) - 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): - 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 - - 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): - dl = ROOT.Experimental.ML.RDataLoader( - [weighted_df1, weighted_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) - 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): - dl = ROOT.Experimental.ML.RDataLoader( - [weighted_df1, weighted_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) - 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): - 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=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) - - 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): - 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=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) - 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): - 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=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, + 300.0, + 3000.0, + 4.0, + 40.0, 0, - 200.0, - 2000.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, @@ -2855,19 +891,6 @@ def test16_vector_padding(self): 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, @@ -2879,7 +902,7 @@ def test16_vector_padding(self): 900.0, 9000.0, ] - results_y_val = [3.0, 4.0, 8.0, 9.0] + results_y_val = [6.0, 7.0, 8.0, 9.0] collected_x_train = [] collected_x_val = [] @@ -2919,25 +942,122 @@ def test16_vector_padding(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test17_dataframes_without_cluster_metadata(self): - # the dataframes are not read from a file, so their entries are counted with Count() + def test17_JAX(self, weighted_df): dl = ROOT.Experimental.ML.RDataLoader( - [squares(5), squares(5, "(int) rdfentry_ + 5")], + 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=True, + load_eager=load_eager, ) + 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] + 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)) -class TestDataLoaderLazyLoadingMultipleDataframes: +class TestDataLoaderMultipleDataframes: # default constants n_train_batch = 2 n_val_batch = 1 @@ -2963,7 +1083,7 @@ 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): + 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_"] @@ -2973,14 +1093,21 @@ def test01_each_element_is_generated_unshuffled(self, df1, df2): 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] + 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 = [] @@ -3027,13 +1154,14 @@ def test01_each_element_is_generated_unshuffled(self, df1, df2): 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): + 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) @@ -3076,13 +1204,14 @@ def test02_each_element_is_generated_shuffled(self, df1, df2): assert len(flat_y_train) == 6 assert len(flat_y_val) == 4 - def test04_dropping_remainder(self, df1, df2): + 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) @@ -3105,7 +1234,7 @@ def test04_dropping_remainder(self, df1, df2): assert len(collected_x) == 3 assert len(collected_y) == 3 - def test05_more_than_one_file(self): + 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")) @@ -3115,14 +1244,21 @@ def test05_more_than_one_file(self): 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] + 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 = [] @@ -3151,7 +1287,7 @@ def test05_more_than_one_file(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test06_multiple_target_columns(self, weighted_df1, weighted_df2): + def test06_multiple_target_columns(self, weighted_df1, weighted_df2, load_eager): dl = ROOT.Experimental.ML.RDataLoader( [weighted_df1, weighted_df2], batch_size=3, @@ -3159,16 +1295,25 @@ def test06_multiple_target_columns(self, weighted_df1, weighted_df2): 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] + 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 = [] @@ -3220,7 +1365,7 @@ def test06_multiple_target_columns(self, weighted_df1, weighted_df2): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test07_multiple_input_columns(self): + def test07_multiple_input_columns(self, load_eager): df1 = from_files( ROOT .RDataFrame(5) @@ -3242,27 +1387,34 @@ def test07_multiple_input_columns(self): 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] + 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 = [] @@ -3302,7 +1454,7 @@ def test07_multiple_input_columns(self): assert results_y_train == flat_y_train assert results_y_val == flat_y_val - def test08_filtered(self, df1, df2): + def test08_filtered(self, df1, df2, load_eager): dff1 = df1.Filter("b1 % 2 == 0", "name") dff2 = df2.Filter("b1 % 2 != 0", "name") @@ -3311,21 +1463,27 @@ def test08_filtered(self, df1, df2): dl = ROOT.Experimental.ML.RDataLoader( [dff1, dff2], - batch_size=4, + batch_size=3 if load_eager else 4, target="b2", shuffle=False, drop_remainder=False, - load_eager=False, + load_eager=load_eager, ) 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] + 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 = [] @@ -3336,8 +1494,8 @@ def test08_filtered(self, df1, df2): val_iter = iter(gen_validation.as_numpy()) x, y = next(train_iter) - assert x.shape == (4, 1) - assert y.shape == (4, 1) + 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()) @@ -3364,7 +1522,7 @@ def test08_filtered(self, df1, df2): 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): + 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")) @@ -3377,10 +1535,15 @@ def test09_filtered_last_chunk(self): target="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) - 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] + 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 = [] @@ -3406,13 +1569,14 @@ def test09_filtered_last_chunk(self): assert results_x_train == flat_x_train assert results_y_train == flat_y_train - def test10_two_epochs_shuffled(self, df1, df2): + 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) @@ -3465,13 +1629,14 @@ def test10_two_epochs_shuffled(self, df1, df2): 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): + 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) @@ -3490,7 +1655,7 @@ def test11_number_of_training_and_validation_batches_remainder(self, df1, df2): 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): + def test12_PyTorch(self, weighted_df1, weighted_df2, load_eager): dl = ROOT.Experimental.ML.RDataLoader( [weighted_df1, weighted_df2], batch_size=3, @@ -3498,16 +1663,25 @@ def test12_PyTorch(self, weighted_df1, weighted_df2): 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] + 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 = [] @@ -3559,7 +1733,7 @@ def test12_PyTorch(self, weighted_df1, weighted_df2): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test13_TensorFlow(self, weighted_df1, weighted_df2): + def test13_TensorFlow(self, weighted_df1, weighted_df2, load_eager): dl = ROOT.Experimental.ML.RDataLoader( [weighted_df1, weighted_df2], batch_size=3, @@ -3567,16 +1741,25 @@ def test13_TensorFlow(self, weighted_df1, weighted_df2): 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] + 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 = [] @@ -3628,7 +1811,7 @@ def test13_TensorFlow(self, weighted_df1, weighted_df2): assert results_z_train == flat_z_train assert results_z_val == flat_z_val - def test14_big_data(self): + def test14_big_data(self, load_eager): entries_in_rdf = randrange(10000, 30000) batch_size = randrange(100, 501) batches_in_memory = randrange(5, 20) @@ -3647,6 +1830,7 @@ def test(size_of_batch, batches_in_memory, num_of_entries): weights="b2", shuffle=False, drop_remainder=False, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.3) @@ -3711,7 +1895,7 @@ def test(size_of_batch, batches_in_memory, num_of_entries): test(batch_size, batches_in_memory, entries_in_rdf) - def test15_two_runs_set_seed(self, df1, df2): + def test15_two_runs_set_seed(self, df1, df2, load_eager): both_runs_collected_x_val = [] both_runs_collected_y_val = [] @@ -3723,6 +1907,7 @@ def test15_two_runs_set_seed(self, df1, df2): shuffle=True, drop_remainder=False, set_seed=42, + load_eager=load_eager, ) gen_train, gen_validation = dl.train_test_split(0.4) @@ -3770,7 +1955,7 @@ def test15_two_runs_set_seed(self, df1, df2): 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): + 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} @@ -3782,66 +1967,125 @@ def test16_vector_padding(self): 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.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] + 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 = [] @@ -3908,6 +2152,23 @@ def test17_shuffled_split_varies_with_seed(self, df1, df2): 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] + class TestDataLoaderRandomUndersampling: # default constants @@ -6674,13 +4935,12 @@ def load_all(self, inputs, **kwargs): y = np.concatenate([y for _, y in batches]).flatten() return x, y - def test01_single_path(self): + def test01_single_path(self, load_eager): self.create_file(self.file_name1) - for load_eager in (False, True): - 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) + 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 From 74e124108968865cb4382547e2934b71d7a73351 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Fri, 23 Oct 2026 18:29:48 +0200 Subject: [PATCH 8/8] [ML] Split the RDataLoader tests into one file per test class Assisted-by: Claude Opus 5.5 --- .../pyroot/pythonizations/test/CMakeLists.txt | 8 +- .../pythonizations/test/ml_dataloader.py | 5029 ----------------- .../test/ml_dataloader_from_files.py | 125 + .../test/ml_dataloader_implicit_mt.py | 25 + .../test/ml_dataloader_multiple_dataframes.py | 1123 ++++ .../test/ml_dataloader_oversampling.py | 1164 ++++ .../pythonizations/test/ml_dataloader_save.py | 210 + .../test/ml_dataloader_single_dataframe.py | 1060 ++++ .../test/ml_dataloader_undersampling.py | 1373 +++++ 9 files changed, 5087 insertions(+), 5030 deletions(-) delete mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_from_files.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_implicit_mt.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_multiple_dataframes.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_oversampling.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_save.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_single_dataframe.py create mode 100644 bindings/pyroot/pythonizations/test/ml_dataloader_undersampling.py diff --git a/bindings/pyroot/pythonizations/test/CMakeLists.txt b/bindings/pyroot/pythonizations/test/CMakeLists.txt index a3aca94a92e6e..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 GENERIC PYTHON_DEPS jax numpy pytest 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/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py deleted file mode 100644 index 415a5404b31e6..0000000000000 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ /dev/null @@ -1,5029 +0,0 @@ -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 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)) - - -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] - - -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 - - -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 - - -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)) - - -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() - - -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_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__]))