test: refactor matrix multiplication tests and improve random matrix generation

Author: jobo322

src/__tests__/index.test.js

@@ -1,3 +1,4 @@
+import { Matrix } from 'ml-matrix';
 import { describe, expect, it } from 'vitest';
 
 import { SparseMatrix } from '../index.js';
@@ -38,33 +39,46 @@ describe('Sparse Matrix', () => {
     expect(m3.cardinality).toBe(1);
 
     expect(m3.get(0, 1)).toBe(2);
-    expect(m3.to2DArray()).toStrictEqual([
+    expectMatrixClose(m3.to2DArray(), [
       [0, 2],
       [0, 0],
     ]);
 
     // Compare with dense multiplication
-    const denseM1 = m1.to2DArray();
-    const denseM2 = m2.to2DArray();
-    const expectedDense = denseMatrixMultiply(denseM1, denseM2);
-    expect(m3.to2DArray()).toStrictEqual(expectedDense);
+    const denseM1 = new Matrix(m1.to2DArray());
+    const denseM2 = new Matrix(m2.to2DArray());
+    const expectedDense = denseM1.mmul(denseM2);
+    expectMatrixClose(m3.to2DArray(), expectedDense.to2DArray());
   });
 
   it('mmul', () => {
     const size = 32;
     const density = 0.1;
-    const m1 = randomSparseMatrix(size, size, density);
-    const m2 = randomSparseMatrix(size, size, density);
-    let m3 = m1.mmul(m2);
-
-    const denseM1 = m1.to2DArray();
-    const denseM2 = m2.to2DArray();
-
-    const newSparse = new SparseMatrix(denseM1);
-    expect(newSparse.to2DArray()).toStrictEqual(denseM1);
-    const expectedDense = denseMatrixMultiply(denseM1, denseM2);
+    const A = randomMatrix(size, size, density * size ** 2);
+    const B = randomMatrix(size, size, density * size ** 2);
+    const m1 = new SparseMatrix(A);
+    const m2 = new SparseMatrix(B);
+    const m3 = m1.mmul(m2);
+
+    const denseM1 = new Matrix(A);
+    const denseM2 = new Matrix(B);
+    const expectedDense = denseM1.mmul(denseM2);
+    expectMatrixClose(m3.to2DArray(), expectedDense.to2DArray());
+  });
 
-    expect(m3.to2DArray()).toStrictEqual(expectedDense);
+  it('mmul with low density', () => {
+    const size = 128;
+    const cardinality = 64;
+    const A = randomMatrix(size, size, cardinality);
+    const B = randomMatrix(size, size, cardinality);
+    const m1 = new SparseMatrix(A);
+    const m2 = new SparseMatrix(B);
+    const m3 = m1.mmul(m2);
+
+    const denseM1 = new Matrix(A);
+    const denseM2 = new Matrix(B);
+    const expectedDense = denseM1.mmul(denseM2);
+    expectMatrixClose(m3.to2DArray(), expectedDense.to2DArray());
   });
 
   it('kronecker', () => {
@@ -168,31 +182,37 @@ describe('Banded matrices', () => {
   });
 });
 
-function denseMatrixMultiply(A, B) {
-  const rowsA = A.length;
-  const colsA = A[0].length;
-  const colsB = B[0].length;
-  const result = Array.from({ length: rowsA }, () => Array(colsB).fill(0));
-  for (let i = 0; i < rowsA; i++) {
-    for (let j = 0; j < colsB; j++) {
-      for (let k = 0; k < colsA; k++) {
-        result[i][j] += A[i][k] * B[k][j];
-      }
+/**
+ * Helper to compare two 2D arrays element-wise using toBeCloseTo.
+ */
+function expectMatrixClose(received, expected, precision = 6) {
+  expect(received.length).toBe(expected.length);
+  for (let i = 0; i < received.length; i++) {
+    expect(received[i].length).toBe(expected[i].length);
+    for (let j = 0; j < received[i].length; j++) {
+      expect(received[i][j]).toBeCloseTo(expected[i][j], precision);
     }
   }
-  return result;
 }
 
-function randomSparseMatrix(rows, cols, density = 0.01) {
-  const matrix = [];
-  for (let i = 0; i < rows; i++) {
-    const row = new Float64Array(cols);
-    for (let j = 0; j < cols; j++) {
-      if (Math.random() < density) {
-        row[j] = Math.random() * 10;
-      }
-    }
-    matrix.push(row);
+function randomMatrix(rows, cols, cardinality) {
+  const total = rows * cols;
+  const positions = new Set();
+
+  // Generate unique random positions
+  while (positions.size < cardinality) {
+    positions.add(Math.floor(Math.random() * total));
   }
-  return new SparseMatrix(matrix);
+
+  // Build the matrix with zeros
+  const matrix = Array.from({ length: rows }, () => new Float64Array(cols));
+
+  // Assign random values to the selected positions
+  for (const pos of positions) {
+    const i = Math.floor(pos / cols);
+    const j = pos % cols;
+    matrix[i][j] = Math.random() * 10;
+  }
+
+  return matrix;
 }

src/index.js

@@ -1,4 +1,5 @@
 import HashTable from 'ml-hash-table';
+
 import { cooToCsr } from './utils/cooToCsr.js';
 
 /** @typedef {(row: number, column: number, value: number) => void} WithEachNonZeroCallback */
@@ -169,13 +170,11 @@ export class SparseMatrix {
         'Number of columns of left matrix are not equal to number of rows of right matrix.',
       );
     }
-
     if (this.cardinality < 42 && other.cardinality < 42) {
       return this._mmulSmall(other);
     } else if (other.rows > 100 && other.cardinality < 100) {
       return this._mmulLowDensity(other);
     }
-
     return this._mmulMediumDensity(other);
   }