chore: refactor getNonZeros method to simplify format handling and improve performance

Author: jobo322

src/__tests__/getNonZeros.test.js

@@ -1,4 +1,5 @@
 import { describe, it, expect } from 'vitest';
+
 import { SparseMatrix } from '../index.js';
 
 describe('Sparse Matrix', () => {
@@ -22,17 +23,10 @@ describe('Sparse Matrix', () => {
     });
 
     // CSR format
-    expect(m2.getNonZeros({ format: 'csr' })).toEqual({
+    expect(m2.getNonZeros({ format: true })).toEqual({
       rows: Float64Array.from([0, 0, 4, 7, 7, 11]),
       columns: Float64Array.from([0, 3, 4, 5, 1, 4, 5, 0, 3, 4, 5]),
       values: Float64Array.from([1, 2, 1, 1, 3, 5, 5, 1, 1, 9, 9]),
     });
-
-    //CSC format
-    expect(m2.getNonZeros({ format: 'csc' })).toEqual({
-      rows: Float64Array.from([1, 4, 2, 1, 4, 1, 2, 4, 1, 2, 4]),
-      columns: Float64Array.from([0, 2, 3, 3, 5, 8, 11]),
-      values: Float64Array.from([1, 1, 3, 2, 1, 1, 5, 9, 1, 5, 9]),
-    });
   });
 });

src/index.js

@@ -163,40 +163,97 @@ export class SparseMatrix {
       );
     }
 
+    if (this.cardinality < 42 && other.cardinality < 42) {
+      return this._mmulSmall(other);
+    } else if (other.rows > 100 && other.cardinality < 110) {
+      return this._mmulLowDensity(other);
+    }
+
+    return this._mmulMediumDensity(other);
+  }
+
+  _mmulSmall(other) {
     const m = this.rows;
     const p = other.columns;
+    const {
+      columns: otherCols,
+      rows: otherRows,
+      values: otherValues,
+    } = other.getNonZeros();
 
-    const result = matrixCreateEmpty(m, p);
+    const nbOtherActive = otherCols.length;
+    const result = new SparseMatrix(m, p);
+    this.withEachNonZero((i, j, v1) => {
+      for (let o = 0; o < nbOtherActive; o++) {
+        if (j === otherRows[o]) {
+          const l = otherCols[o];
+          result.set(i, l, result.get(i, l) + otherValues[o] * v1);
+        }
+      }
+    });
+    return result;
+  }
 
+  _mmulLowDensity(other) {
+    const m = this.rows;
+    const p = other.columns;
     const {
       columns: otherCols,
       rows: otherRows,
       values: otherValues,
-    } = other.getNonZeros({ format: 'csr' });
+    } = other.getNonZeros();
     const {
       columns: thisCols,
       rows: thisRows,
       values: thisValues,
-    } = this.getNonZeros({ format: 'csc' });
-
-    const thisNbCols = this.columns;
-    for (let t = 0; t < thisNbCols; t++) {
-      const tStart = thisCols[t];
-      const tEnd = thisCols[t + 1];
-      const oStart = otherRows[t];
-      const oEnd = otherRows[t + 1];
-      for (let f = tStart; f < tEnd; f++) {
-        for (let k = oStart; k < oEnd; k++) {
-          const i = thisRows[f];
-          const l = otherCols[k];
-          result[i][l] += thisValues[f] * otherValues[k];
+    } = this.getNonZeros();
+
+    const result = new SparseMatrix(m, p);
+    const nbOtherActive = otherCols.length;
+    const nbThisActive = thisCols.length;
+    for (let t = 0; t < nbThisActive; t++) {
+      const i = thisRows[t];
+      const j = thisCols[t];
+      for (let o = 0; o < nbOtherActive; o++) {
+        if (j === otherRows[o]) {
+          const l = otherCols[o];
+          result.set(i, l, result.get(i, l) + otherValues[o] * thisValues[t]);
         }
       }
     }
-
-    return new SparseMatrix(result);
+    // console.log(result.cardinality);
+    return result;
   }
 
+  _mmulMediumDensity(other) {
+    const m = this.rows;
+    const p = other.columns;
+    const {
+      columns: otherCols,
+      rows: otherRows,
+      values: otherValues,
+    } = other.getNonZeros({ format: 'csr' });
+    const {
+      columns: thisCols,
+      rows: thisRows,
+      values: thisValues,
+    } = this.getNonZeros();
+
+    const result = new SparseMatrix(m, p);
+    const nbThisActive = thisCols.length;
+    for (let t = 0; t < nbThisActive; t++) {
+      const i = thisRows[t];
+      const j = thisCols[t];
+      const oStart = otherRows[j];
+      const oEnd = otherRows[j + 1];
+      for (let k = oStart; k < oEnd; k++) {
+        const l = otherCols[k];
+        result.set(i, l, result.get(i, l) + otherValues[k] * thisValues[t]);
+      }
+    }
+
+    return result;
+  }
   /**
    * @param {SparseMatrix} other
    * @returns {SparseMatrix}
@@ -292,11 +349,10 @@ export class SparseMatrix {
    * Returns the non-zero elements of the matrix in coordinate (COO), CSR, or CSC format.
    *
    * @param {Object} [options={}] - Options for output format and sorting.
-   * @param {'csr'|'csc'} [options.format] - If specified, returns the result in CSR or CSC format. Otherwise, returns COO format.
+   * @param {boolean} [options.format] - If specified, returns the result in CSR or CSC format. Otherwise, returns COO format.
    * @param {boolean} [options.sort] - If true, sorts the non-zero elements by their indices.
    * @returns {Object} If no format is specified, returns an object with Float64Array `rows`, `columns`, and `values` (COO format).
-   *                   If format is 'csr', returns { rows, columns, values } in CSR format.
-   *                   If format is 'csc', returns { rows, columns, values } in CSC format.
+   *                   If format is true, returns { rows, columns, values } in CSR format.
    * @throws {Error} If an unsupported format is specified.
    */
   getNonZeros(options = {}) {
@@ -318,16 +374,9 @@ export class SparseMatrix {
       idx++;
     }, sort);
 
-    if (!format) return { rows, columns, values };
-
-    if (!['csr', 'csc'].includes(format)) {
-      throw new Error(`format ${format} is not supported`);
-    }
-
-    const csrMatrix = cooToCsr({ rows, columns, values }, this.rows);
-    return format.toLowerCase() === 'csc'
-      ? csrToCsc(csrMatrix, this.columns)
-      : csrMatrix;
+    return format
+      ? cooToCsr({ rows, columns, values }, this.rows)
+      : { rows, columns, values };
   }
 
   /**
@@ -1469,46 +1518,6 @@ SparseMatrix.prototype.klass = 'Matrix';
 SparseMatrix.identity = SparseMatrix.eye;
 SparseMatrix.prototype.tensorProduct = SparseMatrix.prototype.kroneckerProduct;
 
-/**
- * Converts a matrix from Compressed Sparse Row (CSR) format to Compressed Sparse Column (CSC) format.
- * @param {Object} csrMatrix - The matrix in CSR format with properties: values, columns, rows.
- * @param {number} numCols - The number of columns in the matrix.
- * @returns {{rows: Float64Array, columns: Float64Array, values: Float64Array}} The matrix in CSC format.
- */
-function csrToCsc(csrMatrix, numCols) {
-  const {
-    values: csrValues,
-    columns: csrColIndices,
-    rows: csrRowPtr,
-  } = csrMatrix;
-
-  const cscValues = new Float64Array(csrValues.length);
-  const cscRowIndices = new Float64Array(csrValues.length);
-  const cscColPtr = new Float64Array(numCols + 1);
-
-  for (let i = 0; i < csrColIndices.length; i++) {
-    cscColPtr[csrColIndices[i] + 1]++;
-  }
-
-  for (let i = 1; i <= numCols; i++) {
-    cscColPtr[i] += cscColPtr[i - 1];
-  }
-
-  const next = cscColPtr.slice();
-
-  for (let row = 0; row < csrRowPtr.length - 1; row++) {
-    for (let j = csrRowPtr[row], i = 0; j < csrRowPtr[row + 1]; j++, i++) {
-      const col = csrColIndices[j];
-      const pos = next[col];
-      cscValues[pos] = csrValues[j];
-      cscRowIndices[pos] = row;
-      next[col]++;
-    }
-  }
-
-  return { rows: cscRowIndices, columns: cscColPtr, values: cscValues };
-}
-
 /**
  * Converts a matrix from Coordinate (COO) format to Compressed Sparse Row (CSR) format.
  * @param {Object} cooMatrix - The matrix in COO format with properties: values, columns, rows.
@@ -1518,26 +1527,17 @@ function csrToCsc(csrMatrix, numCols) {
 function cooToCsr(cooMatrix, nbRows) {
   const { values, columns, rows } = cooMatrix;
   const csrRowPtr = new Float64Array(nbRows + 1);
-  const length = values.length;
-  let currentRow = rows[0];
-  for (let index = 0; index < length; ) {
-    while (currentRow === rows[index] && index < length) ++index;
-    csrRowPtr[currentRow + 1] = index;
-    currentRow += 1;
+
+  // Count non-zeros per row
+  const numberOfNonZeros = rows.length;
+  for (let i = 0; i < numberOfNonZeros; i++) {
+    csrRowPtr[rows[i] + 1]++;
   }
-  return { rows: csrRowPtr, columns, values };
-}
 
-/**
- * Creates an empty 2D matrix (array of Float64Array) with the given dimensions.
- * @param {number} nbRows - Number of rows.
- * @param {number} nbColumns - Number of columns.
- * @returns {Float64Array[]} A 2D array representing the matrix.
- */
-function matrixCreateEmpty(nbRows, nbColumns) {
-  const newMatrix = [];
-  for (let row = 0; row < nbRows; row++) {
-    newMatrix.push(new Float64Array(nbColumns));
+  // Compute cumulative sum
+  for (let i = 1; i <= nbRows; i++) {
+    csrRowPtr[i] += csrRowPtr[i - 1];
   }
-  return newMatrix;
+
+  return { rows: csrRowPtr, columns, values };
 }