# HG changeset patch
# User Carnë Draug <carandraug@octave.org>
# Date 1426401009 0
# Node ID 1fae49e34a1af4507a545ad7604562bab6c2a42f
# Parent  45565ecec01912b5e02c17e80db512d53cba05c0
psi: add support for complex numbers.

* libinterp/corefcn/psi.cc: add logic and input check to support complex
numbers (implementation is in lo-specfun.cc).  Add tests.
* liboctave/numeric/lo-specfun.cc, liboctave/numeric/lo-specfun.h: add
template specialization to psi() for std::complex. It is mostly taken from
the implementation in XLiFE++ (also under GPLv3+, see b03c7cccadc2 commit
message for more details).

diff -r 45565ecec019 -r 1fae49e34a1a libinterp/corefcn/psi.cc
--- a/libinterp/corefcn/psi.cc	Sun Mar 15 03:31:16 2015 +0000
+++ b/libinterp/corefcn/psi.cc	Sun Mar 15 06:30:09 2015 +0000
@@ -73,16 +73,28 @@
       E* psi_zv = psi_z.fortran_vec (); \
       const octave_idx_type n = z.numel (); \
       for (octave_idx_type i = 0; i < n; i++) \
-        psi_zv[i] = psi<E> (zv[i]); \
+        psi_zv[i] = psi (zv[i]); \
 \
       retval = psi_z; \
     } 
 
-  FLOAT_BRANCH(double, , , double)
-  else FLOAT_BRANCH(single, Float, float_, float)
+  if (args(0).is_complex_type ())
+    {
+      FLOAT_BRANCH(double, Complex, complex_, Complex)
+      else FLOAT_BRANCH(single, FloatComplex, float_complex_, FloatComplex)
+      else
+        {
+          error ("psi: Z must be a floating point");
+        }
+    }
   else
     {
-      error ("psi: Z must be a floating point");
+      FLOAT_BRANCH(double, , , double)
+      else FLOAT_BRANCH(single, Float, float_, float)
+      else
+        {
+          error ("psi: Z must be a floating point");
+        }
     }
 
 #undef FLOAT_BRANCH
@@ -126,5 +138,22 @@
 %!assert (psi (-2.610720868444144650001537715718724207951074010873480), 0, eps*10)
 %!assert (psi (-3.635293366436901097839181566946017713948423861193530), 0, eps)
 %!assert (psi (-4.653237761743142441714598151148207363719069416133868), 0, eps*100)
+
+## Tests for complex values
+%!shared z
+%! z = [-10:.1:-.1 .1:.1:20]; # drop the 0
+
+## Abramowitz and Stegun, page 259 eq 6.3.10
+%!assert (real (psi (i*z)), real (psi (1 - i*z)))
+
+## Abramowitz and Stegun, page 259 eq 6.3.11
+%!assert (imag (psi (i*z)), 1/2 .* 1./z + 1/2 * pi * coth (pi * z), eps *10)
+
+## Abramowitz and Stegun, page 259 eq 6.3.12
+%!assert (imag (psi (1/2 + i*z)), 1/2 * pi * tanh (pi * z), eps*10)
+
+## Abramowitz and Stegun, page 259 eq 6.3.13
+%!assert (imag (psi (1 + i*z)), - 1./(2*z) + 1/2 * pi * coth (pi * z), eps*10)
+
 */
 
diff -r 45565ecec019 -r 1fae49e34a1a liboctave/numeric/lo-specfun.cc
--- a/liboctave/numeric/lo-specfun.cc	Sun Mar 15 03:31:16 2015 +0000
+++ b/liboctave/numeric/lo-specfun.cc	Sun Mar 15 06:30:09 2015 +0000
@@ -3728,13 +3728,23 @@
     }
 }
 
+static const double euler_mascheroni = 0.577215664901532860606512090082402431042;
+static const double pi = 3.14159265358979323846;
+// Coefficients for C.Lanczos expansion of psi function from XLiFE++ gammaFunctions
+// psi_coef[k] = - (2k+1) * lg_coef[k] (see melina++ gamma functions)
+// -1/12, 3/360,-5/1260, 7/1680,-9/1188, 11*691/360360,-13/156, 15*3617/122400, ? , ?
+static const double psi_coeff[10] = {
+  -0.83333333333333333e-1, 0.83333333333333333e-2,
+  -0.39682539682539683e-2, 0.41666666666666667e-2,
+  -0.75757575757575758e-2, 0.21092796092796093e-1,
+  -0.83333333333333333e-1, 0.4432598039215686,
+  -0.3053954330270122e+1,  0.125318899521531e+2
+};
+
 template<class T>
 T
 psi (const T& z)
 {
-  const T euler_mascheroni = 0.577215664901532860606512090082402431042;
-  const T pi = 3.14159265358979323846;
-
   const bool is_int = (xfloor (z) == z);
 
   T p = 0;
@@ -3768,15 +3778,6 @@
     {
       // adapted from XLiFE++ gammaFunctions
 
-      // Coefficients for C.Lanczos expansion of DiGamma function
-      // dg_coef[k] = - (2k+1) * lg_coef[k] (see melina++ gamma functions)
-      // -1/12, 3/360,-5/1260, 7/1680,-9/1188, 11*691/360360,-13/156, 15*3617/122400, ? , ?
-      const T dg_coeff[10] ={-0.83333333333333333e-1, 0.83333333333333333e-2,
-                             -0.39682539682539683e-2, 0.41666666666666667e-2,
-                             -0.75757575757575758e-2, 0.21092796092796093e-1,
-                             -0.83333333333333333e-1, 0.4432598039215686,
-                             -0.3053954330270122e+1,  0.125318899521531e+2};
-
       T zc = z;
       // Use formula for derivative of LogGamma(z)
       if (z < 10)
@@ -3791,7 +3792,7 @@
 
       p += log (zc) - 0.5 / zc;
       for (octave_idx_type k = 0; k < 10; k++, overz2k *= overz2)
-        p += dg_coeff[k] * overz2k;
+        p += psi_coeff[k] * overz2k;
     }
 
   return p;
@@ -3801,3 +3802,57 @@
 template double psi<double> (const double& z);
 template float  psi<float> (const float& z);
 
+template<class T>
+std::complex<T>
+psi (const std::complex<T>& z)
+{
+  // adapted from XLiFE++ gammaFunctions
+
+  typedef typename std::complex<T>::value_type P;
+
+  P z_r  = z.real ();
+  P z_ra = z_r;
+
+  if (z.imag () == 0)
+    return std::complex<T> (psi (z_r), 0.0);
+  else if (z_r < 0)
+    return psi (P (1.0) - z)- (P (pi) / tan (P (pi) * z));
+  else
+    {
+      // Use formula for derivative of LogGamma(z)
+
+      std::complex<T> dgam = 0.0;
+      std::complex<T> z_p  = z;
+
+      octave_idx_type n = 0;
+      std::complex<T> z_m = z_p;
+      if (z_ra < 8)
+        {
+          n = 8 - octave_idx_type (z_ra);
+          z_m = z_p + std::complex<T> (n, 0.0);
+        }
+
+      // for | Re(z) | > 8, use derivative of C.Lanczos expansion for LogGamma
+      // psi(z) = log(z) - 1/(2z) - 1/12z^2 + 3/360z^4 - 5/1260z^6 + 7/1680z^8 - 9/1188z^10 + ...
+      // (Abramowitz&Stegun, page 259, formula 6.3.18
+      std::complex<T> overz   = P (1.0) / z_m;
+      std::complex<T> overz2  = overz * overz;
+      std::complex<T> overz2k = overz2;
+
+      dgam += log (z_m) - P (0.5) * overz;
+      for (octave_idx_type k = 0; k < 10; k++, overz2k *= overz2)
+        dgam += P (psi_coeff[k]) * overz2k;
+
+      // Recurrence formula
+      // for | Re(z) | < 8 , use recursively DiGamma(z) = DiGamma(z+1) - 1/z
+      for (octave_idx_type k = 0; k < n; k++, z_p += 1.0)
+        dgam -= P (1.0) / z_p;
+
+      return dgam;
+    }
+}
+
+// explicit instantiations
+template Complex psi<double> (const Complex& z);
+template FloatComplex psi<float> (const FloatComplex& z);
+
diff -r 45565ecec019 -r 1fae49e34a1a liboctave/numeric/lo-specfun.h
--- a/liboctave/numeric/lo-specfun.h	Sun Mar 15 03:31:16 2015 +0000
+++ b/liboctave/numeric/lo-specfun.h	Sun Mar 15 06:30:09 2015 +0000
@@ -665,5 +665,7 @@
 
 template<class T>
 extern OCTAVE_API T psi (const T& z);
+template<class T>
+extern OCTAVE_API std::complex<T> psi (const std::complex<T>& z);
 
 #endif