Fix hangs in beta_inc_inv (alternative/extension of JuliaMath#396) (JuliaMath#399)

* Limit number of Newton iterations in beta_inc_inv and allow much
smaller starting values.

* Adjust tolerance instead of number of iterations

Co-authored-by: Andreas Noack <[email protected]>

Author: devmotion

src/beta_inc.jl

@@ -934,7 +934,6 @@ function _beta_inc_inv(a::Float64, b::Float64, p::Float64, q::Float64=1-p)
     r = sqrt(-2*log(p))
     p_approx = r - @horner(r, 2.30753e+00, 0.27061e+00) / @horner(r, 1.0, .99229e+00, .04481e+00)
     fpu = floatmin(Float64)
-    sae = log10(fpu)
     lb = logbeta(a, b)
 
     if a > 1.0 && b > 1.0
@@ -968,12 +967,7 @@ function _beta_inc_inv(a::Float64, b::Float64, p::Float64, q::Float64=1-p)
     sq = 1.0
     prev = 1.0
 
-    if x < 0.0001
-        x = 0.0001
-    end
-    if x > .9999
-        x = .9999
-    end
+    x = clamp(x, fpu, prevfloat(1.0))
 
     # This first argument was proposed in
     #
@@ -997,8 +991,8 @@ function _beta_inc_inv(a::Float64, b::Float64, p::Float64, q::Float64=1-p)
     # The idea with the -5.0/a^2 - 1.0/p^0.2 - 34.0 correction is to
     # use -2r - 2 (for 16 digits) for large values of a and p but use
     # a much smaller tolerance for small values of a and p.
-    iex = max(-5.0/a^2 - 1.0/p^0.2 - 34.0, sae)
-    acu = exp10(iex)
+    iex = -5.0/a^2 - 1.0/p^0.2 - 34.0
+    acu = max(exp10(iex), 10 * fpu) # 10 * fpu instead of fpu avoids hangs for small values
 
     #iterate
     while true
@@ -1012,21 +1006,15 @@ function _beta_inc_inv(a::Float64, b::Float64, p::Float64, q::Float64=1-p)
         if p_approx * p_approx_prev <= 0.0
             prev = max(sq, fpu)
         end
-        g = 1.0
 
-        tx = x - g*p_approx
-        while true
-            adj = g*p_approx
-            sq = adj^2
+	adj = p_approx
+	tx = x - adj
+	while prev <= (sq = adj^2) || tx < 0.0 || tx > 1.0
+            adj /= 3.0
             tx = x - adj
-            if (prev > sq && tx >= 0.0 && tx <= 1.0)
-                break
-            end
-            g /= 3.0
         end
 
         #check if current estimate is acceptable
-
         if prev <= acu || p_approx^2 <= acu
             x = tx
             return (x, 1.0 - x)

test/beta_inc.jl

@@ -299,4 +299,10 @@ end
         @test first(beta_inc_inv(1.0, 1.0, 1e-21)) ≈ 1e-21
         @test beta_inc_inv(1.0e30, 1.0, 0.49) == (1.0, 0.0)
     end
+
+    @testset "Avoid infinite loops" begin
+        # See https://github.com/JuliaStats/StatsFuns.jl/issues/133#issuecomment-1069602721
+        y = 2.0e-280
+        @test beta_inc(2.0, 1.0, beta_inc_inv(2.0, 1.0, y, 1.0)[1])[1] ≈ y
+    end
 end