Merge pull request #236 from pavlin-policar/fix-bh-memory

Fix exploding memory problem in BH for duplicate rows

Author: pavlin-policar

openTSNE/_tsne.pyx

@@ -10,7 +10,7 @@ from cython.parallel import prange, parallel
 from cpython.mem cimport PyMem_Malloc, PyMem_Free
 from libc.stdlib cimport malloc, free
 
-from .quad_tree cimport QuadTree, Node, is_duplicate
+from .quad_tree cimport QuadTree, Node, is_close
 from ._matrix_mul.matrix_mul cimport matrix_multiply_fft_1d, matrix_multiply_fft_2d
 
 
@@ -178,7 +178,7 @@ cpdef double estimate_negative_gradient_bh(
     Py_ssize_t num_threads=1,
     bint pairwise_normalization=True,
 ):
-    """Estimate the negative tSNE gradient using the Barnes Hut approximation.
+    """Estimate the negative t-SNE gradient using the Barnes-Hut approximation.
     
     Notes
     -----
@@ -197,7 +197,7 @@ cpdef double estimate_negative_gradient_bh(
         num_threads = 1
 
     # In order to run gradient estimation in parallel, we need to pass each
-    # worker it's own memory slot to write sum_Qs
+    # worker its own memory slot to write sum_Qs
     for i in prange(num_points, nogil=True, num_threads=num_threads, schedule="guided"):
         _estimate_negative_gradient_single(
             &tree.root, &embedding[i, 0], &gradient[i, 0], &sum_Qi[i], theta, dof
@@ -225,16 +225,16 @@ cdef void _estimate_negative_gradient_single(
     double theta,
     double dof,
 ) nogil:
-    # Make sure that we spend no time on empty nodes or self-interactions
-    if node.num_points == 0 or node.is_leaf and is_duplicate(node, point):
+    # Make sure that we spend no time on empty nodes or simple self-interactions
+    if node.num_points == 0 or node.is_leaf and is_close(node, point, EPSILON):
         return
 
     cdef:
         double distance = EPSILON
         double q_ij, tmp
         Py_ssize_t d
 
-    # Compute the squared euclidean disstance in the embedding space from the
+    # Compute the squared euclidean distance in the embedding space from the
     # new point to the center of mass
     for d in range(node.n_dims):
         tmp = node.center_of_mass[d] - point[d]

openTSNE/quad_tree.pxd

@@ -5,8 +5,6 @@
 # cython: language_level=3
 import numpy as np
 
-cdef double EPSILON = np.finfo(np.float64).eps
-
 ctypedef struct Node:
     Py_ssize_t n_dims
     double *center
@@ -19,7 +17,7 @@ ctypedef struct Node:
     Py_ssize_t num_points
 
 
-cdef bint is_duplicate(Node * node, double * point, double duplicate_eps=*) nogil
+cdef bint is_close(Node * node, double * point, double eps) nogil
 
 
 cdef class QuadTree:

openTSNE/quad_tree.pyx

@@ -71,14 +71,15 @@ cdef Py_ssize_t get_child_idx_for(Node * node, double * point) nogil:
 cdef inline void update_center_of_mass(Node * node, double * point) nogil:
     cdef Py_ssize_t d
     for d in range(node.n_dims):
-        node.center_of_mass[d] = (node.center_of_mass[d] * node.num_points + point[d]) \
-            / (node.num_points + 1)
+        node.center_of_mass[d] = (node.center_of_mass[d] * node.num_points + point[d]) / (node.num_points + 1)
     node.num_points += 1
 
 
 cdef void add_point_to(Node * node, double * point):
-    # If the node is a leaf node and empty, we"re done
-    if node.is_leaf and node.num_points == 0 or is_duplicate(node, point):
+    # If the node is a leaf node and empty, we're done. We'll also limit the
+    # branching of each node to prevent memory explosions. 1e-6 here is
+    # effectively the minimum size that a node can take
+    if node.is_leaf and node.num_points == 0 or is_close(node, point, 1e-6):
         update_center_of_mass(node, point)
         return
 
@@ -124,10 +125,10 @@ cdef void split_node(Node * node):
     PyMem_Free(new_center)
 
 
-cdef inline bint is_duplicate(Node * node, double * point, double duplicate_eps=1e-16) nogil:
+cdef inline bint is_close(Node * node, double * point, double eps) nogil:
     cdef Py_ssize_t d
     for d in range(node.n_dims):
-        if fabs(node.center_of_mass[d] - point[d]) >= duplicate_eps:
+        if fabs(node.center_of_mass[d] - point[d]) >= eps:
             return False
     return True
 

tests/test_correctness.py

@@ -333,3 +333,26 @@ def test_early_exaggeration_does_not_collapse(self):
                 x = np.random.randn(n, d)
                 embedding = openTSNE.TSNE(random_state=42).fit(x)
                 self.assertGreater(np.max(np.abs(embedding)), 1e-8)
+
+
+class TestDataMatricesWithDuplicatedRows(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        from sklearn.preprocessing import KBinsDiscretizer
+
+        # Load up contrived example where we have a large number of duplicated
+        # rows. This is similar to the Titanic data set, which is problematic.
+        np.random.seed(0)
+        iris = datasets.load_iris()
+        x, y = iris.data, iris.target
+
+        discretizer = KBinsDiscretizer(n_bins=2, strategy="uniform")
+        x = discretizer.fit_transform(x).toarray()
+
+        idx = np.random.choice(x.shape[0], size=1000, replace=True)
+        cls.x, cls.y = x[idx], y[idx]
+
+    def test_works_without_error(self):
+        openTSNE.TSNE(
+            early_exaggeration=100, negative_gradient_method="bh", random_state=0
+        ).fit(self.x)