def test_add_train_only():
"""Assert that atom accepts transformers for the train set only."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.add(StandardScaler(), train_only=True)
assert check_scaling(atom.X_train) and not check_scaling(atom.X_test)
len_train, len_test = len(atom.train), len(atom.test)
atom.add(Pruner(), train_only=True)
assert len(atom.train) != len_train and len(atom.test) == len_test
def test_subset_columns():
"""Assert that you can use a subset of the columns."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
# Column indices
cols = atom.columns.copy()
atom.scale(columns=[3, 4])
assert atom.columns == cols # All columns are kept
assert check_scaling(atom.X.iloc[:, [3, 4]])
assert not check_scaling(atom.dataset.iloc[:, [7, 8]])
# Column names
atom.scale(columns=["mean radius", "mean texture"])
assert check_scaling(atom.dataset.iloc[:, [0, 1]])
# Column slice
atom.scale(columns=slice(10, 12))
assert check_scaling(atom.dataset.iloc[:, [10, 11]])
def test_add_complete_dataset():
"""Assert that atom accepts transformers for the complete dataset."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.add(StandardScaler())
assert check_scaling(atom.dataset)
len_dataset = len(atom.dataset)
atom.add(Pruner())
assert len(atom.dataset) != len_dataset
def test_return_scaled_dataset():
"""Assert that the returned dataframe is indeed scaled."""
X = Scaler().fit_transform(X_bin)
assert check_scaling(X)
def test_passthrough_not_scaled():
"""Assert that the features are not scaled when models don't need it."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(["Tree", "RF"])
atom.stacking(passthrough=True)
assert not check_scaling(atom.stack.X.iloc[:, 2:])
def test_scale():
"""Assert that the scale method normalizes the features."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.scale()
assert check_scaling(atom.dataset)
def test_dataset_property():
"""Assert that the dataset property returns scaled data if needed."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["MNB", "LR"])
assert atom.dataset.equals(atom.mnb.dataset)
assert check_scaling(atom.lr.dataset)
def test_X_train_property():
"""Assert that the X_train property returns scaled data if needed."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["MNB", "LR"])
assert atom.X_train.equals(atom.mnb.X_train)
assert check_scaling(atom.lr.X_train)
