linear-regression.js

/**
 * Simple Linear Regression for 2 dimensionnal data
 *
 * @param {Array<Array>} observations - data 
 */
function simpleLinearRegression(observations) {
  // means
  let xSum = 0;
  let ySum = 0;
  const length = observations.length;

  for (let i = 0; i < length; i++) {
    xSum += observations[i][0];
    ySum += observations[i][1];
  }

  const xMean = xSum / length;
  const yMean = ySum / length;

  let sumDiffXMeanSquared = 0; // sum[ pow((x - xMean), 2) ]
  let sumDiffYMeanSquared = 0; // sum[ pow((y - yMean), 2) ]
  let sumDiffXYMean = 0; // sum[ (x - xMean)(y - yMean) ]

  for (let i = 0; i < length; i++) {
    const set = observations[i];
    const diffXMean = set[0] - xMean;
    const diffYMean = set[1] - yMean;

    const diffXMeanSquared = diffXMean * diffXMean;
    const diffYMeanSquared = diffYMean * diffYMean;
    const diffXYMean = diffXMean * diffYMean;

    sumDiffXMeanSquared += diffXMeanSquared;
    sumDiffYMeanSquared += diffYMeanSquared;
    sumDiffXYMean += diffXYMean;
  }

  // Pearson correlation coefficient:
  // cf. https://www.youtube.com/watch?v=2SCg8Kuh0tE
  //
  //                 ∑ [ (x - xMean)(y - yMean) ]
  // r = ------------------------------------------------------
  //     sqrt( ∑ [ pow((x - xMean), 2), pow((y - yMean), 2) ] )
  //
  //
  const r = sumDiffXYMean / Math.sqrt(sumDiffXMeanSquared * sumDiffYMeanSquared);

  // then we have:
  // cf. https://www.youtube.com/watch?v=GhrxgbQnEEU
  //
  // y = a + bx
  // where:
  //         Sy
  // b = r * --
  //         Sx
  //
  // a = yMean - b * xMean
  //
  // S for standard deviation
  //            ∑ [ pow((x - xMean), 2) ]
  // Sx = sqrt( -------------------------  )
  //                      N - 1
  const Sx = Math.sqrt(sumDiffXMeanSquared / (length - 1));
  const Sy = Math.sqrt(sumDiffYMeanSquared / (length - 1));
  const b = r * (Sy / Sx)
  const a = yMean - b * xMean;

  return [a, b];
}