import pickle
import pandas as pd
import numpy as np
from sklearn import svm
import matplotlib.pyplot as plt
from typing import List
[docs]
class ToyModel:
"""Toy model for classification."""
[docs]
def __init__(self, X: np.ndarray, y: np.ndarray):
self.X = X
self.y = y
self.min_x = np.min(self.X[:, 0])
self.max_x = np.max(self.X[:, 0])
self.min_y = np.min(self.X[:, 1])
self.max_y = np.max(self.X[:, 1])
self.estimator = svm.SVC()
@property
def n_samples(self) -> int:
"""The number of samples."""
return self.X.shape[0]
@property
def bbox(self) -> List[float]:
return [self.min_x, self.max_x, self.min_y, self.max_y]
@classmethod
def from_dataframe(
cls, df: pd.DataFrame, predictive_features: list[str], target_feature: str
):
X = np.array(df[predictive_features])
y = np.array([1 if x == "AD" else 0 for x in df[target_feature].values])
return cls(X, y)
def fit(self):
self.estimator.fit(self.X, self.y)
def save(self):
with open("model.pkl", "wb") as fp:
pickle.dump(self.estimator, fp)
def plot(self):
xx, yy = np.meshgrid(
np.linspace(self.min_x, self.max_x, 500),
np.linspace(self.min_y, self.max_y, 500),
)
Z = self.estimator.decision_function(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)
plt.imshow(
Z,
interpolation="nearest",
extent=(xx.min(), xx.max(), yy.min(), yy.max()),
aspect="auto",
origin="lower",
cmap=plt.cm.PuOr_r,
)
contours = plt.contour(xx, yy, Z, levels=[0], linewidths=2, linestyles="dashed")
plt.scatter(
self.X[:, 0],
self.X[:, 1],
s=30,
c=self.y,
cmap=plt.cm.Paired,
edgecolors="k",
)
plt.xticks(())
plt.yticks(())
plt.title(f"Number of sample = {self.n_samples}")
plt.axis(self.bbox)
plt.show()