modified API for multiple dimensions and finite + essential

3 files changed, 96 insertions, 71 deletions

diff --git a/src/python/gudhi/tensorflow/cubical_layer.py b/src/python/gudhi/tensorflow/cubical_layer.py
index b4ff2598..d8177864 100644
--- a/src/python/gudhi/tensorflow/cubical_layer.py
+++ b/src/python/gudhi/tensorflow/cubical_layer.py
@@ -8,41 +8,48 @@ from ..cubical_complex  import CubicalComplex
 
 # The parameters of the model are the pixel values.
 
-def _Cubical(X, dimension):
+def _Cubical(X, dimensions):
     # Parameters: X (image),
-    #             dimension (homology dimension)
+    #             dimensions (homology dimensions)
 
     # Compute the persistence pairs with Gudhi
-    cc = CubicalComplex(dimensions=X.shape, top_dimensional_cells=X.flatten())
+    Xs = X.shape
+    cc = CubicalComplex(dimensions=Xs, top_dimensional_cells=X.flatten())
     cc.persistence()
-    try:
-        cof = cc.cofaces_of_persistence_pairs()[0][dimension]
-    except IndexError:
-        cof = np.array([])
 
-    # Retrieve and ouput image indices/pixels corresponding to positive and negative simplices
-    D = len(Xs) if len(cof) > 0 else 1
-    ocof = np.array([0 for _ in range(D*2*cof.shape[0])])
-    count = 0
-    for idx in range(0,2*cof.shape[0],2):
-        ocof[D*idx:D*(idx+1)]     = np.unravel_index(cof[count,0], Xs)
-        ocof[D*(idx+1):D*(idx+2)] = np.unravel_index(cof[count,1], Xs)
-        count += 1
-    return np.array(ocof, dtype=np.int32)
+    L_cofs = []
+    for dim in dimensions:
+
+        try:
+            cof = cc.cofaces_of_persistence_pairs()[0][dim]
+        except IndexError:
+            cof = np.array([])
+
+        # Retrieve and ouput image indices/pixels corresponding to positive and negative simplices
+        D = len(Xs) if len(cof) > 0 else 1
+        ocof = np.array([0 for _ in range(D*2*cof.shape[0])])
+        count = 0
+        for idx in range(0,2*cof.shape[0],2):
+            ocof[D*idx:D*(idx+1)]     = np.unravel_index(cof[count,0], Xs)
+            ocof[D*(idx+1):D*(idx+2)] = np.unravel_index(cof[count,1], Xs)
+            count += 1
+        L_cofs.append(np.array(ocof, dtype=np.int32))
+
+    return L_cofs
 
 class CubicalLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing cubical persistence out of a cubical complex
     """
-    def __init__(self, dimension=1, **kwargs):
+    def __init__(self, dimensions=[1], **kwargs):
         """
         Constructor for the CubicalLayer class
 
         Parameters:
-            dimension (int): homology dimension
+            dimensions (list of int): homology dimensions
         """
         super().__init__(dynamic=True, **kwargs)
-        self.dimension = dimension
+        self.dimensions = dimensions
 
     def build(self):
         super.build()
@@ -55,11 +62,11 @@ class CubicalLayer(tf.keras.layers.Layer):
             X (TensorFlow variable): pixel values of the cubical complex
 
         Returns:
-            dgm (TensorFlow variable): cubical persistence diagram with shape [num_points, 2]
+            dgms (list of TensorFlow variables): list of cubical persistence diagrams of length self.dimensions, where each element contains a finite persistence diagram of shape [num_finite_points, 2]
         """
         # Compute pixels associated to positive and negative simplices 
         # Don't compute gradient for this operation
-        indices = tf.stop_gradient(_Cubical(X.numpy(), self.dimension))
+        indices = _Cubical(X.numpy(), self.dimensions)
         # Get persistence diagram by simply picking the corresponding entries in the image
-        dgm = tf.reshape(tf.gather_nd(X, tf.reshape(indices, [-1,len(X.shape)])), [-1,2])
-        return dgm
+        self.dgms = [tf.reshape(tf.gather_nd(X, tf.reshape(indice, [-1,len(X.shape)])), [-1,2]) for indice in indices]
+        return self.dgms
diff --git a/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py b/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py
index 4f515386..c509c456 100644
--- a/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py
+++ b/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py
@@ -7,10 +7,10 @@ import tensorflow          as tf
 
 # The parameters of the model are the vertex function values of the simplex tree.
 
-def _LowerStarSimplexTree(simplextree, filtration, dimension):
+def _LowerStarSimplexTree(simplextree, filtration, dimensions):
     # Parameters: simplextree (simplex tree on which to compute persistence)
     #             filtration (function values on the vertices of st),
-    #             dimension (homology dimension),
+    #             dimensions (homology dimensions),
     
     for s,_ in simplextree.get_filtration():
         simplextree.assign_filtration(s, -1e10)
@@ -21,40 +21,38 @@ def _LowerStarSimplexTree(simplextree, filtration, dimension):
     simplextree.make_filtration_non_decreasing()
     
     # Compute persistence diagram
-    dgm = simplextree.compute_persistence()
+    simplextree.compute_persistence()
     
     # Get vertex pairs for optimization. First, get all simplex pairs
-    pairs = simplextree.persistence_pairs()
+    pairs = simplextree.lower_star_persistence_generators()
     
-    # Then, loop over all simplex pairs
-    indices, pers = [], []
-    for s1, s2 in pairs:
-        # Select pairs with good homological dimension and finite lifetime
-        if len(s1) == dimension+1 and len(s2) > 0:
-            # Get IDs of the vertices corresponding to the filtration values of the simplices
-            l1, l2 = np.array(s1), np.array(s2)
-            i1 = l1[np.argmax(filtration[l1])]
-            i2 = l2[np.argmax(filtration[l2])]
-            indices.append(i1)
-            indices.append(i2)
+    L_indices = []
+    for dimension in dimensions:
     
-    indices = np.reshape(indices, [-1,2]).flatten()
-    return np.array(indices, dtype=np.int32)
+        finite_pairs = pairs[0][dimension] if len(pairs[0]) >= dimension+1 else np.empty(shape=[0,2])
+        essential_pairs = pairs[1][dimension] if len(pairs[1]) >= dimension+1 else np.empty(shape=[0,1])
+        
+        finite_indices = np.array(finite_pairs.flatten(), dtype=np.int32)
+        essential_indices = np.array(essential_pairs.flatten(), dtype=np.int32)
+
+        L_indices.append((finite_indices, essential_indices))
+
+    return L_indices
 
 class LowerStarSimplexTreeLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing lower-star persistence out of a simplex tree
     """
-    def __init__(self, simplextree, dimension=0, **kwargs):
+    def __init__(self, simplextree, dimensions=[0], **kwargs):
         """
         Constructor for the LowerStarSimplexTreeLayer class
   
         Parameters:
             simplextree (gudhi.SimplexTree): underlying simplex tree. Its vertices MUST be named with integers from 0 to n = number of vertices
-            dimension (int): homology dimension
+            dimensions (int): homology dimensions
         """
         super().__init__(dynamic=True, **kwargs)
-        self.dimension   = dimension
+        self.dimensions  = dimensions
         self.simplextree = simplextree
     
     def build(self):
@@ -68,11 +66,15 @@ class LowerStarSimplexTreeLayer(tf.keras.layers.Layer):
             F (TensorFlow variable): filter function values over the vertices of the simplex tree. The ith entry of F corresponds to vertex i in self.simplextree
 
         Returns:
-            dgm (TensorFlow variable): lower-star persistence diagram with shape [num_points, 2]
+            dgms (list of tuple of TensorFlow variables): list of lower-star persistence diagrams of length self.dimensions, where each element of the list is a tuple that contains the finite and essential persistence diagrams of shapes [num_finite_points, 2] and [num_essential_points, 1] respectively
         """
         # Don't try to compute gradients for the vertex pairs
-        indices = tf.stop_gradient(_LowerStarSimplexTree(self.simplextree, filtration.numpy(), self.dimension))
-        # Get persistence diagram
-        self.dgm = tf.reshape(tf.gather(filtration, indices), [-1,2]) 
-        return self.dgm
+        indices = _LowerStarSimplexTree(self.simplextree, filtration.numpy(), self.dimensions)
+        # Get persistence diagrams
+        self.dgms = []
+        for idx_dim, dimension in enumerate(self.dimensions):
+            finite_dgm = tf.reshape(tf.gather(filtration, indices[idx_dim][0]), [-1,2])
+            essential_dgm = tf.reshape(tf.gather(filtration, indices[idx_dim][1]), [-1,1])
+            self.dgms.append((finite_dgm, essential_dgm))
+        return self.dgms
 
diff --git a/src/python/gudhi/tensorflow/rips_layer.py b/src/python/gudhi/tensorflow/rips_layer.py
index 6d54871c..83387d21 100644
--- a/src/python/gudhi/tensorflow/rips_layer.py
+++ b/src/python/gudhi/tensorflow/rips_layer.py
@@ -8,38 +8,49 @@ from ..rips_complex     import RipsComplex
 
 # The parameters of the model are the point coordinates.
 
-def _Rips(DX, max_edge, dimension):
+def _Rips(DX, max_edge, dimensions):
     # Parameters: DX (distance matrix), 
     #             max_edge (maximum edge length for Rips filtration), 
-    #             dimension (homology dimension)
+    #             dimensions (homology dimensions)
 
     # Compute the persistence pairs with Gudhi
     rc = RipsComplex(distance_matrix=DX, max_edge_length=max_edge)
-    st = rc.create_simplex_tree(max_dimension=dimension+1)
-    dgm = st.persistence()
-    if dimension == 0:
-        pairs = st.flag_persistence_generators()[0]
-    else:
-        pairs = st.flag_persistence_generators()[1][dimension-1]
+    st = rc.create_simplex_tree(max_dimension=max(dimensions)+1)
+    st.persistence()
+    pairs = st.flag_persistence_generators()
 
-    indices = pairs.flatten()
-    return np.array(indices, dtype=np.int32)
+    L_indices = []
+    for dimension in dimensions:
+
+        if dimension == 0:
+            finite_pairs = pairs[0]
+            essential_pairs = pairs[2]
+        else:
+            finite_pairs = pairs[1][dimension-1] if len(pairs[1]) >= dimension else np.empty(shape=[0,4])
+            essential_pairs = pairs[3][dimension-1] if len(pairs[3]) >= dimension else np.empty(shape=[0,2])
+        
+        finite_indices = np.array(finite_pairs.flatten(), dtype=np.int32)
+        essential_indices = np.array(essential_pairs.flatten(), dtype=np.int32)
+
+        L_indices.append((finite_indices, essential_indices))
+
+    return L_indices
 
 class RipsLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing Rips persistence out of a point cloud
     """
-    def __init__(self, maximum_edge_length=12, dimension=1, **kwargs):
+    def __init__(self, maximum_edge_length=12, dimensions=[1], **kwargs):
         """
         Constructor for the RipsLayer class
 
         Parameters:
             maximum_edge_length (float): maximum edge length for the Rips complex 
-            dimension (int): homology dimension
+            dimensions (int): homology dimensions
         """
         super().__init__(dynamic=True, **kwargs)
         self.max_edge = maximum_edge_length
-        self.dimension = dimension
+        self.dimensions = dimensions
 
     def build(self):
         super.build()
@@ -52,19 +63,24 @@ class RipsLayer(tf.keras.layers.Layer):
             X (TensorFlow variable): point cloud of shape [number of points, number of dimensions]
 
         Returns:
-            dgm (TensorFlow variable): Rips persistence diagram with shape [num_points, 2] with points sorted by 
+            dgms (list of tuple of TensorFlow variables): list of Rips persistence diagrams of length self.dimensions, where each element of the list is a tuple that contains the finite and essential persistence diagrams of shapes [num_finite_points, 2] and [num_essential_points, 1] respectively
         """    
         # Compute distance matrix
         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 = tf.stop_gradient(_Rips(DX.numpy(), self.max_edge, self.dimension))
-        # Get persistence diagram by simply picking the corresponding entries in the distance matrix
-        if self.dimension > 0:
-            dgm = tf.reshape(tf.gather_nd(DX, tf.reshape(indices, [-1,2])), [-1,2])
-        else:
-            #indices = tf.reshape(indices, [-1,2])[1::2,:]
-            indices = indices[:,1:]
-            dgm = tf.concat([tf.zeros([indices.shape[0],1]), tf.reshape(tf.gather_nd(DX, indices), [-1,1])], axis=1)
-        return dgm
+        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):
+            cur_idx = indices[idx_dim]
+            if dimension > 0:
+                finite_dgm = tf.reshape(tf.gather_nd(DX, tf.reshape(cur_idx[0], [-1,2])), [-1,2])
+                essential_dgm = tf.reshape(tf.gather_nd(DX, tf.reshape(cur_idx[1], [-1,2])), [-1,1])
+            else:
+                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))
+        return self.dgms