Feat: Enhance Kronecker product to handle arbitrary dimension (#3461)

Previously `math.kron()` always returned a 2D matrix, and could not handle 3D or greater arrays. Now it always returns an array of the max dimension of its arguments.

Resolves #1753.
---------

Co-authored-by: Delaney Sylvans <delaneysylvans@gmail.com>

Author: gwhitney

AUTHORS

@@ -269,5 +269,6 @@ Don McCurdy <dm@donmccurdy.com>
 Jay Chang <96050090+JayChang4w@users.noreply.github.com>
 mrft <977655+mrft@users.noreply.github.com>
 Kip Robinson <91914404+kiprobinsonknack@users.noreply.github.com>
+Delaney Sylvans <delaneysylvans@gmail.com>
 
 # Generated by tools/update-authors.js

HISTORY.md

@@ -4,6 +4,8 @@
 
 !!! BE CAREFUL: BREAKING CHANGES !!!
 
+- Fix: #1753 Correct dimensionality of Kronecker product on vectors (and
+  extend to arbitrary dimension) (#3455). Thanks @Delaney.
 - Feat: #3349 Decouple precedence of unary percentage operator and binary
   modulus operator (that both use symbol `%`), and raise the former (#3432).
   Thanks @kiprobinsonknack.

src/function/matrix/kron.js

@@ -22,7 +22,7 @@ export const createKron = /* #__PURE__ */ factory(name, dependencies, ({ typed,
      *    // returns [ [ 1, 2, 0, 0 ], [ 3, 4, 0, 0 ], [ 0, 0, 1, 2 ], [ 0, 0, 3, 4 ] ]
      *
      *    math.kron([1,1], [2,3,4])
-     *    // returns [ [ 2, 3, 4, 2, 3, 4 ] ]
+     *    // returns [2, 3, 4, 2, 3, 4]
      *
      * See also:
      *
@@ -49,39 +49,38 @@ export const createKron = /* #__PURE__ */ factory(name, dependencies, ({ typed,
   })
 
   /**
-     * Calculate the Kronecker product of two matrices / vectors
-     * @param {Array} a  First vector
-     * @param {Array} b  Second vector
-     * @returns {Array} Returns the Kronecker product of x and y
-     * @private
+   * Calculate the Kronecker product of two (1-dimensional) vectors,
+   * with no dimension checking
+   * @param {Array} a  First vector
+   * @param {Array} b  Second vector
+   * @returns {Array}  the 1-dimensional Kronecker product of a and b
+   * @private
+   */
+  function _kron1d (a, b) {
+    // TODO in core overhaul: would be faster to see if we can choose a
+    // particular implementation of multiplyScalar at the beginning,
+    // rather than re-dispatch for _every_ ordered pair of entries.
+    return a.flatMap(x => b.map(y => multiplyScalar(x, y)))
+  }
+
+  /**
+   * Calculate the Kronecker product of two possibly multidimensional arrays
+   * @param {Array} a  First array
+   * @param {Array} b  Second array
+   * @param {number} [d]  common dimension; if missing, compute and match args
+   * @returns {Array} Returns the Kronecker product of x and y
+   * @private
      */
-  function _kron (a, b) {
-    // Deal with the dimensions of the matricies.
-    if (size(a).length === 1) {
-      // Wrap it in a 2D Matrix
-      a = [a]
-    }
-    if (size(b).length === 1) {
-      // Wrap it in a 2D Matrix
-      b = [b]
-    }
-    if (size(a).length > 2 || size(b).length > 2) {
-      throw new RangeError('Vectors with dimensions greater then 2 are not supported expected ' +
-            '(Size x = ' + JSON.stringify(a.length) + ', y = ' + JSON.stringify(b.length) + ')')
+  function _kron (a, b, d = -1) {
+    if (d < 0) {
+      let adim = size(a).length
+      let bdim = size(b).length
+      d = Math.max(adim, bdim)
+      while (adim++ < d) a = [a]
+      while (bdim++ < d) b = [b]
     }
-    const t = []
-    let r = []
 
-    return a.map(function (a) {
-      return b.map(function (b) {
-        r = []
-        t.push(r)
-        return a.map(function (y) {
-          return b.map(function (x) {
-            return r.push(multiplyScalar(y, x))
-          })
-        })
-      })
-    }) && t
+    if (d === 1) return _kron1d(a, b)
+    return a.flatMap(aSlice => b.map(bSlice => _kron(aSlice, bSlice, d - 1)))
   }
 })

test/unit-tests/function/matrix/kron.test.js

@@ -5,6 +5,7 @@ import math from '../../../../src/defaultInstance.js'
 
 describe('kron', function () {
   it('should calculate the Kronecker product of two arrays', function () {
+    assert.deepStrictEqual(math.kron([[2]], [[3]]), [[6]])
     assert.deepStrictEqual(math.kron([
       [1, -2, 1],
       [1, 1, 0]
@@ -28,14 +29,27 @@ describe('kron', function () {
     ])
   })
 
+  it('should calculate product for empty 1D Arrays', function () {
+    assert.deepStrictEqual(math.kron([], []), [])
+  })
+
   it('should calculate product for empty 2D Arrays', function () {
     assert.deepStrictEqual(math.kron([[]], [[]]), [[]])
   })
 
   it('should calculate product for 1D Arrays', function () {
+    assert.deepStrictEqual(math.kron([2], [3]), [6])
+    assert.deepStrictEqual(math.kron([1, 2], [3, 4]), [3, 4, 6, 8])
+    assert.deepStrictEqual(math.kron([1, 2, 6, 8], [12, 1, 2, 3]), [12, 1, 2, 3, 24, 2, 4, 6, 72, 6, 12, 18, 96, 8, 16, 24])
+  })
+
+  it('should calculate product for 1D & 2D Arrays', function () {
     assert.deepStrictEqual(math.kron([1, 1], [[1, 0], [0, 1]]), [[1, 0, 1, 0], [0, 1, 0, 1]])
     assert.deepStrictEqual(math.kron([[1, 0], [0, 1]], [1, 1]), [[1, 1, 0, 0], [0, 0, 1, 1]])
-    assert.deepStrictEqual(math.kron([1, 2, 6, 8], [12, 1, 2, 3]), [[12, 1, 2, 3, 24, 2, 4, 6, 72, 6, 12, 18, 96, 8, 16, 24]])
+    assert.deepStrictEqual(math.kron([[1, 2]], [[1, 2, 3]]), [[1, 2, 3, 2, 4, 6]])
+    assert.deepStrictEqual(math.kron([[1], [2]], [[1], [2], [3]]), [[1], [2], [3], [2], [4], [6]])
+    assert.deepStrictEqual(math.kron([[1, 2]], [[1], [2], [3]]), [[1, 2], [2, 4], [3, 6]])
+    assert.deepStrictEqual(math.kron([[1], [2]], [[1, 2, 3]]), [[1, 2, 3], [2, 4, 6]])
   })
 
   it('should support complex numbers', function () {
@@ -55,10 +69,32 @@ describe('kron', function () {
     ])
   })
 
-  it('should throw an error for greater then 2 dimensions', function () {
-    assert.throws(function () {
-      math.kron([[[1, 1], [1, 1]], [[1, 1], [1, 1]]], [[[1, 2, 3], [4, 5, 6]], [[6, 7, 8], [9, 10, 11]]])
-    })
+  it('should calculate a 3D Kronecker product', function () {
+    assert.deepStrictEqual(
+      math.kron([
+        [[1, 2], [2, 3]],
+        [[2, 3], [3, 4]]
+      ], [
+        [[4, 3], [3, 2]],
+        [[3, 2], [2, 1]]
+      ]), [
+        /* eslint-disable no-multi-spaces, array-bracket-spacing */
+        [[4, 3,  8, 6], [3, 2, 6, 4], [ 8, 6, 12,  9], [6, 4,  9, 6]],
+        [[3, 2,  6, 4], [2, 1, 4, 2], [ 6, 4,  9,  6], [4, 2,  6, 3]],
+        [[8, 6, 12, 9], [6, 4, 9, 6], [12, 9, 16, 12], [9, 6, 12, 8]],
+        [[6, 4,  9, 6], [4, 2, 6, 3], [ 9, 6, 12,  8], [6, 3,  8, 4]]
+        /* eslint-enable */
+      ]
+    )
+  })
+
+  it('should allow mixed-dimensional Kronecker products', function () {
+    const b = [[[4, 3], [3, 2]], [[3, 2], [2, 1]]]
+    const a = [1, 2]
+    assert.deepStrictEqual(math.kron(a, b), math.kron([[a]], b))
+    assert.deepStrictEqual(math.kron([a], b), math.kron([[a]], b))
+    assert.deepStrictEqual(math.kron(b, a), math.kron(b, [[a]]))
+    assert.deepStrictEqual(math.kron(b, [a]), math.kron(b, [[a]]))
   })
 
   it('should throw an error if called with an invalid number of arguments', function () {
@@ -81,10 +117,12 @@ describe('kron', function () {
       assert.deepStrictEqual(product.toArray(), [[13, 26, 0, 0], [715, -13, 0, -0], [0, 0, -1, -2], [0, -0, -55, 1]])
     })
 
-    it('should throw an error for invalid Kronecker product of matrix', function () {
-      const y = math.matrix([[[]]])
+    it('should calculate the Kronecker product of 3d matrices', function () {
+      const y = math.matrix([[[3]]])
       const x = math.matrix([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
-      assert.throws(function () { math.kron(y, x) })
+      const product = math.kron(x, y)
+      assert.deepStrictEqual(
+        product.toArray(), [[[3, 3], [3, 3]], [[3, 3], [3, 3]]])
     })
   })