def test_wrap_polynomial_features(): # load data X = fetch_california_housing(as_frame=True).frame # prepare selectors tr = PolynomialFeatures() tr_wrap = SklearnTransformerWrapper(transformer=PolynomialFeatures()) # Test: # When passing variable list varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] tr_wrap.set_params(variables=varlist) Xt = pd.DataFrame(tr.fit_transform(X[varlist]), columns=tr.get_feature_names_out(varlist)) Xw = tr_wrap.fit_transform(X) pd.testing.assert_frame_equal(Xw, pd.concat([X, Xt], axis=1)) assert Xw.shape[1] == len(X.columns) + len( tr.get_feature_names_out(varlist)) # when variable list is None tr_wrap.set_params(variables=None) Xt = pd.DataFrame(tr.fit_transform(X), columns=tr.get_feature_names_out()) Xw = tr_wrap.fit_transform(X) pd.testing.assert_frame_equal(Xw, pd.concat([X, Xt], axis=1)) assert Xw.shape[1] == len(X.columns) + len( tr.get_feature_names_out(X.columns))