feat: add applyAlongAxis to reduce a matrix along a chosen dimension

README.md

@@ -124,6 +124,20 @@ var norm           = A.norm();           // norm = 10.14889156509222 (Frobenius
 var transpose      = A.transpose();      // transpose = Matrix [[1, 10], [1, -1], rows: 2, columns: 2]
 ```
 
+#### Row and column wise reductions
+```js
+var M = new Matrix([
+  [1, 2, 3],
+  [4, 5, 6],
+]);
+
+var sumOf = (vector) => vector.reduce((total, value) => total + value, 0);
+
+var rowSums    = M.applyAlongAxis(sumOf, 'row');    // rowSums    = [6, 15]
+var columnSums = M.applyAlongAxis(sumOf, 'column'); // columnSums = [5, 7, 9]
+```
+The callback receives each row or column as a plain array along with its index, so any reduction can be expressed with it.
+
 #### Instantiation of matrix
 ```js
 var z = Matrix.zeros(3, 2); // z = Matrix [[0, 0], [0, 0], [0, 0], rows: 3, columns: 2]

matrix.d.ts

@@ -292,7 +292,19 @@ export abstract class AbstractMatrix {
    * Applies a callback for each element of the matrix. The function is called in the matrix (this) context.
    * @param callback - Function that will be called for each element in the matrix.
    */
-  apply(callback: (row: number, column: number) => void): this;
+  apply(callback: (this: this, row: number, column: number) => void): this;
+
+  /**
+   * Applies a callback to each row or each column of the matrix and collects the returned values.
+   * The function is called in the matrix (this) context.
+   * @param callback - Function that will be called with each row or column and its index.
+   * @param by - Iterate by 'row' or 'column'.
+   * @returns - An array with the reduced column or row.
+   */
+  applyAlongAxis<ReturnType>(
+    callback: (this: this, vector: number[], index: number) => ReturnType,
+    by: MatrixDimension,
+  ): ReturnType[];
 
   /**
    * Returns a new 1D array filled row by row with the matrix values.

src/__tests__/matrix/applyAlongAxis.test.js

@@ -0,0 +1,157 @@
+import { describe, expect, it } from 'vitest';
+
+import { Matrix, MatrixTransposeView, SymmetricMatrix } from '../..';
+
+describe('applyAlongAxis', () => {
+  const matrix = new Matrix([
+    [1, 2, 3],
+    [4, 5, 6],
+  ]);
+
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('by row returns a column vector', () => {
+    const result = matrix.applyAlongAxis(sum, 'row');
+    expect(result).toStrictEqual([6, 15]);
+  });
+
+  it('by column returns a row vector', () => {
+    const result = matrix.applyAlongAxis(sum, 'column');
+    expect(result).toStrictEqual([5, 7, 9]);
+  });
+
+  it('the callback gets the vector and its index', () => {
+    const seen = [];
+    matrix.applyAlongAxis((vector, index) => {
+      seen.push([index, vector]);
+      return index;
+    }, 'column');
+    expect(seen).toStrictEqual([
+      [0, [1, 4]],
+      [1, [2, 5]],
+      [2, [3, 6]],
+    ]);
+  });
+
+  it('the source matrix is left untouched', () => {
+    matrix.applyAlongAxis(sum, 'row');
+    expect(matrix.to2DArray()).toStrictEqual([
+      [1, 2, 3],
+      [4, 5, 6],
+    ]);
+  });
+
+  it('works with a callback that is not a reduction', () => {
+    const result = matrix.applyAlongAxis(
+      (vector) => Math.max(...vector),
+      'row',
+    );
+    expect(result).toStrictEqual([3, 6]);
+  });
+
+  it('receives the matrix as `this`', () => {
+    const matrix = Matrix.zeros(1, 1);
+    let that;
+    matrix.applyAlongAxis(function cb() {
+      // eslint-disable-next-line no-invalid-this
+      that = this;
+    }, 'row');
+    expect(that).toBe(matrix);
+  });
+
+  it('can return any type', () => {
+    const rows = matrix.applyAlongAxis((row) => row, 'row');
+    expect(rows).toStrictEqual(matrix.to2DArray());
+  });
+});
+
+describe('applyAlongAxis on other matrix kinds', () => {
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('reads a transpose view along the view dimensions', () => {
+    const view = new MatrixTransposeView(
+      new Matrix([
+        [1, 2, 3],
+        [4, 5, 6],
+      ]),
+    );
+    expect(view.applyAlongAxis(sum, 'row')).toStrictEqual([5, 7, 9]);
+  });
+
+  it('works on a symmetric matrix', () => {
+    const symmetric = new SymmetricMatrix([
+      [1, 2],
+      [2, 3],
+    ]);
+    expect(symmetric.applyAlongAxis(sum, 'column')).toStrictEqual([3, 5]);
+  });
+});
+
+describe('applyAlongAxis with degenerate matrices', () => {
+  const emptyMatrix = new Matrix(0, 0);
+  const zeroRowMatrix = new Matrix(0, 2);
+  const zeroColumnMatrix = new Matrix(3, 0);
+
+  function count(vector) {
+    return vector.length;
+  }
+
+  it('by row of a 0x0 matrix', () => {
+    const result = emptyMatrix.applyAlongAxis(count, 'row');
+    expect(result).toStrictEqual([]);
+  });
+
+  it('by column of a 0x0 matrix', () => {
+    const result = emptyMatrix.applyAlongAxis(count, 'column');
+    expect(result).toStrictEqual([]);
+  });
+
+  it('by column of a 0 row matrix', () => {
+    expect(zeroRowMatrix.applyAlongAxis(count, 'column')).toStrictEqual([0, 0]);
+  });
+
+  it('by row of a 0 column matrix', () => {
+    expect(zeroColumnMatrix.applyAlongAxis(count, 'row')).toStrictEqual([
+      0, 0, 0,
+    ]);
+  });
+});
+
+describe('applyAlongAxis error handling', () => {
+  const matrix = new Matrix([
+    [1, 2],
+    [3, 4],
+  ]);
+
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('throws when the callback is missing', () => {
+    expect(() => matrix.applyAlongAxis(undefined, 'row')).toThrow(
+      /^callback must be a function$/,
+    );
+  });
+
+  it('throws when the callback is not a function', () => {
+    expect(() => matrix.applyAlongAxis(42, 'column')).toThrow(
+      /^callback must be a function$/,
+    );
+  });
+
+  it('throws when the dimension is missing', () => {
+    expect(() => matrix.applyAlongAxis(sum)).toThrow(
+      /^invalid option: undefined$/,
+    );
+  });
+
+  it('throws when the dimension is unknown', () => {
+    expect(() => matrix.applyAlongAxis(sum, 'diagonal')).toThrow(
+      /^invalid option: diagonal$/,
+    );
+  });
+});

src/matrix.js

@@ -4,34 +4,34 @@ import rescale from 'ml-array-rescale';
 import { inspectMatrix, inspectMatrixWithOptions } from './inspect';
 import { installMathOperations } from './mathOperations';
 import {
-  sumByRow,
-  sumByColumn,
-  sumAll,
-  productByRow,
-  productByColumn,
-  productAll,
-  varianceByRow,
-  varianceByColumn,
-  varianceAll,
-  centerByRow,
-  centerByColumn,
   centerAll,
-  scaleByRow,
-  scaleByColumn,
-  scaleAll,
-  getScaleByRow,
-  getScaleByColumn,
+  centerByColumn,
+  centerByRow,
   getScaleAll,
+  getScaleByColumn,
+  getScaleByRow,
+  productAll,
+  productByColumn,
+  productByRow,
+  scaleAll,
+  scaleByColumn,
+  scaleByRow,
+  sumAll,
+  sumByColumn,
+  sumByRow,
+  varianceAll,
+  varianceByColumn,
+  varianceByRow,
 } from './stat';
 import {
-  checkRowVector,
-  checkRowIndex,
   checkColumnIndex,
+  checkColumnIndices,
   checkColumnVector,
-  checkRange,
   checkNonEmpty,
+  checkRange,
+  checkRowIndex,
   checkRowIndices,
-  checkColumnIndices,
+  checkRowVector,
 } from './util';
 
 export class AbstractMatrix {
@@ -181,6 +181,30 @@ export class AbstractMatrix {
     return this;
   }
 
+  applyAlongAxis(callback, by) {
+    if (typeof callback !== 'function') {
+      throw new TypeError('callback must be a function');
+    }
+    const result = [];
+    switch (by) {
+      case 'row': {
+        for (let i = 0; i < this.rows; i++) {
+          result.push(callback.call(this, this.getRow(i), i));
+        }
+        break;
+      }
+      case 'column': {
+        for (let i = 0; i < this.columns; i++) {
+          result.push(callback.call(this, this.getColumn(i), i));
+        }
+        break;
+      }
+      default:
+        throw new Error(`invalid option: ${by}`);
+    }
+    return result;
+  }
+
   to1DArray() {
     let array = [];
     for (let i = 0; i < this.rows; i++) {