bun-scikit 0.1.1

package/src/model_selection/StratifiedShuffleSplit.ts

@@ -0,0 +1,211 @@
+import type { FoldIndices } from "./KFold";
+
+export interface StratifiedShuffleSplitOptions {
+  nSplits?: number;
+  testSize?: number;
+  trainSize?: number;
+  randomState?: number;
+}
+
+function mulberry32(seed: number): () => number {
+  let state = seed >>> 0;
+  return () => {
+    state += 0x6d2b79f5;
+    let t = Math.imul(state ^ (state >>> 15), 1 | state);
+    t ^= t + Math.imul(t ^ (t >>> 7), 61 | t);
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+function shuffleInPlace(values: number[], random: () => number): void {
+  for (let i = values.length - 1; i > 0; i -= 1) {
+    const j = Math.floor(random() * (i + 1));
+    const tmp = values[i];
+    values[i] = values[j];
+    values[j] = tmp;
+  }
+}
+
+function resolveSplitCounts(
+  sampleCount: number,
+  testSize: number | undefined,
+  trainSize: number | undefined,
+): { trainCount: number; testCount: number } {
+  function resolveSize(value: number | undefined, label: string): number | null {
+    if (value === undefined) {
+      return null;
+    }
+    if (value > 0 && value < 1) {
+      return Math.max(1, Math.floor(sampleCount * value));
+    }
+    if (Number.isInteger(value) && value >= 1 && value < sampleCount) {
+      return value;
+    }
+    throw new Error(
+      `${label} must be a float in (0, 1) or int in [1, n-1]. Got ${value}.`,
+    );
+  }
+
+  const resolvedTest = resolveSize(testSize, "testSize");
+  const resolvedTrain = resolveSize(trainSize, "trainSize");
+
+  if (resolvedTest === null && resolvedTrain === null) {
+    const defaultTest = Math.max(1, Math.floor(sampleCount * 0.1));
+    return {
+      testCount: defaultTest,
+      trainCount: sampleCount - defaultTest,
+    };
+  }
+
+  if (resolvedTest !== null && resolvedTrain !== null) {
+    if (resolvedTest + resolvedTrain > sampleCount) {
+      throw new Error(
+        `trainSize + testSize must be <= sample count (${sampleCount}). Got ${resolvedTrain + resolvedTest}.`,
+      );
+    }
+    return { trainCount: resolvedTrain, testCount: resolvedTest };
+  }
+
+  if (resolvedTest !== null) {
+    return { testCount: resolvedTest, trainCount: sampleCount - resolvedTest };
+  }
+
+  const trainCount = resolvedTrain!;
+  return { trainCount, testCount: sampleCount - trainCount };
+}
+
+export class StratifiedShuffleSplit {
+  private readonly nSplits: number;
+  private readonly testSize?: number;
+  private readonly trainSize?: number;
+  private readonly randomState: number;
+
+  constructor(options: StratifiedShuffleSplitOptions = {}) {
+    this.nSplits = options.nSplits ?? 10;
+    this.testSize = options.testSize;
+    this.trainSize = options.trainSize;
+    this.randomState = options.randomState ?? 42;
+
+    if (!Number.isInteger(this.nSplits) || this.nSplits < 1) {
+      throw new Error(`nSplits must be an integer >= 1. Got ${this.nSplits}.`);
+    }
+  }
+
+  split<TX>(X: TX[], y: number[]): FoldIndices[] {
+    if (!Array.isArray(X) || X.length === 0) {
+      throw new Error("X must be a non-empty array.");
+    }
+    if (!Array.isArray(y) || y.length !== X.length) {
+      throw new Error(`X and y must have the same length. Got ${X.length} and ${y.length}.`);
+    }
+
+    const { trainCount, testCount } = resolveSplitCounts(X.length, this.testSize, this.trainSize);
+    if (trainCount < 1 || testCount < 1) {
+      throw new Error("Both train and test sets must have at least one sample.");
+    }
+
+    const byClass = new Map<number, number[]>();
+    for (let i = 0; i < y.length; i += 1) {
+      const label = y[i];
+      const bucket = byClass.get(label);
+      if (bucket) {
+        bucket.push(i);
+      } else {
+        byClass.set(label, [i]);
+      }
+    }
+
+    if (byClass.size < 2) {
+      throw new Error("StratifiedShuffleSplit requires at least two classes.");
+    }
+
+    const classEntries = Array.from(byClass.entries());
+    let minClassCount = Number.POSITIVE_INFINITY;
+    for (const [, indices] of classEntries) {
+      if (indices.length < minClassCount) {
+        minClassCount = indices.length;
+      }
+    }
+    if (minClassCount < 2) {
+      throw new Error(
+        "The least populated class has fewer than 2 members, which is not enough for stratified splitting.",
+      );
+    }
+
+    const proportions = classEntries.map(([, indices]) => (indices.length * testCount) / X.length);
+    const testPerClass = proportions.map((value) => Math.floor(value));
+
+    // Respect the global test count while keeping per-class allocations feasible.
+    let allocated = testPerClass.reduce((sum, count) => sum + count, 0);
+    let remaining = testCount - allocated;
+    if (remaining > 0) {
+      const classOrder = proportions
+        .map((target, idx) => ({ idx, frac: target - Math.floor(target) }))
+        .sort((a, b) => b.frac - a.frac)
+        .map((entry) => entry.idx);
+
+      let cursor = 0;
+      while (remaining > 0) {
+        const classIdx = classOrder[cursor % classOrder.length];
+        const classCount = classEntries[classIdx][1].length;
+        if (testPerClass[classIdx] < classCount - 1) {
+          testPerClass[classIdx] += 1;
+          remaining -= 1;
+        }
+        cursor += 1;
+      }
+    }
+
+    for (let i = 0; i < testPerClass.length; i += 1) {
+      const classCount = classEntries[i][1].length;
+      if (testPerClass[i] >= classCount) {
+        testPerClass[i] = classCount - 1;
+      }
+    }
+
+    allocated = testPerClass.reduce((sum, count) => sum + count, 0);
+    if (allocated !== testCount) {
+      throw new Error(
+        `Could not allocate exactly ${testCount} stratified test samples. Allocated ${allocated}.`,
+      );
+    }
+
+    const splits: FoldIndices[] = [];
+    for (let splitIndex = 0; splitIndex < this.nSplits; splitIndex += 1) {
+      const random = mulberry32(this.randomState + splitIndex * 104_729);
+      const testIndices: number[] = [];
+
+      for (let classIdx = 0; classIdx < classEntries.length; classIdx += 1) {
+        const classIndices = classEntries[classIdx][1].slice();
+        shuffleInPlace(classIndices, random);
+        const classTestCount = testPerClass[classIdx];
+        for (let i = 0; i < classTestCount; i += 1) {
+          testIndices.push(classIndices[i]);
+        }
+      }
+
+      testIndices.sort((a, b) => a - b);
+      const testMask = new Uint8Array(X.length);
+      for (let i = 0; i < testIndices.length; i += 1) {
+        testMask[testIndices[i]] = 1;
+      }
+
+      const trainIndices: number[] = [];
+      for (let i = 0; i < X.length; i += 1) {
+        if (testMask[i] === 0) {
+          trainIndices.push(i);
+        }
+      }
+
+      if (trainIndices.length !== trainCount || testIndices.length !== testCount) {
+        throw new Error(
+          `Split sizes mismatch. Expected train/test ${trainCount}/${testCount}, got ${trainIndices.length}/${testIndices.length}.`,
+        );
+      }
+
+      splits.push({ trainIndices, testIndices });
+    }
+
+    return splits;
+  }
+}

package/src/model_selection/crossValScore.ts

@@ -0,0 +1,165 @@
+import type { Matrix, Vector } from "../types";
+import { accuracyScore, f1Score, precisionScore, recallScore } from "../metrics/classification";
+import { meanSquaredError, r2Score } from "../metrics/regression";
+import { assertFiniteMatrix, assertFiniteVector, assertVectorLength } from "../utils/validation";
+import { KFold, type FoldIndices } from "./KFold";
+import { StratifiedKFold } from "./StratifiedKFold";
+
+export type BuiltInScoring =
+  | "accuracy"
+  | "f1"
+  | "precision"
+  | "recall"
+  | "r2"
+  | "mean_squared_error"
+  | "neg_mean_squared_error";
+
+export type ScoringFn = (yTrue: Vector, yPred: Vector) => number;
+
+export interface CrossValEstimator {
+  fit(X: Matrix, y: Vector): unknown;
+  predict(X: Matrix): Vector;
+  score?(X: Matrix, y: Vector): number;
+}
+
+export type CrossValSplitter = {
+  split(X: Matrix, y?: Vector): FoldIndices[];
+};
+
+export interface CrossValScoreOptions {
+  cv?: number | CrossValSplitter;
+  scoring?: BuiltInScoring | ScoringFn;
+}
+
+function isBinaryVector(y: Vector): boolean {
+  for (let i = 0; i < y.length; i += 1) {
+    const value = y[i];
+    if (!(value === 0 || value === 1)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+function subsetMatrix(X: Matrix, indices: number[]): Matrix {
+  const out = new Array(indices.length);
+  for (let i = 0; i < indices.length; i += 1) {
+    out[i] = X[indices[i]];
+  }
+  return out;
+}
+
+function subsetVector(y: Vector, indices: number[]): Vector {
+  const out = new Array(indices.length);
+  for (let i = 0; i < indices.length; i += 1) {
+    out[i] = y[indices[i]];
+  }
+  return out;
+}
+
+function resolveBuiltInScorer(scoring: BuiltInScoring): ScoringFn {
+  switch (scoring) {
+    case "accuracy":
+      return accuracyScore;
+    case "f1":
+      return f1Score;
+    case "precision":
+      return precisionScore;
+    case "recall":
+      return recallScore;
+    case "r2":
+      return r2Score;
+    case "mean_squared_error":
+      return meanSquaredError;
+    case "neg_mean_squared_error":
+      return (yTrue, yPred) => -meanSquaredError(yTrue, yPred);
+    default: {
+      const exhaustive: never = scoring;
+      throw new Error(`Unsupported scoring metric: ${exhaustive}`);
+    }
+  }
+}
+
+function resolveFolds(X: Matrix, y: Vector, cv: number | CrossValSplitter | undefined): FoldIndices[] {
+  if (typeof cv === "number") {
+    if (!Number.isInteger(cv) || cv < 2) {
+      throw new Error(`cv must be an integer >= 2. Got ${cv}.`);
+    }
+    if (isBinaryVector(y)) {
+      return new StratifiedKFold({ nSplits: cv, shuffle: false }).split(X, y);
+    }
+    return new KFold({ nSplits: cv, shuffle: false }).split(X, y);
+  }
+
+  if (cv) {
+    return cv.split(X, y);
+  }
+
+  if (isBinaryVector(y)) {
+    return new StratifiedKFold({ nSplits: 5, shuffle: false }).split(X, y);
+  }
+  return new KFold({ nSplits: 5, shuffle: false }).split(X, y);
+}
+
+export function crossValScore(
+  createEstimator: () => CrossValEstimator,
+  X: Matrix,
+  y: Vector,
+  options: CrossValScoreOptions = {},
+): number[] {
+  if (typeof createEstimator !== "function") {
+    throw new Error("createEstimator must be a function returning a new estimator instance.");
+  }
+
+  if (!Array.isArray(X) || X.length === 0) {
+    throw new Error("X must be a non-empty matrix.");
+  }
+  assertFiniteMatrix(X);
+  assertVectorLength(y, X.length);
+  assertFiniteVector(y);
+
+  const folds = resolveFolds(X, y, options.cv);
+  if (folds.length === 0) {
+    throw new Error("Cross-validation splitter produced no folds.");
+  }
+
+  const explicitScorer =
+    typeof options.scoring === "function"
+      ? options.scoring
+      : options.scoring
+        ? resolveBuiltInScorer(options.scoring)
+        : null;
+
+  const scores = new Array<number>(folds.length);
+  for (let foldIndex = 0; foldIndex < folds.length; foldIndex += 1) {
+    const fold = folds[foldIndex];
+    if (fold.trainIndices.length === 0 || fold.testIndices.length === 0) {
+      throw new Error(`Fold ${foldIndex} must have non-empty train and test indices.`);
+    }
+
+    const XTrain = subsetMatrix(X, fold.trainIndices);
+    const yTrain = subsetVector(y, fold.trainIndices);
+    const XTest = subsetMatrix(X, fold.testIndices);
+    const yTest = subsetVector(y, fold.testIndices);
+
+    const estimator = createEstimator();
+    estimator.fit(XTrain, yTrain);
+
+    if (explicitScorer) {
+      const yPred = estimator.predict(XTest);
+      scores[foldIndex] = explicitScorer(yTest, yPred);
+      continue;
+    }
+
+    if (typeof estimator.score === "function") {
+      scores[foldIndex] = estimator.score(XTest, yTest);
+      continue;
+    }
+
+    throw new Error(
+      "Estimator must implement score() when no explicit scoring function is provided.",
+    );
+  }
+
+  return scores;
+}

package/src/model_selection/trainTestSplit.ts

@@ -0,0 +1,82 @@
+export interface TrainTestSplitOptions {
+  testSize?: number;
+  shuffle?: boolean;
+  randomState?: number;
+}
+
+export interface TrainTestSplitResult<TX, TY> {
+  XTrain: TX[];
+  XTest: TX[];
+  yTrain: TY[];
+  yTest: TY[];
+}
+
+function mulberry32(seed: number): () => number {
+  let state = seed >>> 0;
+  return () => {
+    state += 0x6d2b79f5;
+    let t = Math.imul(state ^ (state >>> 15), 1 | state);
+    t ^= t + Math.imul(t ^ (t >>> 7), 61 | t);
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+export function trainTestSplit<TX, TY>(
+  X: TX[],
+  y: TY[],
+  options: TrainTestSplitOptions = {},
+): TrainTestSplitResult<TX, TY> {
+  if (X.length !== y.length) {
+    throw new Error(`X and y must have the same length. Got ${X.length} and ${y.length}.`);
+  }
+
+  if (X.length < 2) {
+    throw new Error("At least two samples are required for train/test splitting.");
+  }
+
+  const shuffle = options.shuffle ?? true;
+  const randomState = options.randomState ?? 42;
+  const testSize = options.testSize ?? 0.25;
+  const sampleCount = X.length;
+
+  let testCount: number;
+  if (testSize > 0 && testSize < 1) {
+    testCount = Math.max(1, Math.floor(sampleCount * testSize));
+  } else if (Number.isInteger(testSize) && testSize >= 1 && testSize < sampleCount) {
+    testCount = testSize;
+  } else {
+    throw new Error(
+      `testSize must be a float in (0, 1) or int in [1, n-1]. Got ${testSize}.`,
+    );
+  }
+
+  const indices = Array.from({ length: sampleCount }, (_, idx) => idx);
+
+  if (shuffle) {
+    const random = mulberry32(randomState);
+    for (let i = indices.length - 1; i > 0; i -= 1) {
+      const j = Math.floor(random() * (i + 1));
+      const tmp = indices[i];
+      indices[i] = indices[j];
+      indices[j] = tmp;
+    }
+  }
+
+  const testIndices = new Set(indices.slice(0, testCount));
+  const XTrain: TX[] = [];
+  const XTest: TX[] = [];
+  const yTrain: TY[] = [];
+  const yTest: TY[] = [];
+
+  for (let i = 0; i < sampleCount; i += 1) {
+    if (testIndices.has(i)) {
+      XTest.push(X[i]);
+      yTest.push(y[i]);
+    } else {
+      XTrain.push(X[i]);
+      yTrain.push(y[i]);
+    }
+  }
+
+  return { XTrain, XTest, yTrain, yTest };
+}

package/src/naive_bayes/GaussianNB.ts

@@ -0,0 +1,148 @@
+import type { ClassificationModel, Matrix, Vector } from "../types";
+import {
+  assertConsistentRowSize,
+  assertFiniteMatrix,
+  assertFiniteVector,
+  validateClassificationInputs,
+} from "../utils/validation";
+import { accuracyScore } from "../metrics/classification";
+
+const DEFAULT_VAR_SMOOTHING = 1e-9;
+
+export interface GaussianNBOptions {
+  varSmoothing?: number;
+}
+
+export class GaussianNB implements ClassificationModel {
+  classes_: Vector = [0, 1];
+  classPrior_: Vector | null = null;
+  theta_: Matrix | null = null;
+  var_: Matrix | null = null;
+
+  private readonly varSmoothing: number;
+  private fittedFeatureCount = 0;
+  private epsilon = 0;
+
+  constructor(options: GaussianNBOptions = {}) {
+    this.varSmoothing = options.varSmoothing ?? DEFAULT_VAR_SMOOTHING;
+    if (!Number.isFinite(this.varSmoothing) || this.varSmoothing < 0) {
+      throw new Error(`varSmoothing must be finite and >= 0. Got ${this.varSmoothing}.`);
+    }
+  }
+
+  fit(X: Matrix, y: Vector): this {
+    validateClassificationInputs(X, y);
+
+    const sampleCount = X.length;
+    const featureCount = X[0].length;
+    this.fittedFeatureCount = featureCount;
+    this.classes_ = [0, 1];
+
+    let maxVariance = 0;
+    for (let j = 0; j < featureCount; j += 1) {
+      let sum = 0;
+      let sumSquares = 0;
+      for (let i = 0; i < sampleCount; i += 1) {
+        const value = X[i][j];
+        sum += value;
+        sumSquares += value * value;
+      }
+      const mean = sum / sampleCount;
+      const variance = sumSquares / sampleCount - mean * mean;
+      if (variance > maxVariance) {
+        maxVariance = variance;
+      }
+    }
+    this.epsilon = this.varSmoothing * maxVariance;
+
+    const priors = new Array<number>(2).fill(0);
+    const means = Array.from({ length: 2 }, () => new Array<number>(featureCount).fill(0));
+    const variances = Array.from({ length: 2 }, () => new Array<number>(featureCount).fill(0));
+    const counts = new Array<number>(2).fill(0);
+
+    for (let i = 0; i < sampleCount; i += 1) {
+      const label = y[i];
+      counts[label] += 1;
+      for (let j = 0; j < featureCount; j += 1) {
+        means[label][j] += X[i][j];
+      }
+    }
+
+    for (let cls = 0; cls < 2; cls += 1) {
+      if (counts[cls] === 0) {
+        throw new Error(`GaussianNB requires both classes to be present. Missing class ${cls}.`);
+      }
+      priors[cls] = counts[cls] / sampleCount;
+      for (let j = 0; j < featureCount; j += 1) {
+        means[cls][j] /= counts[cls];
+      }
+    }
+
+    for (let i = 0; i < sampleCount; i += 1) {
+      const label = y[i];
+      for (let j = 0; j < featureCount; j += 1) {
+        const diff = X[i][j] - means[label][j];
+        variances[label][j] += diff * diff;
+      }
+    }
+
+    for (let cls = 0; cls < 2; cls += 1) {
+      for (let j = 0; j < featureCount; j += 1) {
+        variances[cls][j] = variances[cls][j] / counts[cls] + this.epsilon;
+      }
+    }
+
+    this.classPrior_ = priors;
+    this.theta_ = means;
+    this.var_ = variances;
+    return this;
+  }
+
+  predictProba(X: Matrix): Matrix {
+    if (!this.classPrior_ || !this.theta_ || !this.var_ || this.fittedFeatureCount === 0) {
+      throw new Error("GaussianNB has not been fitted.");
+    }
+
+    assertConsistentRowSize(X);
+    assertFiniteMatrix(X);
+    if (X[0].length !== this.fittedFeatureCount) {
+      throw new Error(
+        `Feature size mismatch. Expected ${this.fittedFeatureCount}, got ${X[0].length}.`,
+      );
+    }
+
+    const outputs = new Array<number[]>(X.length);
+    for (let i = 0; i < X.length; i += 1) {
+      const row = X[i];
+      const logProb = new Array<number>(2).fill(0);
+      for (let cls = 0; cls < 2; cls += 1) {
+        let sum = Math.log(this.classPrior_[cls]);
+        for (let j = 0; j < this.fittedFeatureCount; j += 1) {
+          const variance = this.var_[cls][j];
+          const mean = this.theta_[cls][j];
+          const diff = row[j] - mean;
+          sum += -0.5 * Math.log(2 * Math.PI * variance) - (diff * diff) / (2 * variance);
+        }
+        logProb[cls] = sum;
+      }
+
+      const maxLog = Math.max(logProb[0], logProb[1]);
+      const exp0 = Math.exp(logProb[0] - maxLog);
+      const exp1 = Math.exp(logProb[1] - maxLog);
+      const denom = exp0 + exp1;
+      outputs[i] = [exp0 / denom, exp1 / denom];
+    }
+
+    return outputs;
+  }
+
+  predict(X: Matrix): Vector {
+    const probabilities = this.predictProba(X);
+    return probabilities.map((pair) => (pair[1] >= 0.5 ? 1 : 0));
+  }
+
+  score(X: Matrix, y: Vector): number {
+    assertFiniteVector(y);
+    return accuracyScore(y, this.predict(X));
+  }
+}