simplex tree skeleton iterator

4 files changed, 29 insertions, 27 deletions

diff --git a/src/python/gudhi/simplex_tree.pxd b/src/python/gudhi/simplex_tree.pxd
index 1b0dc881..66c173a6 100644
--- a/src/python/gudhi/simplex_tree.pxd
+++ b/src/python/gudhi/simplex_tree.pxd
@@ -24,6 +24,13 @@ cdef extern from "Simplex_tree_interface.h" namespace "Gudhi":
     cdef cppclass Simplex_tree_simplex_handle "Gudhi::Simplex_tree_interface<Gudhi::Simplex_tree_options_full_featured>::Simplex_handle":
         pass
 
+    cdef cppclass Simplex_tree_skeleton_iterator "Gudhi::Simplex_tree_interface<Gudhi::Simplex_tree_options_full_featured>::Skeleton_simplex_iterator":
+        Simplex_tree_skeleton_iterator()
+        Simplex_tree_simplex_handle& operator*()
+        Simplex_tree_skeleton_iterator operator++()
+        bint operator!=(Simplex_tree_skeleton_iterator)
+
+
     cdef cppclass Simplex_tree_interface_full_featured "Gudhi::Simplex_tree_interface<Gudhi::Simplex_tree_options_full_featured>":
         Simplex_tree()
         double simplex_filtration(vector[int] simplex)
@@ -37,7 +44,6 @@ cdef extern from "Simplex_tree_interface.h" namespace "Gudhi":
         bool find_simplex(vector[int] simplex)
         bool insert_simplex_and_subfaces(vector[int] simplex,
                                          double filtration)
-        vector[pair[vector[int], double]] get_skeleton(int dimension)
         vector[pair[vector[int], double]] get_star(vector[int] simplex)
         vector[pair[vector[int], double]] get_cofaces(vector[int] simplex,
                                                           int dimension)
@@ -49,6 +55,8 @@ cdef extern from "Simplex_tree_interface.h" namespace "Gudhi":
         pair[vector[int], double] get_simplex_and_filtration(Simplex_tree_simplex_handle f_simplex)
         vector[Simplex_tree_simplex_handle].const_iterator get_filtration_iterator_begin()
         vector[Simplex_tree_simplex_handle].const_iterator get_filtration_iterator_end()
+        Simplex_tree_skeleton_iterator get_skeleton_iterator_begin(int dimension)
+        Simplex_tree_skeleton_iterator get_skeleton_iterator_end(int dimension)
 
 cdef extern from "Persistent_cohomology_interface.h" namespace "Gudhi":
     cdef cppclass Simplex_tree_persistence_interface "Gudhi::Persistent_cohomology_interface<Gudhi::Simplex_tree<Gudhi::Simplex_tree_options_full_featured>>":
diff --git a/src/python/gudhi/simplex_tree.pyx b/src/python/gudhi/simplex_tree.pyx
index 308b3d2d..efac2d80 100644
--- a/src/python/gudhi/simplex_tree.pyx
+++ b/src/python/gudhi/simplex_tree.pyx
@@ -231,15 +231,12 @@ cdef class SimplexTree:
         :returns:  The (simplices of the) skeleton of a maximum dimension.
         :rtype:  list of tuples(simplex, filtration)
         """
-        cdef vector[pair[vector[int], double]] skeleton \
-            = self.get_ptr().get_skeleton(<int>dimension)
-        ct = []
-        for filtered_simplex in skeleton:
-            v = []
-            for vertex in filtered_simplex.first:
-                v.append(vertex)
-            ct.append((v, filtered_simplex.second))
-        return ct
+        cdef Simplex_tree_skeleton_iterator it = self.get_ptr().get_skeleton_iterator_begin(dimension)
+        cdef Simplex_tree_skeleton_iterator end = self.get_ptr().get_skeleton_iterator_end(dimension)
+
+        while it != end:
+            yield self.get_ptr().get_simplex_and_filtration(dereference(it))
+            preincrement(it)
 
     def get_star(self, simplex):
         """This function returns the star of a given N-simplex.
diff --git a/src/python/include/Simplex_tree_interface.h b/src/python/include/Simplex_tree_interface.h
index 878919cc..55d5af97 100644
--- a/src/python/include/Simplex_tree_interface.h
+++ b/src/python/include/Simplex_tree_interface.h
@@ -35,6 +35,7 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
   using Simplex = std::vector<Vertex_handle>;
   using Simplex_and_filtration = std::pair<Simplex, Filtration_value>;
   using Filtered_simplices = std::vector<Simplex_and_filtration>;
+  using Skeleton_simplex_iterator = typename Base::Skeleton_simplex_iterator;
 
  public:
   bool find_simplex(const Simplex& vh) {
@@ -91,18 +92,6 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
     return std::make_pair(std::move(simplex), Base::filtration(f_simplex));
   }
 
-  Filtered_simplices get_skeleton(int dimension) {
-    Filtered_simplices skeletons;
-    for (auto f_simplex : Base::skeleton_simplex_range(dimension)) {
-      Simplex simplex;
-      for (auto vertex : Base::simplex_vertex_range(f_simplex)) {
-        simplex.insert(simplex.begin(), vertex);
-      }
-      skeletons.push_back(std::make_pair(simplex, Base::filtration(f_simplex)));
-    }
-    return skeletons;
-  }
-
   Filtered_simplices get_star(const Simplex& simplex) {
     Filtered_simplices star;
     for (auto f_simplex : Base::star_simplex_range(Base::find(simplex))) {
@@ -134,13 +123,21 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
 
   // Iterator over the simplex tree
   typename std::vector<Simplex_handle>::const_iterator get_filtration_iterator_begin() {
-    Base::initialize_filtration();
+    // Base::initialize_filtration(); already performed in filtration_simplex_range
     return Base::filtration_simplex_range().begin();
   }
 
   typename std::vector<Simplex_handle>::const_iterator get_filtration_iterator_end() {
     return Base::filtration_simplex_range().end();
   }
+
+  Skeleton_simplex_iterator get_skeleton_iterator_begin(int dimension) {
+    return Base::skeleton_simplex_range(dimension).begin();
+  }
+
+  Skeleton_simplex_iterator get_skeleton_iterator_end(int dimension) {
+    return Base::skeleton_simplex_range(dimension).end();
+  }
 };
 
 }  // namespace Gudhi
diff --git a/src/python/test/test_simplex_tree.py b/src/python/test/test_simplex_tree.py
index fa42f2ac..eca3807b 100755
--- a/src/python/test/test_simplex_tree.py
+++ b/src/python/test/test_simplex_tree.py
@@ -56,7 +56,7 @@ def test_insertion():
     assert st.filtration([1]) == 0.0
 
     # skeleton test
-    assert st.get_skeleton(2) == [
+    assert list(st.get_skeleton(2)) == [
         ([0, 1, 2], 4.0),
         ([0, 1], 0.0),
         ([0, 2], 4.0),
@@ -65,7 +65,7 @@ def test_insertion():
         ([1], 0.0),
         ([2], 4.0),
     ]
-    assert st.get_skeleton(1) == [
+    assert list(st.get_skeleton(1)) == [
         ([0, 1], 0.0),
         ([0, 2], 4.0),
         ([0], 0.0),
@@ -73,12 +73,12 @@ def test_insertion():
         ([1], 0.0),
         ([2], 4.0),
     ]
-    assert st.get_skeleton(0) == [([0], 0.0), ([1], 0.0), ([2], 4.0)]
+    assert list(st.get_skeleton(0)) == [([0], 0.0), ([1], 0.0), ([2], 4.0)]
 
     # remove_maximal_simplex test
     assert st.get_cofaces([0, 1, 2], 1) == []
     st.remove_maximal_simplex([0, 1, 2])
-    assert st.get_skeleton(2) == [
+    assert list(st.get_skeleton(2)) == [
         ([0, 1], 0.0),
         ([0, 2], 4.0),
         ([0], 0.0),