forked from Kaggle/docker-python
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_pdpbox.py
26 lines (24 loc) · 1.05 KB
/
test_pdpbox.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import numpy as np
import pandas as pd
from pdpbox import pdp
import unittest
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
class TestPdpbox(unittest.TestCase):
def test_simple_pdp(self):
# set up data
data = pd.read_csv("/input/tests/data/fifa_2018_stats.csv")
y = (data['Man of the Match'] == "Yes")
feature_names = [i for i in data.columns if data[i].dtype in [np.int64]]
X = data[feature_names]
train_X, val_X, train_y, val_y = train_test_split(X, y, random_state=1)
# Build simple model
tree_model = DecisionTreeClassifier(random_state=0,
max_depth=3).fit(train_X, train_y)
# Set up pdp as table
pdp_goals = pdp.pdp_isolate(model=tree_model,
dataset=val_X,
model_features=feature_names,
feature='Goal Scored')
# make plot
pdp.pdp_plot(pdp_goals, 'Goal Scored')