feat: add new matrix addition method and improve performance of mmul and kroneckerProduct method

Author: jobo322

src/__tests__/index.test.js

@@ -3,6 +3,24 @@ import { describe, expect, it } from 'vitest';
 import { SparseMatrix } from '../index.js';
 
 describe('Sparse Matrix', () => {
+  it('add', () => {
+    let m1 = new SparseMatrix([
+      [2, 0, 1],
+      [0, 0, 3],
+      [2, 0, 1],
+    ]);
+    let m2 = new SparseMatrix([
+      [0, 1, 5],
+      [2, 0, 0],
+      [-2, 0, -1],
+    ]);
+    let m3 = m1.add(m2).to2DArray();
+    expect(m3).toStrictEqual([
+      [2, 1, 6],
+      [2, 0, 3],
+      [0, 0, 0],
+    ]);
+  });
   it('mmul', () => {
     let m1 = new SparseMatrix([
       [2, 0, 1],

src/index.js

@@ -18,11 +18,12 @@ export class SparseMatrix {
 
     if (Array.isArray(rows)) {
       const matrix = rows;
-      rows = matrix.length;
-      options = columns || {};
+      const nbRows = matrix.length;
+      const nbColmuns = matrix[0].length;
+      options = columns || { initialCapacity: nbRows * nbColmuns };
       columns = matrix[0].length;
-      this._init(rows, columns, new HashTable(options), options.threshold);
-      for (let i = 0; i < rows; i++) {
+      this._init(nbRows, columns, new HashTable(options), options.threshold);
+      for (let i = 0; i < nbRows; i++) {
         for (let j = 0; j < columns; j++) {
           let value = matrix[i][j];
           if (this.threshold && Math.abs(value) < this.threshold) value = 0;
@@ -31,6 +32,7 @@ export class SparseMatrix {
           }
         }
       }
+      this.elements.maybeShrinkCapacity();
     } else {
       this._init(rows, columns, new HashTable(options), options.threshold);
     }
@@ -169,14 +171,12 @@ export class SparseMatrix {
     const p = other.columns;
 
     const result = new SparseMatrix(m, p);
-    this.forEachNonZero((i, j, v1) => {
-      other.forEachNonZero((k, l, v2) => {
+    this.withEachNonZero((i, j, v1) => {
+      other.withEachNonZero((k, l, v2) => {
         if (j === k) {
           result.set(i, l, result.get(i, l) + v1 * v2);
         }
-        return v2;
       });
-      return v1;
     });
     return result;
   }
@@ -194,16 +194,31 @@ export class SparseMatrix {
     const result = new SparseMatrix(m * p, n * q, {
       initialCapacity: this.cardinality * other.cardinality,
     });
-    this.forEachNonZero((i, j, v1) => {
-      other.forEachNonZero((k, l, v2) => {
-        result.set(p * i + k, q * j + l, v1 * v2);
-        return v2;
+
+    this.withEachNonZero((i, j, v1) => {
+      const pi = p * i;
+      const qj = q * j;
+      other.withEachNonZero((k, l, v2) => {
+        result.set(pi + k, qj + l, v1 * v2);
       });
-      return v1;
     });
     return result;
   }
 
+  withEachNonZero(callback) {
+    const { state, table, values } = this.elements;
+    const nbStates = state.length;
+    const activeIndex = [];
+    for (let i = 0; i < nbStates; i++) {
+      if (state[i] === 1) activeIndex.push(i);
+    }
+    const columns = this.columns;
+    for (const i of activeIndex) {
+      const key = table[i];
+      callback((key / columns) | 0, key % columns, values[i]);
+    }
+  }
+
   /**
    * Calls `callback` for each value in the matrix that is not zero.
    * The callback can return: