Merge pull request #18498 from bashtage/Androp0v-vonmises-fix

BUG: Fixed Von Mises distribution for big values of kappa

7 files changed, 121 insertions, 8 deletions

diff --git a/numpy/random/include/legacy-distributions.h b/numpy/random/include/legacy-distributions.h
index b8ba0841c..f7ccd2cb5 100644
--- a/numpy/random/include/legacy-distributions.h
+++ b/numpy/random/include/legacy-distributions.h
@@ -31,6 +31,7 @@ extern double legacy_f(aug_bitgen_t *aug_state, double dfnum, double dfden);
 extern double legacy_normal(aug_bitgen_t *aug_state, double loc, double scale);
 extern double legacy_standard_gamma(aug_bitgen_t *aug_state, double shape);
 extern double legacy_exponential(aug_bitgen_t *aug_state, double scale);
+extern double legacy_vonmises(bitgen_t *bitgen_state, double mu, double kappa);
 extern int64_t legacy_random_binomial(bitgen_t *bitgen_state, double p,
                                       int64_t n, binomial_t *binomial);
 extern int64_t legacy_negative_binomial(aug_bitgen_t *aug_state, double n,
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
index 8d349b7d8..6f44e271f 100644
--- a/numpy/random/mtrand.pyx
+++ b/numpy/random/mtrand.pyx
@@ -51,7 +51,6 @@ cdef extern from "numpy/random/distributions.h":
     void random_standard_uniform_fill(bitgen_t* bitgen_state, np.npy_intp cnt, double *out) nogil
     int64_t random_positive_int(bitgen_t *bitgen_state) nogil
     double random_uniform(bitgen_t *bitgen_state, double lower, double range) nogil
-    double random_vonmises(bitgen_t *bitgen_state, double mu, double kappa) nogil
     double random_laplace(bitgen_t *bitgen_state, double loc, double scale) nogil
     double random_gumbel(bitgen_t *bitgen_state, double loc, double scale) nogil
     double random_logistic(bitgen_t *bitgen_state, double loc, double scale) nogil
@@ -100,6 +99,7 @@ cdef extern from "include/legacy-distributions.h":
     double legacy_f(aug_bitgen_t *aug_state, double dfnum, double dfden) nogil
     double legacy_exponential(aug_bitgen_t *aug_state, double scale) nogil
     double legacy_power(aug_bitgen_t *state, double a) nogil
+    double legacy_vonmises(bitgen_t *bitgen_state, double mu, double kappa) nogil
 
 np.import_array()
 
@@ -2281,7 +2281,7 @@ cdef class RandomState:
         >>> plt.show()
 
         """
-        return cont(&random_vonmises, &self._bitgen, size, self.lock, 2,
+        return cont(&legacy_vonmises, &self._bitgen, size, self.lock, 2,
                     mu, 'mu', CONS_NONE,
                     kappa, 'kappa', CONS_NON_NEGATIVE,
                     0.0, '', CONS_NONE, None)
diff --git a/numpy/random/src/distributions/distributions.c b/numpy/random/src/distributions/distributions.c
index 4494f860e..f47c54a53 100644
--- a/numpy/random/src/distributions/distributions.c
+++ b/numpy/random/src/distributions/distributions.c
@@ -843,6 +843,7 @@ double random_vonmises(bitgen_t *bitgen_state, double mu, double kappa) {
     return NPY_NAN;
   }
   if (kappa < 1e-8) {
+    /* Use a uniform for very small values of kappa */
     return M_PI * (2 * next_double(bitgen_state) - 1);
   } else {
     /* with double precision rho is zero until 1.4e-8 */
@@ -853,9 +854,23 @@ double random_vonmises(bitgen_t *bitgen_state, double mu, double kappa) {
        */
       s = (1. / kappa + kappa);
     } else {
-      double r = 1 + sqrt(1 + 4 * kappa * kappa);
-      double rho = (r - sqrt(2 * r)) / (2 * kappa);
-      s = (1 + rho * rho) / (2 * rho);
+      if (kappa <= 1e6) {
+        /* Path for 1e-5 <= kappa <= 1e6 */
+        double r = 1 + sqrt(1 + 4 * kappa * kappa);
+        double rho = (r - sqrt(2 * r)) / (2 * kappa);
+        s = (1 + rho * rho) / (2 * rho);
+      } else {
+        /* Fallback to wrapped normal distribution for kappa > 1e6 */
+        result = mu + sqrt(1. / kappa) * random_standard_normal(bitgen_state);
+        /* Ensure result is within bounds */
+        if (result < -M_PI) {
+          result += 2*M_PI;
+        }
+        if (result > M_PI) {
+          result -= 2*M_PI;
+        }
+        return result;
+      }
     }
 
     while (1) {
diff --git a/numpy/random/src/legacy/legacy-distributions.c b/numpy/random/src/legacy/legacy-distributions.c
index fd067fe8d..5b17401dd 100644
--- a/numpy/random/src/legacy/legacy-distributions.c
+++ b/numpy/random/src/legacy/legacy-distributions.c
@@ -1,3 +1,13 @@
+/*
+ * This file contains generation code for distribution that have been modified
+ * since Generator was introduced. These are preserved using identical code
+ * to what was in NumPy 1.16 so that the stream of values generated by
+ * RandomState is not changed when there are changes that affect Generator.
+ *
+ * These functions should not be changed except if they contain code that
+ * cannot be compiled. They should not be changed for bug fixes, performance
+ * improvements that can change the values produced, or enhancements to precision.
+ */
 #include "include/legacy-distributions.h"
 
 
@@ -390,3 +400,61 @@ int64_t legacy_random_logseries(bitgen_t *bitgen_state, double p) {
                                binomial_t *binomial) {
   return random_multinomial(bitgen_state, n, mnix, pix, d, binomial);
 }
+
+double legacy_vonmises(bitgen_t *bitgen_state, double mu, double kappa) {
+  double s;
+  double U, V, W, Y, Z;
+  double result, mod;
+  int neg;
+  if (npy_isnan(kappa)) {
+    return NPY_NAN;
+  }
+  if (kappa < 1e-8) {
+    return M_PI * (2 * next_double(bitgen_state) - 1);
+  } else {
+    /* with double precision rho is zero until 1.4e-8 */
+    if (kappa < 1e-5) {
+      /*
+       * second order taylor expansion around kappa = 0
+       * precise until relatively large kappas as second order is 0
+       */
+      s = (1. / kappa + kappa);
+    } else {
+        /* Path for 1e-5 <= kappa <= 1e6 */
+        double r = 1 + sqrt(1 + 4 * kappa * kappa);
+        double rho = (r - sqrt(2 * r)) / (2 * kappa);
+        s = (1 + rho * rho) / (2 * rho);
+    }
+
+    while (1) {
+      U = next_double(bitgen_state);
+      Z = cos(M_PI * U);
+      W = (1 + s * Z) / (s + Z);
+      Y = kappa * (s - W);
+      V = next_double(bitgen_state);
+      /*
+       * V==0.0 is ok here since Y >= 0 always leads
+       * to accept, while Y < 0 always rejects
+       */
+      if ((Y * (2 - Y) - V >= 0) || (log(Y / V) + 1 - Y >= 0)) {
+        break;
+      }
+    }
+
+    U = next_double(bitgen_state);
+
+    result = acos(W);
+    if (U < 0.5) {
+      result = -result;
+    }
+    result += mu;
+    neg = (result < 0);
+    mod = fabs(result);
+    mod = (fmod(mod + M_PI, 2 * M_PI) - M_PI);
+    if (neg) {
+      mod *= -1;
+    }
+
+    return mod;
+  }
+}
\ No newline at end of file
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
index d68bcd38b..f6bd985b4 100644
--- a/numpy/random/tests/test_generator_mt19937.py
+++ b/numpy/random/tests/test_generator_mt19937.py
@@ -10,7 +10,7 @@ from numpy.testing import (
     assert_warns, assert_no_warnings, assert_array_equal,
     assert_array_almost_equal, suppress_warnings)
 
-from numpy.random import Generator, MT19937, SeedSequence
+from numpy.random import Generator, MT19937, SeedSequence, RandomState
 
 random = Generator(MT19937())
 
@@ -1735,6 +1735,26 @@ class TestRandomDist:
         r = random.vonmises(mu=0., kappa=np.nan)
         assert_(np.isnan(r))
 
+    @pytest.mark.parametrize("kappa", [1e4, 1e15])
+    def test_vonmises_large_kappa(self, kappa):
+        random = Generator(MT19937(self.seed))
+        rs = RandomState(random.bit_generator)
+        state = random.bit_generator.state
+
+        random_state_vals = rs.vonmises(0, kappa, size=10)
+        random.bit_generator.state = state
+        gen_vals = random.vonmises(0, kappa, size=10)
+        if kappa < 1e6:
+            assert_allclose(random_state_vals, gen_vals)
+        else:
+            assert np.all(random_state_vals != gen_vals)
+
+    @pytest.mark.parametrize("mu", [-7., -np.pi, -3.1, np.pi, 3.2])
+    @pytest.mark.parametrize("kappa", [1e-9, 1e-6, 1, 1e3, 1e15])
+    def test_vonmises_large_kappa_range(self, mu, kappa):
+        r = random.vonmises(mu, kappa, 50)
+        assert_(np.all(r > -np.pi) and np.all(r <= np.pi))
+
     def test_wald(self):
         random = Generator(MT19937(self.seed))
         actual = random.wald(mean=1.23, scale=1.54, size=(3, 2))
diff --git a/numpy/random/tests/test_generator_mt19937_regressions.py b/numpy/random/tests/test_generator_mt19937_regressions.py
index 2ef6b0631..9f6dcdc6b 100644
--- a/numpy/random/tests/test_generator_mt19937_regressions.py
+++ b/numpy/random/tests/test_generator_mt19937_regressions.py
@@ -1,14 +1,14 @@
 from numpy.testing import (assert_, assert_array_equal)
 import numpy as np
 import pytest
-from numpy.random import Generator, MT19937
+from numpy.random import Generator, MT19937, RandomState
 
 mt19937 = Generator(MT19937())
 
 
 class TestRegression:
 
-    def test_VonMises_range(self):
+    def test_vonmises_range(self):
         # Make sure generated random variables are in [-pi, pi].
         # Regression test for ticket #986.
         for mu in np.linspace(-7., 7., 5):
diff --git a/numpy/random/tests/test_randomstate.py b/numpy/random/tests/test_randomstate.py
index 7f5f08050..b16275b70 100644
--- a/numpy/random/tests/test_randomstate.py
+++ b/numpy/random/tests/test_randomstate.py
@@ -1238,6 +1238,15 @@ class TestRandomDist:
         r = random.vonmises(mu=0., kappa=1.1e-8, size=10**6)
         assert_(np.isfinite(r).all())
 
+    def test_vonmises_large(self):
+        # guard against changes in RandomState when Generator is fixed
+        random.seed(self.seed)
+        actual = random.vonmises(mu=0., kappa=1e7, size=3)
+        desired = np.array([4.634253748521111e-04,
+                            3.558873596114509e-04,
+                            -2.337119622577433e-04])
+        assert_array_almost_equal(actual, desired, decimal=8)
+
     def test_vonmises_nan(self):
         random.seed(self.seed)
         r = random.vonmises(mu=0., kappa=np.nan)