new fixes

4 files changed, 50 insertions, 17 deletions

diff --git a/src/python/gudhi/simplex_tree.pxd b/src/python/gudhi/simplex_tree.pxd
index ae32eb82..b6284af4 100644
--- a/src/python/gudhi/simplex_tree.pxd
+++ b/src/python/gudhi/simplex_tree.pxd
@@ -57,8 +57,8 @@ 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()
-        void extend_filtration()
-        vector[vector[pair[int, pair[double, double]]]] compute_extended_persistence_subdiagrams(vector[pair[int, pair[double, double]]])
+        void compute_extended_filtration()
+        vector[vector[pair[int, pair[double, double]]]] compute_extended_persistence_subdiagrams(vector[pair[int, pair[double, double]]] dgm)
         # Iterators over Simplex tree
         pair[vector[int], double] get_simplex_and_filtration(Simplex_tree_simplex_handle f_simplex)
         Simplex_tree_simplices_iterator get_simplices_iterator_begin()
diff --git a/src/python/gudhi/simplex_tree.pyx b/src/python/gudhi/simplex_tree.pyx
index 7af44683..3502000a 100644
--- a/src/python/gudhi/simplex_tree.pyx
+++ b/src/python/gudhi/simplex_tree.pyx
@@ -405,7 +405,7 @@ cdef class SimplexTree:
 
             Note that after calling this function, the filtration 
             values are actually modified within the Simplex_tree. 
-            The function :func:`compute_extended_persistence_subdiagrams()<gudhi.SimplexTree.compute_extended_persistence_subdiagrams>`
+            The function :func:`extended_persistence()<gudhi.SimplexTree.extended_persistence>`
             retrieves the original values.
 
         .. note::
@@ -413,21 +413,31 @@ cdef class SimplexTree:
             Note that this code creates an extra vertex internally, so you should make sure that
             the Simplex_tree does not contain a vertex with the largest Vertex_handle. 
         """
-        return self.get_ptr().extend_filtration()
+        return self.get_ptr().compute_extended_filtration()
 
-    def compute_extended_persistence_subdiagrams(self, dgm):
+    def extended_persistence(self, homology_coeff_field=11, min_persistence=0, persistence_dim_max = False):
         """This function retrieves good values for extended persistence, and separate the diagrams 
         into the ordinary, relative, extended+ and extended- subdiagrams.
 
-        :param dgm: Persistence diagram obtained after calling :func:`extend_filtration()<gudhi.SimplexTree.extend_filtration>`, :func:`initialize_filtration()<gudhi.SimplexTree.initialize_filtration>`, and :func:`persistence()<gudhi.SimplexTree.persistence>`.
-
+        :param homology_coeff_field: The homology coefficient field. Must be a
+            prime number. Default value is 11.
+        :type homology_coeff_field: int.
+        :param min_persistence: The minimum persistence value to take into
+            account (strictly greater than min_persistence). Default value is
+            0.0.
+            Sets min_persistence to -1.0 to see all values.
+        :type min_persistence: float.
+        :param persistence_dim_max: If true, the persistent homology for the
+            maximal dimension in the complex is computed. If false, it is
+            ignored. Default is false.
+        :type persistence_dim_max: bool
         :returns: A vector of four persistence diagrams. The first one is Ordinary, the second one is Relative, the third one is Extended+ and the fourth one is Extended-. See section 2.2 in https://link.springer.com/article/10.1007/s10208-017-9370-z for a description of these subtypes.
 
         .. note::
 
             This function should be called only if :func:`extend_filtration()<gudhi.SimplexTree.extend_filtration>`,
             :func:`initialize_filtration()<gudhi.SimplexTree.initialize_filtration>`, 
-            and :func:`persistence()<gudhi.SimplexTree.persistence>` have been called first!
+            and (optionally) :func:`persistence()<gudhi.SimplexTree.persistence>` have been called first!
 
         .. note::
 
@@ -435,7 +445,15 @@ cdef class SimplexTree:
             original filtration values due to the internal transformation (scaling to [-2,-1]) that is 
             performed on these values during the computation of extended persistence.
         """
-        return self.get_ptr().compute_extended_persistence_subdiagrams(dgm)
+        cdef vector[pair[int, pair[double, double]]] persistence_result
+        if self.pcohptr == NULL:
+            self.pcohptr = new Simplex_tree_persistence_interface(self.get_ptr(), persistence_dim_max)
+            if self.pcohptr != NULL:
+                self.pcohptr.get_persistence(homology_coeff_field, min_persistence)
+        if self.pcohptr != NULL:
+            pairs = self.pcohptr.persistence_pairs()
+            persistence_result = [(len(splx1)-1, [self.filtration(splx1), self.filtration(splx2)]) for [splx1, splx2] in pairs]
+        return self.get_ptr().compute_extended_persistence_subdiagrams(persistence_result)
 
 
     def persistence(self, homology_coeff_field=11, min_persistence=0, persistence_dim_max = False):
diff --git a/src/python/include/Simplex_tree_interface.h b/src/python/include/Simplex_tree_interface.h
index 4a7062d6..50ed58d0 100644
--- a/src/python/include/Simplex_tree_interface.h
+++ b/src/python/include/Simplex_tree_interface.h
@@ -37,8 +37,12 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
   using Filtered_simplices = std::vector<Simplex_and_filtration>;
   using Skeleton_simplex_iterator = typename Base::Skeleton_simplex_iterator;
   using Complex_simplex_iterator = typename Base::Complex_simplex_iterator;
+  using Extended_filtration_data = typename Base::Extended_filtration_data;
 
  public:
+
+  Extended_filtration_data efd;
+  
   bool find_simplex(const Simplex& vh) {
     return (Base::find(vh) != Base::null_simplex());
   }
@@ -117,6 +121,15 @@ class Simplex_tree_interface : public Simplex_tree<SimplexTreeOptions> {
     return cofaces;
   }
 
+  void compute_extended_filtration() {
+    this->efd = this->extend_filtration();
+    return;
+  }
+
+  std::vector<std::vector<std::pair<int, std::pair<Filtration_value, Filtration_value>>>> compute_extended_persistence_subdiagrams(const std::vector<std::pair<int, std::pair<Filtration_value, Filtration_value>>>& dgm){
+    return this->extended_persistence_subdiagrams(dgm, this->efd);
+  }
+
   void create_persistence(Gudhi::Persistent_cohomology_interface<Base>* pcoh) {
     Base::initialize_filtration();
     pcoh = new Gudhi::Persistent_cohomology_interface<Base>(*this);
diff --git a/src/python/test/test_simplex_tree.py b/src/python/test/test_simplex_tree.py
index 63eee9a5..20f6aabf 100755
--- a/src/python/test/test_simplex_tree.py
+++ b/src/python/test/test_simplex_tree.py
@@ -9,6 +9,7 @@
 """
 
 from gudhi import SimplexTree
+import pytest
 
 __author__ = "Vincent Rouvreau"
 __copyright__ = "Copyright (C) 2016 Inria"
@@ -322,15 +323,16 @@ def test_extend_filtration():
         ([0, 3, 6], 2.0)
     ]
 
+    dgms = st.extended_persistence()
 
-    dgm = st.persistence()
-    L = st.compute_extended_persistence_subdiagrams(dgm)
-    assert L == [
-        [(0, (1.9999999999999998, 2.9999999999999996))], 
-        [(1, (5.0, 4.0))], 
-        [(0, (1.0, 6.0))], 
-        [(1, (6.0, 1.0))]
-    ]
+    assert dgms[0][0][1][0] == pytest.approx(2.)
+    assert dgms[0][0][1][1] == pytest.approx(3.)
+    assert dgms[1][0][1][0] == pytest.approx(5.)
+    assert dgms[1][0][1][1] == pytest.approx(4.)
+    assert dgms[2][0][1][0] == pytest.approx(1.)
+    assert dgms[2][0][1][1] == pytest.approx(6.)
+    assert dgms[3][0][1][0] == pytest.approx(6.)
+    assert dgms[3][0][1][1] == pytest.approx(1.) 
 
 
 def test_simplices_iterator():