ENH: stats.special_ortho_group: speed up, allow 1x1 and 0x0 ortho and unitary groups (scipy#22304)

* ENH: streamline implementation of stats.special_ortho_group
* ENH: allow 1x1 and 0x0 ortho and unitary groups
* TST: switch from random.seed to random.default_rng

Author: akhmerov

scipy/stats/_multivariate.py

@@ -3620,9 +3620,9 @@ def __call__(self, dim=None, seed=None):
 
     def _process_parameters(self, dim):
         """Dimension N must be specified; it cannot be inferred."""
-        if dim is None or not np.isscalar(dim) or dim <= 1 or dim != int(dim):
+        if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim):
             raise ValueError("""Dimension of rotation must be specified,
-                                and must be a scalar greater than 1.""")
+                                and must be a scalar nonnegative integer.""")
 
         return dim
 
@@ -3644,55 +3644,11 @@ def rvs(self, dim, size=1, random_state=None):
         """
         random_state = self._get_random_state(random_state)
 
-        size = int(size)
-        size = (size,) if size > 1 else ()
-
-        dim = self._process_parameters(dim)
-
-        # H represents a (dim, dim) matrix, while D represents the diagonal of
-        # a (dim, dim) diagonal matrix. The algorithm that follows is
-        # broadcasted on the leading shape in `size` to vectorize along
-        # samples.
-        H = np.empty(size + (dim, dim))
-        H[..., :, :] = np.eye(dim)
-        D = np.empty(size + (dim,))
-
-        for n in range(dim-1):
-
-            # x is a vector with length dim-n, xrow and xcol are views of it as
-            # a row vector and column vector respectively. It's important they
-            # are views and not copies because we are going to modify x
-            # in-place.
-            x = random_state.normal(size=size + (dim-n,))
-            xrow = x[..., None, :]
-            xcol = x[..., :, None]
-
-            # This is the squared norm of x, without vectorization it would be
-            # dot(x, x), to have proper broadcasting we use matmul and squeeze
-            # out (convert to scalar) the resulting 1x1 matrix
-            norm2 = np.matmul(xrow, xcol).squeeze((-2, -1))
-
-            x0 = x[..., 0].copy()
-            D[..., n] = np.where(x0 != 0, np.sign(x0), 1)
-            x[..., 0] += D[..., n]*np.sqrt(norm2)
-
-            # In renormalizing x we have to append an additional axis with
-            # [..., None] to broadcast the scalar against the vector x
-            x /= np.sqrt((norm2 - x0**2 + x[..., 0]**2) / 2.)[..., None]
-
-            # Householder transformation, without vectorization the RHS can be
-            # written as outer(H @ x, x) (apart from the slicing)
-            H[..., :, n:] -= np.matmul(H[..., :, n:], xcol) * xrow
-
-        D[..., -1] = (-1)**(dim-1)*D[..., :-1].prod(axis=-1)
-
-        # Without vectorization this could be written as H = diag(D) @ H,
-        # left-multiplication by a diagonal matrix amounts to multiplying each
-        # row of H by an element of the diagonal, so we add a dummy axis for
-        # the column index
-        H *= D[..., :, None]
-        return H
-
+        q = ortho_group.rvs(dim, size, random_state)
+        dets = np.linalg.det(q)
+        if dim:
+            q[..., 0, :] /= dets[..., np.newaxis]
+        return q
 
 special_ortho_group = special_ortho_group_gen()
 
@@ -3807,9 +3763,9 @@ def __call__(self, dim=None, seed=None):
 
     def _process_parameters(self, dim):
         """Dimension N must be specified; it cannot be inferred."""
-        if dim is None or not np.isscalar(dim) or dim <= 1 or dim != int(dim):
+        if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim):
             raise ValueError("Dimension of rotation must be specified,"
-                             "and must be a scalar greater than 1.")
+                             "and must be a scalar nonnegative integer.")
 
         return dim
 
@@ -4162,7 +4118,7 @@ class unitary_group_gen(multi_rv_generic):
     Parameters
     ----------
     dim : scalar
-        Dimension of matrices, must be greater than 1.
+        Dimension of matrices.
     seed : {None, int, np.random.RandomState, np.random.Generator}, optional
         Used for drawing random variates.
         If `seed` is `None`, the `~np.random.RandomState` singleton is used.
@@ -4219,9 +4175,9 @@ def __call__(self, dim=None, seed=None):
 
     def _process_parameters(self, dim):
         """Dimension N must be specified; it cannot be inferred."""
-        if dim is None or not np.isscalar(dim) or dim <= 1 or dim != int(dim):
+        if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim):
             raise ValueError("Dimension of rotation must be specified,"
-                             "and must be a scalar greater than 1.")
+                             "and must be a scalar nonnegative integer.")
 
         return dim
 

scipy/stats/tests/test_multivariate.py

@@ -1930,21 +1930,16 @@ def test_logpdf_4x4(self):
 
 class TestSpecialOrthoGroup:
     def test_reproducibility(self):
-        np.random.seed(514)
-        x = special_ortho_group.rvs(3)
-        expected = np.array([[-0.99394515, -0.04527879, 0.10011432],
-                             [0.04821555, -0.99846897, 0.02711042],
-                             [0.09873351, 0.03177334, 0.99460653]])
-        assert_array_almost_equal(x, expected)
-
-        random_state = np.random.RandomState(seed=514)
-        x = special_ortho_group.rvs(3, random_state=random_state)
+        x = special_ortho_group.rvs(3, random_state=np.random.default_rng(514))
+        expected = np.array([[-0.93200988, 0.01533561, -0.36210826],
+                             [0.35742128, 0.20446501, -0.91128705],
+                             [0.06006333, -0.97875374, -0.19604469]])
         assert_array_almost_equal(x, expected)
 
     def test_invalid_dim(self):
         assert_raises(ValueError, special_ortho_group.rvs, None)
         assert_raises(ValueError, special_ortho_group.rvs, (2, 2))
-        assert_raises(ValueError, special_ortho_group.rvs, 1)
+        assert_raises(ValueError, special_ortho_group.rvs, -1)
         assert_raises(ValueError, special_ortho_group.rvs, 2.5)
 
     def test_frozen_matrix(self):
@@ -1979,8 +1974,9 @@ def test_haar(self):
         dim = 5
         samples = 1000  # Not too many, or the test takes too long
         ks_prob = .05
-        np.random.seed(514)
-        xs = special_ortho_group.rvs(dim, size=samples)
+        xs = special_ortho_group.rvs(
+            dim, size=samples, random_state=np.random.default_rng(513)
+        )
 
         # Dot a few rows (0, 1, 2) with unit vectors (0, 2, 4, 3),
         #   effectively picking off entries in the matrices of xs.
@@ -1997,6 +1993,14 @@ def test_haar(self):
         ks_tests = [ks_2samp(proj[p0], proj[p1])[1] for (p0, p1) in pairs]
         assert_array_less([ks_prob]*len(pairs), ks_tests)
 
+    def test_one_by_one(self):
+        # Test that the distribution is a delta function at the identity matrix
+        # when dim=1
+        assert_allclose(special_ortho_group.rvs(1, size=1000), 1, rtol=1e-13)
+
+    def test_zero_by_zero(self):
+        assert_equal(special_ortho_group.rvs(0, size=4).shape, (4, 0, 0))
+
 
 class TestOrthoGroup:
     def test_reproducibility(self):
@@ -2015,7 +2019,7 @@ def test_reproducibility(self):
     def test_invalid_dim(self):
         assert_raises(ValueError, ortho_group.rvs, None)
         assert_raises(ValueError, ortho_group.rvs, (2, 2))
-        assert_raises(ValueError, ortho_group.rvs, 1)
+        assert_raises(ValueError, ortho_group.rvs, -1)
         assert_raises(ValueError, ortho_group.rvs, 2.5)
 
     def test_frozen_matrix(self):
@@ -2085,6 +2089,20 @@ def test_haar(self):
         ks_tests = [ks_2samp(proj[p0], proj[p1])[1] for (p0, p1) in pairs]
         assert_array_less([ks_prob]*len(pairs), ks_tests)
 
+    def test_one_by_one(self):
+        # Test that the 1x1 distribution gives ±1 with equal probability.
+        dim = 1
+        xs = ortho_group.rvs(dim, size=5000, random_state=np.random.default_rng(514))
+        assert_allclose(np.abs(xs), 1, rtol=1e-13)
+        k = np.sum(xs > 0)
+        n = len(xs)
+        res = stats.binomtest(k, n)
+        low, high = res.proportion_ci(confidence_level=0.95)
+        assert low < 0.5 < high
+
+    def test_zero_by_zero(self):
+        assert_equal(special_ortho_group.rvs(0, size=4).shape, (4, 0, 0))
+
     @pytest.mark.slow
     def test_pairwise_distances(self):
         # Test that the distribution of pairwise distances is close to correct.
@@ -2290,7 +2308,7 @@ def test_reproducibility(self):
     def test_invalid_dim(self):
         assert_raises(ValueError, unitary_group.rvs, None)
         assert_raises(ValueError, unitary_group.rvs, (2, 2))
-        assert_raises(ValueError, unitary_group.rvs, 1)
+        assert_raises(ValueError, unitary_group.rvs, -1)
         assert_raises(ValueError, unitary_group.rvs, 2.5)
 
     def test_frozen_matrix(self):
@@ -2319,17 +2337,22 @@ def test_haar(self):
         # the complex plane, are uncorrelated.
 
         # Generate samples
-        dim = 5
-        samples = 1000  # Not too many, or the test takes too long
-        np.random.seed(514)  # Note that the test is sensitive to seed too
-        xs = unitary_group.rvs(dim, size=samples)
-
-        # The angles "x" of the eigenvalues should be uniformly distributed
-        # Overall this seems to be a necessary but weak test of the distribution.
-        eigs = np.vstack([scipy.linalg.eigvals(x) for x in xs])
-        x = np.arctan2(eigs.imag, eigs.real)
-        res = kstest(x.ravel(), uniform(-np.pi, 2*np.pi).cdf)
-        assert_(res.pvalue > 0.05)
+        for dim in (1, 5):
+            samples = 1000  # Not too many, or the test takes too long
+            # Note that the test is sensitive to seed too
+            xs = unitary_group.rvs(
+                dim, size=samples, random_state=np.random.default_rng(514)
+            )
+
+            # The angles "x" of the eigenvalues should be uniformly distributed
+            # Overall this seems to be a necessary but weak test of the distribution.
+            eigs = np.vstack([scipy.linalg.eigvals(x) for x in xs])
+            x = np.arctan2(eigs.imag, eigs.real)
+            res = kstest(x.ravel(), uniform(-np.pi, 2*np.pi).cdf)
+            assert_(res.pvalue > 0.05)
+
+    def test_zero_by_zero(self):
+        assert_equal(unitary_group.rvs(0, size=4).shape, (4, 0, 0))
 
 
 class TestMultivariateT: