Merge pull request #10 from DiffusionMapsAcademics/dev_kwargs

Dev kwargs

src/pydiffmap/diffusion_map.py

@@ -28,35 +28,42 @@ class DiffusionMap(object):
         Method for choosing the epsilon.  Currently, the only option is 'fixed' (i.e. don't) or 'bgh' (Berry, Giannakis and Harlim)
     n_evecs : int, optional
         Number of diffusion map eigenvectors to return
+    neighbor_params : dict or None, optional
+        Optional parameters for the nearest Neighbor search. See scikit-learn NearestNeighbors class for details.
     metric : string, optional
         Metric for distances in the kernel. Default is 'euclidean'. The callable should take two arrays as input and return one value indicating the distance between them.
     metric_params : dict or None, optional
         Optional parameters required for the metric given.
-    n_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, optional
-        Algorithm used to compute the nearest neighbors.  See sklearn.neighbors.NearestNeighbors for details
+
+    Examples
+    --------
+    # setup neighbor_params list with as many jobs as CPU cores and kd_tree neighbor search.
+    >>> neighbor_params = {'n_jobs': -1, 'algorithm': 'kd_tree'}
+    # initialize diffusion map object with the top two eigenvalues being computed, epsilon set to 0.1
+    # and alpha set to 1.0.
+    >>> mydmap = DiffusionMap(n_evecs = 2, epsilon = .1, alpha = 1.0, neighbor_params = neighbor_params)
 
     """
 
-    def __init__(self, alpha=0.5, epsilon=1.0, k=64, kernel_type='gaussian', choose_eps='fixed', n_evecs=1, metric='euclidean', metric_params=None, n_algorithm='auto'):
+    def __init__(self, alpha=0.5, epsilon=1.0, k=64, kernel_type='gaussian', choose_eps='fixed', n_evecs=1, neighbor_params=None, metric='euclidean', metric_params=None):
         """
         Initializes Diffusion Map, sets parameters.
         """
         self.alpha = alpha
         self.epsilon = epsilon
+        self.k = k
         self.kernel_type = kernel_type
         self.choose_eps = choose_eps
-        self.k = k
         self.n_evecs = n_evecs
+        self.neighbor_params = neighbor_params
         self.metric = metric
         self.metric_params = metric_params
-        self.n_algorithm = n_algorithm
         return
 
     def _compute_kernel_matrix(self, X):
-        my_kernel = kernel.Kernel(type=self.kernel_type, epsilon=self.epsilon,
-                                  choose_eps=self.choose_eps, k=self.k,
-                                  metric=self.metric, metric_params=self.metric_params,
-                                  n_algorithm=self.n_algorithm)
+        my_kernel = kernel.Kernel(type=self.kernel_type, epsilon=self.epsilon, k = self.k,
+                                  choose_eps=self.choose_eps, neighbor_params=self.neighbor_params,
+                                  metric=self.metric, metric_params=self.metric_params)
         self.local_kernel = my_kernel.fit(X)
         self.epsilon = my_kernel.epsilon
         kernel_matrix = _symmetrize_matrix(my_kernel.compute(X))

src/pydiffmap/kernel.py

@@ -18,26 +18,28 @@ class Kernel(object):
         Type of kernel to construct. Currently the only option is 'gaussian', but more will be implemented.
     epsilon : scalar, optional
         Value of the length-scale parameter.
-    k : int, optional
-        Number of nearest neighbors over which to construct the kernel.
     choose_eps : string, optional
         Method for choosing the epsilon.  Currently, the only option is 'fixed' (i.e. don't), and 'bgh'.
+    k : int, optional
+        Number of nearest neighbors over which to construct the kernel.
+    neighbor_params : dict or None, optional
+        Optional parameters for the nearest Neighbor search. See scikit-learn NearestNeighbors class for details.
     metric : string, optional
         Distance metric to use in constructing the kernel.  This can be selected from any of the scipy.spatial.distance metrics, or a callable function returning the distance.
     metric_params : dict or None, optional
         Optional parameters required for the metric given.
-    n_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, optional
-        Algorithm used to compute the nearest neighbors.  See sklearn.neighbors.NearestNeighbors for details
     """
 
-    def __init__(self, type='gaussian', epsilon=1.0, choose_eps='fixed', k=64, metric='euclidean', metric_params=None, n_algorithm='auto'):
+    def __init__(self, type='gaussian', epsilon=1.0, choose_eps='fixed', k=64, neighbor_params=None, metric='euclidean', metric_params=None):
         self.type = type
         self.epsilon = epsilon
         self.choose_eps = choose_eps
+        self.k = k
         self.metric = metric
         self.metric_params = metric_params
-        self.n_algorithm = n_algorithm
-        self.k = k
+        if neighbor_params is None:
+            neighbor_params = {}
+        self.neighbor_params = neighbor_params
 
     def fit(self, X):
         """
@@ -57,10 +59,10 @@ def fit(self, X):
         # Construct Nearest Neighbor Tree
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="Parameter p is found in metric_params. The corresponding parameter from __init__ is ignored.")
-            self.neigh = NearestNeighbors(n_neighbors=self.k0,
+            self.neigh = NearestNeighbors(n_neighbors = self.k,
                                           metric=self.metric,
                                           metric_params=self.metric_params,
-                                          algorithm=self.n_algorithm)
+                                          **self.neighbor_params)
         self.neigh.fit(X)
         if self.choose_eps != 'fixed':
             self.choose_optimal_epsilon(self.choose_eps)
@@ -84,7 +86,7 @@ def compute(self, Y=None):
         if Y is None:
             Y = self.data
         # perform k nearest neighbour search on X and Y and construct sparse matrix
-        K = self.neigh.kneighbors_graph(Y, n_neighbors=self.k0, mode='distance')
+        K = self.neigh.kneighbors_graph(Y, mode='distance')
         # retrieve all nonzero elements and apply kernel function to it
         v = K.data
         if (self.type == 'gaussian'):

tests/test_kernel.py

@@ -46,8 +46,8 @@ def test_matrix_output(self, y_values, epsilon, metric, metric_params):
         assert(total_error < 1E-8)
 
     @pytest.mark.parametrize('k', np.arange(2, 14, 2))
-    @pytest.mark.parametrize('n_algorithm', ['auto', 'ball_tree'])
-    def test_neighborlists(self, k, n_algorithm):
+    @pytest.mark.parametrize('neighbor_params',[{'algorithm': 'auto'},{'algorithm': 'ball_tree'}])
+    def test_neighborlists(self, k, neighbor_params):
         """
         Test that neighborlisting gives the right number of elements.
         """
@@ -56,7 +56,7 @@ def test_neighborlists(self, k, n_algorithm):
 
         # Construct kernel matrix.
         mykernel = kernel.Kernel(type='gaussian', metric='euclidean',
-                                 epsilon=1., k=k, n_algorithm=n_algorithm)
+                                 epsilon=1., k=k0, neighbor_params=neighbor_params)
         mykernel.fit(x_values)
         K_matrix = mykernel.compute(x_values)