def setUpClass(cls) -> None:
X, y, fi = fetch_creditg_df(preprocess=True)
train_X, test_X, train_y, _ = fair_stratified_train_test_split(X, y, **fi)
cls.train_X = train_X
cls.train_y = train_y
cls.test_X = test_X
cls.fairness_info = fi
def test_fair_stratified_train_test_split(self):
X, y, fairness_info = lale.lib.aif360.fetch_creditg_df(preprocess=False)
z = range(X.shape[0])
(
train_X,
test_X,
train_y,
test_y,
train_z,
test_z,
) = fair_stratified_train_test_split(X, y, z, **fairness_info)
self.assertEqual(train_X.shape[0], train_y.shape[0])
self.assertEqual(train_X.shape[0], len(train_z))
self.assertEqual(test_X.shape[0], test_y.shape[0])
self.assertEqual(test_X.shape[0], len(test_z))
def test_fair_stratified_train_test_split(self):
X = self.creditg_np_num["train_X"]
y = self.creditg_np_num["train_y"]
fairness_info = self.creditg_np_num["fairness_info"]
z = range(X.shape[0])
(
train_X,
test_X,
train_y,
test_y,
train_z,
test_z,
) = fair_stratified_train_test_split(X, y, z, **fairness_info)
self.assertEqual(train_X.shape[0], train_y.shape[0])
self.assertEqual(train_X.shape[0], len(train_z))
self.assertEqual(test_X.shape[0], test_y.shape[0])
self.assertEqual(test_X.shape[0], len(test_z))
self.assertEqual(train_X.shape[0] + test_X.shape[0], X.shape[0])