Add documentation to python torus fonction and apply some modifications according to PR comments

2 files changed, 20 insertions, 6 deletions

diff --git a/src/python/gudhi/datasets/generators/_points.cc b/src/python/gudhi/datasets/generators/_points.cc
index 561fd6d8..003b65a3 100644
--- a/src/python/gudhi/datasets/generators/_points.cc
+++ b/src/python/gudhi/datasets/generators/_points.cc
@@ -56,7 +56,7 @@ PYBIND11_MODULE(_points, m) {
 
     :param n_samples: The number of points to be generated.
     :type n_samples: integer
-    :param dim: The dimension.
+    :param dim: The dimension of the torus on which points would be generated in R^2*dim.
     :type dim: integer
     :param uniform: A flag to define if the points generation is uniform (i.e generated as a grid).
     :type uniform: bool
diff --git a/src/python/gudhi/datasets/generators/points.py b/src/python/gudhi/datasets/generators/points.py
index d5a370ad..a8f5ad54 100644
--- a/src/python/gudhi/datasets/generators/points.py
+++ b/src/python/gudhi/datasets/generators/points.py
@@ -10,7 +10,7 @@
 import numpy as np
 import itertools
 
-def _generate_random_points(n_samples, dim):
+def _generate_random_points_on_torus(n_samples, dim):
 
     # Generate random angles of size n_samples*dim
     alpha = 2*np.pi*np.random.rand(n_samples*dim)
@@ -20,8 +20,9 @@ def _generate_random_points(n_samples, dim):
     
     return array_points
 
-def _generate_grid_points(n_samples, dim):
+def _generate_grid_points_on_torus(n_samples, dim):
     
+    # Generate points on a dim-torus as a grid
     n_samples_grid = int(n_samples**(1./dim))
     alpha = np.linspace(0, 2*np.pi, n_samples_grid, endpoint=False)
     
@@ -31,12 +32,25 @@ def _generate_grid_points(n_samples, dim):
     return array_points
 
 def torus(n_samples, dim, sample='random'):
+    ''' 
+    Generate points on a 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 :
+                if sample is 'random' : (n_samples, 2*dim)
+                if sample is 'grid' : ((int(n_samples**(1./dim)))**dim, 2*dim)
+    '''
     if sample == 'random':
+        # Generate points randomly
         print("Sample is random")
-        return _generate_random_points(n_samples, dim)
+        return _generate_random_points_on_torus(n_samples, dim)
     elif sample == 'grid':
+        # Generate points on a grid
         print("Sample is grid")
-        return _generate_grid_points(n_samples, dim)
+        return _generate_grid_points_on_torus(n_samples, dim)
     else:
-        raise Exception("Sample type '{}' is not supported".format(sample))
+        raise ValueError("Sample type '{}' is not supported".format(sample))
         return