<commit_before>import numpy as np
import scipy as sp
import scipy.special
from scipy.stats import multinomial as dist

def k_categorical(K):
    class Categorical(object):

        problem_type = "classification"
        n_params = K-1

        def __init__(self, params):
            _, N = params.shape
            self.logits = np.zeros((K, N))
            self.logits[1:K,:] = params # default the 0th class logits to 0
            self.probs = sp.special.softmax(self.logits, axis=0)
            # self.dist = dist(n=1, p=self.probs) # scipy doesn't allow vectorized multinomial (!?!?) why allow vectorized versions of the others?
            # this makes me want to refactor all the other code to use lists of distributions, would be more readable imo

        # def __getattr__(self, name):
        #     if name in dir(self.dist):
        #         return getattr(self.dist, name)
        #     return None

        @property
        def params(self):
            names = [f'p{j}' for j in range(self.n_params+1)]
            return {name:p for name, p in zip(names, self.probs)}

        def to_prob(self): 
            return self.probs.T

        def __getitem__(self, key):
            return Categorical(self.logits[1:K,key])

        def __len__(self):
            return self.logits.shape[1]

        def fit(Y):
            _, n = np.unique(Y, return_counts=True)
            p = n/len(Y)
            return np.log(p[1:K]) - np.log(p[0]) 
            # https://math.stackexchange.com/questions/2786600/invert-the-softmax-function

        # log score methods
        def nll(self, Y):
            return -np.log(self.probs[Y, range(len(Y))])

        def D_nll(self, Y):
            return (self.probs.T - np.eye(K)[Y])[:,1:K]

        def fisher_info(self):
            FI = -np.einsum('ji,ki->ijk', self.probs[1:K,:], self.probs[1:K,:])
            d = np.einsum('jii->ij', FI)
            d[:] += self.probs[1:K,:]
            return FI

            # a test:
            # if k==j:
            #     a= FI[i,j,k] == self.probs[k,i] - self.probs[k,i]*self.probs[j,i]
            # else:
            #     a= FI[i,j,k] == -self.probs[k,i]*self.probs[j,i]
            # a

        # crps methods
        def crps(self, Y):
            return np.sum((self.probs - np.eye(K)[Y])**2, axis=1)

        # def D_crps(self, Y):

        # def crps_metric(self):
        #     M = 2 * self.prob ** 2 * np.exp(-2 * self.logit) * (1 + (self.prob / (1 - self.prob)) ** 2)
        #     return M[:, np.newaxis, np.newaxis]

    return Categorical

    Bernoulli = k_categorical(2)
    <commit_msg>Bugfix.<commit_after>import numpy as np
import scipy as sp
import scipy.special
from scipy.stats import multinomial as dist

def k_categorical(K):
    class Categorical(object):

        problem_type = "classification"
        n_params = K-1

        def __init__(self, params):
            _, N = params.shape
            self.logits = np.zeros((K, N))
            self.logits[1:K,:] = params # default the 0th class logits to 0
            self.probs = sp.special.softmax(self.logits, axis=0)
            # self.dist = dist(n=1, p=self.probs) # scipy doesn't allow vectorized multinomial (!?!?) why allow vectorized versions of the others?
            # this makes me want to refactor all the other code to use lists of distributions, would be more readable imo

        # def __getattr__(self, name):
        #     if name in dir(self.dist):
        #         return getattr(self.dist, name)
        #     return None

        @property
        def params(self):
            names = [f'p{j}' for j in range(self.n_params+1)]
            return {name:p for name, p in zip(names, self.probs)}

        def to_prob(self):
            return self.probs.T

        def __getitem__(self, key):
            return Categorical(self.logits[1:K,key])

        def __len__(self):
            return self.logits.shape[1]

        def fit(Y):
            _, n = np.unique(Y, return_counts=True)
            p = n/len(Y)
            return np.log(p[1:K]) - np.log(p[0])
            # https://math.stackexchange.com/questions/2786600/invert-the-softmax-function

        # log score methods
        def nll(self, Y):
            return -np.log(self.probs[Y, range(len(Y))])

        def D_nll(self, Y):
            return (self.probs.T - np.eye(K)[Y])[:,1:K]

        def fisher_info(self):
            FI = -np.einsum('ji,ki->ijk', self.probs[1:K,:], self.probs[1:K,:])
            d = np.einsum('jii->ij', FI)
            d[:] += self.probs[1:K,:]
            return FI

            # a test:
            # if k==j:
            #     a= FI[i,j,k] == self.probs[k,i] - self.probs[k,i]*self.probs[j,i]
            # else:
            #     a= FI[i,j,k] == -self.probs[k,i]*self.probs[j,i]
            # a

        # crps methods
        def crps(self, Y):
            return np.sum((self.probs - np.eye(K)[Y])**2, axis=1)

        # def D_crps(self, Y):

        # def crps_metric(self):
        #     M = 2 * self.prob ** 2 * np.exp(-2 * self.logit) * (1 + (self.prob / (1 - self.prob)) ** 2)
        #     return M[:, np.newaxis, np.newaxis]

    return Categorical

Bernoulli = k_categorical(2)
{'fraction_non_alphanumeric': 0.14208598419124735, 'fraction_numerical': 0.011952959321380375, 'mean_word_length': 1.8303327877795963, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2336098', 'n_tokens_mistral': 1859, 'n_tokens_neox': 1783, 'n_words': 412}