feat(bayes): add CategoricalNB

Author: ObservedObserver

docs/content/docs/apis/bayes/categoricalNB.md

@@ -0,0 +1,24 @@
+---
+title: CategoricalNB
+description: API reference for CategoricalNB
+---
+
+# Bayes.CategoricalNB
+
+```ts
+interface CategoricalNBProps {
+    alpha?: number;
+    forceAlpha?: boolean;
+    fitPrior?: boolean;
+    classPrior?: number[] | null;
+    minCategories?: number | number[] | null;
+}
+constructor(props: CategoricalNBProps = {})
+```
+
+### Example
+```ts
+const clf = new CategoricalNB();
+clf.fit(trainX, trainY);
+const result = clf.predict(testX);
+```

docs/content/docs/apis/bayes/index.md

@@ -4,3 +4,4 @@ description: API reference for Naive Bayes classifiers
 ---
 
 - [BernoulliNB](bernoulliNB.md)
+- [CategoricalNB](categoricalNB.md)

scripts/gen_all.py

@@ -18,6 +18,7 @@ def run(script):
     'gen_optics.py',
     'gen_logistic_regression.py',
     'gen_bernoulli_nb.py',
+    'gen_categorical_nb.py',
     'gen_svc.py',
     'gen_pca.py',
     'gen_spectral_embedding.py',

scripts/gen_categorical_nb.py

@@ -0,0 +1,20 @@
+from sklearn.naive_bayes import CategoricalNB
+import numpy as np
+import json, os
+
+rng = np.random.RandomState(0)
+X = rng.randint(3, size=(50, 5))
+y = rng.randint(2, size=50)
+X_test = rng.randint(3, size=(10, 5))
+clf = CategoricalNB()
+clf.fit(X, y)
+pred = clf.predict(X_test)
+
+os.makedirs('test_data', exist_ok=True)
+with open('test_data/categorical_nb.json', 'w') as f:
+    json.dump({
+        'trainX': X.tolist(),
+        'trainY': y.tolist(),
+        'testX': X_test.tolist(),
+        'expected': pred.tolist()
+    }, f)

src/bayes/__test__/categoricalNB.compare.test.ts

@@ -0,0 +1,12 @@
+import { CategoricalNB } from '../categoricalNB';
+import fs from 'fs';
+import path from 'path';
+
+test('compare with sklearn', () => {
+    const p = path.join(__dirname, '../../../test_data/categorical_nb.json');
+    const data = JSON.parse(fs.readFileSync(p, 'utf8'));
+    const clf = new CategoricalNB();
+    clf.fit(data.trainX, data.trainY);
+    const pred = clf.predict(data.testX);
+    expect(pred).toEqual(data.expected);
+});

src/bayes/__test__/categoricalNB.test.ts

@@ -0,0 +1,20 @@
+import { CategoricalNB } from '../categoricalNB';
+
+test('init', () => {
+    const nb = new CategoricalNB();
+    expect(nb).toBeDefined();
+});
+
+test('simple classification', () => {
+    const X = [
+        [0, 0],
+        [1, 0],
+        [0, 1],
+        [1, 1]
+    ];
+    const Y = [0, 0, 1, 1];
+    const nb = new CategoricalNB();
+    nb.fit(X, Y);
+    const pred = nb.predict(X);
+    expect(pred).toEqual(Y);
+});

src/bayes/categoricalNB.ts

@@ -0,0 +1,133 @@
+import { ClassifierBase } from '../base';
+
+export interface CategoricalNBProps {
+    alpha?: number;
+    forceAlpha?: boolean;
+    fitPrior?: boolean;
+    classPrior?: number[] | null;
+    minCategories?: number | number[] | null;
+}
+
+export class CategoricalNB extends ClassifierBase {
+    private alpha: number;
+    private forceAlpha: boolean;
+    private fitPrior: boolean;
+    private classPrior: number[] | null;
+    private minCategories: number | number[] | null;
+
+    private classes: number[] = [];
+    private classCount: number[] = [];
+    private categoryCount: number[][][] = [];
+    private classLogPrior: number[] = [];
+    private featureLogProb: number[][][] = [];
+    private nCategories: number[] = [];
+
+    constructor(props: CategoricalNBProps = {}) {
+        super();
+        const {
+            alpha = 1.0,
+            forceAlpha = true,
+            fitPrior = true,
+            classPrior = null,
+            minCategories = null
+        } = props;
+        this.alpha = forceAlpha ? alpha : Math.max(alpha, 1e-10);
+        this.forceAlpha = forceAlpha;
+        this.fitPrior = fitPrior;
+        this.classPrior = classPrior;
+        this.minCategories = minCategories;
+    }
+
+    private initCounters(X: number[][]): void {
+        const nFeatures = X[0].length;
+        this.nCategories = new Array(nFeatures).fill(0);
+        for (let j = 0; j < nFeatures; j++) {
+            let maxVal = 0;
+            for (let i = 0; i < X.length; i++) {
+                if (X[i][j] > maxVal) maxVal = X[i][j];
+            }
+            let minCat = 0;
+            if (this.minCategories === null) {
+                minCat = 0;
+            } else if (typeof this.minCategories === 'number') {
+                minCat = this.minCategories;
+            } else {
+                minCat = this.minCategories[j];
+            }
+            this.nCategories[j] = Math.max(maxVal + 1, minCat);
+        }
+        const nClasses = this.classes.length;
+        this.categoryCount = [];
+        this.featureLogProb = [];
+        for (let j = 0; j < nFeatures; j++) {
+            const cats = this.nCategories[j];
+            const mat = Array.from({ length: nClasses }, () => new Array(cats).fill(0));
+            this.categoryCount.push(mat.map(row => row.slice()));
+            this.featureLogProb.push(mat.map(row => row.slice()));
+        }
+        this.classCount = new Array(nClasses).fill(0);
+    }
+
+    public fit(trainX: number[][], trainY: number[]): void {
+        this.classes = Array.from(new Set(trainY)).sort((a, b) => a - b);
+        const classIndex = new Map<number, number>();
+        this.classes.forEach((c, i) => classIndex.set(c, i));
+        this.initCounters(trainX);
+        const nFeatures = trainX[0].length;
+
+        for (let i = 0; i < trainX.length; i++) {
+            const ci = classIndex.get(trainY[i])!;
+            this.classCount[ci] += 1;
+            for (let j = 0; j < nFeatures; j++) {
+                const v = trainX[i][j];
+                if (v >= this.nCategories[j]) continue;
+                this.categoryCount[j][ci][v] += 1;
+            }
+        }
+
+        const nClasses = this.classes.length;
+        if (this.classPrior) {
+            this.classLogPrior = this.classPrior.map(p => Math.log(p));
+        } else if (this.fitPrior) {
+            const totalCount = this.classCount.reduce((a, b) => a + b, 0);
+            this.classLogPrior = this.classCount.map(c => Math.log((c + this.alpha) / (totalCount + nClasses * this.alpha)));
+        } else {
+            this.classLogPrior = new Array(nClasses).fill(Math.log(1 / nClasses));
+        }
+
+        for (let j = 0; j < nFeatures; j++) {
+            for (let c = 0; c < nClasses; c++) {
+                for (let k = 0; k < this.nCategories[j]; k++) {
+                    const count = this.categoryCount[j][c][k];
+                    const denom = this.classCount[c] + this.nCategories[j] * this.alpha;
+                    this.featureLogProb[j][c][k] = Math.log((count + this.alpha) / denom);
+                }
+            }
+        }
+    }
+
+    public predict(testX: number[][]): number[] {
+        const nFeatures = testX[0].length;
+        const nClasses = this.classes.length;
+        const preds: number[] = [];
+        for (const row of testX) {
+            let bestIdx = 0;
+            let bestScore = -Infinity;
+            for (let c = 0; c < nClasses; c++) {
+                let score = this.classLogPrior[c];
+                for (let j = 0; j < nFeatures; j++) {
+                    const v = row[j];
+                    if (v < this.nCategories[j]) {
+                        score += this.featureLogProb[j][c][v];
+                    }
+                }
+                if (score > bestScore) {
+                    bestScore = score;
+                    bestIdx = c;
+                }
+            }
+            preds.push(this.classes[bestIdx]);
+        }
+        return preds;
+    }
+}

src/bayes/index.ts

@@ -1 +1,2 @@
 export { BernoulliNB } from './bernoulliNB';
+export { CategoricalNB } from './categoricalNB';