← back to Watches

services/ml/polynomial-regression.js

82 lines

/**
 * Polynomial Regression Model
 */

export class PolynomialRegression {
  constructor(degree = 2) {
    this.degree = degree;
    this.coefficients = [];
    this.rSquared = 0;
    this.rmse = 0;
  }

  fit(X, y) {
    // Create polynomial features
    const features = X.map(x => Array.from({ length: this.degree + 1 }, (_, i) => Math.pow(x, i)));

    // Normal equation: (X'X)^-1 X'y (simplified for small datasets)
    const n = features.length;
    const m = this.degree + 1;

    // XtX matrix
    const XtX = Array.from({ length: m }, () => Array(m).fill(0));
    const Xty = Array(m).fill(0);

    for (let i = 0; i < n; i++) {
      for (let j = 0; j < m; j++) {
        Xty[j] += features[i][j] * y[i];
        for (let k = 0; k < m; k++) {
          XtX[j][k] += features[i][j] * features[i][k];
        }
      }
    }

    // Solve using Gaussian elimination (simplified)
    this.coefficients = this.solveLinearSystem(XtX, Xty);

    // Calculate metrics
    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 { coefficients: this.coefficients, rSquared: this.rSquared, rmse: this.rmse };
  }

  solveLinearSystem(A, b) {
    const n = b.length;
    const aug = A.map((row, i) => [...row, b[i]]);

    for (let i = 0; i < n; i++) {
      let maxRow = i;
      for (let k = i + 1; k < n; k++) if (Math.abs(aug[k][i]) > Math.abs(aug[maxRow][i])) maxRow = k;
      [aug[i], aug[maxRow]] = [aug[maxRow], aug[i]];

      for (let k = i + 1; k < n; k++) {
        const c = aug[k][i] / aug[i][i];
        for (let j = i; j <= n; j++) aug[k][j] -= c * aug[i][j];
      }
    }

    const x = Array(n).fill(0);
    for (let i = n - 1; i >= 0; i--) {
      x[i] = aug[i][n];
      for (let j = i + 1; j < n; j++) x[i] -= aug[i][j] * x[j];
      x[i] /= aug[i][i];
    }
    return x;
  }

  predict(x) {
    return this.coefficients.reduce((sum, coef, i) => sum + coef * Math.pow(x, i), 0);
  }

  toJSON() {
    return { type: 'polynomial', degree: this.degree, coefficients: this.coefficients, rSquared: this.rSquared, rmse: this.rmse };
  }
}

export default PolynomialRegression;