Make gensim.matutils.normal_factors always return a single factor

gensim/matutils.py

@@ -1422,8 +1422,8 @@ def close(self):
             self.fout.close()
 
 
-def normal_factors(size, rand_obj, loc=0.0, scale=1.0, k=2, iterations=30):
-    r"""Draw random samples from independent k-factors of the normal distribution.
+def normal_factor(size, rand_obj, loc=0.0, scale=1.0, k=2, iterations=30):
+    r"""Draw random samples from a k-factor of the normal distribution.
 
     The product of k independent k-factors is normally distributed. [Pinelis2018]_
     This is useful e.g. for initializing neural network weights that are pointwise multiplied and
@@ -1437,12 +1437,12 @@ def normal_factors(size, rand_obj, loc=0.0, scale=1.0, k=2, iterations=30):
     rand_obj : np.random.Generator
         A random generator object.
     loc : float, optional
-        The mean of the normal distribution. If negative, the sample of the first k-factor will be
-        negated to make the factor product also negative. Default is 0.
+        The mean of the normal distribution. Must be non-negative; if you want to produce a normal
+        distribution with a negative mean, you should negate the first k-factor. Default is 0.
     scale : float, optional
         The standard deviation of the normal distribution. Must be non-negative. Default is 1.
     k : int, optional
-        The number of k-factors. Must be non-negative. Default is 2.
+        The degree of the k-factor. Must be non-negative. Default is 2.
     iterations : int, optional
         The formula of [Pinelis2018]_ contains a sum of a convergent infinite sequence. We approximate
         the sum by a finite sequence of its ``iterations`` initial elements. Increasing the number of
@@ -1451,13 +1451,14 @@ def normal_factors(size, rand_obj, loc=0.0, scale=1.0, k=2, iterations=30):
 
     Returns
     -------
-    factor_samples : tuple of ndarray
-        Random samples drawn from the k independent parametrized k-factors of the normal distribution.
+    out : ndarray
+        Random samples drawn from a k-factor of the normal distribution.
 
     Raises
     ------
     ValueError
-        If the standard deviation, the number of k-factors, or the number of iterations are negative.
+        If the mean, the standard deviation, the number of k-factors, or the number of iterations
+        are negative.
 
     Notes
     -----
@@ -1471,37 +1472,33 @@ def normal_factors(size, rand_obj, loc=0.0, scale=1.0, k=2, iterations=30):
        <https://arxiv.org/abs/1803.09838>`_ *arXiv preprint arXiv:1803.09838*.
 
     """
+    if loc < 0.0:
+        raise ValueError('Scale must be non-negative. If you want to produce a normal distribution '
+                         'with a negative mean, you should negate the first k-factor.')
     if scale < 0.0:
         raise ValueError('Standard deviation must be non-negative')
     if k < 0:
         raise ValueError('The number of k-factors must be non-negative')
     if iterations < 0:
         raise ValueError('The number of iterations must be non-negative')
 
+    def gamma(size):
+        return rand_obj.gamma(1.0 / k, 1.0, size)
+
     if not isinstance(size, tuple):
         size = (size,)
 
     arange = np.arange(iterations) + 1
     exponent = 1.0 / k
-
-    def gamma(size):
-        return rand_obj.gamma(1.0 / k, 1.0, size)
-
-    def factor_sample(sample_number):
-        log_value = np.log(2.0) / (2.0 * k)
-        log_value -= gamma(size)
-        log_value -= np.sum(gamma((*size, iterations)) / (2.0 * arange + 1.0), axis=-1)
-        log_value += np.sum(np.log(1.0 + 1.0 / arange) / (2.0 * k))
-        sample = (rand_obj.uniform(0.0, 1.0, size) < 0.5) * 2.0 - 1.0
-        sample *= np.exp(log_value)
-        sample *= scale**exponent
-        sample += np.abs(loc)**exponent
-        if sample_number == 0 and loc < 0.0:
-            sample *= -1.0
-        return sample
-
-    factor_samples = tuple(factor_sample(i) for i in range(k))
-    return factor_samples
+    log_value = np.log(2.0) / (2.0 * k)
+    log_value -= gamma(size)
+    log_value -= np.sum(gamma((*size, iterations)) / (2.0 * arange + 1.0), axis=-1)
+    log_value += np.sum(np.log(1.0 + 1.0 / arange) / (2.0 * k))
+    sample = (rand_obj.uniform(0.0, 1.0, size) < 0.5) * 2.0 - 1.0
+    sample *= np.exp(log_value)
+    sample *= scale**exponent
+    sample += np.abs(loc)**exponent
+    return sample
 
 
 try:

gensim/test/test_matutils.py

@@ -276,111 +276,99 @@ def setUp(self):
 
     def test_1factor_product_distribution(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, = matutils.normal_factors(self.size, rand_obj, k=1)
+        x = matutils.normal_factor(self.size, rand_obj, k=1)
         factor_product = np.ravel(x)
         _, p_value = stats.kstest(factor_product, 'norm')
         self.assertGreater(p_value, self.alpha)
 
     def test_2factor_product_distribution(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y = matutils.normal_factors(self.size, rand_obj)
+        x = matutils.normal_factor(self.size, rand_obj)
+        y = matutils.normal_factor(self.size, rand_obj)
         factor_product = np.ravel(x * y)
         _, p_value = stats.kstest(factor_product, 'norm')
         self.assertGreater(p_value, self.alpha)
 
     def test_3factor_product_distribution(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y, z = matutils.normal_factors(self.size, rand_obj, k=3)
+        x = matutils.normal_factor(self.size, rand_obj, k=3)
+        y = matutils.normal_factor(self.size, rand_obj, k=3)
+        z = matutils.normal_factor(self.size, rand_obj, k=3)
         factor_product = np.ravel(x * y * z)
         _, p_value = stats.kstest(factor_product, 'norm')
         self.assertGreater(p_value, self.alpha)
 
     def test_2factor_tuple_size(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        factors = matutils.normal_factors(self.size, rand_obj)
+        factor = matutils.normal_factor(self.size, rand_obj)
         expected_shape = self.size
-        self.assertEqual(factors[0].shape, expected_shape)
-        self.assertEqual(factors[1].shape, expected_shape)
+        self.assertEqual(factor.shape, expected_shape)
 
     def test_2factor_int_size(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        factors = matutils.normal_factors(self.size[0], rand_obj)
+        factor = matutils.normal_factor(self.size[0], rand_obj)
         expected_shape = (self.size[0],)
-        self.assertEqual(factors[0].shape, expected_shape)
-        self.assertEqual(factors[1].shape, expected_shape)
+        self.assertEqual(factor.shape, expected_shape)
+
+    def test_negative_loc(self):
+        rand_obj = np.random.default_rng(seed=self.seed)
+        with self.assertRaises(ValueError):
+            matutils.normal_factor(self.size, rand_obj, loc=-self.loc)
 
     def test_1factor_positive_loc(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, = matutils.normal_factors(self.size, rand_obj, loc=self.loc, k=1)
+        x = matutils.normal_factor(self.size, rand_obj, loc=self.loc, k=1)
         factor_product = np.ravel(x)
         expected_loc = self.loc
         actual_loc = np.mean(factor_product)
         self.assertAlmostEqual(expected_loc, actual_loc, places=0)
 
     def test_2factor_positive_loc(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y = matutils.normal_factors(self.size, rand_obj, loc=self.loc)
+        x = matutils.normal_factor(self.size, rand_obj, loc=self.loc)
+        y = matutils.normal_factor(self.size, rand_obj, loc=self.loc)
         factor_product = np.ravel(x * y)
         expected_loc = self.loc
         actual_loc = np.mean(factor_product)
         self.assertAlmostEqual(expected_loc, actual_loc, places=0)
 
     def test_3factor_positive_loc(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y, z = matutils.normal_factors(self.size, rand_obj, loc=self.loc, k=3)
+        x = matutils.normal_factor(self.size, rand_obj, loc=self.loc, k=3)
+        y = matutils.normal_factor(self.size, rand_obj, loc=self.loc, k=3)
+        z = matutils.normal_factor(self.size, rand_obj, loc=self.loc, k=3)
         factor_product = np.ravel(x * y * z)
         expected_loc = self.loc
         actual_loc = np.mean(factor_product)
         self.assertAlmostEqual(expected_loc, actual_loc, places=0)
 
-    def test_1factor_negative_loc(self):
-        rand_obj = np.random.default_rng(seed=self.seed)
-        x, = matutils.normal_factors(self.size, rand_obj, loc=-self.loc, k=1)
-        factor_product = np.ravel(x)
-        expected_loc = -self.loc
-        actual_loc = np.mean(factor_product)
-        self.assertAlmostEqual(expected_loc, actual_loc, places=0)
-
-    def test_2factor_negative_loc(self):
-        rand_obj = np.random.default_rng(seed=self.seed)
-        x, y = matutils.normal_factors(self.size, rand_obj, loc=-self.loc)
-        factor_product = np.ravel(x * y)
-        expected_loc = -self.loc
-        actual_loc = np.mean(factor_product)
-        self.assertAlmostEqual(expected_loc, actual_loc, places=0)
-
-    def test_3factor_negative_loc(self):
-        rand_obj = np.random.default_rng(seed=self.seed)
-        x, y, z = matutils.normal_factors(self.size, rand_obj, loc=-self.loc, k=3)
-        factor_product = np.ravel(x * y * z)
-        expected_loc = -self.loc
-        actual_loc = np.mean(factor_product)
-        self.assertAlmostEqual(expected_loc, actual_loc, places=0)
-
     def test_negative_scale(self):
         rand_obj = np.random.default_rng(seed=self.seed)
         with self.assertRaises(ValueError):
-            matutils.normal_factors(self.size, rand_obj, scale=-self.scale, k=1)
+            matutils.normal_factor(self.size, rand_obj, scale=-self.scale)
 
     def test_1factor_positive_scale(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, = matutils.normal_factors(self.size, rand_obj, scale=self.scale, k=1)
+        x = matutils.normal_factor(self.size, rand_obj, scale=self.scale, k=1)
         factor_product = np.ravel(x)
         expected_scale = self.scale
         actual_scale = np.std(factor_product)
         self.assertAlmostEqual(expected_scale, actual_scale, places=0)
 
     def test_2factor_positive_scale(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y = matutils.normal_factors(self.size, rand_obj, scale=self.scale)
+        x = matutils.normal_factor(self.size, rand_obj, scale=self.scale)
+        y = matutils.normal_factor(self.size, rand_obj, scale=self.scale)
         factor_product = np.ravel(x * y)
         expected_scale = self.scale
         actual_scale = np.std(factor_product)
         self.assertAlmostEqual(expected_scale, actual_scale, places=0)
 
     def test_3factor_positive_scale(self):
         rand_obj = np.random.default_rng(seed=self.seed)
-        x, y, z = matutils.normal_factors(self.size, rand_obj, scale=self.scale, k=3)
+        x = matutils.normal_factor(self.size, rand_obj, scale=self.scale, k=3)
+        y = matutils.normal_factor(self.size, rand_obj, scale=self.scale, k=3)
+        z = matutils.normal_factor(self.size, rand_obj, scale=self.scale, k=3)
         factor_product = np.ravel(x * y * z)
         expected_scale = self.scale
         actual_scale = np.std(factor_product)