Issue #17141: random.vonmisesvariate() no more hangs for large kappas.

2 files changed, 38 insertions, 10 deletions

diff --git a/Lib/random.py b/Lib/random.py
index c275071..7d8d4f3 100644
--- a/Lib/random.py
+++ b/Lib/random.py
@@ -431,22 +431,20 @@ class Random(_random.Random):
         if kappa <= 1e-6:
             return TWOPI * random()
 
-        a = 1.0 + _sqrt(1.0 + 4.0 * kappa * kappa)
-        b = (a - _sqrt(2.0 * a))/(2.0 * kappa)
-        r = (1.0 + b * b)/(2.0 * b)
+        s = 0.5 / kappa
+        r = s + _sqrt(1.0 + s * s)
 
         while 1:
             u1 = random()
-
             z = _cos(_pi * u1)
-            f = (1.0 + r * z)/(r + z)
-            c = kappa * (r - f)
 
+            d = z / (r + z)
             u2 = random()
-
-            if u2 < c * (2.0 - c) or u2 <= c * _exp(1.0 - c):
+            if u2 < 1.0 - d * d or u2 <= (1.0 - d) * _exp(d):
                 break
 
+        q = 1.0 / r
+        f = (q + z) / (1.0 + q * z)
         u3 = random()
         if u3 > 0.5:
             theta = (mu + _acos(f)) % TWOPI
diff --git a/Lib/test/test_random.py b/Lib/test/test_random.py
index b176ea1..c3ab7d2 100644
--- a/Lib/test/test_random.py
+++ b/Lib/test/test_random.py
@@ -440,6 +440,7 @@ class TestDistributions(unittest.TestCase):
         g.random = x[:].pop; g.paretovariate(1.0)
         g.random = x[:].pop; g.expovariate(1.0)
         g.random = x[:].pop; g.weibullvariate(1.0, 1.0)
+        g.random = x[:].pop; g.vonmisesvariate(1.0, 1.0)
         g.random = x[:].pop; g.normalvariate(0.0, 1.0)
         g.random = x[:].pop; g.gauss(0.0, 1.0)
         g.random = x[:].pop; g.lognormvariate(0.0, 1.0)
@@ -460,6 +461,7 @@ class TestDistributions(unittest.TestCase):
                 (g.uniform, (1.0,10.0), (10.0+1.0)/2, (10.0-1.0)**2/12),
                 (g.triangular, (0.0, 1.0, 1.0/3.0), 4.0/9.0, 7.0/9.0/18.0),
                 (g.expovariate, (1.5,), 1/1.5, 1/1.5**2),
+                (g.vonmisesvariate, (1.23, 0), pi, pi**2/3),
                 (g.paretovariate, (5.0,), 5.0/(5.0-1),
                                   5.0/((5.0-1)**2*(5.0-2))),
                 (g.weibullvariate, (1.0, 3.0), gamma(1+1/3.0),
@@ -476,8 +478,30 @@ class TestDistributions(unittest.TestCase):
                 s1 += e
                 s2 += (e - mu) ** 2
             N = len(y)
-            self.assertAlmostEqual(s1/N, mu, places=2)
-            self.assertAlmostEqual(s2/(N-1), sigmasqrd, places=2)
+            self.assertAlmostEqual(s1/N, mu, places=2,
+                                   msg='%s%r' % (variate.__name__, args))
+            self.assertAlmostEqual(s2/(N-1), sigmasqrd, places=2,
+                                   msg='%s%r' % (variate.__name__, args))
+
+    def test_constant(self):
+        g = random.Random()
+        N = 100
+        for variate, args, expected in [
+                (g.uniform, (10.0, 10.0), 10.0),
+                (g.triangular, (10.0, 10.0), 10.0),
+                #(g.triangular, (10.0, 10.0, 10.0), 10.0),
+                (g.expovariate, (float('inf'),), 0.0),
+                (g.vonmisesvariate, (3.0, float('inf')), 3.0),
+                (g.gauss, (10.0, 0.0), 10.0),
+                (g.lognormvariate, (0.0, 0.0), 1.0),
+                (g.lognormvariate, (-float('inf'), 0.0), 0.0),
+                (g.normalvariate, (10.0, 0.0), 10.0),
+                (g.paretovariate, (float('inf'),), 1.0),
+                (g.weibullvariate, (10.0, float('inf')), 10.0),
+                (g.weibullvariate, (0.0, 10.0), 0.0),
+            ]:
+            for i in range(N):
+                self.assertEqual(variate(*args), expected)
 
     def test_von_mises_range(self):
         # Issue 17149: von mises variates were not consistently in the
@@ -493,6 +517,12 @@ class TestDistributions(unittest.TestCase):
                         msg=("vonmisesvariate({}, {}) produced a result {} out"
                              " of range [0, 2*pi]").format(mu, kappa, sample))
 
+    def test_von_mises_large_kappa(self):
+        # Issue #17141: vonmisesvariate() was hang for large kappas
+        random.vonmisesvariate(0, 1e15)
+        random.vonmisesvariate(0, 1e100)
+
+
 class TestModule(unittest.TestCase):
     def testMagicConstants(self):
         self.assertAlmostEqual(random.NV_MAGICCONST, 1.71552776992141)