feat: add transposeMultiply (thisᵀ · other) without materializing the transpose (#211)

Author: lpatiny

benchmark/transposeMultiply.js

@@ -0,0 +1,68 @@
+'use strict';
+
+// Compares `A.transpose().mmul(B)` (materializes the transpose) with the fused
+// `A.transposeMultiply(B)` (streams rows, never materializes the transpose).
+// Run after `npm run compile`: node benchmark/transposeMultiply.js
+
+let benchmark = require('benchmark');
+
+let { Matrix } = require('..');
+let { SparseMatrix } = require('ml-sparse-matrix');
+
+// (m x n)ᵀ · (m x p) -> (n x p)
+const cases = [
+  { m: 256, n: 256, p: 256, label: 'dense 256' },
+  { m: 512, n: 512, p: 512, label: 'dense 512' },
+  { m: 252, n: 252, p: 210, label: 'dense 252x210 (NMR block)' },
+];
+
+for (const { m, n, p, label } of cases) {
+  const A = Matrix.rand(m, n);
+  const B = Matrix.rand(m, p);
+
+  const reference = A.transpose().mmul(B);
+  const fused = A.transposeMultiply(B);
+  if (!reference.to2DArray().flat().every((v, i) => v === fused.to2DArray().flat()[i])) {
+    throw new Error('transposeMultiply result differs from transpose().mmul()');
+  }
+
+  new benchmark.Suite(label)
+    .add(`${label}: transpose().mmul()`, () => {
+      A.transpose().mmul(B);
+    })
+    .add(`${label}: transposeMultiply()`, () => {
+      A.transposeMultiply(B);
+    })
+    .on('cycle', (event) => {
+      console.log(String(event.target));
+    })
+    .on('complete', function onComplete() {
+      console.log(`  fastest: ${this.filter('fastest').map('name')}\n`);
+    })
+    .run();
+}
+
+// Sparse `this`: the zero-skip in transposeMultiply avoids most of the work.
+// Build the sparse operand with the real ml-sparse-matrix package (a few
+// non-zeros per row, as in the NMR spin-simulation operators) and densify it,
+// the path such callers take before multiplying.
+const builder = new SparseMatrix(512, 512);
+for (let i = 0; i < 512; i++) {
+  for (let k = 0; k < 9; k++) builder.set(i, (i * 31 + k * 101) % 512, k + 1);
+}
+const sparse = new Matrix(builder.to2DArray());
+const dense = Matrix.rand(512, 256);
+new benchmark.Suite('sparse')
+  .add('sparse 512x512 (9 nnz/row): transpose().mmul()', () => {
+    sparse.transpose().mmul(dense);
+  })
+  .add('sparse 512x512 (9 nnz/row): transposeMultiply()', () => {
+    sparse.transposeMultiply(dense);
+  })
+  .on('cycle', (event) => {
+    console.log(String(event.target));
+  })
+  .on('complete', function onComplete() {
+    console.log(`  fastest: ${this.filter('fastest').map('name')}`);
+  })
+  .run();

matrix.d.ts

@@ -626,6 +626,16 @@ export abstract class AbstractMatrix {
    */
   gram(): Matrix;
 
+  /**
+   * Returns the matrix product `thisᵀ · other` without materializing the
+   * transpose: each shared row is streamed once and applied as a rank-1 update,
+   * so both operands are read row-major (contiguously) and zero entries of
+   * `this` are skipped. As fast as `this.transpose().mmul(other)` on dense
+   * matrices and much faster on sparse ones, with an identical result.
+   * @param other - Other matrix, with the same number of rows as `this`.
+   */
+  transposeMultiply(other: MaybeMatrix): Matrix;
+
   /**
    * Returns the matrix product between `this` and its transpose (`this · thisᵀ`),
    * optionally weighting each column by `scale` (`this · diag(scale) · thisᵀ`).

package.json

@@ -75,6 +75,7 @@
     "jest-matcher-deep-close-to": "^3.0.2",
     "mathjs": "^14.5.2",
     "ml-dataset-iris": "^1.2.1",
+    "ml-sparse-matrix": "^3.1.0",
     "ml-xsadd": "^3.0.1",
     "numeric": "^1.2.6",
     "prettier": "^3.5.3",

src/__tests__/matrix/utility.test.js

@@ -1,3 +1,4 @@
+import { SparseMatrix } from 'ml-sparse-matrix';
 import { describe, it, beforeEach, expect } from 'vitest';
 
 import { Matrix, pseudoInverse, determinant } from '../..';
@@ -443,6 +444,75 @@ describe('utility methods', () => {
     );
   });
 
+  it('transposeMultiply', () => {
+    let a = new Matrix([
+      [1, 2, 3],
+      [4, 5, 6],
+    ]);
+    let b = new Matrix([
+      [7, 8],
+      [9, 10],
+    ]);
+    // aᵀ · b, equivalent to a.transpose().mmul(b) without materializing the transpose
+    expect(a.transposeMultiply(b).to2DArray()).toStrictEqual(
+      a.transpose().mmul(b).to2DArray(),
+    );
+  });
+
+  it('transposeMultiply stays identical on a sparse matrix', () => {
+    let a = new Matrix([
+      [1, 0, 2],
+      [0, 3, 0],
+      [4, 0, 0],
+      [0, 0, 5],
+    ]);
+    let b = new Matrix([
+      [1, 2, 3],
+      [4, 5, 6],
+      [7, 8, 9],
+      [10, 11, 12],
+    ]);
+    expect(a.transposeMultiply(b).to2DArray()).toStrictEqual(
+      a.transpose().mmul(b).to2DArray(),
+    );
+  });
+
+  it('transposeMultiply throws on row mismatch', () => {
+    let a = new Matrix([[1, 2, 3]]);
+    let b = new Matrix([
+      [1, 2],
+      [3, 4],
+    ]);
+    expect(() => a.transposeMultiply(b)).toThrow(
+      'the number of rows of the two matrices must be equal',
+    );
+  });
+
+  it('transposeMultiply on a matrix built with ml-sparse-matrix', () => {
+    // A genuinely sparse operand, built with our ml-sparse-matrix package.
+    let sparse = new SparseMatrix(6, 4);
+    sparse.set(0, 1, 2);
+    sparse.set(2, 0, 3);
+    sparse.set(3, 3, 5);
+    sparse.set(5, 2, 7);
+    expect(sparse.cardinality).toBe(4);
+
+    // Densify into a Matrix, the path a caller such as the NMR spin simulation
+    // takes before multiplying, and check `sparseᵀ · b`.
+    let a = new Matrix(sparse.to2DArray());
+    let b = new Matrix([
+      [1, 2, 3],
+      [4, 5, 6],
+      [7, 8, 9],
+      [10, 11, 12],
+      [13, 14, 15],
+      [16, 17, 18],
+    ]);
+    expect(a.transposeMultiply(b).to2DArray()).toStrictEqual(
+      a.transpose().mmul(b).to2DArray(),
+    );
+  });
+
   it('pseudoinverse', () => {
     // Actual values calculated by the Numpy library
 

src/correlation.js

@@ -39,7 +39,7 @@ export function correlation(xMatrix, yMatrix = xMatrix, options = {}) {
     ? sdx
     : yMatrix.standardDeviation('column', { unbiased: true });
 
-  const corr = xMatrix.transpose().mmul(yMatrix);
+  const corr = xMatrix.transposeMultiply(yMatrix);
   for (let i = 0; i < corr.rows; i++) {
     for (let j = 0; j < corr.columns; j++) {
       corr.set(

src/covariance.js

@@ -26,7 +26,7 @@ export function covariance(xMatrix, yMatrix = xMatrix, options = {}) {
       yMatrix = yMatrix.center('column');
     }
   }
-  const cov = xMatrix.transpose().mmul(yMatrix);
+  const cov = xMatrix.transposeMultiply(yMatrix);
   for (let i = 0; i < cov.rows; i++) {
     for (let j = 0; j < cov.columns; j++) {
       cov.set(i, j, cov.get(i, j) * (1 / (xMatrix.rows - 1)));

src/matrix.js

@@ -908,6 +908,34 @@ export class AbstractMatrix {
     return result;
   }
 
+  transposeMultiply(other) {
+    other = Matrix.checkMatrix(other);
+    if (this.rows !== other.rows) {
+      throw new RangeError(
+        'the number of rows of the two matrices must be equal',
+      );
+    }
+    const n = this.columns;
+    const p = other.columns;
+
+    const result = new Matrix(n, p);
+    const otherRow = new Float64Array(p);
+    for (let r = 0; r < this.rows; r++) {
+      for (let j = 0; j < p; j++) {
+        otherRow[j] = other.get(r, j);
+      }
+      for (let i = 0; i < n; i++) {
+        const value = this.get(r, i);
+        if (value === 0) continue;
+        const resultRow = result.data[i];
+        for (let j = 0; j < p; j++) {
+          resultRow[j] += value * otherRow[j];
+        }
+      }
+    }
+    return result;
+  }
+
   mmulByTranspose(scale) {
     let m = this.rows;
     let n = this.columns;