@wlearn/ensemble 0.1.0

package/package.json

@@ -0,0 +1,26 @@
+{
+  "name": "@wlearn/ensemble",
+  "version": "0.1.0",
+  "description": "Ensemble methods for wlearn: voting, stacking, Caruana selection",
+  "type": "module",
+  "main": "src/index.js",
+  "exports": {
+    ".": "./src/index.js"
+  },
+  "files": [
+    "src/"
+  ],
+  "sideEffects": false,
+  "dependencies": {
+    "@wlearn/core": "0.1.0",
+    "@wlearn/types": "0.1.0"
+  },
+  "scripts": {
+    "test": "node --test test/*.js"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "author": "Anton Zemlyansky",
+  "license": "MIT"
+}

package/src/bagging.js

@@ -0,0 +1,410 @@
+import {
+  encodeBundle, decodeBundle, register, load as registryLoad,
+  normalizeX, normalizeY, accuracy, r2Score,
+  stratifiedKFold, kFold,
+  ValidationError, NotFittedError, DisposedError,
+  lift
+} from '@wlearn/core'
+
+const TYPE_ID_CLS = 'wlearn.ensemble.bagged.classifier@1'
+const TYPE_ID_REG = 'wlearn.ensemble.bagged.regressor@1'
+let _registered = false
+
+/**
+ * K-fold bagged estimator with out-of-fold prediction storage.
+ *
+ * Trains K * nRepeats copies of a base model. Each repeat uses a different
+ * seed for fold assignment. OOF predictions are accumulated (sum + count)
+ * and averaged, matching AutoGluon's BaggedEnsembleModel pattern.
+ */
+export class BaggedEstimator {
+  #spec         // [name, Class, params]
+  #kFold
+  #nRepeats
+  #task
+  #seed
+  #foldModels   // fitted model instances, length K * nRepeats
+  #classes
+  #nClasses = 0
+  #nSamples = 0
+  #oofAccum     // Float64Array: accumulated OOF predictions (sum)
+  #oofCounts    // Uint8Array: per-sample prediction count
+  #fitted = false
+  #disposed = false
+
+  constructor(params = {}) {
+    this.#spec = params.estimator || null
+    this.#kFold = params.kFold || 5
+    this.#nRepeats = params.nRepeats || 1
+    this.#task = params.task || 'classification'
+    this.#seed = params.seed ?? 42
+    this.#foldModels = null
+    this.#classes = null
+    this.#oofAccum = null
+    this.#oofCounts = null
+    BaggedEstimator._register()
+  }
+
+  static async create(params = {}) {
+    return new BaggedEstimator(params)
+  }
+
+  #ensureAlive() {
+    if (this.#disposed) throw new DisposedError('BaggedEstimator has been disposed.')
+  }
+
+  #ensureFitted() {
+    this.#ensureAlive()
+    if (!this.#fitted) throw new NotFittedError('BaggedEstimator is not fitted. Call fit() first.')
+  }
+
+  async fit(X, y) {
+    this.#ensureAlive()
+    const Xn = normalizeX(X)
+    const yn = normalizeY(y)
+    const n = Xn.rows
+    this.#nSamples = n
+
+    if (this.#task === 'classification') {
+      const labelSet = new Set()
+      for (let i = 0; i < yn.length; i++) labelSet.add(yn[i])
+      this.#classes = new Int32Array([...labelSet].sort((a, b) => a - b))
+      this.#nClasses = this.#classes.length
+    }
+
+    // Initialize OOF accumulation
+    if (this.#task === 'classification') {
+      this.#oofAccum = new Float64Array(n * this.#nClasses)
+    } else {
+      this.#oofAccum = new Float64Array(n)
+    }
+    this.#oofCounts = new Uint8Array(n)
+
+    const [, EstClass, params] = this.#spec
+    this.#foldModels = []
+
+    for (let repeat = 0; repeat < this.#nRepeats; repeat++) {
+      const repeatSeed = this.#seed + repeat
+
+      const folds = this.#task === 'classification'
+        ? stratifiedKFold(yn, this.#kFold, { shuffle: true, seed: repeatSeed })
+        : kFold(n, this.#kFold, { shuffle: true, seed: repeatSeed })
+
+      for (const { train, test } of folds) {
+        const Xtrain = _subsetX(Xn, train)
+        const ytrain = _subsetY(yn, train)
+        const Xtest = _subsetX(Xn, test)
+
+        const model = await EstClass.create(params || {})
+        model.fit(Xtrain, ytrain)
+
+        // Accumulate OOF predictions
+        if (this.#task === 'classification') {
+          const proba = await model.predictProba(Xtest)
+          const nc = this.#nClasses
+          for (let i = 0; i < test.length; i++) {
+            const row = test[i]
+            for (let c = 0; c < nc; c++) {
+              this.#oofAccum[row * nc + c] += proba[i * nc + c]
+            }
+          }
+        } else {
+          const preds = await model.predict(Xtest)
+          for (let i = 0; i < test.length; i++) {
+            this.#oofAccum[test[i]] += preds[i]
+          }
+        }
+
+        for (let i = 0; i < test.length; i++) {
+          this.#oofCounts[test[i]] += 1
+        }
+
+        this.#foldModels.push(model)
+      }
+    }
+
+    this.#fitted = true
+    return this
+  }
+
+  predict(X) {
+    this.#ensureFitted()
+    const Xn = normalizeX(X)
+    const n = Xn.rows
+
+    if (this.#task === 'regression') {
+      return this.#averagePredictions(Xn, n)
+    }
+
+    // Classification: average probabilities, then argmax
+    const proba = this.predictProba(Xn)
+    return lift(proba, p => {
+      const nc = this.#nClasses
+      const out = new Float64Array(n)
+      for (let i = 0; i < n; i++) {
+        let bestC = 0, bestV = -Infinity
+        for (let c = 0; c < nc; c++) {
+          if (p[i * nc + c] > bestV) {
+            bestV = p[i * nc + c]
+            bestC = c
+          }
+        }
+        out[i] = this.#classes[bestC]
+      }
+      return out
+    })
+  }
+
+  predictProba(X) {
+    this.#ensureFitted()
+    if (this.#task !== 'classification') {
+      throw new ValidationError('predictProba is only available for classification')
+    }
+
+    const Xn = normalizeX(X)
+    const n = Xn.rows
+    const nc = this.#nClasses
+    const nModels = this.#foldModels.length
+
+    const rawOutputs = []
+    let hasPromise = false
+    for (const model of this.#foldModels) {
+      const out = model.predictProba(Xn)
+      if (out != null && typeof out.then === 'function') hasPromise = true
+      rawOutputs.push(out)
+    }
+
+    const assemble = (outputs) => {
+      const result = new Float64Array(n * nc)
+      for (const proba of outputs) {
+        for (let i = 0; i < n * nc; i++) result[i] += proba[i]
+      }
+      for (let i = 0; i < n * nc; i++) result[i] /= nModels
+      return result
+    }
+
+    return hasPromise ? Promise.all(rawOutputs).then(assemble) : assemble(rawOutputs)
+  }
+
+  score(X, y) {
+    this.#ensureFitted()
+    const preds = this.predict(X)
+    const yn = normalizeY(y)
+    const scorer = this.#task === 'classification' ? accuracy : r2Score
+    return lift(preds, p => scorer(yn, p))
+  }
+
+  #averagePredictions(Xn, n) {
+    const nModels = this.#foldModels.length
+    const rawOutputs = []
+    let hasPromise = false
+    for (const model of this.#foldModels) {
+      const out = model.predict(Xn)
+      if (out != null && typeof out.then === 'function') hasPromise = true
+      rawOutputs.push(out)
+    }
+    const assemble = (outputs) => {
+      const result = new Float64Array(n)
+      for (const preds of outputs) {
+        for (let i = 0; i < n; i++) result[i] += preds[i]
+      }
+      for (let i = 0; i < n; i++) result[i] /= nModels
+      return result
+    }
+    return hasPromise ? Promise.all(rawOutputs).then(assemble) : assemble(rawOutputs)
+  }
+
+  /**
+   * Averaged OOF predictions.
+   * Classification: flat (n * nClasses) row-major probabilities.
+   * Regression: flat (n) predictions.
+   */
+  get oofPredictions() {
+    this.#ensureFitted()
+    const counts = new Uint8Array(this.#oofCounts)
+    for (let i = 0; i < counts.length; i++) {
+      if (counts[i] === 0) counts[i] = 1
+    }
+
+    if (this.#task === 'classification') {
+      const nc = this.#nClasses
+      const oof = new Float64Array(this.#oofAccum)
+      for (let i = 0; i < this.#nSamples; i++) {
+        const c = counts[i]
+        for (let j = 0; j < nc; j++) {
+          oof[i * nc + j] /= c
+        }
+      }
+      return oof
+    }
+
+    const oof = new Float64Array(this.#oofAccum)
+    for (let i = 0; i < this.#nSamples; i++) {
+      oof[i] /= counts[i]
+    }
+    return oof
+  }
+
+  save() {
+    this.#ensureFitted()
+    const typeId = this.#task === 'classification' ? TYPE_ID_CLS : TYPE_ID_REG
+
+    const manifest = {
+      typeId,
+      params: {
+        task: this.#task,
+        kFold: this.#kFold,
+        nRepeats: this.#nRepeats,
+        seed: this.#seed,
+        estimatorName: this.#spec[0],
+        classes: this.#classes ? [...this.#classes] : null,
+        nClasses: this.#nClasses,
+        nSamples: this.#nSamples,
+      },
+    }
+
+    const artifacts = this.#foldModels.map((model, i) => ({
+      id: `fold_${i}`,
+      data: model.save(),
+      mediaType: 'application/x-wlearn-bundle',
+    }))
+
+    // Store OOF data as raw float64 LE bytes
+    const oof = this.oofPredictions
+    const oofBytes = new Uint8Array(oof.buffer, oof.byteOffset, oof.byteLength)
+    artifacts.push({
+      id: 'oof',
+      data: oofBytes,
+      mediaType: 'application/octet-stream',
+    })
+
+    return encodeBundle(manifest, artifacts)
+  }
+
+  static async load(bytes) {
+    const { manifest, toc, blobs } = decodeBundle(bytes)
+    return BaggedEstimator._loadFromParts(manifest, toc, blobs)
+  }
+
+  dispose() {
+    if (this.#disposed) return
+    this.#disposed = true
+    if (this.#foldModels) {
+      for (const m of this.#foldModels) m.dispose()
+    }
+    this.#foldModels = null
+    this.#oofAccum = null
+    this.#oofCounts = null
+  }
+
+  getParams() {
+    return {
+      task: this.#task,
+      kFold: this.#kFold,
+      nRepeats: this.#nRepeats,
+      seed: this.#seed,
+      estimatorName: this.#spec ? this.#spec[0] : null,
+    }
+  }
+
+  setParams(p) {
+    this.#ensureAlive()
+    if (p.kFold !== undefined) this.#kFold = p.kFold
+    if (p.nRepeats !== undefined) this.#nRepeats = p.nRepeats
+    if (p.seed !== undefined) this.#seed = p.seed
+    return this
+  }
+
+  get capabilities() {
+    return {
+      classifier: this.#task === 'classification',
+      regressor: this.#task === 'regression',
+      predictProba: this.#task === 'classification',
+      decisionFunction: false,
+      sampleWeight: false,
+      csr: false,
+      earlyStopping: false,
+    }
+  }
+
+  get isFitted() { return this.#fitted && !this.#disposed }
+  get classes() { return this.#classes }
+
+  // --- Static internals ---
+
+  static _register() {
+    if (_registered) return
+    _registered = true
+    const loader = (manifest, toc, blobs) =>
+      BaggedEstimator._loadFromParts(manifest, toc, blobs)
+    register(TYPE_ID_CLS, loader)
+    register(TYPE_ID_REG, loader)
+  }
+
+  static async _loadFromParts(manifest, toc, blobs) {
+    const p = manifest.params
+    const bag = new BaggedEstimator({
+      task: p.task,
+      kFold: p.kFold || 5,
+      nRepeats: p.nRepeats || 1,
+      seed: p.seed ?? 42,
+    })
+    bag.#classes = p.classes ? new Int32Array(p.classes) : null
+    bag.#nClasses = p.nClasses || 0
+    bag.#nSamples = p.nSamples || 0
+    bag.#spec = [p.estimatorName || 'base', null, null]
+
+    // Load fold models
+    const nFoldModels = bag.#kFold * bag.#nRepeats
+    bag.#foldModels = []
+    for (let i = 0; i < nFoldModels; i++) {
+      const foldId = `fold_${i}`
+      const entry = toc.find(t => t.id === foldId)
+      if (!entry) throw new ValidationError(`No artifact for "${foldId}"`)
+      const blob = blobs.subarray(entry.offset, entry.offset + entry.length)
+      bag.#foldModels.push(await registryLoad(blob))
+    }
+
+    // Load OOF data
+    const oofEntry = toc.find(t => t.id === 'oof')
+    if (oofEntry) {
+      const oofBlob = blobs.subarray(oofEntry.offset, oofEntry.offset + oofEntry.length)
+      const oof = new Float64Array(
+        oofBlob.buffer.slice(oofBlob.byteOffset, oofBlob.byteOffset + oofBlob.byteLength)
+      )
+      bag.#oofAccum = oof
+      bag.#oofCounts = new Uint8Array(bag.#nSamples).fill(1)
+    } else {
+      if (bag.#task === 'classification') {
+        bag.#oofAccum = new Float64Array(bag.#nSamples * bag.#nClasses)
+      } else {
+        bag.#oofAccum = new Float64Array(bag.#nSamples)
+      }
+      bag.#oofCounts = new Uint8Array(bag.#nSamples)
+    }
+
+    bag.#fitted = true
+    return bag
+  }
+}
+
+// --- Subset helpers ---
+
+function _subsetX(X, indices) {
+  const { data, cols } = X
+  const rows = indices.length
+  const out = new Float64Array(rows * cols)
+  for (let i = 0; i < rows; i++) {
+    const srcOff = indices[i] * cols
+    out.set(data.subarray(srcOff, srcOff + cols), i * cols)
+  }
+  return { data: out, rows, cols }
+}
+
+function _subsetY(y, indices) {
+  const out = new (y.constructor)(indices.length)
+  for (let i = 0; i < indices.length; i++) {
+    out[i] = y[indices[i]]
+  }
+  return out
+}

package/src/index.js

@@ -0,0 +1,6 @@
+export { VotingEnsemble } from './voting.js'
+export { StackingEnsemble } from './stacking.js'
+export { BaggedEstimator } from './bagging.js'
+export { caruanaSelect } from './selection.js'
+export { getOofPredictions } from './oof.js'
+export { optimizeWeights, projectSimplex } from './weights.js'

package/src/oof.js

@@ -0,0 +1,96 @@
+import { stratifiedKFold, kFold, normalizeX, normalizeY, ValidationError } from '@wlearn/core'
+
+/**
+ * Generate out-of-fold predictions for a list of estimator specs.
+ *
+ * @param {Array<[string, Function, Object]>} estimatorSpecs - [name, EstimatorClass, params]
+ * @param {Object|number[][]} X - feature matrix
+ * @param {TypedArray|number[]} y - labels
+ * @param {Object} opts
+ * @returns {{ oofPreds: Float64Array[], classes: Int32Array|null }}
+ */
+export async function getOofPredictions(estimatorSpecs, X, y, {
+  cv = 5,
+  seed = 42,
+  task = 'classification',
+} = {}) {
+  const Xn = normalizeX(X)
+  const yn = normalizeY(y)
+  const n = Xn.rows
+
+  const folds = task === 'classification'
+    ? stratifiedKFold(yn, cv, { shuffle: true, seed })
+    : kFold(n, cv, { shuffle: true, seed })
+
+  // Discover classes for classification
+  let classes = null
+  let nClasses = 0
+  if (task === 'classification') {
+    const labelSet = new Set()
+    for (let i = 0; i < yn.length; i++) labelSet.add(yn[i])
+    classes = new Int32Array([...labelSet].sort((a, b) => a - b))
+    nClasses = classes.length
+  }
+
+  const oofPreds = []
+
+  for (const [name, EstimatorClass, params] of estimatorSpecs) {
+    let oof
+    if (task === 'classification') {
+      oof = new Float64Array(n * nClasses)
+    } else {
+      oof = new Float64Array(n)
+    }
+
+    for (const { train, test } of folds) {
+      const Xtrain = _subsetX(Xn, train)
+      const ytrain = _subsetY(yn, train)
+      const Xtest = _subsetX(Xn, test)
+
+      const model = await EstimatorClass.create(params || {})
+      try {
+        model.fit(Xtrain, ytrain)
+        if (task === 'classification') {
+          const proba = await model.predictProba(Xtest)
+          for (let i = 0; i < test.length; i++) {
+            const row = test[i]
+            for (let c = 0; c < nClasses; c++) {
+              oof[row * nClasses + c] = proba[i * nClasses + c]
+            }
+          }
+        } else {
+          const preds = await model.predict(Xtest)
+          for (let i = 0; i < test.length; i++) {
+            oof[test[i]] = preds[i]
+          }
+        }
+      } finally {
+        model.dispose()
+      }
+    }
+    oofPreds.push(oof)
+  }
+
+  return { oofPreds, classes }
+}
+
+// --- Internal helpers (same as core/cv.js) ---
+
+function _subsetX(X, indices) {
+  const { data, cols } = X
+  const rows = indices.length
+  const out = new Float64Array(rows * cols)
+  for (let i = 0; i < rows; i++) {
+    const srcOff = indices[i] * cols
+    out.set(data.subarray(srcOff, srcOff + cols), i * cols)
+  }
+  return { data: out, rows, cols }
+}
+
+function _subsetY(y, indices) {
+  const out = new (y.constructor)(indices.length)
+  for (let i = 0; i < indices.length; i++) {
+    out[i] = y[indices[i]]
+  }
+  return out
+}

package/src/selection.js

@@ -0,0 +1,127 @@
+import { getScorer, normalizeY, ValidationError } from '@wlearn/core'
+import { optimizeWeights } from './weights.js'
+
+/**
+ * Caruana greedy ensemble selection (Caruana et al., 2004).
+ *
+ * Selects a weighted subset from a pool of OOF predictions by greedily
+ * adding the candidate that most improves the ensemble score at each step.
+ * Candidates can be selected multiple times (with replacement).
+ *
+ * @param {Float64Array[]} oofPredictions - per-candidate OOF predictions
+ *   Classification: each is n_samples * n_classes (flat row-major proba)
+ *   Regression: each is n_samples
+ * @param {TypedArray|number[]} yTrue - true labels
+ * @param {Object} opts
+ * @param {boolean} opts.refineWeights - if true, optimize weights after selection
+ * @returns {{ indices: Int32Array, weights: Float64Array, scores: Float64Array }}
+ */
+export function caruanaSelect(oofPredictions, yTrue, {
+  maxSize = 20,
+  scoring = 'accuracy',
+  task = 'classification',
+  nClasses = 0,
+  refineWeights = true,
+} = {}) {
+  const yn = normalizeY(yTrue)
+  const n = yn.length
+  const nCandidates = oofPredictions.length
+
+  if (nCandidates === 0) {
+    throw new ValidationError('caruanaSelect: need at least 1 candidate')
+  }
+
+  const scorerFn = getScorer(scoring)
+
+  // Determine prediction size per sample
+  const predSize = oofPredictions[0].length / n
+  if (predSize !== Math.floor(predSize)) {
+    throw new ValidationError('caruanaSelect: oofPredictions[0].length must be divisible by n')
+  }
+
+  if (task === 'classification' && nClasses === 0) {
+    nClasses = predSize
+  }
+
+  // Current ensemble prediction (running weighted average)
+  const current = new Float64Array(oofPredictions[0].length)
+  const selected = []
+  const scores = []
+
+  for (let t = 0; t < maxSize; t++) {
+    let bestIdx = -1
+    let bestScore = -Infinity
+
+    for (let i = 0; i < nCandidates; i++) {
+      // Trial: ((t) * current + P[i]) / (t + 1)
+      const trial = _trialPredictions(current, oofPredictions[i], t, t + 1)
+      const trialScore = _score(trial, yn, scorerFn, task, nClasses, n)
+      if (trialScore > bestScore) {
+        bestScore = trialScore
+        bestIdx = i
+      }
+    }
+
+    selected.push(bestIdx)
+    scores.push(bestScore)
+
+    // Update running ensemble: current = (t * current + P[bestIdx]) / (t + 1)
+    const P = oofPredictions[bestIdx]
+    for (let j = 0; j < current.length; j++) {
+      current[j] = (t * current[j] + P[j]) / (t + 1)
+    }
+  }
+
+  // Compute weights from selection counts
+  const counts = new Map()
+  for (const idx of selected) {
+    counts.set(idx, (counts.get(idx) || 0) + 1)
+  }
+  const uniqueIndices = new Int32Array([...counts.keys()].sort((a, b) => a - b))
+  const weights = new Float64Array(uniqueIndices.length)
+  for (let i = 0; i < uniqueIndices.length; i++) {
+    weights[i] = counts.get(uniqueIndices[i]) / maxSize
+  }
+
+  const result = {
+    indices: uniqueIndices,
+    weights,
+    scores: new Float64Array(scores),
+  }
+
+  if (refineWeights && uniqueIndices.length > 1) {
+    const selectedOofs = Array.from(uniqueIndices, idx => oofPredictions[idx])
+    result.weights = optimizeWeights(selectedOofs, yn, weights, { task })
+  }
+
+  return result
+}
+
+// --- Internal helpers ---
+
+function _trialPredictions(current, candidate, tCount, tTotal) {
+  const trial = new Float64Array(current.length)
+  for (let j = 0; j < current.length; j++) {
+    trial[j] = (tCount * current[j] + candidate[j]) / tTotal
+  }
+  return trial
+}
+
+function _score(preds, yTrue, scorerFn, task, nClasses, n) {
+  if (task === 'regression') {
+    return scorerFn(yTrue, preds)
+  }
+  // Classification: convert proba to hard predictions via argmax
+  const hardPreds = new Float64Array(n)
+  for (let i = 0; i < n; i++) {
+    let bestC = 0, bestV = -Infinity
+    for (let c = 0; c < nClasses; c++) {
+      if (preds[i * nClasses + c] > bestV) {
+        bestV = preds[i * nClasses + c]
+        bestC = c
+      }
+    }
+    hardPreds[i] = bestC
+  }
+  return scorerFn(yTrue, hardPreds)
+}