diff --git a/matrix.d.ts b/matrix.d.ts index c6f6153..6e086dc 100644 --- a/matrix.d.ts +++ b/matrix.d.ts @@ -668,6 +668,12 @@ export abstract class AbstractMatrix { strassen3x3(other: MaybeMatrix): Matrix; + /** + * Returns the same product as {@link AbstractMatrix.mmul}, computed with the Strassen algorithm. + * The recursive path only kicks in once both dimensions of the padded operands pass 512, + * below that the call is handed to the plain multiplication. + * @param y - The right operand. + */ mmulStrassen(y: MaybeMatrix): Matrix; /** diff --git a/src/__tests__/matrix/mmulStrassen.test.js b/src/__tests__/matrix/mmulStrassen.test.js new file mode 100644 index 0000000..754d2b2 --- /dev/null +++ b/src/__tests__/matrix/mmulStrassen.test.js @@ -0,0 +1,135 @@ +import { afterEach, describe, it, expect, vi } from 'vitest'; + +import { AbstractMatrix, Matrix } from '../..'; + +afterEach(() => { + vi.restoreAllMocks(); +}); + +// https://github.com/mljs/matrix/issues/114 +describe('mmulStrassen agrees with mmul', () => { + const shapes = [ + [1, 1, 1, 1], + [1, 2, 2, 2], + [1, 3, 3, 1], + [3, 1, 1, 3], + [3, 2, 2, 4], + [4, 2, 2, 3], + [2, 5, 5, 2], + [5, 5, 5, 5], + [7, 3, 3, 6], + ]; + + for (const [r1, c1, r2, c2] of shapes) { + it(`${r1}x${c1} by ${r2}x${c2}`, () => { + const a = Matrix.randInt(r1, c1, { min: 0, max: 9 }); + const b = Matrix.randInt(r2, c2, { min: 0, max: 9 }); + const expected = a.mmul(b); + const result = a.mmulStrassen(b); + expect(result.rows).toBe(r1); + expect(result.columns).toBe(c2); + expect(result.to2DArray()).toStrictEqual(expected.to2DArray()); + }); + } + + it('keeps the operands untouched', () => { + const a = new Matrix([ + [1, 2], + [3, 4], + ]); + const b = new Matrix([ + [5, 6], + [7, 8], + ]); + a.mmulStrassen(b); + expect(a.to2DArray()).toStrictEqual([ + [1, 2], + [3, 4], + ]); + expect(b.to2DArray()).toStrictEqual([ + [5, 6], + [7, 8], + ]); + }); + + it('accepts a 2D array', () => { + const a = new Matrix([[1, 2]]); + expect( + a + .mmulStrassen([ + [1, 2], + [3, 4], + ]) + .to2DArray(), + ).toStrictEqual([[7, 10]]); + }); +}); + +describe('mmulStrassen above the recursion threshold', () => { + // the recursive path only runs once both dimensions pass 512, so anything + // smaller was delegating to mmul and never exercised the block split. + // these two cases cover the odd padded size along with the even one. + // small integers keep every intermediate exact in a float64. + const cases = [ + { r1: 513, c1: 513, c2: 513 }, + { r1: 514, c1: 514, c2: 514 }, + ]; + + for (const { r1, c1, c2 } of cases) { + it(`${r1}x${c1} by ${c1}x${c2} matches mmul exactly`, () => { + const a = Matrix.randInt(r1, c1, { min: 0, max: 3 }); + const b = Matrix.randInt(c1, c2, { min: 0, max: 3 }); + const expected = a.mmul(b); + const result = a.mmulStrassen(b); + expect(result.rows).toBe(r1); + expect(result.columns).toBe(c2); + + let differing = 0; + for (let i = 0; i < r1; i++) { + for (let j = 0; j < c2; j++) { + if (result.get(i, j) !== expected.get(i, j)) differing++; + } + } + expect(differing).toBe(0); + }, 120000); + } + + const rectangularCases = [ + { r1: 1, c1: 513, c2: 1 }, + { r1: 10, c1: 600, c2: 10 }, + ]; + + for (const { r1, c1, c2 } of rectangularCases) { + it(`${r1}x${c1} by ${c1}x${c2} delegates to mmul`, () => { + const a = Matrix.randInt(r1, c1, { min: 0, max: 3 }); + const b = Matrix.randInt(c1, c2, { min: 0, max: 3 }); + const expected = a.mmul(b); + const multiply = vi.spyOn(AbstractMatrix.prototype, 'mmul'); + + const result = a.mmulStrassen(b); + + expect(multiply).toHaveBeenCalledTimes(1); + expect(result.to2DArray()).toStrictEqual(expected.to2DArray()); + }); + } +}); + +describe('mmulStrassen with degenerate matrices', () => { + it('a matrix without rows', () => { + const result = new Matrix(0, 2).mmulStrassen(new Matrix(2, 3)); + expect(result.rows).toBe(0); + expect(result.columns).toBe(3); + }); + + it('a matrix without columns', () => { + const result = new Matrix(2, 3).mmulStrassen(new Matrix(3, 0)); + expect(result.rows).toBe(2); + expect(result.columns).toBe(0); + }); + + it('two 0x0 matrices', () => { + const result = new Matrix(0, 0).mmulStrassen(new Matrix(0, 0)); + expect(result.rows).toBe(0); + expect(result.columns).toBe(0); + }); +}); diff --git a/src/__tests__/matrix/utility.test.js b/src/__tests__/matrix/utility.test.js index b1893a3..52fa6ee 100644 --- a/src/__tests__/matrix/utility.test.js +++ b/src/__tests__/matrix/utility.test.js @@ -374,14 +374,10 @@ describe('utility methods', () => { it('mmul strassen on empty matrices', () => { // https://github.com/mljs/matrix/issues/114 - // while the mathematically correct result is 0x0, we assert a 2x2 padded result that the current implementation produces - // (this call is actually just delegated to standard multiplication in mmul()) - expect( - new Matrix(0, 2).mmulStrassen(new Matrix(2, 0)).to2DArray(), - ).toStrictEqual([ - [0, 0], - [0, 0], - ]); + const result = new Matrix(0, 2).mmulStrassen(new Matrix(2, 0)); + expect(result.rows).toBe(0); + expect(result.columns).toBe(0); + expect(result.to2DArray()).toStrictEqual([]); }); it('mmul 2x2 and 3x3', () => { diff --git a/src/matrix.js b/src/matrix.js index 887b59b..acd90c3 100644 --- a/src/matrix.js +++ b/src/matrix.js @@ -1135,6 +1135,17 @@ export class AbstractMatrix { `Multiplying ${r1} x ${c1} and ${r2} x ${c2} matrix: dimensions do not match.`, ); } + if (r1 === 0 || c2 === 0) { + return new Matrix(r1, c2); + } + + // Padding a rectangular product into a square whose side is the largest + // dimension can turn a cheap multiplication into an enormous allocation. + // The recursive implementation is only useful for equally sized square + // operands; mmul already handles every other compatible shape directly. + if (c1 === r2 && (r1 !== c1 || r2 !== c2 || r1 !== r2)) { + return x.mmul(y); + } // Put a matrix into the top left of a matrix of zeros. // `rows` and `cols` are the dimensions of the output matrix. @@ -1150,14 +1161,12 @@ export class AbstractMatrix { } } - // Make sure both matrices are the same size. - // This is exclusively for simplicity: - // this algorithm can be implemented with matrices of different sizes. - - let r = Math.max(r1, r2); - let c = Math.max(c1, c2); - x = embed(x, r, c); - y = embed(y, r, c); + // pad both operands into the same square so that the block split lines up. + // zeros never reach the top left r1 x c2 corner of the product, which is + // the part that gets returned. + let n = Math.max(r1, c1, r2, c2); + x = embed(x, n, n); + y = embed(y, n, n); // Our recursive multiplication function. function blockMult(a, b, rows, cols) { @@ -1230,13 +1239,17 @@ export class AbstractMatrix { // Crop output to the desired size (undo dynamic padding). let result = AbstractMatrix.zeros(2 * c11.rows, 2 * c11.columns); result = result.setSubMatrix(c11, 0, 0); - result = result.setSubMatrix(c12, c11.rows, 0); - result = result.setSubMatrix(c21, 0, c11.columns); + result = result.setSubMatrix(c12, 0, c11.columns); + result = result.setSubMatrix(c21, c11.rows, 0); result = result.setSubMatrix(c22, c11.rows, c11.columns); return result.subMatrix(0, rows - 1, 0, cols - 1); } - return blockMult(x, y, r, c); + const product = blockMult(x, y, n, n); + if (product.rows === r1 && product.columns === c2) { + return product; + } + return product.subMatrix(0, r1 - 1, 0, c2 - 1); } scaleRows(options = {}) {