Durand-Kerner-Aberth method (well-known algorithm to compute all roots of a polynominal)

dka.ml

open Complex

(** [roots_initvals ?r [c(n); ...; c(2); c(1); c(0)]] computes
 initial values for [roots] by using Aberth's method.
 *)
let roots_initvals ?(r=1000.0) coeff_list =
  match coeff_list with
  | a::b::_ ->
     let n = List.length coeff_list - 1 in
     let s = 3.14159265358979 /. (float n) in
     let t = div b (mul a {re = ~-.(float n); im = 0.}) in
     let rec loop acc i =
       let angle = s *. (2. *. (float i) +. 0.5) in
       let zi = add t (mul {re=r; im=0.} (exp {re=0.; im=angle})) in
       let acc' = zi :: acc in
       if i = 0 then acc' else loop acc' (i - 1)
     in
     loop [] (n - 1)
  | _ -> invalid_arg "roots_aberth_initvals: # of coefficients < 2"

(** [roots ?epsilon ?init [c(n); ...; c(2); c(1); c(0)]] computes roots of a
 polynominal [c(n)*x**n + ... + c(2)*x**2 + c(1)*x + c(0)] by using the
 Durand-Kerner method.
 *)
let roots ?(epsilon=1e-6) ?init coeff_list =
  let z_list = match init with (* initial values of roots *)
    | None -> roots_initvals coeff_list
    | Some zz -> zz in
  let calc_poly x = (* calculate a polynomial *)
    List.fold_left (fun acc ci -> add (mul acc x) ci) zero coeff_list in
  let cn = List.hd coeff_list in (* c(n) *)
  let rec update_z z_list = (* update z(0), ..., z(n-1) until they converge *)
    let update_zi z_list i zi = (* update z(i) *)
      let f (j, acc) zj = (j + 1, if i = j then acc else mul acc (sub zi zj)) in
      let (_, deno) = List.fold_left f (0, cn) z_list in
      sub zi (div (calc_poly zi) deno) (* new z(i) *)
    in
    let z_list' = List.mapi (update_zi z_list) z_list in (*new z(0),...,z(n-1)*)
    if List.for_all2 (fun zi zi' -> norm2 (sub zi zi') < epsilon) z_list z_list'
    then z_list' (* converged! *)
    else update_z z_list'
  in
  update_z z_list

(* -300 + 320 x - 59 x^2 - 26 x^3 + 5 x^4 + 2 x^5 = 0 *)
let c = [{re=2.;im=0.}; {re=5.;im=0.}; {re=(-26.);im=0.};
         {re=(-59.);im=0.}; {re=(320.);im=0.}; {re=(-300.);im=0.}];;
let z = roots c;;