重み付き線形回帰

regression.test.ts

import { linearRegression } from "simple-statistics"

/**
 * 重み付き線形回帰
 * https://chemstat.hatenablog.com/entry/2021/01/23/194405
 * @param pricesY
 * @param pricesX
 * @param weights
 * @returns expectedY = intercept + slope * x
 */
export const weightLinearRegression = (
  pricesY: number[],
  pricesX: number[],
  weights?: number[],
) => {
  const w = weights ? weights : pricesY.map(() => 1)
  const wSum = w.reduce((a, b) => a + b, 0)
  // const xSum = pricesX.reduce((a, b) => a + b, 0)
  // const ySum = pricesY.reduce((a, b) => a + b, 0)
  // const x2 = pricesX.map((x) => x * x).reduce((a, b) => a + b, 0)
  // const y2 = pricesY.map((y) => y * y).reduce((a, b) => a + b, 0)
  // const xy = pricesX.map((x, i) => x * pricesY[i]).reduce((a, b) => a + b, 0)
  const wx2 = pricesX.map((x, i) => x * x * w[i]).reduce((a, b) => a + b, 0)
  const wxy = pricesX
    .map((x, i) => x * pricesY[i] * w[i])
    .reduce((a, b) => a + b, 0)
  const wx = pricesX.map((x, i) => x * w[i]).reduce((a, b) => a + b, 0)
  const wy = pricesY.map((y, i) => y * w[i]).reduce((a, b) => a + b, 0)
  const denominator = wSum * wx2 - wx ** 2
  const slope = (wSum * wxy - wx * wy) / denominator
  const intercept = (wx2 * wy - wx * wxy) / denominator

  return {
    intercept,
    slope,
  }
}

describe("lineRegression", () => {
  test("重みなし", () => {
    // [x, y][x, y]...
    const data = [
      [1, 1],
      [2, 3],
      [3, 2],
      [4, 4],
      [5, 9],
    ]
    const correctResult = linearRegression(data)
    expect(correctResult).toEqual({
      b: -1.2999999999999998,
      m: 1.7,
    })
    const result = weightLinearRegression(
      data.map((d) => d[1]),
      data.map((d) => d[0]),
    )
    expect(result).toEqual({
      intercept: -1.3,
      slope: 1.7,
    })
  })
  test("重みあり", () => {
    // https://chemstat.hatenablog.com/entry/2021/01/23/194405
    const [pricesY, pricesX, weights] = [
      [2, 6, 6, 9, 6],
      [1, 2, 3, 4, 5],
      [1, 0.5, 1, 0.5, 1],
    ]
    const result = weightLinearRegression(pricesY, pricesX, weights)
    expect(result).toEqual({
      slope: 1.0555555555555556,
      intercept: 2.2083333333333335,
    })
  })
  test("重みあり(0をちゃんと無視するかどうか)", () => {
    const [pricesY, pricesX, weights] = [
      [100, 100, 2, 4, 6],
      [100, 100, 3, 4, 5],
      [0, 0, 100, 10, 1],
    ]
    const result = weightLinearRegression(pricesY, pricesX, weights)
    expect(result).toEqual({
      slope: 2,
      intercept: -4,
    })
  })
})