KSG Conditional MI Estimation

KSG Conditional MI Estimation#

The Kraskov–Grassberger–Stögbauer (KSG) method was originally developed for estimating MI (Kraskov-Stoegbauer-Grassberger (KSG) MI Estimation) [KSG11]. Ever since it has been adapted to estimate other information measures such as Conditional MI (CMI) [FP07]. The CMI in between two variables \(X\) and \(Y\) by conditioning the third variable \(Z\) is obtained by the following expression:

\[ I(X; Y \mid Z) = \psi(k) + \langle \psi(n_z(i) + 1) - \psi(n_{xz}(i) + 1) - \psi(n_{yz}(i) + 1) \rangle \]

where

  • \(k\) is the number of nearest neighbors,

  • \(n_x(\cdot)\) refers to the number of neighbors which are with in a hypercube that defines the search range around a statevector. The size of the hypercube in each of the marginal spaces is defined based on the distance to the \(k\)-th nearest neighbor in the joint space.

  • \(\psi(\cdot)\) denotes the digamma function,

  • \(\langle \cdot \rangle\) represents the expectation operator.

Similarly, the local conditional MI estimator is:

\[ i(x; y \mid z) = \psi(k) + \psi(n_z(i) + 1) - \psi(n_{xz}(i) + 1) - \psi(n_{yz}(i) + 1) \]
import infomeasure as im

x = [-2.49, 1.64, -3.05, 7.95, -5.96, 1.77, -5.24, 7.03, -0.11, -1.86]
y = [7.95, -5.96, 7.03, -0.11, -1.86, 1.77, -2.49, 1.64, -3.05, -5.24]
z = [-8.59, 8.41, 3.76, 3.77, 5.69, 1.75, -3.2, -4.0, -4.0, 6.85]

im.cmi(x, y, cond=z, approach='ksg')
np.float64(-0.061309523809523855)

The Local Conditional MI is calculated as follows:

est = im.estimator(
    x, y, cond=z,
    measure='cmi',  # or 'conditional_mutual_information'
    approach='ksg'
)
est.local_vals()
array([ 0.        ,  0.        , -0.18214286,  0.        , -0.05      ,
       -0.10714286, -0.14047619,  0.        , -0.08333333, -0.05      ])