update code and doc

1 files changed, 13 insertions, 6 deletions

diff --git a/src/python/gudhi/tensorflow/rips_layer.py b/src/python/gudhi/tensorflow/rips_layer.py
index 97f28d74..a5f212e3 100644
--- a/src/python/gudhi/tensorflow/rips_layer.py
+++ b/src/python/gudhi/tensorflow/rips_layer.py
@@ -8,7 +8,7 @@ from ..rips_complex     import RipsComplex
 
 # The parameters of the model are the point coordinates.
 
-def _Rips(DX, max_edge, dimensions, min_persistence):
+def _Rips(DX, max_edge, dimensions):
     # Parameters: DX (distance matrix), 
     #             max_edge (maximum edge length for Rips filtration), 
     #             dimensions (homology dimensions)
@@ -16,7 +16,7 @@ def _Rips(DX, max_edge, dimensions, min_persistence):
     # Compute the persistence pairs with Gudhi
     rc = RipsComplex(distance_matrix=DX, max_edge_length=max_edge)
     st = rc.create_simplex_tree(max_dimension=max(dimensions)+1)
-    st.compute_persistence(min_persistence=min_persistence)
+    st.compute_persistence()
     pairs = st.flag_persistence_generators()
 
     L_indices = []
@@ -40,18 +40,20 @@ class RipsLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing Rips persistence out of a point cloud
     """
-    def __init__(self, dimensions, maximum_edge_length=12, min_persistence=0., **kwargs):
+    def __init__(self, dimensions, maximum_edge_length=12, min_persistence=None, **kwargs):
         """
         Constructor for the RipsLayer class
 
         Parameters:
             maximum_edge_length (float): maximum edge length for the Rips complex 
             dimensions (List[int]): homology dimensions
+            min_persistence (List[float]): minimum distance-to-diagonal of the points in the output persistence diagrams (default None, in which case 0. is used for all dimensions)
         """
         super().__init__(dynamic=True, **kwargs)
         self.max_edge = maximum_edge_length
         self.dimensions = dimensions
-        self.min_persistence = min_persistence
+        self.min_persistence = min_persistence if min_persistence != None else [0. for _ in range(len(self.dimensions))]
+        assert len(self.min_persistence) == len(self.dimensions)
         
     def call(self, X):
         """
@@ -67,7 +69,7 @@ class RipsLayer(tf.keras.layers.Layer):
         DX = tf.math.sqrt(tf.reduce_sum((tf.expand_dims(X, 1)-tf.expand_dims(X, 0))**2, 2))
         # Compute vertices associated to positive and negative simplices 
         # Don't compute gradient for this operation
-        indices = _Rips(DX.numpy(), self.max_edge, self.dimensions, self.min_persistence)
+        indices = _Rips(DX.numpy(), self.max_edge, self.dimensions)
         # Get persistence diagrams by simply picking the corresponding entries in the distance matrix
         self.dgms = []
         for idx_dim, dimension in enumerate(self.dimensions):
@@ -79,6 +81,11 @@ class RipsLayer(tf.keras.layers.Layer):
                 reshaped_cur_idx = tf.reshape(cur_idx[0], [-1,3])
                 finite_dgm = tf.concat([tf.zeros([reshaped_cur_idx.shape[0],1]), tf.reshape(tf.gather_nd(DX, reshaped_cur_idx[:,1:]), [-1,1])], axis=1)
                 essential_dgm = tf.zeros([cur_idx[1].shape[0],1])
-            self.dgms.append((finite_dgm, essential_dgm))
+            min_pers = self.min_persistence[idx_dim]
+            if min_pers >= 0:
+                persistent_indices = np.argwhere(np.abs(finite_dgm[:,1]-finite_dgm[:,0]) > min_pers).ravel()
+                self.dgms.append((tf.gather(finite_dgm, indices=persistent_indices), essential_dgm))
+            else:
+                self.dgms.append((finite_dgm, essential_dgm))
         return self.dgms