Add comments and some minor changes following code review

4 files changed, 13 insertions, 10 deletions

diff --git a/src/python/doc/datasets_generators.rst b/src/python/doc/datasets_generators.rst
index c0bbb973..3700b8a2 100644
--- a/src/python/doc/datasets_generators.rst
+++ b/src/python/doc/datasets_generators.rst
@@ -48,22 +48,23 @@ You can also generate points on a torus.
 
 Two functions are available and give the same output: the first one depends on **CGAL** and the second does not and consists of full python code.
 
-On another hand, two sample types are provided : you can either generate i.i.d. points on a d-torus in :math:`R^{2d}` *randomly* or on a *grid*.
+On another hand, two sample types are provided: you can either generate i.i.d. points on a d-torus in :math:`R^{2d}` *randomly* or on a *grid*.
 
-First function : **ctorus**
+First function: **ctorus**
 """""""""""""""""""""""""""
 
 The user should provide the number of points to be generated on the torus :code:`n_samples`, and the dimension :code:`dim` of the torus on which points would be generated in :math:`R^{2dim}`.
 The :code:`sample` argument is optional and is set to **'random'** by default.
 In this case, the returned generated points would be an array of shape :math:`(n\_samples, 2*dim)`.
-Otherwise, if set to **'grid'**, the points are generated on a grid and would be given as an array of shape :
+Otherwise, if set to **'grid'**, the points are generated on a grid and would be given as an array of shape:
 
 .. math::
 
-   ( [n\_samples^{1 \over {dim}}]^{dim}, 2*dim )
+   ( ⌊n\_samples^{1 \over {dim}}⌋^{dim}, 2*dim )
 
+**Note 1:** The output array first shape is rounded down to the closest perfect :math:`dim^{th}` power.
 
-**Note:** This version is recommended when the user wishes to use **'grid'** as sample type, or **'random'** with a relatively small number of samples (~ less than 150).
+**Note 2:** This version is recommended when the user wishes to use **'grid'** as sample type, or **'random'** with a relatively small number of samples (~ less than 150).
 
 Example
 """""""
@@ -79,7 +80,7 @@ Example
 
 .. autofunction:: gudhi.datasets.generators.points.ctorus
 
-Second function : **torus**
+Second function: **torus**
 """""""""""""""""""""""""""
 
 The user should provide the number of points to be generated on the torus :code:`n_samples` and the dimension :code:`dim` of the torus on which points would be generated in :math:`R^{2dim}`.
diff --git a/src/python/gudhi/datasets/generators/_points.cc b/src/python/gudhi/datasets/generators/_points.cc
index 536fa949..5d675930 100644
--- a/src/python/gudhi/datasets/generators/_points.cc
+++ b/src/python/gudhi/datasets/generators/_points.cc
@@ -110,7 +110,7 @@ PYBIND11_MODULE(_points, m) {
           :rtype: numpy array of float.
           The shape of returned numpy array is :
               if sample is 'random' : (n_samples, 2*dim).
-              if sample is 'grid' : ([n_samples**(1./dim)]**dim, 2*dim).
+              if sample is 'grid' : (⌊n_samples**(1./dim)⌋**dim, 2*dim), where shape[0] is rounded down to the closest perfect 'dim'th power.
           :returns: the generated points on a torus.
           )pbdoc");
 }
diff --git a/src/python/gudhi/datasets/generators/points.py b/src/python/gudhi/datasets/generators/points.py
index 1995f769..7f4667af 100644
--- a/src/python/gudhi/datasets/generators/points.py
+++ b/src/python/gudhi/datasets/generators/points.py
@@ -36,15 +36,15 @@ def _generate_grid_points_on_torus(n_samples, dim):
 
 def torus(n_samples, dim, sample='random'):
     """ 
-    Generate points on a dim-torus in R^2dim either randomly or on a grid
+    Generate points on a flat dim-torus in R^2dim either randomly or on a grid
     
     :param n_samples: The number of points to be generated.
     :param dim: The dimension of the torus on which points would be generated in R^2*dim.
     :param sample: The sample type of the generated points. Can be 'random' or 'grid'.
     :returns: numpy array containing the generated points on a torus.
-        The shape of returned numpy array is :
+        The shape of returned numpy array is:
         if sample is 'random' : (n_samples, 2*dim).
-        if sample is 'grid' : ([n_samples**(1./dim)]**dim, 2*dim).
+        if sample is 'grid' : (⌊n_samples**(1./dim)⌋**dim, 2*dim), where shape[0] is rounded down to the closest perfect 'dim'th power.
     """
     if sample == 'random':
         # Generate points randomly
diff --git a/src/python/test/test_datasets_generators.py b/src/python/test/test_datasets_generators.py
index e2d300e0..933a763e 100755
--- a/src/python/test/test_datasets_generators.py
+++ b/src/python/test/test_datasets_generators.py
@@ -23,6 +23,8 @@ def _basic_torus(impl):
     assert impl(n_samples = 64, dim = 3, sample = 'grid').shape == (64, 6)
 
     assert impl(n_samples = 10, dim = 4, sample = 'random').shape == (10, 8)
+
+    # Here 1**dim < n_samples < 2**dim, the output shape is therefore (1, 2*dim) = (1, 8), where shape[0] is rounded down to the closest perfect 'dim'th power
     assert impl(n_samples = 10, dim = 4, sample = 'grid').shape == (1, 8)
 
     with pytest.raises(ValueError):