feat: Indexing with an array of booleans (#2994)

* Included math to syntax when missing

* Included solveODE

* renamed initialStep as firstStep

* Included tests for solveODE

* Test the full state instead of the final state

* Fixed issue with tolerance

* Indexing with an array of booleans

* Indexing with booleans and with empty

* Changed index embedded docs

* removed solveODE

* typos on tests

* included config.predictable

* Throws an error if the size doesn't match

* Included config predictable to get subset

* Can do replacement by broadcasting

* DenseMatrix set can broadcast first

* Added tests for broadcasted subset in the parser

* Faster cloning of deep arrays

* Included docs and better test coverage

* Test coverage for `subset`

* Removed config predictable from subset

* Removed config from index and sparseMatrix

* Redaction and typos

* Cleanup unnecesary changes

* fixed issue when there is no need to broadcast

* Inline ifs

* Included specific broadcasting test

* Reduced repetition

---------

Co-authored-by: David Contreras <[email protected]>
Co-authored-by: Jos de Jong <[email protected]>

docs/datatypes/matrices.md

@@ -331,6 +331,49 @@ method `.set()`, the matrix will be resized. By default, new items will be
 initialized with zero, but it is possible to specify an alternative value using
 the optional third argument `defaultValue`.
 
+## Advanced Indexing
+
+Boolean array indexing is a technique that allows you to filter, replace, and set values in an array based on logical conditions. This can be done by creating a boolean array that represents the desired conditions, and then using that array as an index to select the elements of the original array that meet those conditions.
+
+For example, a boolean array can be created to represent all the even numbers in an array, and then used to filter the original array to only include the even numbers. Alternatively, a boolean array can be created to represent all the elements of an array that are greater than a certain value, and then used to replace all the elements of the original array that are greater than that value with a new value.
+
+
+```js
+const q = [1, 2, 3, 4]
+math.subset(q, math.index([true, false, true, false]))      // Array [1, 3]
+
+// filtering
+math.subset(q, math.index(math.larger(q, 2)))               // Array [3, 4]
+
+// filtering with no matches
+math.subset(q, math.index(math.larger(q, 5)))               // Array []
+
+// setting specific values, please note that the replacement value is broadcasted
+q = math.subset(q, math.index(math.smaller(q, 3)), 0)       // q = [0, 0, 3, 4]
+
+// replacing specific values
+math.subset(q, math.index(math.equal(q, 0)), [1, 2])        // q = [1, 2, 3, 4]
+```
+
+The same can be accomplished in the parser in a much more compact manner. Please note that everything after `#` are comments.
+```js
+math.evaluate(`
+q = [1, 2, 3, 4]
+q[[true, false, true, false]] # Matrix [1, 3]
+q[q>2]                        # Matrix [3, 4]
+q[q>5]                        # Matrix []
+q[q <3] = 0                   # q = [0, 0, 3, 4]
+q[q==0] = [1, 2]              # q = [1, 2, 3, 4]
+`)
+```
+The expression inside the index can be as complex as needed as long it evaluates to an array of booleans of the same size.
+```js
+math.evaluate(`
+q = [1, 2, 3, 4]
+r = [6, 5, 4, 3]
+q[q > 3 and r < 4]     # [4]
+`)
+```
 
 ## Iterating
 

src/expression/embeddedDocs/construction/index.js

@@ -12,11 +12,12 @@ export const indexDocs = {
   description:
       'Create an index to get or replace a subset of a matrix',
   examples: [
-    '[1, 2, 3]',
     'A = [1, 2, 3; 4, 5, 6]',
     'A[1, :]',
     'A[1, 2] = 50',
-    'A[1:2, 1:2] = ones(2, 2)'
+    'A[1:2, 1:2] = 1',
+    'B = [1, 2, 3]',
+    'B[B>1 and B<3]'
   ],
   seealso: [
     'bignumber', 'boolean', 'complex', 'matrix,', 'number', 'range', 'string', 'unit'

src/expression/transform/index.transform.js

@@ -2,9 +2,9 @@ import { isArray, isBigNumber, isMatrix, isNumber, isRange } from '../../utils/i
 import { factory } from '../../utils/factory.js'
 
 const name = 'index'
-const dependencies = ['Index']
+const dependencies = ['Index', 'getMatrixDataType']
 
-export const createIndexTransform = /* #__PURE__ */ factory(name, dependencies, ({ Index }) => {
+export const createIndexTransform = /* #__PURE__ */ factory(name, dependencies, ({ Index, getMatrixDataType }) => {
   /**
    * Attach a transform function to math.index
    * Adds a property transform containing the transform function.
@@ -16,20 +16,22 @@ export const createIndexTransform = /* #__PURE__ */ factory(name, dependencies,
     for (let i = 0, ii = arguments.length; i < ii; i++) {
       let arg = arguments[i]
 
-      // change from one-based to zero based, and convert BigNumber to number
+      // change from one-based to zero based, convert BigNumber to number and leave Array of Booleans as is
       if (isRange(arg)) {
         arg.start--
         arg.end -= (arg.step > 0 ? 0 : 2)
       } else if (arg && arg.isSet === true) {
         arg = arg.map(function (v) { return v - 1 })
       } else if (isArray(arg) || isMatrix(arg)) {
-        arg = arg.map(function (v) { return v - 1 })
+        if (getMatrixDataType(arg) !== 'boolean') {
+          arg = arg.map(function (v) { return v - 1 })
+        }
       } else if (isNumber(arg)) {
         arg--
       } else if (isBigNumber(arg)) {
         arg = arg.toNumber() - 1
       } else if (typeof arg === 'string') {
-        // leave as is
+      // leave as is
       } else {
         throw new TypeError('Dimension must be an Array, Matrix, number, string, or Range')
       }

src/expression/transform/subset.transform.js

@@ -3,10 +3,10 @@ import { errorTransform } from './utils/errorTransform.js'
 import { createSubset } from '../../function/matrix/subset.js'
 
 const name = 'subset'
-const dependencies = ['typed', 'matrix']
+const dependencies = ['typed', 'matrix', 'zeros', 'add']
 
-export const createSubsetTransform = /* #__PURE__ */ factory(name, dependencies, ({ typed, matrix }) => {
-  const subset = createSubset({ typed, matrix })
+export const createSubsetTransform = /* #__PURE__ */ factory(name, dependencies, ({ typed, matrix, zeros, add }) => {
+  const subset = createSubset({ typed, matrix, zeros, add })
 
   /**
    * Attach a transform function to math.subset

src/function/matrix/concat.js

@@ -1,6 +1,6 @@
 import { isBigNumber, isMatrix, isNumber } from '../../utils/is.js'
 import { clone } from '../../utils/object.js'
-import { arraySize } from '../../utils/array.js'
+import { arraySize, concat as _concat } from '../../utils/array.js'
 import { IndexError } from '../../error/IndexError.js'
 import { DimensionError } from '../../error/DimensionError.js'
 import { factory } from '../../utils/factory.js'
@@ -94,7 +94,7 @@ export const createConcat = /* #__PURE__ */ factory(name, dependencies, ({ typed
 
       let res = matrices.shift()
       while (matrices.length) {
-        res = _concat(res, matrices.shift(), dim, 0)
+        res = _concat(res, matrices.shift(), dim)
       }
 
       return asMatrix ? matrix(res) : res
@@ -105,31 +105,3 @@ export const createConcat = /* #__PURE__ */ factory(name, dependencies, ({ typed
     }
   })
 })
-
-/**
- * Recursively concatenate two matrices.
- * The contents of the matrices is not cloned.
- * @param {Array} a             Multi dimensional array
- * @param {Array} b             Multi dimensional array
- * @param {number} concatDim    The dimension on which to concatenate (zero-based)
- * @param {number} dim          The current dim (zero-based)
- * @return {Array} c            The concatenated matrix
- * @private
- */
-function _concat (a, b, concatDim, dim) {
-  if (dim < concatDim) {
-    // recurse into next dimension
-    if (a.length !== b.length) {
-      throw new DimensionError(a.length, b.length)
-    }
-
-    const c = []
-    for (let i = 0; i < a.length; i++) {
-      c[i] = _concat(a[i], b[i], concatDim, dim + 1)
-    }
-    return c
-  } else {
-    // concatenate this dimension
-    return a.concat(b)
-  }
-}

src/function/matrix/subset.js

@@ -1,14 +1,14 @@
 import { isIndex } from '../../utils/is.js'
 import { clone } from '../../utils/object.js'
-import { validateIndex } from '../../utils/array.js'
+import { isEmptyIndex, validateIndex, validateIndexSourceSize } from '../../utils/array.js'
 import { getSafeProperty, setSafeProperty } from '../../utils/customs.js'
 import { DimensionError } from '../../error/DimensionError.js'
 import { factory } from '../../utils/factory.js'
 
 const name = 'subset'
-const dependencies = ['typed', 'matrix']
+const dependencies = ['typed', 'matrix', 'zeros', 'add']
 
-export const createSubset = /* #__PURE__ */ factory(name, dependencies, ({ typed, matrix }) => {
+export const createSubset = /* #__PURE__ */ factory(name, dependencies, ({ typed, matrix, zeros, add }) => {
   /**
    * Get or set a subset of a matrix or string.
    *
@@ -20,21 +20,23 @@ export const createSubset = /* #__PURE__ */ factory(name, dependencies, ({ typed
    *
    *     // get a subset
    *     const d = [[1, 2], [3, 4]]
-   *     math.subset(d, math.index(1, 0))        // returns 3
-   *     math.subset(d, math.index([0, 1], 1))   // returns [[2], [4]]
+   *     math.subset(d, math.index(1, 0))             // returns 3
+   *     math.subset(d, math.index([0, 1], 1))        // returns [[2], [4]]
+   *     math.subset(d, math.index([false, true], 0)) // returns [[3]]
    *
    *     // replace a subset
    *     const e = []
-   *     const f = math.subset(e, math.index(0, [0, 2]), [5, 6])  // f = [[5, 6]]
-   *     const g = math.subset(f, math.index(1, 1), 7, 0)         // g = [[5, 6], [0, 7]]
+   *     const f = math.subset(e, math.index(0, [0, 2]), [5, 6])  // f = [[5, 0, 6]]
+   *     const g = math.subset(f, math.index(1, 1), 7, 0)         // g = [[5, 0, 6], [0, 7, 0]]
+   *     math.subset(g, math.index([false, true], 1), 8)          // returns [[5, 0, 6], [0, 8, 0]]
    *
    *     // get submatrix using ranges
    *     const M = [
    *       [1,2,3],
    *       [4,5,6],
    *       [7,8,9]
    *     ]
-   *     math.subset(M, math.index(math.range(0,2), math.range(0,3))) // [[1,2,3],[4,5,6]]
+   *     math.subset(M, math.index(math.range(0,2), math.range(0,3))) // [[1, 2, 3], [4, 5, 6]]
    *
    * See also:
    *
@@ -53,49 +55,81 @@ export const createSubset = /* #__PURE__ */ factory(name, dependencies, ({ typed
    *                                          math.matrix elements will be left undefined.
    * @return {Array | Matrix | string} Either the retrieved subset or the updated matrix.
    */
+
   return typed(name, {
     // get subset
-    'Array, Index': function (value, index) {
-      const m = matrix(value)
-      const subset = m.subset(index) // returns a Matrix
-      return index.isScalar()
-        ? subset
-        : subset.valueOf() // return an Array (like the input)
-    },
-
     'Matrix, Index': function (value, index) {
+      if (isEmptyIndex(index)) { return matrix() }
+      validateIndexSourceSize(value, index)
       return value.subset(index)
     },
 
+    'Array, Index': typed.referTo('Matrix, Index', function (subsetRef) {
+      return function (value, index) {
+        const subsetResult = subsetRef(matrix(value), index)
+        return index.isScalar() ? subsetResult : subsetResult.valueOf()
+      }
+    }),
+
     'Object, Index': _getObjectProperty,
 
     'string, Index': _getSubstring,
 
     // set subset
-    'Array, Index, any': function (value, index, replacement) {
-      return matrix(clone(value))
-        .subset(index, replacement, undefined)
-        .valueOf()
+    'Matrix, Index, any, any': function (value, index, replacement, defaultValue) {
+      if (isEmptyIndex(index)) { return value }
+      validateIndexSourceSize(value, index)
+      return value.clone().subset(index, _broadcastReplacement(replacement, index), defaultValue)
     },
 
-    'Array, Index, any, any': function (value, index, replacement, defaultValue) {
-      return matrix(clone(value))
-        .subset(index, replacement, defaultValue)
-        .valueOf()
-    },
+    'Array, Index, any, any': typed.referTo('Matrix, Index, any, any', function (subsetRef) {
+      return function (value, index, replacement, defaultValue) {
+        const subsetResult = subsetRef(matrix(value), index, replacement, defaultValue)
+        return subsetResult.isMatrix ? subsetResult.valueOf() : subsetResult
+      }
+    }),
 
-    'Matrix, Index, any': function (value, index, replacement) {
-      return value.clone().subset(index, replacement)
-    },
+    'Array, Index, any': typed.referTo('Matrix, Index, any, any', function (subsetRef) {
+      return function (value, index, replacement) {
+        return subsetRef(matrix(value), index, replacement, undefined).valueOf()
+      }
+    }),
 
-    'Matrix, Index, any, any': function (value, index, replacement, defaultValue) {
-      return value.clone().subset(index, replacement, defaultValue)
-    },
+    'Matrix, Index, any': typed.referTo('Matrix, Index, any, any', function (subsetRef) {
+      return function (value, index, replacement) { return subsetRef(value, index, replacement, undefined) }
+    }),
 
     'string, Index, string': _setSubstring,
     'string, Index, string, string': _setSubstring,
     'Object, Index, any': _setObjectProperty
   })
+
+  /**
+   * Broadcasts a replacment value to be the same size as index
+   * @param {number | BigNumber | Array | Matrix} replacement Replacement value to try to broadcast
+   * @param {*} index Index value
+   * @returns broadcasted replacement that matches the size of index
+   */
+
+  function _broadcastReplacement (replacement, index) {
+    if (typeof replacement === 'string') {
+      throw new Error('can\'t boradcast a string')
+    }
+    if (index._isScalar) {
+      return replacement
+    }
+
+    const indexSize = index.size()
+    if (indexSize.every(d => d > 0)) {
+      try {
+        return add(replacement, zeros(indexSize))
+      } catch (error) {
+        return replacement
+      }
+    } else {
+      return replacement
+    }
+  }
 })
 
 /**
@@ -110,6 +144,10 @@ function _getSubstring (str, index) {
     // TODO: better error message
     throw new TypeError('Index expected')
   }
+
+  if (isEmptyIndex(index)) { return '' }
+  validateIndexSourceSize(Array.from(str), index)
+
   if (index.size().length !== 1) {
     throw new DimensionError(index.size().length, 1)
   }
@@ -134,7 +172,7 @@ function _getSubstring (str, index) {
  * @param {string} str            string to be replaced
  * @param {Index} index           An index or list of indices (character positions)
  * @param {string} replacement    Replacement string
- * @param {string} [defaultValue] Default value to be uses when resizing
+ * @param {string} [defaultValue] Default value to be used when resizing
  *                                the string. is ' ' by default
  * @returns {string} result
  * @private
@@ -144,6 +182,8 @@ function _setSubstring (str, index, replacement, defaultValue) {
     // TODO: better error message
     throw new TypeError('Index expected')
   }
+  if (isEmptyIndex(index)) { return str }
+  validateIndexSourceSize(Array.from(str), index)
   if (index.size().length !== 1) {
     throw new DimensionError(index.size().length, 1)
   }
@@ -197,6 +237,8 @@ function _setSubstring (str, index, replacement, defaultValue) {
  * @private
  */
 function _getObjectProperty (object, index) {
+  if (isEmptyIndex(index)) { return undefined }
+
   if (index.size().length !== 1) {
     throw new DimensionError(index.size(), 1)
   }
@@ -218,6 +260,7 @@ function _getObjectProperty (object, index) {
  * @private
  */
 function _setObjectProperty (object, index, replacement) {
+  if (isEmptyIndex(index)) { return object }
   if (index.size().length !== 1) {
     throw new DimensionError(index.size(), 1)
   }

src/type/matrix/DenseMatrix.js

@@ -1,5 +1,5 @@
 import { isArray, isBigNumber, isCollection, isIndex, isMatrix, isNumber, isString, typeOf } from '../../utils/is.js'
-import { arraySize, getArrayDataType, processSizesWildcard, reshape, resize, unsqueeze, validate, validateIndex } from '../../utils/array.js'
+import { arraySize, getArrayDataType, processSizesWildcard, reshape, resize, unsqueeze, validate, validateIndex, broadcastTo } from '../../utils/array.js'
 import { format } from '../../utils/string.js'
 import { isInteger } from '../../utils/number.js'
 import { clone, deepStrictEqual } from '../../utils/object.js'
@@ -321,11 +321,23 @@ export const createDenseMatrixClass = /* #__PURE__ */ factory(name, dependencies
       if (sSize.length !== 0) {
         throw new TypeError('Scalar expected')
       }
-
       matrix.set(index.min(), submatrix, defaultValue)
     } else {
       // set a submatrix
 
+      // broadcast submatrix
+      if (!deepStrictEqual(sSize, iSize)) {
+        try {
+          if (sSize.length === 0) {
+            submatrix = broadcastTo([submatrix], iSize)
+          } else {
+            submatrix = broadcastTo(submatrix, iSize)
+          }
+          sSize = arraySize(submatrix)
+        } catch {
+        }
+      }
+
       // validate dimensions
       if (iSize.length < matrix._size.length) {
         throw new DimensionError(iSize.length, matrix._size.length, '<')