regression2d-ndarray.js

import pool from "ndarray-scratch"
import pack from "ndarray-pack"
import concatRows from "ndarray-concat-rows"
import concatCols from "ndarray-concat-cols"
import qr from "ndarray-householder-qr"

export function regression2d(
  y0,
  y1,
  x0,
  x1
) {
  const n = x0.length
  const z = pool.zeros([n])
  const o = pool.ones([n])
  const d = pool.zeros([6])
  const x = pack(x0)
  const y = pack(x1)
  const X = pack(y0)
  const Y = pack(y1)
  const A = concatRows([
    concatCols([x, y, o, z, z, z]),
    concatCols([z, z, z, x, y, o]),
  ])

  const B = concatRows([X, Y])
  B.shape = [n * 2, 1]

  qr.factor(A, d)
  qr.solve(A, d, B)

  return [...Array(6)].map((_, i) => B.get(i))
}