Haskell implementation of Alberth method for finding all roots of a polynomial

alberth_method.hs

-- Alberth method for finding all roots of a (complex) polynomial
--   https://en.wikipedia.org/wiki/Aberth_method
-- The calculations have to use complex numbers.
-- Some optimizations are possible if the coefficients are real,
-- This is not implemented here.
-- Implementation was prompted by
--   https://byorgey.wordpress.com/2019/02/13/finding-roots-of-polynomials-in-haskell/
-- This is not meant to be robust implementation, just a proof of concept.
-- For analysis and implementation of this specific method,
-- see the work by D Bini and others.
--
-- Jan Van lent
-- Started: 2019-02-20 Wed
-- Last updated: 2019-02-21 Thu

import Data.Complex

-- |Given [p0, p1, ..., pn] x, evaluate polynomial p0 + p1*x + ... + pn*x^n
polyval :: Num a => [a] -> a -> a
polyval p x = rec p 1 0
  where rec [] z y = y
        rec (p0:p) z y = rec p (x*z) (y + p0*z)
-- alternative: Horner scheme, sparse polynomials, FFT, etc.

diff :: Num a => [a] -> [a]
diff p = zipWith (*) (map fromIntegral [1..]) (tail p)

scales :: Fractional a => [a] -> [a]
scales zs = [ sum [ 1/(zk-zj) | (zj, j) <- zip zs [1..], j /= k ] |
              (zk, k) <- zip zs [1..] ]

offset :: Fractional a => a -> a -> a
offset q s = -q/(1-q*s)

offsets :: Fractional a => [a] -> [a] -> [a] -> [a]
offsets fs gs ss = ws
  where
      qs = zipWith (/) fs gs
      ws = zipWith offset qs ss

solve' :: RealFloat a => [Complex a] -> [Complex a] -> [Complex a] -> Int -> a -> [Complex a]
solve' p p' zs maxit tol
  | maxit <= 0 = zs
  -- | maxf <= tol = zs
  | maxw <= tol = newzs
  | otherwise = solve' p p' newzs (maxit-1) tol
    where
      fs = map (polyval p) zs
      maxf = maximum (map magnitude fs)
      gs = map (polyval p') zs
      ss = scales zs
      ws = offsets fs gs ss
      maxw = maximum (map magnitude ws)
      newzs = zipWith (+) zs ws

solve :: RealFloat a => [Complex a] -> [Complex a] -> Int -> a -> [Complex a]
solve p zs maxit tol = solve' p (diff p) zs maxit tol

rootBound :: RealFloat a => [Complex a] -> a
rootBound p = r
  where
    revp = reverse p
    num = maximum (map magnitude (tail revp))
    den = magnitude (head revp)
    r = 1 + num / den

-- choose initial points equally spaced on a circle with radius based on bounds
-- for the absolute values of the roots
-- for a polynomial with real coefficients it would make sense to start
-- with conjugate pairs, since this is preserved by the iteration
-- I suspect the literature on this method will have much more to say
-- about how to choose good initial values
circleInit :: RealFloat a => [Complex a] -> [Complex a]
circleInit p = [ mkPolar (r/2) (((fromIntegral i)+0.5)*step) | i <- [0..n-1] ]
  where
    step = 2*pi/(fromIntegral n)
    r = rootBound p
    n = (length p) - 1

phi = ((sqrt 5) + 1)/2
goldenAngle = 2 * pi / phi**2

-- similar to circleInit
-- the golden angle is used to generate a quasi-uniform distribution
-- this is still deterministic, but a bit more like random
goldenCircleInit :: [Complex Double] -> [Complex Double]
goldenCircleInit p = [ mkPolar (r/2) ((fromIntegral i)*goldenAngle) | i <- [0..n-1] ]
  where
    r = rootBound p
    n = (length p) - 1

roots :: [Complex Double] -> Int -> Double -> [Complex Double]
--roots p tol = solve p (circleInit p) tol
roots p maxit tol = solve p (goldenCircleInit p) maxit tol

rootpoly' :: Num a => [a] -> [a] -> [a]
rootpoly' [] p = p
rootpoly' (z:zs) p = rootpoly' zs (zipWith (-) (0:p) ((map (z*) p)++[0]))

-- |find the coefficients of a polynomial with given roots
rootpoly :: Num a => [a] -> [a]
rootpoly zs = rootpoly' zs [1]



test1 = (z0s, ss, ws)
  where
    --p = [3:+0, 2:+0, 1:+0]
    p = [3, 2, 1]
    z0s = goldenCircleInit p
    p' = diff p
    ss = scales z0s
    fs = map (polyval p) z0s
    gs = map (polyval p') z0s
    ws = offsets fs gs ss

test2 = zs
  where
    p = rootpoly [1:+0, 2:+0, 3:+0]
    zs = roots p 30 1e-3

test3 n = zs
  where
    p = rootpoly [ i:+i^2 | i <- [1..n] ]
    zs = roots p 30 1e-3