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

Author: jobo322

src/__tests__/test.js

@@ -1,6 +1,24 @@
-import { SparseMatrix } from '..';
+import { SparseMatrix } from '../index';
 
 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

@@ -17,11 +17,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;
@@ -30,6 +31,7 @@ export class SparseMatrix {
           }
         }
       }
+      this.elements.maybeShrinkCapacity();
     } else {
       this._init(rows, columns, new HashTable(options), options.threshold);
     }
@@ -132,6 +134,19 @@ export class SparseMatrix {
     return this;
   }
 
+  add(other) {
+    if (typeof other === 'number') return this.addS(other);
+    if (this.rows !== other.rows || this.columns !== other.columns) {
+      throw new RangeError('Matrices dimensions must be equal');
+    }
+
+    other.withEachNonZero((i, j, v) => {
+      this.set(i, j, v + this.get(i, j));
+    });
+
+    return this;
+  }
+
   mmul(other) {
     if (this.columns !== other.rows) {
       // eslint-disable-next-line no-console
@@ -144,14 +159,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;
   }
@@ -165,16 +178,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]);
+    }
+  }
+
   forEachNonZero(callback) {
     this.elements.forEachPair((key, value) => {
       const i = (key / this.columns) | 0;