A first naive iterator implementation with yield

4 files changed, 37 insertions, 21 deletions

diff --git a/src/Simplex_tree/example/simple_simplex_tree.cpp b/src/Simplex_tree/example/simple_simplex_tree.cpp
index 4353939f..92ab923b 100644
--- a/src/Simplex_tree/example/simple_simplex_tree.cpp
+++ b/src/Simplex_tree/example/simple_simplex_tree.cpp
@@ -165,6 +165,10 @@ int main(int argc, char* const argv[]) {
 
   // ++ GENERAL VARIABLE SET
 
+  //std::vector<Simplex_tree::Simplex_handle>::const_iterator
+  std::vector<Simplex_tree::Simplex_handle>::const_iterator begin = simplexTree.filtration_simplex_range().begin();
+  auto end = simplexTree.filtration_simplex_range().end();
+  
   std::cout << "********************************************************************\n";
   // Display the Simplex_tree - Can not be done in the middle of 2 inserts
   std::cout << "* The complex contains " << simplexTree.num_simplices() << " simplices\n";
diff --git a/src/python/gudhi/simplex_tree.pxd b/src/python/gudhi/simplex_tree.pxd
index 96d14079..caf3c459 100644
--- a/src/python/gudhi/simplex_tree.pxd
+++ b/src/python/gudhi/simplex_tree.pxd
@@ -21,6 +21,9 @@ cdef extern from "Simplex_tree_interface.h" namespace "Gudhi":
     cdef cppclass Simplex_tree_options_full_featured:
         pass
 
+    cdef cppclass Simplex_tree_simplex_handle "Gudhi::Simplex_tree_interface<Gudhi::Simplex_tree_options_full_featured>::Simplex_handle":
+        pass
+
     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)
@@ -34,7 +37,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_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,
@@ -43,6 +45,10 @@ cdef extern from "Simplex_tree_interface.h" namespace "Gudhi":
         void remove_maximal_simplex(vector[int] simplex)
         bool prune_above_filtration(double filtration)
         bool make_filtration_non_decreasing()
+        # Iterators over Simplex tree
+        pair[vector[int], double] get_simplex_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()
 
 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 b18627c4..478139de 100644
--- a/src/python/gudhi/simplex_tree.pyx
+++ b/src/python/gudhi/simplex_tree.pyx
@@ -7,6 +7,7 @@
 # Modification(s):
 #   - YYYY/MM Author: Description of the modification
 
+from cython.operator import dereference, preincrement
 from libc.stdint cimport intptr_t
 from numpy import array as np_array
 cimport simplex_tree
@@ -214,15 +215,14 @@ cdef class SimplexTree:
         :returns:  The simplices sorted by increasing filtration values.
         :rtype:  list of tuples(simplex, filtration)
         """
-        cdef vector[pair[vector[int], double]] filtration \
-            = self.get_ptr().get_filtration()
-        ct = []
-        for filtered_complex in filtration:
-            v = []
-            for vertex in filtered_complex.first:
-                v.append(vertex)
-            ct.append((v, filtered_complex.second))
-        return ct
+        cdef vector[Simplex_tree_simplex_handle].const_iterator it = self.get_ptr().get_filtration_iterator_begin()
+        cdef vector[Simplex_tree_simplex_handle].const_iterator end = self.get_ptr().get_filtration_iterator_end()
+
+        while True:
+            yield(self.get_ptr().get_simplex_filtration(dereference(it)))
+            preincrement(it)
+            if it == end:
+                raise StopIteration
 
     def get_skeleton(self, dimension):
         """This function returns the (simplices of the) skeleton of a maximum
diff --git a/src/python/include/Simplex_tree_interface.h b/src/python/include/Simplex_tree_interface.h
index 06f31341..843966cd 100644
--- a/src/python/include/Simplex_tree_interface.h
+++ b/src/python/include/Simplex_tree_interface.h
@@ -33,7 +33,8 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
   using Simplex_handle = typename Base::Simplex_handle;
   using Insertion_result = typename std::pair<Simplex_handle, bool>;
   using Simplex = std::vector<Vertex_handle>;
-  using Filtered_simplices = std::vector<std::pair<Simplex, Filtration_value>>;
+  using Filtered_simplex = std::pair<Simplex, Filtration_value>;
+  using Filtered_simplices = std::vector<Filtered_simplex>;
 
  public:
   bool find_simplex(const Simplex& vh) {
@@ -82,17 +83,12 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
     Base::initialize_filtration();
   }
 
-  Filtered_simplices get_filtration() {
-    Base::initialize_filtration();
-    Filtered_simplices filtrations;
-    for (auto f_simplex : Base::filtration_simplex_range()) {
-      Simplex simplex;
-      for (auto vertex : Base::simplex_vertex_range(f_simplex)) {
-        simplex.insert(simplex.begin(), vertex);
-      }
-      filtrations.push_back(std::make_pair(simplex, Base::filtration(f_simplex)));
+  Filtered_simplex get_simplex_filtration(Simplex_handle f_simplex) {
+    Simplex simplex;
+    for (auto vertex : Base::simplex_vertex_range(f_simplex)) {
+      simplex.insert(simplex.begin(), vertex);
     }
-    return filtrations;
+    return std::make_pair(simplex, Base::filtration(f_simplex));
   }
 
   Filtered_simplices get_skeleton(int dimension) {
@@ -135,6 +131,16 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
     Base::initialize_filtration();
     pcoh = new Gudhi::Persistent_cohomology_interface<Base>(*this);
   }
+
+  // Iterator over the simplex tree
+  typename std::vector<Simplex_handle>::const_iterator get_filtration_iterator_begin() {
+    Base::initialize_filtration();
+    return Base::filtration_simplex_range().begin();
+  }
+
+  typename std::vector<Simplex_handle>::const_iterator get_filtration_iterator_end() {
+    return Base::filtration_simplex_range().end();
+  }
 };
 
 }  // namespace Gudhi