From dbfd8e4c39e75a5d344c1058f94f0fc21dbd8713 Mon Sep 17 00:00:00 2001
From: Dan Rose <danoftheroses@gmail.com>
Date: Tue, 15 Oct 2024 09:16:28 -0500
Subject: [PATCH] feat: add `mpow` to compute matrix power using exponentiation
 by squaring (#193)

---
 matrix.d.ts                       |  6 ++++
 src/__tests__/matrix/mpow.test.js | 53 +++++++++++++++++++++++++++++++
 src/matrix.js                     | 20 ++++++++++++
 3 files changed, 79 insertions(+)
 create mode 100644 src/__tests__/matrix/mpow.test.js

diff --git a/matrix.d.ts b/matrix.d.ts
index 4ac3fe4..5025952 100644
--- a/matrix.d.ts
+++ b/matrix.d.ts
@@ -616,6 +616,12 @@ export abstract class AbstractMatrix {
    */
   mmul(other: MaybeMatrix): Matrix;
 
+  /**
+   * Returns the square matrix raised to the given power
+   * @param scalar - the non-negative integer power to raise this matrix to
+   */
+  mpow(scalar: number): Matrix;
+
   strassen2x2(other: MaybeMatrix): Matrix;
 
   strassen3x3(other: MaybeMatrix): Matrix;
diff --git a/src/__tests__/matrix/mpow.test.js b/src/__tests__/matrix/mpow.test.js
new file mode 100644
index 0000000..940b1d8
--- /dev/null
+++ b/src/__tests__/matrix/mpow.test.js
@@ -0,0 +1,53 @@
+import { describe, it, expect } from 'vitest';
+
+import { Matrix } from '../..';
+
+describe('matrix power', () => {
+  it('power of a non-square matrix is not defined', () => {
+    let x = new Matrix(3, 5);
+    expect(() => x.mpow(2)).toThrowError();
+  });
+  it('negative power is rejected', () => {
+    let x = new Matrix(3, 3);
+    expect(() => x.mpow(-2)).toThrowError();
+  });
+  it('A matrix to the power 0 is identity', () => {
+    expect(
+      new Matrix([
+        [1, 2, 3],
+        [4, 5, 6],
+        [7, 8, 9],
+      ])
+        .mpow(0)
+        .to2DArray(),
+    ).toStrictEqual([
+      [1, 0, 0],
+      [0, 1, 0],
+      [0, 0, 1],
+    ]);
+  });
+  it('Matches scalar exponentiation', () => {
+    let x = new Matrix([[1.5]]).mpow(50);
+    expect(x.size).toBe(1);
+    expect(x.get(0, 0)).toBeCloseTo(1.5 ** 50);
+  });
+  it('Handles high integer exponent', () => {
+    let e = Number.MAX_SAFE_INTEGER;
+    let x = new Matrix([[-1]]).mpow(e);
+    expect(x.size).toBe(1);
+    expect(x.get(0, 0)).toBe(-1);
+  });
+  it('Exponentiates a small matrix correctly', () => {
+    let u = new Matrix([
+      [1, 2],
+      [3, -1],
+    ]);
+    let e = 21;
+    let resultLoop = Matrix.eye(2);
+    for (let i = 0; i < e; ++i) {
+      resultLoop = resultLoop.mmul(u);
+    }
+    let resultMpow = u.mpow(e);
+    expect(resultLoop.to2DArray()).toStrictEqual(resultMpow.to2DArray());
+  });
+});
diff --git a/src/matrix.js b/src/matrix.js
index afb63a7..ef9feda 100644
--- a/src/matrix.js
+++ b/src/matrix.js
@@ -874,6 +874,26 @@ export class AbstractMatrix {
     return result;
   }
 
+  mpow(scalar) {
+    if (!this.isSquare()) {
+      throw new RangeError('Matrix must be square');
+    }
+    if (!Number.isInteger(scalar) || scalar < 0) {
+      throw new RangeError('Exponent must be a non-negative integer');
+    }
+    // Russian Peasant exponentiation, i.e. exponentiation by squaring
+    let result = Matrix.eye(this.rows);
+    let bb = this;
+    // Note: Don't bit shift. In JS, that would truncate at 32 bits
+    for (let e = scalar; e > 1; e /= 2) {
+      if ((e & 1) !== 0) {
+        result = result.mmul(bb);
+      }
+      bb = bb.mmul(bb);
+    }
+    return result;
+  }
+
   strassen2x2(other) {
     other = Matrix.checkMatrix(other);
     let result = new Matrix(2, 2);