← back to Watches

services/ml/random-forest.js

83 lines

/**
 * Simplified Random Forest for Price Prediction
 * Uses ensemble of decision stumps
 */

export class RandomForest {
  constructor(nTrees = 10, maxDepth = 3) {
    this.nTrees = nTrees;
    this.maxDepth = maxDepth;
    this.trees = [];
    this.rSquared = 0;
    this.rmse = 0;
  }

  fit(X, y) {
    // X is array of feature arrays, y is target
    const n = X.length;

    for (let t = 0; t < this.nTrees; t++) {
      // Bootstrap sample
      const indices = Array.from({ length: n }, () => Math.floor(Math.random() * n));
      const sampleX = indices.map(i => X[i]);
      const sampleY = indices.map(i => y[i]);

      // Build simple decision tree (stump)
      const tree = this.buildTree(sampleX, sampleY, 0);
      this.trees.push(tree);
    }

    // Calculate metrics on full dataset
    const predictions = X.map(x => this.predict(x));
    const yMean = y.reduce((a, b) => a + b, 0) / n;
    const ssTotal = y.reduce((acc, yi) => acc + Math.pow(yi - yMean, 2), 0);
    const ssRes = y.reduce((acc, yi, i) => acc + Math.pow(yi - predictions[i], 2), 0);
    this.rSquared = 1 - (ssRes / ssTotal);
    this.rmse = Math.sqrt(ssRes / n);

    return { nTrees: this.nTrees, rSquared: this.rSquared, rmse: this.rmse };
  }

  buildTree(X, y, depth) {
    if (depth >= this.maxDepth || X.length < 5) {
      return { leaf: true, value: y.reduce((a, b) => a + b, 0) / y.length };
    }

    // Find best split (simplified: use first feature, median split)
    const featureIdx = Math.floor(Math.random() * X[0].length);
    const values = X.map(x => x[featureIdx]).sort((a, b) => a - b);
    const threshold = values[Math.floor(values.length / 2)];

    const leftIdx = X.map((x, i) => x[featureIdx] <= threshold ? i : -1).filter(i => i >= 0);
    const rightIdx = X.map((x, i) => x[featureIdx] > threshold ? i : -1).filter(i => i >= 0);

    if (leftIdx.length === 0 || rightIdx.length === 0) {
      return { leaf: true, value: y.reduce((a, b) => a + b, 0) / y.length };
    }

    return {
      leaf: false,
      featureIdx,
      threshold,
      left: this.buildTree(leftIdx.map(i => X[i]), leftIdx.map(i => y[i]), depth + 1),
      right: this.buildTree(rightIdx.map(i => X[i]), rightIdx.map(i => y[i]), depth + 1)
    };
  }

  predictTree(tree, x) {
    if (tree.leaf) return tree.value;
    return x[tree.featureIdx] <= tree.threshold ? this.predictTree(tree.left, x) : this.predictTree(tree.right, x);
  }

  predict(x) {
    const predictions = this.trees.map(tree => this.predictTree(tree, x));
    return predictions.reduce((a, b) => a + b, 0) / predictions.length;
  }

  toJSON() {
    return { type: 'random_forest', nTrees: this.nTrees, maxDepth: this.maxDepth, rSquared: this.rSquared, rmse: this.rmse };
  }
}

export default RandomForest;