diffusion_map.pyx

#cython: wraparound=False, boundscheck=False, cdivision=True, overflowcheck=False
#cython: profile=False, nonecheck=False, cdivision_warnings=True

import numpy as np
cimport numpy as np
from libc.math import sqrt, exp
from sklearn.utils.extmath import fast_dot
from sklearn.metrics.pairwise import pairwise_distances

def diffusion_map(X, n_components=2, metric="euclidean"):
	distance_matrix = pairwise_distances(X, metric=metric) 
	P = calc_markov_matrix(distance_matrix)	
	
	evals, evecs = np.eigsh(P, n_components + 1, which="LA")
	evals, evecs = evals[::-1], evecs[:, ::-1] # largest to smallest
	evals, evecs = evals[1:], evecs[:, 1:] # discard 1st bc it will always be 1
	
	proj = fast_dot(P, evecs[:, 1:])	
	return proj

def calc_markov_matrix(distance_matrix, epsilon=0.2):
	N = distance_matrix.shape[0]
	return np.asarray(_calc_markov_matrix(distance_matrix, epsilon, N))

cdef double[:,:] _calc_markov_matrix(double[:,:] distance_matrix, double epsilon, int N):
	cdef:
		int i, j
		double[:] D = np.zeros(N)
		double[:] Dtilda = np.zeros(N)
		double[:,:] K = np.zeros((N,N))
		double[:,:] Ktilda = np.zeros((N,N))
		double[:,:] P = np.zeros((N,N))

	with nogil:
		for i in range(N):
			for j in range(N):
				K[i, j] = exp( (-distance_matrix[i, j]**2) / (2*epsilon**2)
				D[i] += K[i, j]

		for i in range(N):
			for j in range(N):
				Ktilda[i, j] = K[i, j] / sqrt(D[i]**2)
				Dtilda[i] += Ktilda[i, j]

		for i in range(N):
			for j in range(N):
				P[i, j] = Ktilda[i, j] / Dtilda[i]

	return P