path: root/stdlib/source/lux/data/number/complex.lux

(.module: {#.doc "Complex arithmetic."}
  [lux #*
   [math]
   [control
    [equivalence (#+ Equivalence)]
    number
    codec
    ["M" monad (#+ do Monad)]
    ["p" parser]]
   [data
    [number ("frac/" Number<Frac>)]
    [text ("text/" Monoid<Text>)]
    [maybe]
    [collection [list ("list/" Functor<List>)]]]
   ["." macro
    [code]
    ["s" syntax (#+ syntax: Syntax)]]])

(type: #export Complex
  {#real Frac
   #imaginary Frac})

(syntax: #export (complex real {?imaginary (p.maybe s.any)})
  {#.doc (doc "Complex literals."
              (complex real imaginary)
              "The imaginary part can be omitted if it's 0."
              (complex real))}
  (wrap (list (` {#..real (~ real)
                  #..imaginary (~ (maybe.default (' 0.0)
                                                 ?imaginary))}))))

(def: #export i Complex (complex 0.0 1.0))

(def: #export one Complex (complex 1.0 0.0))

(def: #export zero Complex (complex 0.0 0.0))

(def: #export (not-a-number? complex)
  (or (number.not-a-number? (get@ #real complex))
      (number.not-a-number? (get@ #imaginary complex))))

(def: #export (= param input)
  (-> Complex Complex Bit)
  (and (f/= (get@ #real param)
            (get@ #real input))
       (f/= (get@ #imaginary param)
            (get@ #imaginary input))))

(do-template [<name> <op>]
  [(def: #export (<name> param input)
     (-> Complex Complex Complex)
     {#real (<op> (get@ #real param)
                  (get@ #real input))
      #imaginary (<op> (get@ #imaginary param)
                       (get@ #imaginary input))})]

  [+ f/+]
  [- f/-]
  )

(structure: #export _ (Equivalence Complex)
  (def: = ..=))

(def: #export negate
  (-> Complex Complex)
  (|>> (update@ #real frac/negate)
       (update@ #imaginary frac/negate)))

(def: #export signum
  (-> Complex Complex)
  (|>> (update@ #real frac/signum)
       (update@ #imaginary frac/signum)))

(def: #export conjugate
  (-> Complex Complex)
  (update@ #imaginary frac/negate))

(def: #export (*' param input)
  (-> Frac Complex Complex)
  {#real (f/* param
              (get@ #real input))
   #imaginary (f/* param
                   (get@ #imaginary input))})

(def: #export (* param input)
  (-> Complex Complex Complex)
  {#real (f/- (f/* (get@ #imaginary param)
                   (get@ #imaginary input))
              (f/* (get@ #real param)
                   (get@ #real input)))
   #imaginary (f/+ (f/* (get@ #real param)
                        (get@ #imaginary input))
                   (f/* (get@ #imaginary param)
                        (get@ #real input)))})

(def: #export (/ param input)
  (-> Complex Complex Complex)
  (let [(^slots [#real #imaginary]) param]
    (if (f/< (frac/abs imaginary)
             (frac/abs real))
      (let [quot (f// imaginary real)
            denom (|> real (f/* quot) (f/+ imaginary))]
        {#real (|> (get@ #real input) (f/* quot) (f/+ (get@ #imaginary input)) (f// denom))
         #imaginary (|> (get@ #imaginary input) (f/* quot) (f/- (get@ #real input)) (f// denom))})
      (let [quot (f// real imaginary)
            denom (|> imaginary (f/* quot) (f/+ real))]
        {#real (|> (get@ #imaginary input) (f/* quot) (f/+ (get@ #real input)) (f// denom))
         #imaginary (|> (get@ #imaginary input) (f/- (f/* quot (get@ #real input))) (f// denom))}))))

(def: #export (/' param subject)
  (-> Frac Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (f// param real)
     #imaginary (f// param imaginary)}))

(def: #export (% param input)
  (-> Complex Complex Complex)
  (let [scaled (/ param input)
        quotient (|> scaled
                     (update@ #real math.floor)
                     (update@ #imaginary math.floor))]
    (- (* quotient param)
       input)))

(def: #export (cos subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (f/* (math.cosh imaginary)
                (math.cos real))
     #imaginary (frac/negate (f/* (math.sinh imaginary)
                                  (math.sin real)))}))

(def: #export (cosh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (f/* (math.cos imaginary)
                (math.cosh real))
     #imaginary (f/* (math.sin imaginary)
                     (math.sinh real))}))

(def: #export (sin subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (f/* (math.cosh imaginary)
                (math.sin real))
     #imaginary (f/* (math.sinh imaginary)
                     (math.cos real))}))

(def: #export (sinh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (f/* (math.cos imaginary)
                (math.sinh real))
     #imaginary (f/* (math.sin imaginary)
                     (math.cosh real))}))

(def: #export (tan subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r2 (f/* 2.0 real)
        i2 (f/* 2.0 imaginary)
        d (f/+ (math.cos r2) (math.cosh i2))]
    {#real (f// d (math.sin r2))
     #imaginary (f// d (math.sinh i2))}))

(def: #export (tanh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r2 (f/* 2.0 real)
        i2 (f/* 2.0 imaginary)
        d (f/+ (math.cosh r2) (math.cos i2))]
    {#real (f// d (math.sinh r2))
     #imaginary (f// d (math.sin i2))}))

(def: #export (abs subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    (complex (if (f/< (frac/abs imaginary)
                      (frac/abs real))
               (if (f/= 0.0 imaginary)
                 (frac/abs real)
                 (let [q (f// imaginary real)]
                   (f/* (math.pow 0.5 (f/+ 1.0 (f/* q q)))
                        (frac/abs imaginary))))
               (if (f/= 0.0 real)
                 (frac/abs imaginary)
                 (let [q (f// real imaginary)]
                   (f/* (math.pow 0.5 (f/+ 1.0 (f/* q q)))
                        (frac/abs real))))
               ))))

(structure: #export _ (Number Complex)
  (def: + ..+)
  (def: - ..-)
  (def: * ..*)
  (def: / ../)
  (def: % ..%)
  (def: (negate x)
    (|> x
        (update@ #real frac/negate)
        (update@ #imaginary frac/negate)))
  (def: abs ..abs)
  (def: (signum x)
    (|> x
        (update@ #real frac/signum)
        (update@ #imaginary frac/signum))))

(def: #export (exp subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r-exp (math.exp real)]
    {#real (f/* r-exp (math.cos imaginary))
     #imaginary (f/* r-exp (math.sin imaginary))}))

(def: #export (log subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (|> subject ..abs (get@ #real) math.log)
     #imaginary (math.atan2 real imaginary)}))

(do-template [<name> <type> <op>]
  [(def: #export (<name> param input)
     (-> <type> Complex Complex)
     (|> input log (<op> param) exp))]

  [pow  Complex ..*]
  [pow' Frac    ..*']
  )

(def: (copy-sign sign magnitude)
  (-> Frac Frac Frac)
  (f/* (frac/signum sign) magnitude))

(def: #export (root2 (^@ input (^slots [#real #imaginary])))
  (-> Complex Complex)
  (let [t (|> input ..abs (get@ #real) (f/+ (frac/abs real)) (f// 2.0) (math.pow 0.5))]
    (if (f/>= 0.0 real)
      {#real t
       #imaginary (f// (f/* 2.0 t)
                       imaginary)}
      {#real (f// (f/* 2.0 t)
                  (frac/abs imaginary))
       #imaginary (f/* t (copy-sign imaginary 1.0))})))

(def: #export (root2-1z input)
  (-> Complex Complex)
  (|> (complex 1.0) (- (* input input)) root2))

(def: #export (reciprocal (^slots [#real #imaginary]))
  (-> Complex Complex)
  (if (f/< (frac/abs imaginary)
           (frac/abs real))
    (let [q (f// imaginary real)
          scale (f// (|> real (f/* q) (f/+ imaginary))
                     1.0)]
      {#real (f/* q scale)
       #imaginary (frac/negate scale)})
    (let [q (f// real imaginary)
          scale (f// (|> imaginary (f/* q) (f/+ real))
                     1.0)]
      {#real scale
       #imaginary (|> scale frac/negate (f/* q))})))

(def: #export (acos input)
  (-> Complex Complex)
  (|> input
      (+ (|> input root2-1z (* i)))
      log
      (* (negate i))))

(def: #export (asin input)
  (-> Complex Complex)
  (|> input
      root2-1z
      (+ (* i input))
      log
      (* (negate i))))

(def: #export (atan input)
  (-> Complex Complex)
  (|> input
      (+ i)
      (/ (- input i))
      log
      (* (/ (complex 2.0) i))))

(def: #export (argument (^slots [#real #imaginary]))
  (-> Complex Frac)
  (math.atan2 real imaginary))

(def: #export (roots nth input)
  (-> Nat Complex (List Complex))
  (if (n/= +0 nth)
    (list)
    (let [r-nth (|> nth .int int-to-frac)
          nth-root-of-abs (|> input abs (get@ #real) (math.pow (f// r-nth 1.0)))
          nth-phi (|> input argument (f// r-nth))
          slice (|> math.pi (f/* 2.0) (f// r-nth))]
      (|> (list.n/range +0 (dec nth))
          (list/map (function (_ nth')
                      (let [inner (|> nth' .int int-to-frac
                                      (f/* slice)
                                      (f/+ nth-phi))
                            real (f/* nth-root-of-abs
                                      (math.cos inner))
                            imaginary (f/* nth-root-of-abs
                                           (math.sin inner))]
                        {#real real
                         #imaginary imaginary})))))))