Adding the logarithmic derivative of the gamma function (digamma) to stdlib_specialfunctions_gamma about stdlib HOT 1 OPEN

program polygamy

  use iso_fortran_env, only: sp => real32, dp => real64, qp => real128, int64

  implicit none

  integer :: n
  complex(dp) :: p
  complex(dp) :: z
  complex(dp), parameter :: ci = (0.0_dp, 1.0_dp)

  n = 1
  z = 0.5_dp
  p = polygamma_cdp(n, z)
  print*, p
  p = digamma_cdp(z)
  print*, p


  contains

    recursive impure elemental function polygamma_cdp(n, z) result (res)
      !! Returns the polygamma function of order n evaluated at z. See the
      !! Digitial Library of Mathematical Functions (DLMF), section 5.15
      !! for reference.
      !!
      !! References for the nth derivative of the cotangent function, used in the reflection formula of the 
      !! polygamma function :
      !!
      !! (1) The Polygamma Function and the Derivatives of the Cotangent Function for Rational Arguments, K. S. Kölbig
      !!       CERN, https://cds.cern.ch/record/298844
      !!
      !! (2) Derivative Plynomials for Tangent and Secant, Michael E. Hoffman
      !!       The American Mathematical Monthly, Vol. 102, No. 1, pp. 23–30, https://doi.org/10.2307/2974853

      integer, intent(in) :: n
      complex(dp), intent(in) :: z
      complex(dp) :: res

      integer :: j, k, m
      integer, parameter :: nbernoulli = 8
      integer(int64) :: factorial
      integer(int64) :: binom
      real(dp), parameter :: zero_k1 = 0.0_dp
      real(dp), parameter :: z_limit = 10.0_dp
      real(qp), parameter :: one = 1.0_qp
      real(qp), parameter :: two = 2.0_qp
      real(qp), parameter :: pi = acos(-one)
      real(qp), parameter :: B(nbernoulli) = [        &
                                one / 6.0_qp,         &
                              - one / 30.0_qp,        &
                                one / 42.0_qp,        &
                              - one / 30.0_qp,        &
                                5.0_qp / 66.0_qp,     &
                              - 691.0_qp / 2730.0_qp, &
                                7.0_qp / 6.0_qp,      &
                              - 3617.0_qp / 510.0_qp  &
                             ]
        !! The Bernoulli numbers B(1) = B2, B(2) = B4, B(3) = B6, etc.

      complex(qp) :: P(0:n)
      complex(qp) :: dcotan
      complex(qp) :: series
      complex(qp) :: factor
      complex(qp) :: z2, zr, zr2

      factorial(m)   = gamma(real(m + 1, kind = qp))
      binom(k, j) = factorial(k) / (factorial(j) * factorial(k - j))

      if(n .eq. 0) then
        res = digamma_cdp(z)
        return
      endif

      z2 = z * one
      zr = one / z2
      zr2 = zr * zr

      if( z % re <= zero_k1 ) then
        ! -- reflection (-1)^n ψ^(n)(1 - z) - ψ^(n)(z) = π (d/dx)^n cot(πz), including the imaginary axis

        P(0) = cotan(pi * z2)
        P(1) = cotan(pi * z2) ** 2 + 1
        do m = 1, n - 1
          P(m + 1) = sum( [( binom(m, j) * P(j) * (P(m - j)), j = 0, m )] )
        enddo

        ! res = pi * (d/dz)^n cot(pi z) = pi^{n+1} (-1)^n P(cot(piz))
        ! res = pi * (-pi) ** n * P(n) / sin(pi * z) ** (n + 1)
        dcotan =  pi ** (n + 1)  * (-1) ** n * P(n)
        res = ( (-1) ** n * polygamma_cdp(n, 1 - z) ) - cmplx(dcotan % re, dcotan % im, kind = dp)

        return

      elseif( z % re > z_limit ) then
        ! -- Poincaré asymptotic expansion

        series = factorial(n - 1) * zr ** n + factorial(n) * zr ** (n + 1) / two
        series = series + sum( [( factorial(2 * k + n - 1) / factorial(2 * k) * B(k) * zr ** (2 * k + n), k = 1, nbernoulli )] )

        res = (-1) ** (n - 1) * cmplx(series % re, series % im, kind = dp)

        return

      endif

      !-- recurrence relation ψ^n(z + 1) = ψ^n(z) + (-1)^n n! / z^{n + 1}
      factor = (-1) ** n * factorial(n) * zr ** (n + 1)
      res = polygamma_cdp(n, z + 1) - cmplx(factor % re, factor % im, kind = dp)

    end function polygamma_cdp

    recursive impure elemental function digamma_cdp(z) result (res)
      !!
      !! Returns the digamma function for any complex number, excluding negative 
      !! whole numbers, by reflection (z < 0), upward recurence (z % re < z_limit), or a 
      !! truncated Stirling / de Moivre series
      !!
        complex(dp), intent(in) :: z
        complex(dp) :: res
        integer, parameter :: n = 7
        integer :: k
        complex(qp) :: z2, zr, zr2, series
        complex(qp) :: res2
        real(dp), parameter :: z_limit = 10.0_dp
        real(dp), parameter :: zero_k1 = 0.0_dp
        real(qp), parameter :: one = 1.0_qp, two = 2.0_qp, pi = acos(-one)
        real(qp), parameter :: a(n) = [                     &
                          -one / 12.0_qp,             &
                           one / 120.0_qp,            &
                          -one / 252.0_qp,            &
                           one / 240.0_qp,            &
                          -one / 132.0_qp,            &
                           691.0_qp / 32760.0_qp, &
                          -one / 12.0_qp]

        z2 = z * one

        if( z % re <= zero_k1 ) then

          ! -- reflection ψ(1 - x) - ψ(x) = π cot(πx), including the imaginary axis
          res = digamma_cdp(1.0_dp - z)
          res2 = - pi * cotan(pi * z2)
          res = res + cmplx(res2 % re, res2 % im, kind = dp)

          return

        elseif( z % re > z_limit ) then

          zr  = one / z2
          zr2 = zr * zr
          series = log(z2) - zr / two  + sum( [(  a(k) * (zr2 ** k), k = 1, n )] )
          res = cmplx(series % re, series % im, kind = dp)

          return

        endif

        !-- recurrence relation ψ(z + 1) = ψ(z) + 1 / z
        res = digamma_cdp(z + 1) - 1 / z


    end function digamma_cdp


end program polygamy