def test_transform_data():
"""Assert that the data is transformed correctly."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.prune(columns=slice(3, 10))
atom.apply(lambda x: x + 2, column="mean radius")
atom.feature_generation(strategy="dfs", n_features=5)
atom.feature_selection(strategy="sfm", solver="lgb", n_features=10)
atom.save(FILE_DIR + "atom", save_data=False)
atom2 = ATOMLoader(FILE_DIR + "atom",
data=(X_bin, y_bin),
transform_data=True)
assert atom2.dataset.shape == atom.dataset.shape
atom3 = ATOMLoader(FILE_DIR + "atom",
data=(X_bin, y_bin),
transform_data=False)
assert atom3.dataset.shape == merge(X_bin, y_bin).shape
def test_getattr_column():
"""Assert that the columns can be accessed as attributes."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.apply(lambda x: np.log(x["mean radius"]), column="log_column")
assert isinstance(atom.log_column, pd.Series)
def test_apply_new_column():
"""Assert that apply can create a new column."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.apply(lambda x: 1, column="new column")
assert atom["new column"].sum() == atom.shape[0]
def test_apply_same_column():
"""Assert that apply can transform an existing column."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.apply(lambda x: 1, column=0)
assert atom["mean radius"].sum() == atom.shape[0]
assert str(atom.pipeline[0]).startswith("FuncTransformer(func=<lambda>")