feat: Fully define complex cube root (#39)

Also extends add and multiply to allow multiple arguments.

Resolves #29.

Reviewed-on: #39
Co-authored-by: Glen Whitney <glen@studioinfinity.org>
Co-committed-by: Glen Whitney <glen@studioinfinity.org>

src/complex/__test__/arithmetic.spec.js

@@ -132,4 +132,15 @@ describe('complex arithmetic operations', () => {
       assert.deepStrictEqual(
          sub(cplx(z,z), cplx(10,20)), cplx(cplx(-7, 4), cplx(-17, 4)))
    })
+   it('cube roots complex numbers and quaternions', () => {
+      assert(math.equal(math.cbrt(cplx(0, 8))[0], cplx(Math.sqrt(3), 1)))
+      assert.deepStrictEqual(
+         math.cbrt(cplx(2, 3)),
+         [cplx(1.4518566183526649, 0.49340353410400467),
+            cplx(-1.1532283040274218, 1.0106429470939742),
+            cplx(-0.29862831432524256, -1.5040464811979786)])
+      const quat = cplx(cplx(1, 2), cplx(2, 1))
+      const root = math.cbrt(quat)[0]
+      assert(math.equal(math.multiply(root, root, root), quat))
+   })
 })

src/complex/__test__/type.spec.js

@@ -19,6 +19,13 @@ describe('complex type operations', () => {
       assert.strictEqual(math.arg(cplx(1, Math.sqrt(3))), Math.PI/3)
       assert.strictEqual(math.arg(cplx(true, true)), Math.PI/4)
    })
+   it('calculates the real parts of complex numbers and quaternions', () => {
+      assert.strictEqual(math.re(cplx(2, -3)), 2)
+      assert.strictEqual(math.re(cplx(cplx(0.5), cplx(8, -7))), 0.5)
+   })
+   it('calculates the argument of a quaternion', () => {
+      assert(math.equal(math.arg(cplx(cplx(1, 1), cplx(1, 1))), Math.acos(0.5)))
+   })
    it('detects associates of a complex number', () => {
       const z = cplx(3, 4)
       const assoc = math.associate

src/complex/arithmetic.js

@@ -4,6 +4,7 @@ import {ResolutionError} from '#core/helpers.js'
 import {Returns, ReturnType, ReturnTyping} from '#core/Type.js'
 import {match} from '#core/TypePatterns.js'
 import {ReturnsAs} from '#generic/helpers.js'
+import {NumberT} from '#number/NumberT.js'
 
 const {conservative, full, free} = ReturnTyping
 
@@ -141,4 +142,44 @@ export const sqrt = match(Complex, (math, C, strategy) => {
    return ReturnsAs(prune, z => prune(sqrtImp(z)))
 })
 
+const TAU3 = 2 * Math.PI / 3
+
+export const cbrt = match(Complex, (math, C) => {
+   const arg = math.arg.resolve(C)
+   const divArg = math.divide.resolve([ReturnType(arg), NumberT])
+   const absC = math.abs.resolve(C)
+   const cbrtR = math.cbrt.resolve(ReturnType(absC), conservative)
+   const im = math.im.resolve(C)
+   const absIm = math.abs.resolve(ReturnType(im))
+   const divIm = math.divide.resolve([ReturnType(im), ReturnType(absIm)])
+   // TODO: replace with nanomath cos and sin when available
+   // const cos = math.cos.resolve(ReturnType(divArg))
+   // const CosType = ReturnType(cos)  // and similarly for sin
+   const cos = Math.cos
+   const CosType = NumberT
+   const sin = Math.sin
+   const SinType = NumberT
+   const mulIm = math.multiply.resolve([ReturnType(divIm), SinType])
+   const addUp = math.add.resolve([CosType, ReturnType(mulIm)])
+   const addAngle = math.add.resolve([ReturnType(divArg), math.typeOf(TAU3)])
+   const subAngle = math.subtract.resolve(
+      [ReturnType(divArg), math.typeOf(TAU3)])
+   const scale = math.multiply.resolve([ReturnType(cbrtR), ReturnType(addUp)])
+   const Cout = ReturnType(scale)
+   const vec = math.vector.resolve([Cout, Cout, Cout])
+   return ReturnsAs(vec, z => {
+      const arg3 = divArg(arg(z), 3)
+      const mag = absC(z)
+      const r = cbrtR(mag)
+      const imz = im(z)
+      const unit = divIm(imz, absIm(imz))
+      // At this point, z = mag*(cos(arg) + unit * sin(arg))
+      // so one cube root is r*(cos(arg3) + unit * sin(arg3))
+      // and we get the other two by adjusting arg3 by +- TAU3
+      const cus = theta => scale(r, addUp(cos(theta), mulIm(unit, sin(theta))))
+      return vec(
+         cus(arg3), cus(addAngle(arg3, TAU3)), cus(subAngle(arg3, TAU3)))
+   })
+})
+
 export const subtract = promoteBinary('subtract')

src/complex/type.js

@@ -31,20 +31,26 @@ export const complex = [
    })
 ]
 
-export const arg = // [  // enable when we have atan2 in mathjs
-//   match(Complex, (math, C) => {
-//      const re = math.re.resolve(C)
-//      const R = ReturnType(re)
-//      const im = math.im.resolve(C)
-//      const abs = math.abs.resolve(C)
-//      const atan2 = math.atan2.resolve([R, R], conservative)
-//      return Returns(R, z => atan2(abs(im(z)), re(z)))
-//   }),  // note always between 0 and tau/2; need to use in conjunction
-//        // with a complex unit function that gives you the proper
-//        // imaginary unit, ±i in the simple complex case, to restore the
-//        // full circle of values for the direction of a complex number
-   match(Complex(NumberT), Returns(NumberT, z => Math.atan2(z.im, z.re)))
-//]
+export const arg = match(Complex, (math, C) => {
+   const re = math.re.resolve(C)
+   const atan2 = Math.atan2
+   const ArgType = NumberT
+   if (ReturnType(re) === C.Component) {
+      // "plain" Complex (not nested, i.e. not quaternions)
+      // FIXME: use math.atan2 once available
+      // const atan2 = math.atan2.resolve(
+      //   [C.Component, C.Component], conservative)
+      // const argType = ReturnType(atan2)
+      return Returns(ArgType, z => atan2(z.im, z.re))
+   }
+   const im = math.im.resolve(C)
+   const abs = math.abs.resolve(ReturnType(im))
+   // FIXME: use math.atan2 once available
+   // const atan2 = math.atan2.resolve(
+   //    [ReturnType(abs), ReturnType(re)], conservative)
+   // const ArgType = ReturnType(atan2)
+   return Returns(ArgType, z => atan2(abs(im(z)), re(z)))
+})
 
 /* Returns true if w is z multiplied by a complex unit */
 export const associate = match(

src/core/TypeDispatcher.js

@@ -275,7 +275,8 @@ export class TypeDispatcher {
          throw new ReferenceError(`no method or value for key '${key}'`)
       }
       if (!Array.isArray(types)) types = [types]
-      if (types.some(T => T === Unknown)) {
+      // Unknown and union, i.e., OneOf(...), stymie resolution:
+      if (types.some(T => T === Unknown || T.unifies)) {
          if (!strategy) return this[key]
          const thisTypeOf = whichType(this.types)
          return (...args) => {

src/generic/__test__/arithmetic.spec.js

@@ -1,18 +1,53 @@
 import assert from 'assert'
 import math from '#nanomath'
-import {ReturnType, ReturnTyping} from '#core/Type.js'
+import {ReturnType, ReturnTyping, Unknown} from '#core/Type.js'
 
 const {Complex, NumberT} = math.types
 
 describe('generic arithmetic', () => {
+   const cplx = math.complex
    it('squares anything', () => {
       const sq = math.square
       assert.strictEqual(sq(7), 49)
       assert.strictEqual(ReturnType(math.square.resolve([NumberT])), NumberT)
-      assert.deepStrictEqual(sq(math.complex(3, 4)), math.complex(-7, 24))
-      const eyes = math.complex(0, 2)
+      assert.deepStrictEqual(sq(cplx(3, 4)), cplx(-7, 24))
+      const eyes = cplx(0, 2)
       assert.strictEqual(sq(eyes), -4)
       const sqFull = math.square.resolve(Complex(NumberT), ReturnTyping.full)
-      assert.deepStrictEqual(sqFull(eyes), math.complex(-4, 0))
+      assert.deepStrictEqual(sqFull(eyes), cplx(-4, 0))
+   })
+
+   it('adds up multiple arguments', () => {
+      const add = math.add
+      assert.strictEqual(add(), 0)
+      assert.strictEqual(add(1), 1)
+      assert.deepStrictEqual(add(cplx(2)), cplx(2))
+      assert.deepStrictEqual(
+         ReturnType(add.resolve([Complex(NumberT)])), Complex(NumberT))
+      assert.strictEqual(add(3, 4, 5), 12)
+      assert.deepStrictEqual(add(6, cplx(0, 7), 8), cplx(14, 7))
+      assert.strictEqual(
+         ReturnType(add.resolve([NumberT, Complex(NumberT), NumberT])),
+         Unknown) // loses track of whether it comes out real or complex
+      assert.deepStrictEqual(
+         add(9, cplx(0, 10), cplx(11), cplx(0, -10)), 20)
+      const saveTyping = math.config.returnTyping
+      math.config.returnTyping = ReturnTyping.full
+      assert.strictEqual(
+         ReturnType(add.resolve([NumberT, Complex(NumberT), NumberT])),
+         Complex(NumberT)) // now definite
+      assert.deepStrictEqual(
+         add(9, cplx(0, 10), cplx(11), cplx(0, -10)), cplx(20))
+      math.config.returnTyping = saveTyping
+   })
+
+   it('multiplies multiple arguments', () => {
+      const mult = math.multiply
+      assert.strictEqual(mult(), 1)
+      assert.strictEqual(mult(2), 2)
+      assert.deepStrictEqual(mult(cplx(3)), cplx(3))
+      assert.strictEqual(mult(4, 5, 6), 120)
+      assert.deepStrictEqual(mult(7, cplx(0, 8), 9), cplx(0, 504))
+      assert.deepStrictEqual(mult(10, cplx(0, 1), cplx(0, 2), 3), -60)
    })
 })

src/generic/arithmetic.js

@@ -1,4 +1,5 @@
-import {ReturnsAs} from './helpers.js'
+import {iterateBinary, ReturnsAs} from './helpers.js'
+import {plain} from "#number/helpers.js"
 
 import {Returns, ReturnType, ReturnTyping} from '#core/Type.js'
 import {match, Any} from '#core/TypePatterns.js'
@@ -18,3 +19,9 @@ export const square = match(Any, (math, T, strategy) => {
    const mult = math.multiply.resolve([T, T], strategy)
    return ReturnsAs(mult, t => mult(t, t))
 })
+
+export const add = iterateBinary('add')
+add.push(plain(() => 0))
+
+export const multiply = iterateBinary('multiply')
+multiply.push(plain(() => 1))

src/generic/helpers.js

@@ -1,3 +1,24 @@
 import {Returns, ReturnType} from '#core/Type.js'
+import {match, Any, Multiple} from '#core/TypePatterns.js'
 
 export const ReturnsAs = (g, f) => Returns(ReturnType(g), f)
+export const iterateBinary = operation => [
+   match(Any, (_math, T) => Returns(T, t => t)),
+   match([Any, Any, Any, Multiple(Any)], (math, [T, U, V, Rest], strategy) => {
+      const op1 = math[operation].resolve([T, U], strategy)
+      const op2 = math[operation].resolve([ReturnType(op1), V])
+      let finalOp = op2
+      let restOps = []
+      for (const Typ of Rest) {
+         finalOp = math[operation].resolve([ReturnType(finalOp), Typ])
+         restOps.push(finalOp)
+      }
+      return ReturnsAs(finalOp, (t, u, v, rest) => {
+         let result = op2(op1(t, u), v)
+         for (let i = 0; i < rest.length; ++i) {
+            result = restOps[i](result, rest[i])
+         }
+         return result
+      })
+   })
+]

src/number/__test__/arithmetic.spec.js

@@ -30,4 +30,26 @@ describe('number arithmetic', () => {
       assert.deepStrictEqual(math.sqrt(-4), math.complex(0, 2))
       math.config.returnTyping = ReturnTyping.free
    })
+   it('takes cube root of numbers appropriately', () => {
+      assert(isNaN(math.cbrt(NaN)))
+      assert.strictEqual(math.cbrt(8), 2)
+      assert.strictEqual(math.cbrt(-8), -2)
+      const cbrt10 = Math.cbrt(10)
+      assert.strictEqual(math.cbrt(10), cbrt10)
+      const saveTyping = math.config.returnTyping
+      math.config.returnTyping = ReturnTyping.full
+      assert(isNaN(math.cbrt(NaN)))
+      const cbrtIm = Math.sqrt(3) / 2
+      const cplx = math.complex
+      const rt8 = math.cbrt(8)
+      assert.deepStrictEqual(
+         rt8, [cplx(2), cplx(-1, 2 * cbrtIm), cplx(-1, -2 * cbrtIm)])
+      assert(math.equal(math.multiply.apply(null, rt8), cplx(8)))
+      assert.deepStrictEqual(
+         math.cbrt(10),
+         [cplx(cbrt10),
+            cplx(-cbrt10 / 2, cbrt10 * cbrtIm),
+            cplx(-cbrt10 / 2, -cbrt10 * cbrtIm)])
+      math.config.returnTyping = saveTyping
+   })
 })

src/number/arithmetic.js

@@ -3,6 +3,7 @@ import {NumberT} from './NumberT.js'
 import {OneOf, Returns, ReturnTyping, Undefined} from '#core/Type.js'
 import {match} from '#core/TypePatterns.js'
 import {Complex} from '#complex/Complex.js'
+import {ReturnsAs} from '#generic/helpers.js'
 
 const {conservative, full} = ReturnTyping
 
@@ -15,7 +16,7 @@ export const add = [
    match([NumberT, Undefined], Returns(NumberT, () => NaN))
 ]
 export const divide = plain((a, b) => a / b)
-export const cbrt = plain(a => {
+const numberCbrt = a => {
    if (a === 0) return a
    const negate = a < 0 
    if (negate) a = -a
@@ -25,7 +26,29 @@ export const cbrt = plain(a => {
       result = (a / (result * result) + (2 * result)) / 3
    }
    return negate ? -result : result
+}
+const cbrtIm = Math.sqrt(3) / 2
+export const cbrt = match(NumberT, (math, _N, strategy) => {
+   if (strategy === full && math.types.Complex && math.types.Vector) {
+      // return vector of all three complex roots, real one first
+      const C = Complex(NumberT)
+      const vec = math.vector.resolve([C, C, C])
+      const promote = math.complex.resolve([NumberT])
+      const cplx = math.complex.resolve([NumberT, NumberT])
+      return ReturnsAs(vec, x => {
+         const realroot = numberCbrt(x)
+         if (!isFinite(realroot)) {
+            const full = cplx(realroot, realroot)
+            return(promote(realroot), full, full)
+         }
+         return vec(promote(realroot),
+            cplx(-realroot / 2, realroot * cbrtIm),
+            cplx(-realroot / 2, -realroot * cbrtIm))
+      })
+   }
+   return Returns(NumberT, numberCbrt)
 })
+
 export const invert = plain(a => 1/a)
 export const multiply = [
    plain((a, b) => a * b),