from sklearn.feature_selection import RFE
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt

cancer = load_breast_cancer()

rng = np.random.RandomState(42)
noise = rng.normal(size=(len(cancer.data), 50))
X_w_noise = np.hstack([cancer.data, noise])

X_train, X_test, y_train, y_test = train_test_split(X_w_noise, cancer.target, random_state=0, test_size=.5)
select = RFE(RandomForestClassifier(n_estimators=100, random_state=42),
             n_features_to_select=40)

select.fit(X_train, y_train)
mask = select.get_support()
plt.matshow(mask.reshape(1, -1), cmap='gray_r')
plt.xlabel("attr num")
plt.yticks([0])
plt.show()

X_train_rfe = select.transform(X_train)
X_test_rfe = select.transform(X_test)

lr = LogisticRegression()
lr.fit(X_train, y_train)
print("x data: {:.3f}".format(lr.score(X_test, y_test)))
score = LogisticRegression().fit(X_train_rfe, y_train).score(X_test_rfe, y_test)
print("rfe data: {:.3f}".format(score))

