[racket] ATTN: Neil T. Racket/Math digamma (psi) function implementation

From: Ray Racine (ray.racine at gmail.com)
Date: Mon May 6 15:25:50 EDT 2013

Neil,

Here in is an implementation of the digamma function.  Initially I tried
Haskell's, then a Java impl from Mallet  (LDA library) but they didn't
align (BTW I think the Haskell impl is flat wrong for digamma (2)) and
sought alternate sources.

I finally found
http://code.google.com/p/pmtk3/source/browse/trunk/util/digamma.m?r=227&spec=svn227
from pmtk2 a probabilistic modeling toolkit for Matlab/Octave, which is MIT
LICENSED so I think we are good there.

My impl is based (well ok copied) from that impl.

Octave also appears to have a Gnu Science Library FFI which has an impl for
psi (digamma) as well which I've used as a sanity check for positive x
values.

 apt-get install octave
 apt-get install octave-gls
 apt-get install octave-gsl
octave:6> format long
octave:7> psi (10)
ans =  2.25175258906672

> (digamma 10.0)
2.251752589068216

Since it works for x = 10 to a few decimal places, I'm sure it works
everywhere ... perfectly.

So far the Racket impl appears to be aligned to Octave+GSL, but I have no
intuition as what is acceptable epsilon variance in floating point vs what
is true error in implementation.  The impl handles z < 0, yet Octave+GSL
blows up.

I added addition poly coefficients up to 12 but the impl currently only
uses 7.

impl follows:

#lang typed/racket/base

#|
 Reference:

    J Bernardo,
    Psi ( Digamma ) Function,
    Algorithm AS 103,
    Applied Statistics,
    Volume 25, Number 3, pages 315-317, 1976.

 From http://www.psc.edu/~burkardt/src/dirichlet/dirichlet.f
|#

(provide:
 [digamma (Float -> Float)])

(require
 (only-in racket/flonum
      fltan
      fllog)
 (only-in racket/math
      pi))

(define: PI : Float pi)
(define: NEG-PI : Float (- PI))
(define: PI-OVER-TWO : Float (/ PI 2.0))
(define: PI-SQUARED-OVER-SIX : Float (/ (* PI PI) 6))
(define: EULER-MASCHERONI : Float -0.5772156649015328606065121)

(define: DIGAMMA-1 : Float EULER-MASCHERONI)     ;; digamma(1)
(define: DIGAMMA-2 : Float PI-SQUARED-OVER-SIX)

(define: DIGAMMA-COEF-3 : Float (/ 1.0 12.0))
(define: DIGAMMA-COEF-4 : Float (/ 1.0 120.0))
(define: DIGAMMA-COEF-5 : Float (/ 1.0 252.0))
(define: DIGAMMA-COEF-6 : Float (/ 1.0 240.0))
(define: DIGAMMA-COEF-7 : Float (/ 1.0 132.0))
;;(define: DIGAMMA-COEF-8 : Float (/ 691.0 32760.0))
;;(define: DIGAMMA-COEF-9 : Float (/ 1.0 12.0))
;;(define: DIGAMMA-COEF-10 : Float (/ 3617.0 8160.0))
;;(define: DIGAMMA-COEF-11 : Float (/ 43867.0 14364.0))
;;(define: DIGAMMA-COEF-12 : Float (/ 174611.0 6600.0))

(define: DIGAMMA-LARGE : Float 9.5)
(define: DIGAMMA-SMALL : Float 0.000001)

;;anxn + … + a2x2 + a1x + a0 = a0 + x(a1 + x(a2 + x(… an) …) )
(: digamma-poly (Float -> Float))
(define (digamma-poly z)
  (* z (- DIGAMMA-COEF-3
      (* z (- DIGAMMA-COEF-4
          (* z (- DIGAMMA-COEF-5
              (* z (+ DIGAMMA-COEF-6
                  (* z (+ DIGAMMA-COEF-7)))))))))))

;; Assumes z < 0
;;  Use the reflection formula (Jeffrey 11.1.6):
;;  digamma(-x) = digamma(x+1) + pi*cot(pi*x)
;;   y = digamma(-x+1) + pi*cot(-pi*x);
;;  This is related to the identity
;;  digamma(-x) = digamma(x+1) - digamma(z) + digamma(1-z)
;;  where z is the fractional part of x
;;  For example:
;;  digamma(-3.1) = 1/3.1 + 1/2.1 + 1/1.1 + 1/0.1 + digamma(1-0.1)
;;                = digamma(4.1) - digamma(0.1) + digamma(1-0.1)
;;  Then we use
;;  digamma(1-z) - digamma(z) = pi*cot(pi*z)
;;  cot(x) = tan (pi/2 - x)

(: cotan (Float -> Float))
(define (cotan x)
  (fltan (- PI-OVER-TWO x)))

(: digamma-negative (Float -> Float))
(define (digamma-negative z)
  ;; (assert (< z 0.0))
  (+ (digamma (+ (- z) 1.0))
     (* PI (cotan (* z NEG-PI)))))

(: digamma (Float -> Float))
(define (digamma z)
  (cond
   ((eq? z 0.0)
    -inf.0)
   ((< z 0.0)
    (digamma-negative z))
   ((eq? z 1.0)
    DIGAMMA-1)
   ((<= z DIGAMMA-SMALL)
    (/ (- DIGAMMA-1 1.0)
       (+ z (* DIGAMMA-2 z))))
   ((< z DIGAMMA-LARGE)             ;; Reduce to digamma(x+n) where (x+n)
>= large
    (- (digamma (+ z 1.0)) (/ 1 z)))
   (else                            ;; Use de Moivre's expansion if
argument >= large.
    (let ((r (/ 1.0 z)))
      (let ((A (- (fllog z) (* 0.5 r))))
    (- A (digamma-poly (* r r))))))))
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.racket-lang.org/users/archive/attachments/20130506/be332bed/attachment-0001.html>

Posted on the users mailing list.