support 2D or 3D tsne

.gitignore

@@ -36,6 +36,7 @@ nosetests.xml
 
 # Project specific
 tsne/bh_sne.cpp
+tsne/bh_sne_3d.cpp
 tsne/bh_sne_src/bh_tsne
 *.pkl.gz
 .cache

setup.py

@@ -38,14 +38,24 @@
                    extra_compile_args=extra_compile_args + ['-ffast-math', '-O3'],
                    extra_link_args=['-Wl,-framework', '-Wl,Accelerate', '-lcblas'],
                    language='c++')]
+
 else:
     # LINUX
+
     ext_modules = [Extension(name='bh_sne',
                    sources=['tsne/bh_sne_src/sptree.cpp', 'tsne/bh_sne_src/tsne.cpp', 'tsne/bh_sne.pyx'],
                    include_dirs=[numpy.get_include(), '/usr/local/include', 'tsne/bh_sne_src/'],
                    library_dirs=['/usr/local/lib', '/usr/lib64/atlas'],
                    extra_compile_args=['-msse2', '-O3', '-fPIC', '-w', '-ffast-math'],
                    extra_link_args=['-Wl,-Bstatic', '-lcblas', '-Wl,-Bdynamic'],
+                   language='c++'),
+
+                   Extension(name='bh_sne_3d',
+                   sources=['tsne/bh_sne_src/sptree.cpp', 'tsne/bh_sne_src/tsne.cpp', 'tsne/bh_sne_3d.pyx'],
+                   include_dirs=[numpy.get_include(), '/usr/local/include', 'tsne/bh_sne_src/'],
+                   library_dirs=['/usr/local/lib', '/usr/lib64/atlas'],
+                   extra_compile_args=['-msse2', '-O3', '-fPIC', '-w', '-ffast-math', '-DTSNE3D'],
+                   extra_link_args=['-Wl,-Bstatic', '-lcblas', '-Wl,-Bdynamic'],
                    language='c++')]
 
 ext_modules = cythonize(ext_modules)

tsne/__init__.py

@@ -4,6 +4,7 @@
 import scipy.linalg as la
 import sys
 from bh_sne import BH_SNE
+from bh_sne_3d import BH_SNE_3D
 
 def bh_sne(data, pca_d=None, d=2, perplexity=30., theta=0.5,
            random_state=None, copy_data=False, init=None,
@@ -16,8 +17,7 @@ def bh_sne(data, pca_d=None, d=2, perplexity=30., theta=0.5,
     @param pca_d        The dimensionality of data is reduced via PCA
                         to this dimensionality.
 
-    @param d            The embedding dimensionality. Must be fixed to
-                        2.
+    @param d            The embedding dimensionality. Must be 2 or 3.
 
     @param perplexity   The perplexity controls the effective number of
                         neighbors.
@@ -79,7 +79,13 @@ def bh_sne(data, pca_d=None, d=2, perplexity=30., theta=0.5,
     if mom_switch_iter is None:
         mom_switch_iter = 250
 
-    tsne = BH_SNE()
+    if d == 2:
+        tsne = BH_SNE()
+    elif d == 3:
+        tsne = BH_SNE_3D()
+    else:
+        raise Exception("TSNE dimensions must be 2 or 3")
+
     Y = tsne.run(X, N, X.shape[1], d, perplexity, theta, seed, init=init, use_init=use_init,
                  max_iter=max_iter, stop_lying_iter=stop_lying_iter, mom_switch_iter=mom_switch_iter)
     return Y

tsne/bh_sne.pyx

@@ -9,7 +9,6 @@ cdef extern from "tsne.h":
         TSNE()
         void run(double* X, int N, int D, double* Y, int no_dims, double perplexity, double theta, int rand_seed, bool skip_random_init, double *init, bool use_init, int max_iter, int stop_lying_iter, int mom_switch_iter)
 
-
 cdef class BH_SNE:
     cdef TSNE* thisptr # hold a C++ instance
 

tsne/bh_sne_3d.pyx

@@ -0,0 +1,28 @@
+# distutils: language = c++
+import numpy as np
+cimport numpy as np
+cimport cython
+from libcpp cimport bool
+
+cdef extern from "tsne.h":
+    cdef cppclass TSNE:
+        TSNE()
+        void run(double* X, int N, int D, double* Y, int no_dims, double perplexity, double theta, int rand_seed, bool skip_random_init, double *init, bool use_init, int max_iter, int stop_lying_iter, int mom_switch_iter)
+
+cdef class BH_SNE_3D:
+    cdef TSNE* thisptr # hold a C++ instance
+
+    def __cinit__(self):
+        self.thisptr = new TSNE()
+
+    def __dealloc__(self):
+        del self.thisptr
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def run(self, X, N, D, d, perplexity, theta, seed, init, use_init, max_iter, stop_lying_iter, mom_switch_iter):
+        cdef np.ndarray[np.float64_t, ndim=2, mode='c'] _X = np.ascontiguousarray(X)
+        cdef np.ndarray[np.float64_t, ndim=2, mode='c'] _init = np.ascontiguousarray(init)
+        cdef np.ndarray[np.float64_t, ndim=2, mode='c'] Y = np.zeros((N, d), dtype=np.float64)
+        self.thisptr.run(&_X[0,0], N, D, &Y[0,0], d, perplexity, theta, seed, False, &_init[0,0], use_init, max_iter, stop_lying_iter, mom_switch_iter)
+        return Y

tsne/bh_sne_src/Makefile

@@ -5,16 +5,22 @@
 CXX = g++
 CFLAGS = -ffast-math -O3
 
-all: bh_tsne
+all: bh_tsne bh_tsne_3d
 
 bh_tsne: tsne.o sptree.o
 	$(CXX) $(CFLAGS) tsne.o sptree.o -o bh_tsne
 
+bh_tsne_3d: tsne_3d.o sptree.o
+	$(CXX) $(CFLAGS) tsne_3d.o sptree.o -o bh_tsne_3d
+
 sptree.o: sptree.cpp sptree.h
 	$(CXX) $(CFLAGS) -c sptree.cpp
 
 tsne.o: tsne.cpp tsne.h sptree.h vptree.h
 	$(CXX) $(CFLAGS) -c tsne.cpp
 
+tsne_3d.o: tsne.cpp tsne.h sptree.h vptree.h
+	$(CXX) $(CFLAGS) -DTSNE3D -c tsne.cpp
+
 clean:
-	rm -Rf *.o bh_tsne
+	rm -Rf *.o bh_tsne bh_tsne_3d

tsne/bh_sne_src/tsne.cpp

@@ -30,7 +30,6 @@
  *
  */
 
-
 #include <iostream>
 #include <fstream>
 #include <string>
@@ -45,9 +44,14 @@
 #include "sptree.h"
 #include "tsne.h"
 
-
 using namespace std;
 
+#ifdef TSNE3D
+#define NDIMS 3
+#else
+#define NDIMS 2
+#endif
+
 // Perform t-SNE
 void TSNE::run(double* X, int N, int D, double* Y, int no_dims, double perplexity, double theta, int rand_seed,
                bool skip_random_init, double *init, bool use_init,
@@ -217,7 +221,7 @@ void TSNE::computeGradient(double* P, unsigned int* inp_row_P, unsigned int* inp
 {
 
     // Construct space-partitioning tree on current map
-    SPTree<2>* tree = new SPTree<2>(Y, N);
+    SPTree<NDIMS>* tree = new SPTree<NDIMS>(Y, N);
 
     // Compute all terms required for t-SNE gradient
     double sum_Q = .0;
@@ -327,7 +331,7 @@ double TSNE::evaluateError(unsigned int* row_P, unsigned int* col_P, double* val
 {
 
     // Get estimate of normalization term
-    SPTree<2>* tree = new SPTree<2>(Y, N);
+    SPTree<NDIMS>* tree = new SPTree<NDIMS>(Y, N);
     double* buff = (double*) calloc(D, sizeof(double));
     double sum_Q = .0;
     for(int n = 0; n < N; n++) tree->computeNonEdgeForces(n, theta, buff, &sum_Q);
@@ -761,9 +765,6 @@ int main(int argc, char *argv[]) {
     // Read the parameters and the dataset
 	if(tsne->load_data(dat_file_c, &data, &origN, &D, &no_dims, &theta, &perplexity, &rand_seed, &max_iter)) {
 
-        // assert that no_dims = 2. if not, the template hack won't work.
-        if (no_dims != 2) { fprintf(stderr,"Number of dims must be 2!\n"); exit(1); }
-
 		// Make dummy landmarks
         N = origN;
         int* landmarks = (int*) malloc(N * sizeof(int));