homemade matrix library experiment draft

matrix.js

//@ts-check
'use strict'
/**
@typedef {number|bigint} numeric
@template {numeric} T
@typedef {T[][]} Matrix<T>
*/

// x*A = (x*I)*A ?
// (x^2)(A^2) = (xA)^2 ?

/**
@template T
@param {T} x
@return {T extends numeric ? true : false}
*/
const is_numeric = x => typeof x == 'number' || typeof x == 'bigint'

/**
@template {numeric} T
@param {Matrix<T>} a
@return {boolean}
*/
const is_nxn = a => {
	const size = a.length
	for (const row of a)
		if (row.length !== size)
			return false
	return true
}

/**
@template {numeric} T
@param {Matrix<T>} a
@return {boolean}
*/
const is_scalarizable = a => {
	if (!is_nxn(a) || a[0].length == 0)
		return false

	const a00 = a[0][0]

	for (const [i, row] of a.entries())
		for (const [j, x] of row.entries())
			if (x != (i == j ? a00 : 0))
				return false
	return true
}

/**
@template {boolean} B
@param {number} n
@param {B} bigint
@return {Matrix<B extends true ? bigint : number>}
*/
const matrix_iden = (n = 2, bigint) =>
	Array.from({ length: n }, (_, i) =>
		Array.from({ length: n }, (_, j) =>
			(bigint ? BigInt : Number)(i == j)
		)
	)

/**
@template {numeric} T
@param {Matrix<T>} a
@return {T[]}
*/
const get_column = (a, j) => a.map((a_i, i) => a_i[j])

/**
@template {numeric} T
@param {Matrix<T>} a
@return {T[]}
*/
const get_diag = a => {
	if (!is_nxn(a))
		throw new RangeError('non-square matrix')
	return a.map((a_i, i) => a_i[i])
}

/**
@template {numeric} T
@param {Matrix<T>} a
@return {T}
*/
const trace = a => sum(get_diag(a))

/**
@template {numeric} T
@param {Matrix<T>} a
@param {((v: T, i: number, j: number, matrix: Matrix<T>) => T)} cb
@param {unknown?} thisArg
@return {Matrix<T>}
*/
const matrix_map_entries = (a, cb, thisArg) =>
	a.map((row, i) =>
		row.map((v, j) => cb.call(thisArg, v, i, j, a))
	)

/**
@template {numeric} T
@param {T} x
@param {number} n
@return {Matrix<T>}
*/
const vectorize = (x, n) =>
	mul(x, matrix_iden(n, typeof x == 'bigint'))

/**
@template {numeric} T
@param {Matrix<T>} a
@return {T}
*/
const scalarize = a => {
	if (is_scalarizable(a))
		return a[0][0]
	throw new RangeError('cannot scalarize')
}

/**
@template {numeric} T
@template {T|Matrix<T>} A
@template {T|Matrix<T>} B
@param {A} a
@param {B} b
@return {A extends T ? B extends T ? T : Matrix<T> : Matrix<T>}
*/
const add = (a, b) => {
	if (is_numeric(a) && is_numeric(b))
		return a + b

	if (is_numeric(a) != is_numeric(b))
		is_numeric(a)
			? (a = vectorize(a, b.length))
			: (b = vectorize(b, a.length))

	if (a.length !== b.length)
		throw new RangeError('not same-size')

	return matrix_map_entries(a, (a_i_j, i, j) => a_i_j + b[i][j])
}

// to-do: define comparison

/**
@template {numeric} T
@template {T|Matrix<T>} A
@param {A} a
@return {A extends T ?  T : Matrix<T>}
*/
const neg = a =>
	is_numeric(a) ? -a : mul(typeof a[0][0] == 'bigint' ? -1n : -1, a)

/**
@template {numeric} T
@template {T|Matrix<T>} A
@template {T|Matrix<T>} B
@param {A} a
@param {B} b
@return {A extends T ? B extends T ? T : Matrix<T> : Matrix<T>}
*/
const sub = (a, b) => add(a, neg(b))

/**
@template {numeric} T
@param {T[]} a
@return {T|0}
*/
const sum = a => a.reduce((a, x) => a + x, 0)

/**
@template {numeric} T
@param {T[]} a
@param {T[]} b
@return {T|0}
*/
const dot = (a, b) => sum(a.map((a_i, i) => a_i * b[i]))

/**
@template {numeric} T
@template {T|Matrix<T>} A
@template {T|Matrix<T>} B
@param {A} a
@param {B} b
@return {A extends T ? B extends T ? T : Matrix<T> : Matrix<T>}
*/
const mul = (a, b) => {
	if (is_numeric(a) && is_numeric(b))
		return a * b

	if (is_numeric(a) != is_numeric(b)) {
		if (is_numeric(b))
			[a, b] = [b, a]
		// now `a` is scalar and `b` is matrix
		return matrix_map_entries(b, x => a * x)
	}

	if (a[0].length !== b.length)
		throw new RangeError('cannot multiply')

	return a.map((a_i, i_a) => b.map((b_i, i_b) => b_i))
}