removed unraveling in cubical

3 files changed, 7 insertions, 14 deletions

diff --git a/src/python/gudhi/tensorflow/cubical_layer.py b/src/python/gudhi/tensorflow/cubical_layer.py
index 55bd2685..70528f98 100644
--- a/src/python/gudhi/tensorflow/cubical_layer.py
+++ b/src/python/gudhi/tensorflow/cubical_layer.py
@@ -17,6 +17,7 @@ def _Cubical(X, dimensions):
     cc = CubicalComplex(top_dimensional_cells=X)
     cc.compute_persistence()
 
+    # Retrieve and ouput image indices/pixels corresponding to positive and negative simplices    
     cof_pp = cc.cofaces_of_persistence_pairs()
     
     L_cofs = []
@@ -27,15 +28,7 @@ def _Cubical(X, dimensions):
         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.zeros(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, order='F')
-            ocof[D*(idx+1):D*(idx+2)] = np.unravel_index(cof[count,1], Xs, order='F')
-            count += 1
-        L_cofs.append(np.array(ocof, dtype=np.int32))
+        L_cofs.append(np.array(cof, dtype=np.int32))
 
     return L_cofs
 
@@ -43,7 +36,7 @@ class CubicalLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing cubical persistence out of a cubical complex
     """
-    def __init__(self, dimensions=[0], **kwargs):
+    def __init__(self, dimensions, **kwargs):
         """
         Constructor for the CubicalLayer class
 
@@ -70,5 +63,5 @@ class CubicalLayer(tf.keras.layers.Layer):
         # Don't compute gradient for this operation
         indices = _Cubical(X.numpy(), self.dimensions)
         # Get persistence diagram by simply picking the corresponding entries in the image
-        self.dgms = [tf.reshape(tf.gather_nd(X, tf.reshape(indice, [-1,len(X.shape)])), [-1,2]) for indice in indices]
+        self.dgms = [tf.reshape(tf.gather(  tf.reshape(tf.transpose(X), [-1]), indice  ), [-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 c509c456..fd1698ea 100644
--- a/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py
+++ b/src/python/gudhi/tensorflow/lower_star_simplex_tree_layer.py
@@ -43,7 +43,7 @@ class LowerStarSimplexTreeLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing lower-star persistence out of a simplex tree
     """
-    def __init__(self, simplextree, dimensions=[0], **kwargs):
+    def __init__(self, simplextree, dimensions, **kwargs):
         """
         Constructor for the LowerStarSimplexTreeLayer class
   
diff --git a/src/python/gudhi/tensorflow/rips_layer.py b/src/python/gudhi/tensorflow/rips_layer.py
index 7735db67..a314229b 100644
--- a/src/python/gudhi/tensorflow/rips_layer.py
+++ b/src/python/gudhi/tensorflow/rips_layer.py
@@ -16,7 +16,7 @@ def _Rips(DX, max_edge, dimensions):
     # 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.persistence()
+    st.compute_persistence()
     pairs = st.flag_persistence_generators()
 
     L_indices = []
@@ -40,7 +40,7 @@ class RipsLayer(tf.keras.layers.Layer):
     """
     TensorFlow layer for computing Rips persistence out of a point cloud
     """
-    def __init__(self, maximum_edge_length=12, dimensions=[0], **kwargs):
+    def __init__(self, maximum_edge_length=12, dimensions, **kwargs):
         """
         Constructor for the RipsLayer class