PRD algorithm constant calculator (Haskell)

prd_constant_calculator.hs

--
-- Pseudo-random distribution constant calculator.
--
-- The constants are calculated by numerically approximating the roots of
-- the function f(p) = p - p(c) via Newton-Raphson / bisection.
-- Calculating p(c) is easy: evaluate a polynomial and invert it.
--
-- For an explanation and visualization of the PRD algorithm, please see
-- https://observablehq.com/@manuelblanc/pseudo-random-distribution
--
-- This is free and unencumbered software released into the public domain.
-- For more information, please refer to <http://unlicense.org/>

import Control.Monad (forM_, liftM2)
import Text.Printf (printf)

config_TOLERANCE = 1e-15 :: Double
config_STEPS = 100 :: Int


-- Main entrypoint. Outputs a lookup formatted as a Lua table.
main :: IO ()
main = do
  putStrLn "local prd_constant_lookup = {"
  forM_ [1 .. config_STEPS-1] $ liftM2 (printf "\t[%d] = %.16f,\n") id calc
  putStrLn "}"
  where
    step_size = recip $ fromIntegral config_STEPS
    calc i = calculate_c config_TOLERANCE (step_size * fromIntegral i)


-- Given a PRD constant c, calculates p(c) and p'(c).
nominal_p_and_p' :: RealFrac a => a -> (a, a)
nominal_p_and_p' c = (p, p')
  where
    p = recip q
    p' = negate q' * p*p
    n = subtract 1 $ ceiling $ recip c
    (q, q') = foldl eval_polynomial (1, 0) $ map fromInteger [n, n-1 .. 1]
    eval_polynomial (a, a') k = (1 + (1 - k*c)*a, (1 - k*c)*a' - k*a)


-- Given a nominal probability and an initial guess `c0`, calculates a better
-- approximation of the real value `c = c(p)` using Newton-Raphson.
approximate_c :: RealFrac a => a -> a -> a
approximate_c p c0
  | c1 < 0    = c0 / 2 -- Degrade to bisection if underflowing.
  | otherwise = c1
  where
    c1 = c0 + step
    step = (p - p_c) / p_c'
    (p_c, p_c') = nominal_p_and_p' c0


-- Calculates the PRD constant for a given p, up to the given tolerance.
calculate_c tolerance p
  | p >= 1     = 1
  | p <= 0     = 0
  | otherwise  = c
  where
    c = converge initial $ drop 1 $ iterate (approximate_c p) initial
    initial = p*p
    converge y (x:xs)
      | (y - x) < tolerance = x
      | otherwise           = converge x xs