Naive Bayes Classifier From Scratch With PyTorch

students_mat = pd.read_csv(os.path.join('data', 'student-mat.csv'))
students_por = pd.read_csv(os.path.join('data', 'student-por.csv'))

# Concatenating students data from math and portuguese class
students = pd.concat([students_mat, students_por], axis=0)

# Averaging three grades into one single grade
students['grade'] = (students['G1'] + students['G1'] + students['G3']) / 3

# Combining weekly and weekend alcohol consumption into a single attribute
students['alc'] = students['Walc'] + students['Dalc']

# Drop the combined columns
students = students.drop(columns=['G1', 'G2', 'G3', 'school'])

students.head(5)

categorical_dict = {}
for col in students.columns:
    # For each column of type object, use sklearn label encoder and add the mapping to a dictionary
    if students[col].dtype == 'object':
        le = LabelEncoder() 
        students[col] = le.fit_transform(students[col])
        categorical_dict[col] = dict(zip(le.classes_, le.transform(le.classes_)))

pprint(categorical_dict)

sns.countplot(students['alc'])

# Converting weekly consumption to classes
students.loc[students['alc'] <= 3, 'alc'] = 0                            # little 
students.loc[(students['alc'] > 3) & (students['alc'] <= 5), 'alc'] = 1  # moderate
students.loc[students['alc'] > 5, 'alc'] = 2                             # severe

# We need to distinguish categorical from numeric values to fit different distributions
# when we will fit the model
numeric_cols = ['age',
                'traveltime',
                'studytime',
                'failures',
                'famrel',
                'freetime',
                'goout',
                'health',
                'absences',
                'grade',
                'alc']

students = students.drop(columns=['Walc', 'Dalc'])

is_categorical = []
for col in students.columns:
    if col in numeric_cols:
        is_categorical.append(0)
    else:
        is_categorical.append(1)
        
# Convert data to torch tensor
X = torch.from_numpy(students.iloc[:, :-1].values).float()
y = torch.from_numpy(students.iloc[:, -1].values).float()

class NaiveBayesClassifier(BaseEstimator):
    """Class for the naive bayes classifier
    
    Inherits from sklearn BaseEstimator class to use cross validation.
    
    Attributes:
        offset: An integer to increment the conditional probabilities
            in order to smooth probabilities to avoid that a posterior
            probability be 0.
        is_categorical: A list containing 0 and 1 for indicating if a
            feature is categorical or numerical.
        nb_features: An integer for the numbers of feature of the data.
        nb_class: An integer for the number of classes in the labels.
        class_probs: A torch tensor for the proportion of each class.
        cond_probs: A torch tensor for the conditional probability of
            having a given value on a certain feature in the population 
            of each class.
    """
    
    def __init__(self, offset=1):
        """Init function for the naive bayes class"""
        self.offset = offset
    
    def fit(self, X, y, **kwargs):
        """Fits the model given data and labels as input
        
        Args:
            X: A torch tensor for the data.
            y: A torch tensor for the labels.
        """
        # It is mandatory to pass a list describing if each feature is categorical or numerical
        if 'is_categorical' not in kwargs:
            raise ValueError('must pass \'is_categorical\' to fit through **kwargs')
        self.is_categorical = kwargs['is_categorical']
        size = X.shape[0]
        
        self.nb_features = X.shape[1]
        y_uvals = y.unique()
        self.nb_class = len(y_uvals)
        # Probability of each class in the training set
        self.class_probs = y.int().bincount().float() / size
        
        features_maxvals = torch.zeros((self.nb_features,), dtype=torch.int32)
        for j in range(self.nb_features):
            features_maxvals[j] = X[:, j].max()
        
        # All the posterior probabilites
        cond_probs = [] 
        for i in range(self.nb_class):
            cond_probs.append([])
            # Group samples by class
            idx = torch.where(y == y_uvals[i])[0]
            elts = X[idx]
            
            size_class = elts.shape[0]
            for j in range(self.nb_features):
                cond_probs[i].append([])
                if self.is_categorical[j]:
                    # If categorical
                    # For each features
                    for k in range(features_maxvals[j] + 1):
                        # Count the number of occurence of each value in this feature given the group class
                        # Divided by the number of samples in the class
                        p_x_k = (torch.where(elts[:, j] == k)[0].shape[0] + self.offset) / size_class
                        # Append to posteriors probabilities
                        cond_probs[i][j].append(p_x_k)
                else:
                    # If numerical
                    features_class = elts[:, j]
                    # Compute mean and std
                    mean = features_class.mean()
                    std = (features_class - mean).pow(2).mean().sqrt()
                    # Store these value to use them for the gaussian likelihood
                    cond_probs[i][j] = [mean, std]
        self.cond_probs = cond_probs
        return 0
    
    def gaussian_likelihood(self, X, mean, std):
        """Computes the gaussian likelihood
        
        Args:
            X: A torch tensor for the data.
            mean: A float for the mean of the gaussian.
            std: A flot for the standard deviation of the gaussian.
        """
        return (1 / (2 * math.pi * std.pow(2))) * torch.exp(-0.5 * ((X - mean) / std).pow(2))
    
    def predict(self, X):
        """Predicts labels given an input
        
        Args:
            X: A torch tensor containing a batch of data.
        """
        if len(X.shape) == 1:
            X = X.unsqueeze(0)
        
        nb_samples = X.shape[0]
        pred_probs = torch.zeros((nb_samples, self.nb_class), dtype=torch.float32)
        for k in range(nb_samples):
            elt = X[k]
            for i in range(self.nb_class):
                # Set probability by the prior (class probability)
                pred_probs[k][i] = self.class_probs[i]
                prob_feature_per_class = self.cond_probs[i]
                for j in range(self.nb_features):
                    if self.is_categorical[j]:
                        # If categorical get the probability of drawing the value of the input on feature j
                        # inside class i
                        pred_probs[k][i] *= prob_feature_per_class[j][elt[j].int()]
                    else:
                        # If numerical, multiply by the gaussian likelihood with parameters
                        # mean and std of the class i on feature j
                        mean, std = prob_feature_per_class[j]
                        pred_probs[k][i] *= self.gaussian_likelihood(elt[j], mean, std)
        # Return the highest probability among all classes
        return pred_probs.argmax(dim=1)

nbc = NaiveBayesClassifier()

fit_params = {'is_categorical' : is_categorical}
cross_val_score(nbc, X, y, cv=5, scoring='accuracy', fit_params=fit_params).mean()