feat: add polynomialRoot function (#2839)

* feat: add polynomialRoot function

  This is intended as a benchmark for general arithmetic/basic algebra
  functionality of mathjs, but was chosen to be something of potential
  independent utility as well, worth adding to mathjs in its own right.

  Currently ol=nly computes the numerical roots for real or complex
  polynomials of degree three or less.

  As usual, adds documentation, embedded documentation, tests, TypeScript
  declaration, and TypeScript tests.

  Also updates doc.test.js to make it easier to specify an array of complex
  numbers as the expected output, and comapres with the appropriate fuzz
  in such cases.

  Finally, adds a benchmark that computes the roots of all cubics with
  nonnegative integer coefficients no larger than five.

* doc: fix typo in polynomialRoot embedded docs

Thanks, SamuelTLG

* style: avoid slightly cryptic Boolean flag stand-in

src/expression/embeddedDocs/embeddedDocs.js

@@ -36,6 +36,7 @@ import { lsolveDocs } from './function/algebra/lsolve.js'
 import { lsolveAllDocs } from './function/algebra/lsolveAll.js'
 import { lupDocs } from './function/algebra/lup.js'
 import { lusolveDocs } from './function/algebra/lusolve.js'
+import { polynomialRootDocs } from './function/algebra/polynomialRoot.js'
 import { qrDocs } from './function/algebra/qr.js'
 import { rationalizeDocs } from './function/algebra/rationalize.js'
 import { resolveDocs } from './function/algebra/resolve.js'
@@ -340,6 +341,7 @@ export const embeddedDocs = {
   lup: lupDocs,
   lusolve: lusolveDocs,
   leafCount: leafCountDocs,
+  polynomialRoot: polynomialRootDocs,
   resolve: resolveDocs,
   simplify: simplifyDocs,
   simplifyConstant: simplifyConstantDocs,

src/expression/embeddedDocs/function/algebra/polynomialRoot.js

@@ -0,0 +1,16 @@
+export const polynomialRootDocs = {
+  name: 'polynomialRoot',
+  category: 'Algebra',
+  syntax: [
+    'x=polynomialRoot(-6, 3)',
+    'x=polynomialRoot(4, -4, 1)',
+    'x=polynomialRoot(-8, 12, -6, 1)'
+  ],
+  description: 'Finds the roots of a univariate polynomial given by its coefficients starting from constant, linear, and so on, increasing in degree.',
+  examples: [
+    'a = polynomialRoot(-6, 11, -6 1)'
+  ],
+  seealso: [
+    'cbrt', 'sqrt'
+  ]
+}

src/factoriesAny.js

@@ -210,6 +210,7 @@ export { createLup } from './function/algebra/decomposition/lup.js'
 export { createQr } from './function/algebra/decomposition/qr.js'
 export { createSlu } from './function/algebra/decomposition/slu.js'
 export { createLusolve } from './function/algebra/solver/lusolve.js'
+export { createPolynomialRoot } from './function/algebra/polynomialRoot.js'
 export { createHelpClass } from './expression/Help.js'
 export { createChainClass } from './type/chain/Chain.js'
 export { createHelp } from './expression/function/help.js'

src/function/algebra/polynomialRoot.js

@@ -0,0 +1,153 @@
+import { factory } from '../../utils/factory.js'
+
+const name = 'polynomialRoot'
+const dependencies = [
+  'typed',
+  'isZero',
+  'equalScalar',
+  'add',
+  'subtract',
+  'multiply',
+  'divide',
+  'sqrt',
+  'unaryMinus',
+  'cbrt',
+  'typeOf',
+  'im',
+  're'
+]
+
+export const createPolynomialRoot = /* #__PURE__ */ factory(name, dependencies, ({
+  typed,
+  isZero,
+  equalScalar,
+  add,
+  subtract,
+  multiply,
+  divide,
+  sqrt,
+  unaryMinus,
+  cbrt,
+  typeOf,
+  im,
+  re
+}) => {
+  /**
+   * Finds the numerical values of the distinct roots of a polynomial with real or complex coefficients.
+   * Currently operates only on linear, quadratic, and cubic polynomials using the standard
+   * formulas for the roots.
+   *
+   * Syntax:
+   *
+   *     polynomialRoot(constant, linearCoeff, quadraticCoeff, cubicCoeff)
+   *
+   * Examples:
+   *     // linear
+   *     math.polynomialRoot(6, 3)                                        // [-2]
+   *     math.polynomialRoot(math.complex(6,3), 3)                        // [-2 - i]
+   *     math.polynomialRoot(math.complex(6,3), math.complex(2,1))        // [-3 + 0i]
+   *     // quadratic
+   *     math.polynomialRoot(2, -3, 1)                                    // [2, 1]
+   *     math.polynomialRoot(8, 8, 2)                                     // [-2]
+   *     math.polynomialRoot(-2, 0, 1)                                    // [1.4142135623730951, -1.4142135623730951]
+   *     math.polynomialRoot(2, -2, 1)                                    // [1 + i, 1 - i]
+   *     math.polynomialRoot(math.complex(1,3), math.complex(-3, -2), 1)  // [2 + i, 1 + i]
+   *     // cubic
+   *     math.polynomialRoot(-6, 11, -6, 1)                               // [1, 3, 2]
+   *     math.polynomialRoot(-8, 0, 0, 1)                                 // [-1 - 1.7320508075688774i, 2, -1 + 1.7320508075688774i]
+   *     math.polynomialRoot(0, 8, 8, 2)                                  // [0, -2]
+   *     math.polynomialRoot(1, 1, 1, 1)                                  // [-1 + 0i, 0 - i, 0 + i]
+   *
+   * See also:
+   *     cbrt, sqrt
+   *
+   * @param {... number | Complex} coeffs
+   *     The coefficients of the polynomial, starting with with the constant coefficent, followed
+   *     by the linear coefficient and subsequent coefficients of increasing powers.
+   * @return {Array} The distinct roots of the polynomial
+   */
+
+  return typed(name, {
+    'number|Complex, ...number|Complex': (constant, restCoeffs) => {
+      const coeffs = [constant, ...restCoeffs]
+      while (coeffs.length > 0 && isZero(coeffs[coeffs.length - 1])) {
+        coeffs.pop()
+      }
+      if (coeffs.length < 2) {
+        throw new RangeError(
+          `Polynomial [${constant}, ${restCoeffs}] must have a non-zero non-constant coefficient`)
+      }
+      switch (coeffs.length) {
+        case 2: // linear
+          return [unaryMinus(divide(coeffs[0], coeffs[1]))]
+        case 3: { // quadratic
+          const [c, b, a] = coeffs
+          const denom = multiply(2, a)
+          const d1 = multiply(b, b)
+          const d2 = multiply(4, a, c)
+          if (equalScalar(d1, d2)) return [divide(unaryMinus(b), denom)]
+          const discriminant = sqrt(subtract(d1, d2))
+          return [
+            divide(subtract(discriminant, b), denom),
+            divide(subtract(unaryMinus(discriminant), b), denom)
+          ]
+        }
+        case 4: { // cubic, cf. https://en.wikipedia.org/wiki/Cubic_equation
+          const [d, c, b, a] = coeffs
+          const denom = unaryMinus(multiply(3, a))
+          const D0_1 = multiply(b, b)
+          const D0_2 = multiply(3, a, c)
+          const D1_1 = add(multiply(2, b, b, b), multiply(27, a, a, d))
+          const D1_2 = multiply(9, a, b, c)
+          if (equalScalar(D0_1, D0_2) && equalScalar(D1_1, D1_2)) {
+            return [divide(b, denom)]
+          }
+          const Delta0 = subtract(D0_1, D0_2)
+          const Delta1 = subtract(D1_1, D1_2)
+          const discriminant1 = add(
+            multiply(18, a, b, c, d), multiply(b, b, c, c))
+          const discriminant2 = add(
+            multiply(4, b, b, b, d),
+            multiply(4, a, c, c, c),
+            multiply(27, a, a, d, d))
+          if (equalScalar(discriminant1, discriminant2)) {
+            return [
+              divide(
+                subtract(
+                  multiply(4, a, b, c),
+                  add(multiply(9, a, a, d), multiply(b, b, b))),
+                multiply(a, Delta0)), // simple root
+              divide(
+                subtract(multiply(9, a, d), multiply(b, c)),
+                multiply(2, Delta0)) // double root
+            ]
+          }
+          // OK, we have three distinct roots
+          let Ccubed
+          if (equalScalar(D0_1, D0_2)) {
+            Ccubed = Delta1
+          } else {
+            Ccubed = divide(
+              add(
+                Delta1,
+                sqrt(subtract(
+                  multiply(Delta1, Delta1), multiply(4, Delta0, Delta0, Delta0)))
+              ),
+              2)
+          }
+          const allRoots = true
+          const rawRoots = cbrt(Ccubed, allRoots).toArray().map(
+            C => divide(add(b, C, divide(Delta0, C)), denom))
+          return rawRoots.map(r => {
+            if (typeOf(r) === 'Complex' && equalScalar(re(r), re(r) + im(r))) {
+              return re(r)
+            }
+            return r
+          })
+        }
+        default:
+          throw new RangeError(`only implemented for cubic or lower-order polynomials, not ${coeffs}`)
+      }
+    }
+  })
+})

test/benchmark/index.js

@@ -1,6 +1,7 @@
 // run all benchmarks
 require('./expression_parser')
 require('./algebra')
+require('./roots')
 require('./matrix_operations')
 require('./prime')
 require('./load')

test/benchmark/roots.js

@@ -0,0 +1,45 @@
+// test performance of the expression parser in node.js
+
+const Benchmark = require('benchmark')
+const padRight = require('pad-right')
+const math = require('../..')
+
+function pad (text) {
+  return padRight(text, 40, ' ')
+}
+
+const maxCoeff = 5
+function countRoots () {
+  let polys = 0
+  let roots = 0
+  for (let d = 0; d <= maxCoeff; ++d) {
+    for (let c = 0; c <= maxCoeff; ++c) {
+      for (let b = 0; b <= maxCoeff; ++b) {
+        for (let a = 1; a <= maxCoeff; ++a) {
+          polys += 1
+          roots += math.polynomialRoot(d, c, b, a).length
+        }
+      }
+    }
+  }
+  return [polys, roots]
+}
+
+const test = countRoots()
+console.log('There are', test[1], 'roots of the', test[0], 'integer cubic')
+console.log('polynomials (with coefficients <=', maxCoeff, ')')
+
+const results = []
+
+const suite = new Benchmark.Suite()
+suite
+  .add(pad('count roots'), function () {
+    const res = countRoots()
+    results.push(res)
+  })
+  .on('cycle', function (event) {
+    console.log(String(event.target))
+  })
+  .on('complete', function () {
+  })
+  .run()

test/node-tests/doc.test.js

@@ -63,7 +63,14 @@ function extractValue (spec) {
   try {
     value = eval(spec) // eslint-disable-line no-eval
   } catch (err) {
-    if (err instanceof SyntaxError || err instanceof ReferenceError) {
+    if (spec[0] === '[') {
+      // maybe it was an array with mathjs expressions in it
+      try {
+        value = math.evaluate(spec).toArray()
+      } catch (newError) {
+        value = spec
+      }
+    } else if (err instanceof SyntaxError || err instanceof ReferenceError) {
       value = spec
     } else {
       throw err
@@ -110,17 +117,21 @@ function maybeCheckExpectation (name, expected, expectedFrom, got, gotFrom) {
 }
 
 function checkExpectation (want, got) {
-  if (Array.isArray(want) && !Array.isArray(got)) {
-    approx.deepEqual(got, math.matrix(want), 1e-9)
-  } else if (want instanceof math.Unit && got instanceof math.Unit) {
-    approx.deepEqual(got, want, 1e-9)
-  } else if (want instanceof math.Complex && got instanceof math.Complex) {
-    approx.deepEqual(got, want, 1e-9)
-  } else if (typeof want === 'number' &&
-             typeof got === 'number' &&
-             want !== got) {
-    approx.equal(got, want, 1e-9)
+  if (Array.isArray(want)) {
+    if (!Array.isArray(got)) {
+      want = math.matrix(want)
+    }
+    return approx.deepEqual(got, want, 1e-9)
+  }
+  if (want instanceof math.Unit && got instanceof math.Unit) {
+    return approx.deepEqual(got, want, 1e-9)
+  }
+  if (want instanceof math.Complex && got instanceof math.Complex) {
+    return approx.deepEqual(got, want, 1e-9)
+  }
+  if (typeof want === 'number' && typeof got === 'number' && want !== got) {
     console.log(`  Note: return value ${got} not exactly as expected: ${want}`)
+    return approx.equal(got, want, 1e-9)
   } else {
     assert.deepEqual(got, want)
   }

test/typescript-tests/testTypes.ts

@@ -183,6 +183,18 @@ Bignumbers examples
   // const _eanother: number = math.exp(true)
 }
 
+/*
+  Algebra function examples
+*/
+{
+  const math = create(all, {})
+  const derivVal = math.derivative('x^3 + x^2', 'x').evaluate({ x: 2 })
+  assert.strictEqual(derivVal, 16)
+  const roots = math.polynomialRoot(9, 6, 1)
+  assert.strictEqual(roots.length, 1) // double root, so only one of them
+  assert.strictEqual(roots[0], -3)
+}
+
 /*
 Chaining examples
 */

test/unit-tests/function/algebra/polynomialRoot.test.js

@@ -0,0 +1,53 @@
+import approx from '../../../../tools/approx.js'
+import math from '../../../../src/defaultInstance.js'
+
+const complex = math.complex
+const pRoot = math.polynomialRoot
+
+describe('polynomialRoot', function () {
+  it('should solve a linear equation with real or complex coefficients',
+    function () {
+      approx.deepEqual(pRoot(6, 3), [-2])
+      approx.deepEqual(pRoot(complex(-3, 2), 2), [complex(1.5, -1)])
+      approx.deepEqual(pRoot(complex(3, 1), complex(-1, -1)), [complex(2, -1)])
+    })
+  it('should solve a quadratic equation with a double root', function () {
+    approx.deepEqual(pRoot(4, 4, 1), [-2])
+    approx.deepEqual(pRoot(complex(0, 2), complex(2, 2), 1), [complex(-1, -1)])
+  })
+  it('should solve a quadratic with two distinct roots', function () {
+    approx.deepEqual(pRoot(-3, 2, 1), [1, -3])
+    approx.deepEqual(pRoot(-2, 0, 1), [math.sqrt(2), -math.sqrt(2)])
+    approx.deepEqual(
+      pRoot(4, 2, 1), [complex(-1, math.sqrt(3)), complex(-1, -math.sqrt(3))])
+    approx.deepEqual(
+      pRoot(complex(3, 1), -3, 1), [complex(2, -1), complex(1, 1)])
+  })
+  it('should solve a cubic with a triple root', function () {
+    approx.deepEqual(pRoot(8, 12, 6, 1), [-2])
+    approx.deepEqual(
+      pRoot(complex(-2, 11), complex(9, -12), complex(-6, 3), 1),
+      [complex(2, -1)])
+  })
+  it('should solve a cubic with one simple and one double root', function () {
+    approx.deepEqual(pRoot(4, 0, -3, 1), [-1, 2])
+    approx.deepEqual(
+      pRoot(complex(9, 9), complex(15, 6), complex(7, 1), 1),
+      [complex(-1, -1), -3])
+    approx.deepEqual(
+      pRoot(complex(0, 6), complex(6, 8), complex(5, 2), 1),
+      [-3, complex(-1, -1)])
+    approx.deepEqual(
+      pRoot(complex(2, 6), complex(8, 6), complex(5, 1), 1),
+      [complex(-3, 1), complex(-1, -1)])
+  })
+  it('should solve a cubic with three distinct roots', function () {
+    approx.deepEqual(pRoot(6, 11, 6, 1), [-3, -1, -2])
+    approx.deepEqual(
+      pRoot(-1, -2, 0, 1), [-1, (1 + math.sqrt(5)) / 2, (1 - math.sqrt(5)) / 2])
+    approx.deepEqual(pRoot(1, 1, 1, 1), [-1, complex(0, -1), complex(0, 1)])
+    approx.deepEqual(
+      pRoot(complex(0, -10), complex(8, 12), complex(-6, -3), 1),
+      [complex(1, 1), complex(3, 1), complex(2, 1)])
+  })
+})

types/index.d.ts

@@ -899,6 +899,22 @@ declare namespace math {
       threshold?: number
     ): MathArray
 
+    /* Finds the roots of a polynomial of degree three or less. Coefficients are given constant first
+     * followed by linear and higher powers in order; coefficients beyond the degree of the polynomial
+     * need not be specified.
+     * @param {number|Complex} constantCoeff
+     * @param {number|Complex} linearCoeff
+     * @param {number|Complex} quadraticCoeff
+     * @param {number|Complex} cubicCoeff
+     * @returns {Array<number|Complex>} array of roots of specified polynomial
+     */
+    polynomialRoot(
+      constantCoeff: number | Complex,
+      linearCoeff: number | Complex,
+      quadraticCoeff?: number | Complex,
+      cubicCoeff?: number | Complex
+    ): (number | Complex)[]
+
     /**
      * Calculate the Matrix QR decomposition. Matrix A is decomposed in two
      * matrices (Q, R) where Q is an orthogonal matrix and R is an upper