mutual information 2

def envelope_ecog(ecog, band, flp, fs):
    bband, aband = sg.butter(4, np.array(band)/(fs/2), btype='bandpass')
    blow, alow = sg.butter(4, flp/(fs/2), btype='low')
    ecog_band = sg.filtfilt(bband, aband, ecog, axis=0)
    ecog_abs = np.abs(ecog_band)
    ecog_env = sg.filtfilt(blow, alow, ecog_abs, axis=0)
    return ecog_env

def envelope_sound(sound, fhp, flp, fs):
    bband, aband = sg.butter(4, fhp/(fs/2), btype='high')
    blow, alow = sg.butter(4, flp/(fs/2), btype='low')
    sound_band = sg.filtfilt(bband, aband, sound, axis=0)
    sound_logabs = np.log(np.abs(sound_band)+1)
    sound_env = sg.filtfilt(blow, alow, sound_logabs, axis=0)
    return sound_env

def _mutual_information(x, y, bins):
    c_xy, b1, b2 = np.histogram2d(x, y, bins)
    mi = mutual_info_score(None, None, contingency=c_xy)
    return mi

def mutual_information(ecog_envs_pictures, sound_envs_pictures, clip_alpha=0.99, bins=10):
    # clip
    q = np.quantile(ecog_envs_pictures, clip_alpha, axis=0)
    sound_entropy = _mutual_information(sound_envs_pictures, sound_envs_pictures, bins)
    sound_entropy = sound_entropy if sound_entropy > 1e-7 else 1000
    mi = []
    for i in range(ecog_envs_pictures.shape[1]):
        ecog_env_in_channel = ecog_envs_pictures[:,i]
        ecog_env_in_channel[ecog_env_in_channel > q[i]] = q[i]
        ecog_entropy = _mutual_information(ecog_env_in_channel, ecog_env_in_channel, bins)
        ecog_entropy = ecog_entropy if ecog_entropy > 1e-7 else 1000
        mi_channel = _mutual_information(sound_envs_pictures, ecog_env_in_channel, bins) / np.sqrt(sound_entropy*ecog_entropy)
        mi.append(mi_channel)
    return np.array(mi)