├── README.md ├── LICENSE └── knn_entropy.py /README.md: -------------------------------------------------------------------------------- 1 | # A set of functions to compute entropy and mutual information 2 | using k-nearest neighbor estimators. 3 | ## Dependencies 4 | numpy, scipy 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Blake 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /knn_entropy.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.special import gamma,psi 3 | from scipy.spatial.distance import cdist 4 | 5 | # Get the k nearest neighbors 6 | # between points in a random variable/vector. 7 | # Uses Euclidean style distances. 8 | def k_nearest_neighbors(X, k=1): 9 | #length 10 | nX = len(X) 11 | #initialize knn dict 12 | knn = {key: [] for key in xrange(nX)} 13 | #make sure X has the right shape for the cdist function 14 | X = np.reshape(X, (nX,-1)) 15 | dists_arr = cdist(X, X) 16 | distances = [[i,j,dists_arr[i,j]] for i in xrange(nX-1) for j in xrange(i+1,nX)] 17 | #sort distances 18 | distances.sort(key=lambda x: x[2]) 19 | #pick up the k nearest 20 | for d in distances: 21 | i = d[0] 22 | j = d[1] 23 | dist = d[2] 24 | if len(knn[i]) < k: 25 | knn[i].append([j, dist]) 26 | if len(knn[j]) < k: 27 | knn[j].append([i, dist]) 28 | return knn 29 | 30 | 31 | # knn kth neighbor distances for entropy calcs. 32 | def kth_nearest_neighbor_distances(X, k=1): 33 | 34 | #length 35 | nX = len(X) 36 | #make sure X has the right shape for the cdist function 37 | X = np.reshape(X, (nX,-1)) 38 | dists_arr = cdist(X, X) 39 | #sorts each row 40 | dists_arr.sort() 41 | return [dists_arr[i][k] for i in xrange(nX)] 42 | 43 | def shannon_entropy(X, k=1, kth_dists=None): 44 | # KL entropy estimator from 45 | # https://arxiv.org/pdf/1506.06501v1.pdf 46 | # Also see: 47 | # https://www.cs.tut.fi/~timhome/tim/tim/core/differential_entropy_kl_details.htm 48 | # and 49 | # Kozachenko, L. F. & Leonenko, N. N. 1987 Sample estimate of entropy 50 | # of a random vector. Probl. Inf. Transm. 23, 95-101. 51 | # the kth nearest neighbor distances 52 | r_k = kth_dists 53 | if kth_dists is None: 54 | r_k = kth_nearest_neighbor_distances(X, k=k) 55 | #length 56 | n = len(X) 57 | #dimension 58 | d = 1 59 | if len(X.shape) == 2: 60 | d = X.shape[1] 61 | # volume of unit ball in d^n 62 | v_unit_ball = np.pi**(0.5*d)/gamma(0.5*d + 1.0) 63 | # log distances 64 | lr_k = np.log(r_k) 65 | # Shannon entropy estimate 66 | H = psi(n) - psi(k) + np.log(v_unit_ball) + (np.float(d)/np.float(n))*( lr_k.sum()) 67 | 68 | return H 69 | 70 | # Not entirely sure this one is correct. 71 | def shannon_entropy_pc(X, k=1, kth_dists=None): 72 | # entropy estimator from 73 | # F. Perez-Cruz, (2008). Estimation of Information Theoretic Measures 74 | # for Continuous Random Variables. Advances in Neural Information 75 | # Processing Systems 21 (NIPS). Vancouver (Canada), December. 76 | # https://papers.nips.cc/paper/3417-estimation-of-information-theoretic-measures-for-continuous-random-variables.pdf 77 | # 78 | # 79 | r_k = kth_dists 80 | if kth_dists is None: 81 | r_k = np.array(kth_nearest_neighbor_distances(X, k=k)) 82 | n = len(X) 83 | d = 1 84 | if len(X.shape) == 2: 85 | d = X.shape[1] 86 | #volume of the unit ball 87 | v_unit_ball = np.pi**(0.5*d)/gamma(0.5*d + 1.0) 88 | # probability estimator using knn distances 89 | p_k_hat = (k / (n -1.0)) * (1.0/v_unit_ball) * (1.0/r_k**d) 90 | # log probability 91 | log_p_k_hat = np.log(p_k_hat) 92 | #entropy estimator 93 | h_k_hat = log_p_k_hat.sum() / (-1.0*n) 94 | return h_k_hat 95 | 96 | #mutual information 97 | def mutual_information(var_tuple, k=2): 98 | nvar = len(var_tuple) 99 | #make sure the input arrays are properly shaped for hstacking 100 | var_tuple = tuple( var_tuple[i].reshape(len(var_tuple[i]),-1) for i in xrange(nvar) ) 101 | #compute the individual entropies of each variable 102 | Hx = [shannon_entropy(var_tuple[i],k=k) for i in xrange(nvar)] 103 | Hx = np.array(Hx) 104 | # and get the sum 105 | Hxtot = Hx.sum() 106 | #now get the entropy of the joint distribution 107 | joint = np.hstack(var_tuple) 108 | Hjoint = shannon_entropy(joint, k=k) 109 | #get the mutual information 110 | MI = Hxtot - Hjoint 111 | #set to zero if value is negative 112 | if MI < 0.0: MI = 0.0 113 | #return 114 | return MI 115 | 116 | # conditional mutual information 117 | def conditional_mutual_information(var_tuple, cond_tuple, k=2): 118 | nvar = len(var_tuple) 119 | ncon = len(cond_tuple) 120 | #make sure the input arrays are properly shaped for hstacking 121 | var_tuple = tuple( var_tuple[i].reshape(len(var_tuple[i]),-1) for i in xrange(nvar) ) 122 | cond_tuple = tuple( cond_tuple[i].reshape(len(cond_tuple[i]),-1) for i in xrange(ncon) ) 123 | # compute pair joint entropies 124 | Hxz = [shannon_entropy(np.hstack(var_tuple[i]+cond_tuple), k=k) for i in xrange(nvar)] 125 | Hxz = np.array(Hxz) 126 | jtup = var_tuple + cond_tuple 127 | joint = np.hstack( jtup ) 128 | Hj = shannon_entropy(joint, k=k) 129 | Hz = 0.0 130 | if len(cond_tuple) > 1: 131 | joint = np.hstack(cond_tuple) 132 | Hz = shannon_entropy(joint, k=k) 133 | else: 134 | Hz = shannon_entropy(cond_tuple[0], k=k) 135 | Hxzsum = Hxz.sum() 136 | # print "Hxzsum: ",Hxzsum, " Hj: ",Hj, " Hz: ",Hz 137 | MIc = Hxzsum - Hj - Hz 138 | # print "MIc: ",MIc 139 | if MIc < 0.0: MIc = 0.0 140 | #print MIc 141 | return MIc 142 | --------------------------------------------------------------------------------