feat: support multiple inputs in function map (josdejong#3228)

alenLym · Aug 22, 2024 · bcf0da4 · bcf0da4
1 parent 88a4b35
commit bcf0da4
Show file tree

Hide file tree

Showing 12 changed files with 542 additions and 88 deletions.
diff --git a/docs/datatypes/matrices.md b/docs/datatypes/matrices.md
@@ -414,6 +414,39 @@ const cum = a.map(function (value, index, matrix) {
 console.log(cum.toString())  // [[0, 1], [3, 6], [10, 15]]
 ```
 
+### Iterating over multiple Matrixes or Arrays
+
+You can iterate over multiple matrices or arrays by using the `map` function. Mapping allows to perform element-wise operations on matrices by automatically adjusting their sizes to match each other.
+
+To iterate over multiple matrices, you can use the `map` function. The `map` function applies a given function to each element of the matrices and returns a new matrix with the results.
+
+Here's an example of iterating over two matrices and adding their corresponding elements:
+
+```js
+const a = math.matrix([[1, 2], [3, 4]]);
+const b = math.matrix([[5, 6], [7, 8]]);
+
+const result = math.map(a, b, (x, y) => x + y);
+
+console.log(result); // [[6, 8], [10, 12]]
+```
+
+In this example, the `map` function takes matrices as the first two arguments and a callback function `(x, y) => x + y` as the third argument. The callback function is applied to each element of the matrices, where `x` represents the corresponding element from matrix `a` and `y` represents the corresponding element from matrix `b`. The result is a new matrix with the element-wise sum of the two matrices.
+
+By using broadcasting and the `map` function, you can easily iterate over multiple matrices and perform element-wise operations.
+
+```js
+const a = math.matrix([10, 20])
+const b = math.matrix([[3, 4], [5, 6]])
+
+const result = math.map(a, b, (x, y) => x + y)
+console.log(result); // [[13, 24], [15, 26]]
+```
+
+It's also possible to provide a callback with an index and the broadcasted arrays. Like `(valueA, valueB, index)` or even `(valueA, valueB, index, broadcastedMatrixA, broadcastedMatrixB)`. There is no specific limit for the number of matrices `N` that can be mapped. Thus, the callback can have `N` arguments, `N+1` arguments in the case of including the index, or `2N+1` arguments in the case of including the index and the broadcasted matrices in the callback. 
+
+At this moment `forEach` doesn't include the same functionality.
+
 ## Storage types
 
 Math.js supports both dense matrices as well as sparse matrices. Sparse matrices are efficient for matrices largely containing zeros. In that case they save a lot of memory, and calculations can be much faster than for dense matrices.

diff --git a/src/expression/embeddedDocs/function/matrix/map.js b/src/expression/embeddedDocs/function/matrix/map.js
@@ -2,11 +2,13 @@ export const mapDocs = {
   name: 'map',
   category: 'Matrix',
   syntax: [
-    'map(x, callback)'
+    'map(x, callback)',
+    'map(x, y, ..., callback)'
   ],
-  description: 'Create a new matrix or array with the results of the callback function executed on each entry of the matrix/array.',
+  description: 'Create a new matrix or array with the results of the callback function executed on each entry of the matrix/array or the matrices/arrays.',
   examples: [
-    'map([1, 2, 3], square)'
+    'map([1, 2, 3], square)',
+    'map([1, 2], [3, 4], f(a,b) = a + b)'
   ],
   seealso: ['filter', 'forEach']
 }
diff --git a/src/expression/transform/map.transform.js b/src/expression/transform/map.transform.js
@@ -1,7 +1,6 @@
-import { applyCallback } from '../../utils/applyCallback.js'
-import { map } from '../../utils/array.js'
 import { factory } from '../../utils/factory.js'
 import { isFunctionAssignmentNode, isSymbolNode } from '../../utils/is.js'
+import { createMap } from '../../function/matrix/map.js'
 import { compileInlineExpression } from './utils/compileInlineExpression.js'
 
 const name = 'map'
@@ -14,61 +13,123 @@ export const createMapTransform = /* #__PURE__ */ factory(name, dependencies, ({
    *
    * This transform creates a one-based index instead of a zero-based index
    */
+  const map = createMap({ typed })
+
   function mapTransform (args, math, scope) {
-    let x, callback
+    if (args.length === 0) {
+      return map()
+    }
 
-    if (args[0]) {
-      x = args[0].compile().evaluate(scope)
+    if (args.length === 1) {
+      return map(args[0])
     }
+    const N = args.length - 1
+    let X, callback
+    callback = args[N]
+    X = args.slice(0, N)
+    X = X.map(arg => _compileAndEvaluate(arg, scope))
 
-    if (args[1]) {
-      if (isSymbolNode(args[1]) || isFunctionAssignmentNode(args[1])) {
+    if (callback) {
+      if (isSymbolNode(callback) || isFunctionAssignmentNode(callback)) {
         // a function pointer, like filter([3, -2, 5], myTestFunction)
-        callback = args[1].compile().evaluate(scope)
+        callback = _compileAndEvaluate(callback, scope)
       } else {
         // an expression like filter([3, -2, 5], x > 0)
-        callback = compileInlineExpression(args[1], math, scope)
+        callback = compileInlineExpression(callback, math, scope)
       }
     }
+    return map(...X, _transformCallback(callback, N))
 
-    return map(x, callback)
+    function _compileAndEvaluate (arg, scope) {
+      return arg.compile().evaluate(scope)
+    }
   }
   mapTransform.rawArgs = true
 
-  // one-based version of map function
-  const map = typed('map', {
-    'Array, function': function (x, callback) {
-      return _map(x, callback, x)
-    },
+  return mapTransform
 
-    'Matrix, function': function (x, callback) {
-      return x.create(_map(x.valueOf(), callback, x), x.datatype())
+  /**
+   * Transforms the given callback function based on its type and number of arrays.
+   *
+   * @param {Function} callback - The callback function to transform.
+   * @param {number} numberOfArrays - The number of arrays to pass to the callback function.
+   * @returns {*} - The transformed callback function.
+   */
+  function _transformCallback (callback, numberOfArrays) {
+    if (typed.isTypedFunction(callback)) {
+      return _transformTypedCallbackFunction(callback, numberOfArrays)
+    } else {
+      return _transformCallbackFunction(callback, callback.length, numberOfArrays)
     }
-  })
+  }
 
-  return mapTransform
+  /**
+   * Transforms the given typed callback function based on the number of arrays.
+   *
+   * @param {Function} typedFunction - The typed callback function to transform.
+   * @param {number} numberOfArrays - The number of arrays to pass to the callback function.
+   * @returns {*} - The transformed typed callback function.
+   */
+  function _transformTypedCallbackFunction (typedFunction, numberOfArrays) {
+    const signatures = Object.fromEntries(
+      Object.entries(typedFunction.signatures)
+        .map(([signature, callbackFunction]) => {
+          const numberOfCallbackInputs = signature.split(',').length
+          if (typed.isTypedFunction(callbackFunction)) {
+            return [signature, _transformTypedCallbackFunction(callbackFunction, numberOfArrays)]
+          } else {
+            return [signature, _transformCallbackFunction(callbackFunction, numberOfCallbackInputs, numberOfArrays)]
+          }
+        })
+    )
+
+    if (typeof typedFunction.name === 'string') {
+      return typed(typedFunction.name, signatures)
+    } else {
+      return typed(signatures)
+    }
+  }
 }, { isTransformFunction: true })
 
 /**
- * Map for a multidimensional array. One-based indexes
- * @param {Array} array
- * @param {function} callback
- * @param {Array} orig
- * @return {Array}
- * @private
+ * Transforms the callback function based on the number of callback inputs and arrays.
+ * There are three cases:
+ * 1. The callback function has N arguments.
+ * 2. The callback function has N+1 arguments.
+ * 3. The callback function has 2N+1 arguments.
+ *
+ * @param {Function} callbackFunction - The callback function to transform.
+ * @param {number} numberOfCallbackInputs - The number of callback inputs.
+ * @param {number} numberOfArrays - The number of arrays.
+ * @returns {Function} The transformed callback function.
  */
-function _map (array, callback, orig) {
-  function recurse (value, index) {
-    if (Array.isArray(value)) {
-      return map(value, function (child, i) {
-        // we create a copy of the index array and append the new index value
-        return recurse(child, index.concat(i + 1)) // one based index, hence i + 1
-      })
-    } else {
-      // invoke the (typed) callback function with the right number of arguments
-      return applyCallback(callback, value, index, orig, 'map')
+function _transformCallbackFunction (callbackFunction, numberOfCallbackInputs, numberOfArrays) {
+  if (numberOfCallbackInputs === numberOfArrays) {
+    return callbackFunction
+  } else if (numberOfCallbackInputs === numberOfArrays + 1) {
+    return function (...args) {
+      const vals = args.slice(0, numberOfArrays)
+      const idx = _transformDims(args[numberOfArrays])
+      return callbackFunction(...vals, idx)
+    }
+  } else if (numberOfCallbackInputs > numberOfArrays + 1) {
+    return function (...args) {
+      const vals = args.slice(0, numberOfArrays)
+      const idx = _transformDims(args[numberOfArrays])
+      const rest = args.slice(numberOfArrays + 1)
+      return callbackFunction(...vals, idx, ...rest)
     }
+  } else {
+    return callbackFunction
   }
+}
 
-  return recurse(array, [])
+/**
+ * Transforms the dimensions by adding 1 to each dimension.
+ *
+ * @param {Array} dims - The dimensions to transform.
+ * @returns {Array} The transformed dimensions.
+ */
+function _transformDims (dims) {
+  return dims.map(dim => dim.isBigNumber ? dim.plus(1) : dim + 1)
 }
diff --git a/src/expression/transform/utils/dimToZeroBase.js b/src/expression/transform/utils/dimToZeroBase.js
@@ -0,0 +1,17 @@
+import { isNumber, isBigNumber } from '../../../utils/is.js'
+/**
+ * Change last argument dim from one-based to zero-based.
+ */
+export function dimToZeroBase (dim) {
+  if (isNumber(dim)) {
+    return dim - 1
+  } else if (isBigNumber(dim)) {
+    return dim.minus(1)
+  } else {
+    return dim
+  }
+}
+
+export function isNumberOrBigNumber (n) {
+  return isNumber(n) || isBigNumber(n)
+}
diff --git a/src/expression/transform/utils/lastDimToZeroBase.js b/src/expression/transform/utils/lastDimToZeroBase.js
@@ -1,16 +1,14 @@
-import { isBigNumber, isCollection, isNumber } from '../../../utils/is.js'
-
+import { isCollection } from '../../../utils/is.js'
+import { dimToZeroBase, isNumberOrBigNumber } from './dimToZeroBase.js'
 /**
  * Change last argument dim from one-based to zero-based.
  */
 export function lastDimToZeroBase (args) {
   if (args.length === 2 && isCollection(args[0])) {
     args = args.slice()
     const dim = args[1]
-    if (isNumber(dim)) {
-      args[1] = dim - 1
-    } else if (isBigNumber(dim)) {
-      args[1] = dim.minus(1)
+    if (isNumberOrBigNumber(dim)) {
+      args[1] = dimToZeroBase(dim)
     }
   }
   return args

diff --git a/src/function/matrix/flatten.js b/src/function/matrix/flatten.js
@@ -2,9 +2,9 @@ import { flatten as flattenArray } from '../../utils/array.js'
 import { factory } from '../../utils/factory.js'
 
 const name = 'flatten'
-const dependencies = ['typed', 'matrix']
+const dependencies = ['typed']
 
-export const createFlatten = /* #__PURE__ */ factory(name, dependencies, ({ typed, matrix }) => {
+export const createFlatten = /* #__PURE__ */ factory(name, dependencies, ({ typed }) => {
   /**
    * Flatten a multidimensional matrix into a single dimensional matrix.
    * A new matrix is returned, the original matrix is left untouched.
@@ -30,9 +30,9 @@ export const createFlatten = /* #__PURE__ */ factory(name, dependencies, ({ type
     },
 
     Matrix: function (x) {
-      const flat = flattenArray(x.toArray())
-      // TODO: return the same matrix type as x (Dense or Sparse Matrix)
-      return matrix(flat)
+      // Return the same matrix type as x (Dense or Sparse Matrix)
+      // Return the same data type as x
+      return x.create(flattenArray(x.toArray()), x.datatype())
     }
   })
 })