MetOffice / XBTs_classification

Project for the classification of eXpendable Bathy Thermographs
BSD 3-Clause "New" or "Revised" License
4 stars 2 forks source link

Decision tree visualisation #105

Open stevehadd opened 3 years ago

stevehadd commented 3 years ago

For the paper, it would be really useful to visualise the output tree. Apparently this ias actually really easy with scikit-learn:

from sklearn.tree import DecisionTreeClassifier
from sklearn import tree
from matplotlib import pyplot as plt 

dataset=np.concatenate([mat_original,np.expand_dims(y_km,axis=1)],axis=1)
dt = DecisionTreeClassifier(max_depth=15, min_samples_leaf=30)
classifier=dt.fit(dataset[:,:-1], dataset[:,-1])
fig, axes=plt.subplots(1, 1, figsize=(400, 200), sharex=True, sharey=True)
tree.plot_tree(classifier)
fig.savefig('tree.png')

It would be interesting to combine this code with the idea to make visualisation in a notebook easy. We could create a wrapper of the class that has a _to_html() method so that when you execute a cell with just the tree object name, it prints out a graph, like the nice options you get with iris cubes or pandas dataframes in a jupyter notebook. This could potentially be fed back to scikit learn in some way.