Skip to content

tetragamma function, implemented using Horner approach in Julia's trigamma - #541

Open
essandess wants to merge 1 commit into
JuliaMath:masterfrom
essandess:tetragamma
Open

tetragamma function, implemented using Horner approach in Julia's trigamma#541
essandess wants to merge 1 commit into
JuliaMath:masterfrom
essandess:tetragamma

Conversation

@essandess

Copy link
Copy Markdown

tetragamma function, implemented using the Horner approach in Julia's trigamma function.

Polynomial order to achieve complex128 accuracy determined using Lentz's method near $|z| = 10$.

Asymptotic coefficients computed using $B_{2k}(2k+1)$, $k=1…10$ (A&S 6.4.11, $m=2$). Closed-form tests use Apéry's constant $\zeta(3)$.

Reference issue

@stevengj

stevengj commented Aug 11, 2026

Copy link
Copy Markdown
Member

How do the performance and accuracy compare to the existing polygamma?

(What kind of applications need higher performance for tetragamma but not even higher order? That is, at what point do we stop adding specialized versions of polygamma for particular orders?)

@essandess

essandess commented Aug 11, 2026

Copy link
Copy Markdown
Author

How do the performance and accuracy compare to the existing polygamma?

(What kind of applications need higher performance for tetragamma but not even higher order? That is, at what point do we stop adding specialized versions of polygamma for particular orders?)

trigamma and tetragamma are famously the last two polygamma functions that have simple rational continued fraction/Bernoulli × rational asymptotic coefficients. See:

  1. Cuyt et al., Handbook of Continued Fractions for Special Functions, Springer, 2008.
  2. Catherine M. Bonan-Hamada and William B. Jones, Stieltjes continued fractions for polygamma functions; speed of convergence, J. Computational and Appl. Math. 179(1–2) : 47–55 (2005), https://www.sciencedirect.com/science/article/pii/S037704270400442X

There are no nice, simple, and fast expansions for the polygamma after the tetragamma, and we need to rely on more complicated methods to compute Hurwitz zeta ζ(s, z), as polygamma necessarily does in in this package. You can see in the code that it'll be a lot more efficient to use the Bernoulli expansion for ψ₂(z) than to compute ζ(2, z).

There are plenty of mathematical and other applications for ψ₂(z), given its relationship to ζ(2, z), Apéry's constant ζ(3), and as a derivative of the gamma function. But aside from applications it's a good to have an implementation of the very nice computational properties that ψ₂(z) possesses.

@stevengj

stevengj commented Aug 11, 2026

Copy link
Copy Markdown
Member

trigamma and tetragamma are famously the last two polygamma functions that have simple rational continued fraction/Bernoulli × rational asymptotic coefficients. See:

The second paper you cite says that polygamma has a continued fraction for all nonnegative integer orders $k$:

image

Indeed, IIRC, the polygamma implementation in Julia is mostly based on simply taking $k$ derivatives of the digamma implementation's formulas (the tricky part being adjusting the cutoffs and some of the corner cases).

@essandess

essandess commented Aug 11, 2026

Copy link
Copy Markdown
Author

Yes, Stieltges continued fractions exist for polygammas of all orders. But only cases m = 0, 1, and 2 (digamma, trigamma, and tetragamma) have simple expressions where it’s worth the trouble to derive and code specific fast algorithms. digamma and tetragamma already have specific implementations—why not just delete these and use polygamma for everything?

That would be a bad idea for many reasons:

  1. The specialized digamma/trigamma/tetragamma functions all perform about 10–20 floating-point operations total per call (shifts, a fixed-degree Horner evaluation, a reciprocal), and are fixed algorithms with pre-derived parameters that will run fast in compiled JIT functions. The _zeta function has several order-specific operations that are unknown at compile or JIT time, especially the order-dependent @pg_horner macro, which takes about double the number of operations of the specific cases to handle the general case. This is still small but this generality adds a constant overhead that one expects makes the general case 2+× slower than the specific ones.
  2. The specialized digamma/trigamma are used a lot in statistical applications especially because of LDA, and if the tetragamma isn’t used a lot, it’s still used. In LDA, the distribution is the gamma, its log-derivative is the digamma, and its Hessian is the trigamma. Trigamma is the Fisher information / score-variance term for the entire exponential family that includes the Gamma, Beta, and Dirichlet distributions. It's the single hottest special function inside variational-EM inference for LDA, where the variational Dirichlet parameter updates require Newton-Raphson steps built on digamma and trigamma. You want fast digamma and trigamma functions because calling a general polygamma many times in an LDA application would slow down the algorithm a lot.
  3. The tetragamma arises as the next derivative when computing statistical curvature using e.g. Amari’s α-connection where the Riemannian curvature in LDA is given by a trigamma, and its Christoffel symbols are given by the tetragamma. The tetragamma also arises in methods that use the derivative of Fisher information, like Halley's method for accelerated ML. You want a fast specialized tetragamma for all these applications, not a slower general polygamma method.

We want a specialized tetragamma function for the same reasons we want specialized digamma and trigamma ones: these appear all the time in theory and applications. If the pentagamma were ever relevant, yes, we could use Stieltjes existence of CF to derive its specific algorithm, but I am unaware of any demand for this. But the tetragamma does appear, and is worth implementing as a specific function. Furthermore, we cannot predict how this function will appear in other applications, but if it’s useful for these applications, it’s likely to be useful for unknown others too.

@essandess

essandess commented Aug 11, 2026

Copy link
Copy Markdown
Author

Here’s a simple benchmark that illustrates the algorithm analysis between a specialized tetragamma and the general polygamma that uses zeta. As expected, tetragamma is more than twice as fast as a polygamma and a lot more memory efficient because Julia’s JIT is given a set of fixed operations and doesn’t have to garbage collect after general macro calls.

Using Distributions
Using BenchmarkTools

# edit out case: s == 3 && return tetragamma(z)
include("Downloads/gamma.jl")

args = rand(Exponential(), 1_000_000) .* exp.(im*2*pi*rand(Uniform(), 1_000_000));
function tgloop(args)
    for x in args
        tetragamma(x)
    end
end
function ztloop(args)
    for x in args
        zeta(3, x) 
    end
end

@btime tgloop($args)
  40.288 ms (0 allocations: 0 bytes)

@btime ztloop($args)
  90.557 ms ([1000000](tel:1000000) allocations: 30.52 MiB)

@essandess

Copy link
Copy Markdown
Author

trigamma and tetragamma are famously the last two polygamma functions that have simple rational continued fraction/Bernoulli × rational asymptotic coefficients. See:

The second paper you cite says that polygamma has a continued fraction for all nonnegative integer orders k :

For completeness, the relevant part of this paper is here, where it is stated that that the trigamma and tetragamma are “the only two polygamma functions for which there are known closed expressions for the coefficients $a_m^{(k)}$”:
IMG_1330

@stevengj

stevengj commented Aug 11, 2026

Copy link
Copy Markdown
Member

I don't see how the comment you cited, which is about the continued-fraction coefficients, is relevant here? We currently don't use the continued-fraction expansion, we use the (Stirling) asymptotic series, with the shift+reflection formulas to push the argument into the regime where the asymptotic series is accurate. This is the Kölbig (1972) algorithm (and derivatives thereof).

It looks like your tetragamma implementation also uses the asymptotic series, and indeed a comment in your code indicates that it is the 2nd derivative of the Kölbig digamma method.

We could therefore do the same thing for any desired order, and hence the question is where to stop. digamma is used so frequently that it is worth having a specialized implementation. trigamma is used less frequently but is present in many special functions libraries. I haven't seen any widely used library that has a specialized tetragamma implementation, am I missing one?

Definitely a specialized implementation for particular orders can be faster than the general polygamma code, but it trades off code, testing, and maintenance costs. (In your benchmark, as you point out, the performance benefit is only a factor of ≈ 2. Also something is weird in your benchmark because ztloop is not allocation-free … if that is a problem in the library, it should be fixed, and then the performance advantage will be even less.)

@stevengj

stevengj commented Aug 11, 2026

Copy link
Copy Markdown
Member

I can't reproduce the allocations in your benchmark (using SpecialFunctions v2.7.2):

julia> @btime ztloop($args);
  168.567 ms (0 allocations: 0 bytes)

@essandess

essandess commented Aug 11, 2026

Copy link
Copy Markdown
Author

I can't reproduce the allocations in your benchmark (using SpecialFunctions v2.7.2):

julia> @btime ztloop($args);
  168.567 ms (0 allocations: 0 bytes)

The memory issue I observed is from redefining _zeta in the include("Downloads/gamma.jl"), so ignore that. If I use the the existing SpecialFunctions.jl and use polygamma directly, the benchmarks show about a 3× speedup, which about is what you'd expect just looking at the code:

using SpecialFunctions
using Distributions
using BenchmarkTools

args = rand(Exponential(), 1_000_000) .* exp.(im*2*pi*rand(Uniform(), 1_000_000));
function tgloop(args)
    for x in args
        tetragamma(x)
    end
end
function pgloop(args)
   for x in args
      polygamma(2, x)
   end
end

@btime tgloop($args)
  40.2224 ms (0 allocations: 0 bytes)

@btime pgloop($args)
  116.516 ms (0 allocations: 0 bytes)

@essandess

Copy link
Copy Markdown
Author

I don't see how the comment you cited, which is about the continued-fraction coefficients, is relevant here? We currently don't use the continued-fraction expansion, we use the (Stirling) asymptotic series, with the shift+reflection formulas to push the argument into the regime where the asymptotic series is accurate. This is the Kölbig (1972) algorithm (and derivatives thereof).

It looks like your tetragamma implementation also uses the asymptotic series, and indeed a comment in your code indicates that it is the 2nd derivative of the Kölbig digamma method.

We could therefore do the same thing for any desired order, and hence the question is where to stop. digamma is used so frequently that it is worth having a specialized implementation. trigamma is used less frequently but is present in many special functions libraries. I haven't seen any widely used library that has a specialized tetragamma implementation, am I missing one?

Definitely a specialized implementation for particular orders can be faster than the general polygamma code, but it trades off code, testing, and maintenance costs. (In your benchmark, as you point out, the performance benefit is only a factor of ≈ 2. Also something is weird in your benchmark because ztloop is not allocation-free … if that is a problem in the library, it should be fixed, and then the performance advantage will be even less.)

I addressed your questions about performance: about 2–3×. And "what kind of applications need higher performance for tetragamma but not even higher order": statistical analysis that relies upon exponential families, statistical curvature, and accelerated ML. I'm not aware of an application for pentagamma or higher, and haven't argued to extend beyond the tetragamma. If you'd prefer not to include tetragamma because of code, testing, and maintenance costs, or because other widely used libraries don't have it, that's fine, but please just make that call for those reasons.

@essandess

Copy link
Copy Markdown
Author

I'm not aware of an application for pentagamma

Correcting myself: Riemannian curvature of an LDA statistical manifold involves first derivatives of the Christoffel symbols, which are tetragammas, so the curvature of LDA will be pentagammas. And if you look up the Fisher–Rao geometry of the Dirichlet, sure enough there’s a pentagamma in Proposition 15 in that paper about sectional curvature. And of course any paper like this will be filled with tetragammas, as this one is. The same sort of results would exist for nearly all distributions in the exponential family. This isn’t an argument to include a pentagamma, just to point out that the “low” higher derivatives of the log-gamma are important in their own right.

Geometrically, manifolds are isometric iff there’s a correspondence of the tensors g, R, ∇R, ∇²R, …., ∇ᵏR so in principle some sequence k of the polygammas are relevant to the exponential family, but actual applications pretty much end at the pentagamma, and those are rare. I personally believe that the tetragamma is interesting and useful enough to warrant a specific implementation.

@essandess

Copy link
Copy Markdown
Author

I haven't seen any widely used library that has a specialized tetragamma implementation, am I missing one?

The widely cited R package limma uses a nice convex Newton's method for its trigammaInverse function, used to compute dof from moments. Having a tetragamma allows a nice implementation of this, the same way it's nice to have the trigamma for Minka's invdigamma in the existing code. I implemented Smyth and Phipson's algorithm in the PR.

…gamma

Asymptotic coefficients computed using B_{2k}*(2k+1), k=1..10 (A&S 6.4.11, m=2)

Polynomial order to achieve complex128 accuracy determined using Lentz's method near |z| = 10

Reference issue and PRs:
* scipy/scipy#7410
* scipy/scipy#17933
* scipy/xsf#239
@stevengj

stevengj commented Aug 12, 2026

Copy link
Copy Markdown
Member

Again, we have tetragamma via polygamma, so you can easily implement the same method in Julia. (Perhaps even more easily than in R, because we have an AD rule for trigamma.)

The question is whether it is performance sensitive enough in real applications to implement a specialized version to gain a factor of 2–3. The R limma package is not much of an argument for this because it is implemented in pure (non-vectorized) R and hence accepts a big performance penalty right out of the gate... and yet, despite that, it is widely used and cited (≈ 7500 times on google scholar). If the performance of tetragamma and/or trigammaInverse were critical, wouldn't someone have implemented a specialized compiled version by now?

(That being said, I'm sympathetic to the argument that derivatives of critical functions are critical nowadays too.)

@essandess

Copy link
Copy Markdown
Author

Again, we have tetragamma via polygamma, so you can easily implement the same method in Julia. (Perhaps even more easily than in R, because we have an AD rule for trigamma.)

The question is whether it is performance sensitive enough in real applications to implement a specialized version to gain a factor of 2–3. The R limma package is not much of an argument for this because it is implemented in pure (non-vectorized) R and hence accepts a big performance penalty right out of the gate... and yet, despite that, it is widely used and cited (≈ 7500 times on google scholar). If the performance of tetragamma and/or trigammaInverse were critical, wouldn't someone have implemented a specialized compiled version by now?

(That being said, I'm sympathetic to the argument that derivatives of critical functions are critical nowadays too.)

All your arguments against inclusion of a tetragamma apply to the trigamma too and yet the trigamma is already included.

I mentioned the limma package directly in response to your question about widely used packages that use the tetragamma, and because it has a very nice inverse trigamma algorithm that would enhance SpecialFunctions.jl, not because limma represents the best possible way to implement these algorithms. If you're simply opposed to inclusion of a specfic tetragamma function, that’s your call, and I’d prefer that to a back and forth with moving goal posts.

@stevengj

stevengj commented Aug 12, 2026

Copy link
Copy Markdown
Member

All your arguments against inclusion of a tetragamma apply to the trigamma too and yet the trigamma is already included.

First, the bar for adding a new function implementation, trading code size for performance without adding functionality, is higher than the bar for keeping an existing function.

Second, there are specialized compiled implementations of trigamma such as this one, which has been around since 1978.

So far, I can't find any such examples for tetragamma. I'm not saying that it's useless, I'm saying that there is a tradeoff here and it would be nice to have some outside evidence that the tradeoff is worthwhile, which often takes the form of finding other examples where people have optimized this particular function or where it is unambiguously performance critical.

@essandess

Copy link
Copy Markdown
Author

All your arguments against inclusion of a tetragamma apply to the trigamma too and yet the trigamma is already included.

First, the bar for adding a new function implementation, trading code size for performance without adding functionality, is higher than the bar for keeping an existing function.

Second, there are specialized compiled implementations of trigamma such as this one, which has been around since 1978.

So far, I can't find any such examples for tetragamma. I'm not saying that it's useless, I'm saying that there is a tradeoff here and it would be nice to have some outside evidence that the tradeoff is worthwhile, which often takes the form of finding other examples where people have optimized this particular function or where it is unambiguously performance critical.

Old Fortran implementations appearing in 1970s volumes of the RSS is now the goalpost? Seriously?! What about Cleve’s implementation of the tetragamma: The Tetragamma Function and Numerical Craftsmanship. Or this is a 21st century reference in Matlab so it doesn’t count? Cleve added it to Matlab 25 years ago for the same reasons I’ve been advocating: new statistical and machine learning tools demand new functions.

Again, if you're simply opposed to inclusion of a specific tetragamma function, that’s your call, and I’d prefer that to a back and forth with moving goal posts.

@stevengj

stevengj commented Aug 12, 2026

Copy link
Copy Markdown
Member

Most of the special-function libraries (SLATEC, Cephes, Amos) in widespread use today are decades old, so mere age of the code does not indicate lack of present interest. (On the contrary, if there are papers dating back decades on implementing a particular function, and multiple implementations of that algorithm still online today, that is a good indication that the function is of critical interest.)

Although Cleve's 2002 article was about the tetragamma example, what was actually added to Matlab in 2002 was the psi function that computes polygamma for arbitrary order. (Of course, we don't know for sure how many special cases Matlab has optimized internally, although I guess one could benchmark it.)

Again I'm simply asking for some evidence that tetragamma is performance-critical in real applications, to the point that it is worth the tradeoff of specializing it rather than using a generic polygamma implementation. (It's not just a question of whether this function is useful, because we already have it in polygamma.) One simple way to demonstrate this is to point to someone who has bothered to optimize this special case. I've been asking for this consistently. I don't know what you mean by "moving the goalposts".

Again, we have to draw the line somewhere on how many orders we optimize, and so far I haven't seen a clear argument for why tetragamma should be that line. I'm not trying to be hard on you, just trying to be practical about the maintenance tradeoffs we face here.

Let's not get combative, please.

@essandess

essandess commented Aug 12, 2026

Copy link
Copy Markdown
Author

Most of the special-function libraries (SLATEC, Cephes, Amos) in widespread use today are decades old, so mere age of the code does not indicate lack of present interest. (On the contrary, if there are papers dating back decades on implementing a particular function, and multiple implementations of that algorithm still online today, that is a good indication that the function is of critical interest.)

Although Cleve's 2002 article was about the tetragamma example, what was actually added to Matlab in 2002 was the psi function that computes polygamma for arbitrary order. (Of course, we don't know for sure how many special cases Matlab has optimized internally, although I guess one could benchmark it.)

Again I'm simply asking for some evidence that tetragamma is performance-critical in real applications, to the point that it is worth the tradeoff of specializing it rather than using a generic polygamma implementation. (It's not just a question of whether this function is useful, because we already have it in polygamma.) One simple way to demonstrate this is to point to someone who has bothered to optimize this special case. I've been asking for this consistently. I don't know what you mean by "moving the goalposts".

Again, we have to draw the line somewhere on how many orders we optimize, and so far I haven't seen a clear argument for why tetragamma should be that line. I'm not trying to be hard on you, just trying to be practical about the maintenance tradeoffs we face here.

Let's not get combative, please.

Matlab’s psi is almost certainly Amos’s dpsifn.

Here’s a benchmark that skits on real arguments, the existing trigamma is as fast as polygamma, but there’s a 9x speedup for tetragamma. I’ve also benched a Julia port of Amos’s code, which is almost as fast as polygamma. The Rmath library uses an R port of dpsifn that’s limited up to pentagamma.

Matlab is great, but a lot of its code has remained unoptimized, which is why Julia and PyTorch are often better, faster options. Here’s a chance to optimize at least one function in this package.

using SpecialFunctions
using Distributions
using BenchmarkTools

# real arguments, χ distributed
args = abs.(randn(1_000_000) + im*randn(1_000_000));

benchmarks = [
   "trigamma"      => args -> foreach(trigamma, args),
   "polygamma(1)"  => args -> foreach(x -> polygamma(1, x), args),
   "psifn(1)"      => args -> foreach(x -> psifn(x, 1, 1, 1), args),
   "tetragamma"    => args -> foreach(tetragamma, args),
   "polygamma(2)"  => args -> foreach(x -> polygamma(2, x), args),
   "psifn(2)"      => args -> foreach(x -> psifn(x, 2, 1, 1), args)
]
for (name, fn) in benchmarks
   print("Benchmarking $name: ")
   @btime $fn($args)
end

Benchmarking trigamma:   11.378 ms (0 allocations: 0 bytes)
Benchmarking polygamma(1):   12.474 ms (0 allocations: 0 bytes)
Benchmarking psifn(1):   106.335 ms (6000000 allocations: 1.15 GiB)
Benchmarking tetragamma:   11.945 ms (0 allocations: 0 bytes)
Benchmarking polygamma(2):   94.223 ms (0 allocations: 0 bytes)
Benchmarking psifn(2):   118.465 ms (6000000 allocations: 1.15 GiB)
tetragamma
const ComplexOrReal{T} = Union{T,Complex{T}}

"""
   tetragamma(x)

Compute the tetragamma function of `x` (the logarithmic third derivative of `gamma(x)`).
"""
tetragamma(x::Number) = _tetragamma(float(x))

function _tetragamma(z::ComplexOrReal{Float64})
   # via the second derivative of the Kölbig digamma formulation
   x = real(z)
   if z isa Complex && imag(z) == 0
       # also catches 0.0+0.0im → -Inf because π / 0. is Inf in IEEE 754,
       # but π+0im / 0+0im is Nan+Nan*im in complex
       return complex(_tetragamma(x))
   end
   if x <= 0 # reflection formula
       if !(z isa Complex) || abs(imag(z)) < log(floatmax(Float64))
           return tetragamma(1 - z) - 2*π^3*cospi(z)*inv(sinpi(z))^3
       else
           # omit numerically 0 term and avoid overflow at large imag(z)
           return tetragamma(1 - z)
       end
   end
   ψ = zero(z)
   N = 10
   if x < N
       # shift using recurrence formula
       n = N - floor(Int,x)
       ψ -= 2*inv(z)^3
       for ν = 1:n-1
           ψ -= 2*inv(z + ν)^3
       end
       z += n
   end
   t = inv(z)
   w = t * t # 1/z^2
   ψ += -w * (1.0 + t)
   # the coefficients here are Float64(-(2*(1:10) .+ 1) .* bernoulli[2:11])
   # order determined by Lentz's method near |z| = 10
   ψ += w*w * @evalpoly(w,-0.5,0.16666666666666666,-0.16666666666666666,0.3,-0.8333333333333334,3.2904761904761907,-17.5,120.56666666666666,-1044.452380952381,11111.609090909091)
end
Amos's `DPSIFN` Fortran to Julia port
const ComplexOrReal{T} = Union{T,Complex{T}}

const _dpsifn_nmax = 100

const _dpsifn_B = (
   1.00000000000000000e+00, -5.00000000000000000e-01, 1.66666666666666667e-01,
   -3.33333333333333333e-02, 2.38095238095238095e-02, -3.33333333333333333e-02,
   7.57575757575757576e-02, -2.53113553113553114e-01, 1.16666666666666667e+00,
   -7.09215686274509804e+00, 5.49711779448621554e+01, -5.29124242424242424e+02,
   6.19212318840579710e+03, -8.65802531135531136e+04, 1.42551716666666667e+06,
   -2.72982310678160920e+07, 6.01580873900642368e+08, -1.51163157670921569e+10,
   4.29614643061166667e+11, -1.37116552050883328e+13, 4.88332318973593167e+14,
   -1.92965793419400681e+16,
)

const _dpsifn_wdtol = max(eps(Float64) / 2, 0.5e-18)
const _dpsifn_elim = 0.99 * log(floatmax(Float64))
const _dpsifn_rln = min(log10(2.0) * 53, 18.06)

function _dpsifn_n0_recur(s::T, z::T, nx::Integer) where {T<:ComplexOrReal{Float64}}
   for i = 1:nx
       s += 1.0 / (z + float(nx - i))
   end
   return s
end

function _dpsifn_finalize_n0(s::T, zdmln::T, kode::Integer, zdmy::T, z::T) where {T<:ComplexOrReal{Float64}}
   kode != 2 && return s - zdmln
   zdmy == z && return s
   return s - log(zdmy / z)
end

"""
   psifn(z::Number, n::Integer, kode::Integer, m::Integer) -> Vector

Port of D. E. Amos's `DPSIFN` (ACM Algorithm 610), computing `m` member
sequences of scaled derivatives `W(k,z) = (-1)^(k+1) ψ^(k)(z)/k!`,
`k = n,...,n+m-1`, with `kode=2` returning `-ψ(z)+log(z)` instead of
`-ψ(z)` for `k=0`.

Extended to complex `z` and to `Re(z) ≤ 0` — the original Fortran
requires `x > 0` and errors otherwise — via the same reflection
identity already used by `_polygamma` above, expressed in terms of the
scaled `W` convention as

   W(k,z) = signflip(k, zeta(k+1, 1-z)) + cotderiv(k, z),  k ≥ 2,

with `k=0,1` handled directly by `digamma`/`trigamma` (which already
have their own complex-capable reflection formulas), avoiding the pole
of `zeta(1,·)` the same way `_polygamma` avoids it for `m=0`.

For `Re(z) > 0`, the original Amos asymptotic-expansion / backward-shift
algorithm is used essentially unchanged (it only assumed `z` real, not
that it's positive along a cut — the derivation holds for complex `z`
with large positive real part equally well).
"""
psifn(z::Number, n::Integer, kode::Integer, m::Integer) = _psifn(float(z), n, kode, m)

function _psifn(z::T, n::Integer, kode::Integer, m::Integer) where {T<:ComplexOrReal{Float64}}
   n < 0 && throw(DomainError(n, "`n` must be nonnegative."))
   (kode < 1 || kode > 2) && throw(DomainError(kode, "`kode` must be 1 or 2."))
   m < 1 && throw(DomainError(m, "`m` must be positive."))
   iszero(z) && throw(DomainError(z, "`z` must be nonzero."))

   real(z) <= 0 && return _psifn_reflect(z, n, kode, m)
   return _psifn_asym(z, n, kode, m)
end

# New: reflection branch for Re(z) ≤ 0. Not present in the original
# Fortran at all (DPSIFN requires x > 0); built entirely from
# digamma/trigamma/zeta/cotderiv/signflip, exactly as _polygamma does.
function _psifn_reflect(z::T, n::Integer, kode::Integer, m::Integer) where {T<:ComplexOrReal{Float64}}
   ans = Vector{T}(undef, m)
   @inbounds for j = 1:m
       k = n + j - 1
       if k == 0
           w = -digamma(z)
           kode == 2 && (w += log(z))
           ans[j] = w
       elseif k == 1
           ans[j] = trigamma(z)
       else
           ans[j] = signflip(k, zeta(Float64(k + 1), 1 - z)) + cotderiv(k, z)
       end
   end
   return ans
end

# Amos's asymptotic-expansion / backward-shift algorithm, generalized
# from real x to any z::ComplexOrReal{Float64} with real(z) > 0.
function _psifn_asym(z::T, n::Integer, kode::Integer, m::Integer) where {T<:ComplexOrReal{Float64}}
   nmax = _dpsifn_nmax
   B = _dpsifn_B
   wdtol = _dpsifn_wdtol
   elim = _dpsifn_elim
   x = real(z)

   ans = zeros(T, m)
   trm = zeros(T, 22)

   nn = n + m - 1
   fn = float(nn)
   fnp = fn + 1.0

   zln = log(z)
   t = fnp * zln
   if abs(real(t)) > elim
       throw(ArgumentError(real(t) > 0 ?
           "psifn: underflow, z too large or n+m-1 too large" :
           "psifn: overflow, z too small or n+m-1 too large"))
   end

   if abs(z) < wdtol
       # small |z|
       ans[1] = z^(-n - 1)
       if m > 1
           @inbounds for i = 2:m
               ans[i] = ans[i-1] / z
           end
       end
       n == 0 && kode == 2 && (ans[1] += zln)
       return ans
   end

   rln = _dpsifn_rln
   fln = max(rln, 3.0) - 3.0
   yint = 3.50 + 0.40 * fln
   slope = 0.21 + fln * (0.0006038 * fln + 0.008677)
   xm = yint + slope * fn
   xmin = float(floor(Int, xm) + 1)

   use_series = false
   fns = 0.0
   if n != 0
       xm2 = -2.302 * rln - min(0.0, real(zln))
       fns = float(n)
       arg = min(0.0, xm2 / fns)
       epsx = exp(arg)
       xm2 = abs(arg) < 1.0e-3 ? -arg : 1.0 - epsx
       fln = x * xm2 / epsx
       use_series = (xmin - x) > 7.0 && fln < 15.0
   end

   if use_series
       # series (z+k)^(-(n+1)), k=0,1,2,...
       nn2 = floor(Int, fln) + 1
       np = n + 1
       t1 = (fns + 1.0) * zln
       t = exp(-t1)
       s = t
       den = z
       trmv = Vector{T}(undef, nn2)  # sized to what's actually needed
       @inbounds for i = 1:nn2
           den += 1.0
           trmv[i] = den^(-np)
           s += trmv[i]
       end
       ans[1] = (n == 0 && kode == 2) ? s + zln : s
       if m > 1
           tol = wdtol / 5.0
           for j = 2:m
               t /= z
               s = t
               tols = abs(t) * tol
               den = z
               @inbounds for i = 1:nn2
                   den += 1.0
                   trmv[i] /= den
                   s += trmv[i]
                   abs(trmv[i]) < tols && break
               end
               ans[j] = s
           end
       end
       return ans
   end

   # asymptotic expansion for w(n+m-1,z)
   zdmy = z
   zdmln = zln
   xinc = 0.0
   if x < xmin
       xinc = xmin - float(floor(Int, x))
       zdmy = z + xinc
       zdmln = log(zdmy)
   end

   t = fn * zdmln
   t1 = zdmln + zdmln
   t2 = t + zdmln
   tk = max(abs(real(t)), abs(real(t1)), abs(real(t2)))
   tk > elim && throw(ArgumentError("psifn: underflow, z too large or n+m-1 too large"))

   tss = exp(-t)
   tt = 0.5 / zdmy
   t1 = nn != 0 ? tt + 1.0 / fn : tt
   tst = wdtol * abs(tt)
   rzsq = 1.0 / (zdmy * zdmy)
   ta = 0.5 * rzsq
   t = fnp * ta
   s = t * B[3]
   if abs(s) >= tst
       tk = 2.0
       @inbounds for k = 4:22
           t *= ((tk + fn + 1.0) / (tk + 1.0)) * ((tk + fn) / (tk + 2.0)) * rzsq
           trm[k] = t * B[k]
           abs(trm[k]) < tst && break
           s += trm[k]
           tk += 2.0
       end
   end
   s = (s + t1) * tss

   nx = 0
   trmr = T[]
   if xinc != 0.0
       nx = floor(Int, xinc)
       np = nn + 1
       nx > nmax && throw(ArgumentError("psifn: shift $nx exceeds internal limit ($nmax)"))
       trmr = Vector{T}(undef, nx)  # was a fixed-size-100 buffer in the Fortran
       if nn == 0
           s = _dpsifn_n0_recur(s, z, nx)
           ans[1] = _dpsifn_finalize_n0(s, zdmln, kode, zdmy, z)
           return ans
       end
       xm = xinc - 1.0
       fx = z + xm
       @inbounds for i = 1:nx
           trmr[i] = fx^(-np)
           s += trmr[i]
           xm -= 1.0
           fx = z + xm
       end
   end
   ans[m] = s

   if fn == 0.0
       ans[1] = _dpsifn_finalize_n0(s, zdmln, kode, zdmy, z)
       return ans
   end
   m == 1 && return ans

   # lower derivatives, j < n+m-1
   for j = 2:m
       fnp = fn
       fn -= 1.0
       tss *= zdmy
       t1 = fn != 0.0 ? tt + 1.0 / fn : tt
       t = fnp * ta
       s = t * B[3]
       if abs(s) >= tst
           tk = 3.0 + fnp
           @inbounds for k = 4:22
               trm[k] *= fnp / tk
               abs(trm[k]) < tst && break
               s += trm[k]
               tk += 2.0
           end
       end
       s = (s + t1) * tss

       if xinc != 0.0
           if fn == 0.0
               s = _dpsifn_n0_recur(s, z, nx)
               ans[1] = _dpsifn_finalize_n0(s, zdmln, kode, zdmy, z)
               return ans
           end
           xm = xinc - 1.0
           fx = z + xm
           @inbounds for i = 1:nx
               trmr[i] *= fx
               s += trmr[i]
               xm -= 1.0
               fx = z + xm
           end
       end

       ans[m-j+1] = s
       fn == 0.0 && (ans[1] = _dpsifn_finalize_n0(s, zdmln, kode, zdmy, z); return ans)
   end
   return ans
end

@essandess

Copy link
Copy Markdown
Author

One simple way to demonstrate this is to point to someone who has bothered to optimize this special case.

These are honestly very small, well understood, and well tested functions and would not take a great deal of effort to maintain.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants