SimplexTree.insert_batch for multiple simplices of the same dimension

2 files changed, 123 insertions, 46 deletions

diff --git a/src/python/gudhi/simplex_tree.pyx b/src/python/gudhi/simplex_tree.pyx
index 3eebfff7..3646d659 100644
--- a/src/python/gudhi/simplex_tree.pyx
+++ b/src/python/gudhi/simplex_tree.pyx
@@ -8,14 +8,23 @@
 #   - YYYY/MM Author: Description of the modification
 
 from cython.operator import dereference, preincrement
-from libc.stdint cimport intptr_t
+from libc.stdint cimport intptr_t, int32_t, int64_t
 import numpy as np
 cimport gudhi.simplex_tree
+cimport cython
 
 __author__ = "Vincent Rouvreau"
 __copyright__ = "Copyright (C) 2016 Inria"
 __license__ = "MIT"
 
+ctypedef fused some_int:
+    int32_t
+    int64_t
+
+ctypedef fused some_float:
+    float
+    double
+
 # SimplexTree python interface
 cdef class SimplexTree:
     """The simplex tree is an efficient and flexible data structure for
@@ -226,6 +235,7 @@ cdef class SimplexTree:
         return self.get_ptr().insert(simplex, <double>filtration)
 
     @staticmethod
+    @cython.boundscheck(False)
     def create_from_array(filtrations, double max_filtration=np.inf):
         """Creates a new, empty complex and inserts vertices and edges. The vertices are numbered from 0 to n-1, and
         the filtration values are encoded in the array, with the diagonal representing the vertices. It is the
@@ -265,10 +275,38 @@ cdef class SimplexTree:
         :param edges: the edges to insert and their filtration values.
         :type edges: scipy.sparse.coo_matrix of shape (n,n)
         """
-        # TODO: optimize this
+        # TODO: optimize this?
         for edge in zip(edges.row, edges.col, edges.data):
             self.get_ptr().insert((edge[0], edge[1]), edge[2])
 
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def insert_batch(self, some_int[:,:] vertex_array, some_float[:] filtrations):
+        """Inserts k-simplices given by a sparse array in a format similar
+        to `torch.sparse <https://pytorch.org/docs/stable/sparse.html>`_.
+        Duplicate entries are not allowed. Missing entries are not inserted.
+        Simplices with a repeated vertex are currently interpreted as lower
+        dimensional simplices, but we do not guarantee this behavior in the
+        future. Any time a simplex is inserted, its faces are inserted as well
+        if needed to preserve a simplicial complex.
+
+        :param vertex_array: the k-simplices to insert.
+        :type vertex_array: numpy.array of shape (k+1,n)
+        :param filtrations: the filtration values.
+        :type filtrations: numpy.array of shape (n,)
+        """
+        cdef Py_ssize_t k = vertex_array.shape[0]
+        cdef Py_ssize_t n = vertex_array.shape[1]
+        assert filtrations.shape[0] == n
+        cdef Py_ssize_t i
+        cdef Py_ssize_t j
+        cdef vector[int] v
+        for i in range(n):
+            for j in range(k):
+                v.push_back(vertex_array[j, i])
+            self.get_ptr().insert(v, filtrations[i])
+            v.clear()
+
     def get_simplices(self):
         """This function returns a generator with simplices and their given
         filtration values.
diff --git a/src/python/test/test_simplex_tree.py b/src/python/test/test_simplex_tree.py
index 0436d891..a5b8ffe0 100755
--- a/src/python/test/test_simplex_tree.py
+++ b/src/python/test/test_simplex_tree.py
@@ -345,9 +345,10 @@ def test_simplices_iterator():
         print("filtration is: ", simplex[1])
         assert st.filtration(simplex[0]) == simplex[1]
 
+
 def test_collapse_edges():
     st = SimplexTree()
-    
+
     assert st.insert([0, 1], filtration=1.0) == True
     assert st.insert([1, 2], filtration=1.0) == True
     assert st.insert([2, 3], filtration=1.0) == True
@@ -362,33 +363,35 @@ def test_collapse_edges():
         assert st.num_simplices() == 9
         assert st.find([1, 3]) == False
         for simplex in st.get_skeleton(0):
-            assert simplex[1] == 1.
+            assert simplex[1] == 1.0
     else:
         # If no Eigen3, collapse_edges throws an exception
         with pytest.raises(RuntimeError):
             st.collapse_edges()
 
+
 def test_reset_filtration():
     st = SimplexTree()
-    
-    assert st.insert([0, 1, 2], 3.) == True
-    assert st.insert([0, 3], 2.) == True
-    assert st.insert([3, 4, 5], 3.) == True
-    assert st.insert([0, 1, 6, 7], 4.) == True
+
+    assert st.insert([0, 1, 2], 3.0) == True
+    assert st.insert([0, 3], 2.0) == True
+    assert st.insert([3, 4, 5], 3.0) == True
+    assert st.insert([0, 1, 6, 7], 4.0) == True
 
     # Guaranteed by construction
     for simplex in st.get_simplices():
-        assert st.filtration(simplex[0]) >= 2.
-    
+        assert st.filtration(simplex[0]) >= 2.0
+
     # dimension until 5 even if simplex tree is of dimension 3 to test the limits
     for dimension in range(5, -1, -1):
-        st.reset_filtration(0., dimension)
+        st.reset_filtration(0.0, dimension)
         for simplex in st.get_skeleton(3):
             print(simplex)
             if len(simplex[0]) < (dimension) + 1:
-                assert st.filtration(simplex[0]) >= 2.
+                assert st.filtration(simplex[0]) >= 2.0
             else:
-                assert st.filtration(simplex[0]) == 0.
+                assert st.filtration(simplex[0]) == 0.0
+
 
 def test_boundaries_iterator():
     st = SimplexTree()
@@ -404,16 +407,17 @@ def test_boundaries_iterator():
         list(st.get_boundaries([]))
 
     with pytest.raises(RuntimeError):
-        list(st.get_boundaries([0, 4])) # (0, 4) does not exist
+        list(st.get_boundaries([0, 4]))  # (0, 4) does not exist
 
     with pytest.raises(RuntimeError):
-        list(st.get_boundaries([6])) # (6) does not exist
+        list(st.get_boundaries([6]))  # (6) does not exist
+
 
 def test_persistence_intervals_in_dimension():
     # Here is our triangulation of a 2-torus - taken from https://dioscuri-tda.org/Paris_TDA_Tutorial_2021.html
     #   0-----3-----4-----0
     #   | \   | \   | \   | \   |
-    #   |   \ |   \ |    \|   \ | 
+    #   |   \ |   \ |    \|   \ |
     #   1-----8-----7-----1
     #   | \   | \   | \   | \   |
     #   |   \ |   \ |   \ |   \ |
@@ -422,50 +426,52 @@ def test_persistence_intervals_in_dimension():
     #   |   \ |   \ |   \ |   \ |
     #   0-----3-----4-----0
     st = SimplexTree()
-    st.insert([0,1,8])
-    st.insert([0,3,8])
-    st.insert([3,7,8])
-    st.insert([3,4,7])
-    st.insert([1,4,7])
-    st.insert([0,1,4])
-    st.insert([1,2,5])
-    st.insert([1,5,8])
-    st.insert([5,6,8])
-    st.insert([6,7,8])
-    st.insert([2,6,7])
-    st.insert([1,2,7])
-    st.insert([0,2,3])
-    st.insert([2,3,5])
-    st.insert([3,4,5])
-    st.insert([4,5,6])
-    st.insert([0,4,6])
-    st.insert([0,2,6])
+    st.insert([0, 1, 8])
+    st.insert([0, 3, 8])
+    st.insert([3, 7, 8])
+    st.insert([3, 4, 7])
+    st.insert([1, 4, 7])
+    st.insert([0, 1, 4])
+    st.insert([1, 2, 5])
+    st.insert([1, 5, 8])
+    st.insert([5, 6, 8])
+    st.insert([6, 7, 8])
+    st.insert([2, 6, 7])
+    st.insert([1, 2, 7])
+    st.insert([0, 2, 3])
+    st.insert([2, 3, 5])
+    st.insert([3, 4, 5])
+    st.insert([4, 5, 6])
+    st.insert([0, 4, 6])
+    st.insert([0, 2, 6])
     st.compute_persistence(persistence_dim_max=True)
-    
+
     H0 = st.persistence_intervals_in_dimension(0)
-    assert np.array_equal(H0, np.array([[ 0., float("inf")]]))
+    assert np.array_equal(H0, np.array([[0.0, float("inf")]]))
     H1 = st.persistence_intervals_in_dimension(1)
-    assert np.array_equal(H1, np.array([[ 0., float("inf")], [ 0., float("inf")]]))
+    assert np.array_equal(H1, np.array([[0.0, float("inf")], [0.0, float("inf")]]))
     H2 = st.persistence_intervals_in_dimension(2)
-    assert np.array_equal(H2, np.array([[ 0., float("inf")]]))
+    assert np.array_equal(H2, np.array([[0.0, float("inf")]]))
     # Test empty case
     assert st.persistence_intervals_in_dimension(3).shape == (0, 2)
 
+
 def test_equality_operator():
     st1 = SimplexTree()
     st2 = SimplexTree()
 
     assert st1 == st2
 
-    st1.insert([1,2,3], 4.)
+    st1.insert([1, 2, 3], 4.0)
     assert st1 != st2
 
-    st2.insert([1,2,3], 4.)
+    st2.insert([1, 2, 3], 4.0)
     assert st1 == st2
 
+
 def test_simplex_tree_deep_copy():
     st = SimplexTree()
-    st.insert([1, 2, 3], 0.)
+    st.insert([1, 2, 3], 0.0)
     # compute persistence only on the original
     st.compute_persistence()
 
@@ -484,14 +490,15 @@ def test_simplex_tree_deep_copy():
 
     for a_splx in a_filt_list:
         assert a_splx in st_filt_list
-    
+
     # test double free
     del st
     del st_copy
 
+
 def test_simplex_tree_deep_copy_constructor():
     st = SimplexTree()
-    st.insert([1, 2, 3], 0.)
+    st.insert([1, 2, 3], 0.0)
     # compute persistence only on the original
     st.compute_persistence()
 
@@ -510,20 +517,23 @@ def test_simplex_tree_deep_copy_constructor():
 
     for a_splx in a_filt_list:
         assert a_splx in st_filt_list
-    
+
     # test double free
     del st
     del st_copy
 
+
 def test_simplex_tree_constructor_exception():
     with pytest.raises(TypeError):
-        st = SimplexTree(other = "Construction from a string shall raise an exception")
+        st = SimplexTree(other="Construction from a string shall raise an exception")
+
 
 def test_create_from_array():
     a = np.array([[1, 4, 13, 6], [4, 3, 11, 5], [13, 11, 10, 12], [6, 5, 12, 2]])
     st = SimplexTree.create_from_array(a, max_filtration=5)
     assert list(st.get_filtration()) == [([0], 1.0), ([3], 2.0), ([1], 3.0), ([0, 1], 4.0), ([1, 3], 5.0)]
 
+
 def test_insert_edges_from_coo_matrix():
     try:
         from scipy.sparse import coo_matrix
@@ -561,3 +571,32 @@ def test_insert_edges_from_coo_matrix():
     st = SimplexTree()
     st.insert_edges_from_coo_matrix(edges)
     assert 100 < st.num_simplices() < 1000
+
+
+def test_insert_batch():
+    st = SimplexTree()
+    # vertices
+    st.insert_batch(np.array([[6, 1, 5]]), np.array([-5.0, 2.0, -3.0]))
+    # triangles
+    st.insert_batch(np.array([[2, 10], [5, 0], [6, 11]]), np.array([4.0, 0.0]))
+    # edges
+    st.insert_batch(np.array([[1, 5], [2, 5]]), np.array([1.0, 3.0]))
+
+    assert list(st.get_filtration()) == [
+        ([6], -5.0),
+        ([5], -3.0),
+        ([0], 0.0),
+        ([10], 0.0),
+        ([0, 10], 0.0),
+        ([11], 0.0),
+        ([0, 11], 0.0),
+        ([10, 11], 0.0),
+        ([0, 10, 11], 0.0),
+        ([1], 1.0),
+        ([2], 1.0),
+        ([1, 2], 1.0),
+        ([2, 5], 4.0),
+        ([2, 6], 4.0),
+        ([5, 6], 4.0),
+        ([2, 5, 6], 4.0),
+    ]