persistence generators for lower-star and flag filtrations

1 files changed, 55 insertions, 0 deletions

diff --git a/src/python/gudhi/simplex_tree.pyx b/src/python/gudhi/simplex_tree.pyx
index b18627c4..1c9b9cf1 100644
--- a/src/python/gudhi/simplex_tree.pyx
+++ b/src/python/gudhi/simplex_tree.pyx
@@ -514,3 +514,58 @@ cdef class SimplexTree:
         else:
             print("intervals_in_dim function requires persistence function"
                   " to be launched first.")
+
+    def lower_star_persistence_generators(self):
+        """Assuming this is a lower-star filtration, this function returns the persistence pairs,
+        where each simplex is replaced with the vertex that gave it its filtration value.
+
+        :returns: first the regular persistence pairs, grouped by dimension, with one vertex per extremity,
+            and second the essential features, grouped by dimension, with one vertex each
+        :rtype: Tuple[List[numpy.array[int] of shape (n,2)], List[numpy.array[int] of shape (m,)]]
+
+        :note: intervals_in_dim function requires
+            :func:`persistence()<gudhi.SimplexTree.persistence>`
+            function to be launched first.
+        """
+        if self.pcohptr != NULL:
+            gen = self.pcohptr.lower_star_generators()
+            normal = [np_array(d).reshape(-1,2) for d in gen.first]
+            infinite = [np_array(d) for d in gen.second]
+            return (normal, infinite)
+        else:
+            print("lower_star_persistence_generators() requires that persistence() be called first.")
+
+    def flag_persistence_generators(self):
+        """Assuming this is a flag complex, this function returns the persistence pairs,
+        where each simplex is replaced with the vertices of the edges that gave it its filtration value.
+
+        :returns: first the regular persistence pairs of dimension 0, with one vertex for birth and two for death;
+            then the other regular persistence pairs, grouped by dimension, with 2 vertices per extremity;
+            then the connected components, with one vertex each;
+            finally the other essential features, grouped by dimension, with 2 vertices for birth.
+        :rtype: Tuple[List[numpy.array[int] of shape (n,3)], List[numpy.array[int] of shape (m,4)], List[numpy.array[int] of shape (l,)], List[numpy.array[int] of shape (k,2)]]
+
+        :note: intervals_in_dim function requires
+            :func:`persistence()<gudhi.SimplexTree.persistence>`
+            function to be launched first.
+        """
+        if self.pcohptr != NULL:
+            gen = self.pcohptr.flag_generators()
+            if len(gen.first) == 0:
+                normal0 = np_array([])
+                normals = np_array([])
+            else:
+                l = iter(gen.first)
+                normal0 = np_array(next(l)).reshape(-1,3)
+                normals = [np_array(d).reshape(-1,4) for d in l]
+            if len(gen.second) == 0:
+                infinite0 = np_array([])
+                infinites = np_array([])
+            else:
+                l = iter(gen.second)
+                infinite0 = np_array(next(l))
+                infinites = [np_array(d).reshape(-1,3) for d in l]
+
+            return (normal0, normals, infinite0, infinites)
+        else:
+            print("lower_star_persistence_generators() requires that persistence() be called first.")