#!/usr/bin/env python
# -*- coding: latin-1 -*-
#
#   Copyright 2016-2019 Blaise Frederick
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
from __future__ import print_function, division

import numpy as np
import pylab as plt
import rapidtide.io as tide_io

from rapidtide.correlate import calc_MI


def test_calc_MI(display=False):
    inlen = 1000
    offset = 100
    filename1 = "testdata/lforegressor.txt"
    filename2 = "testdata/lforegressor.txt"
    sig1 = tide_io.readvec(filename1)
    sig2 = np.power(sig1, 2.0)
    sig3 = np.power(sig1, 3.0)
    
    kstart = 3
    kend = 100
    linmivals = []
    sqmivals = []
    cubemivals = []
    for clustersize in range(kstart, kend, 2):
        linmivals.append(calc_MI(sig1, sig1, clustersize) / np.log(clustersize))
        sqmivals.append(calc_MI(sig2, sig1, clustersize) / np.log(clustersize))
        cubemivals.append(calc_MI(sig3, sig1, clustersize) / np.log(clustersize))

    if display:
        plt.figure()
        #plt.ylim([-1.0, 3.0])
        plt.plot(np.array(range(kstart, kend, 2)), np.array(linmivals), 'r')
        plt.plot(np.array(range(kstart, kend, 2)), np.array(sqmivals), 'g')
        plt.plot(np.array(range(kstart, kend, 2)), np.array(cubemivals), 'b')
        #print('maximum occurs at offset', np.argmax(stdcorrelate_result) - midpoint + 1)
        plt.legend(['Mutual information', 'Squared mutual information', 'Cubed mutual information'])
        plt.show()

    aethresh = 10
    np.testing.assert_almost_equal(1.0, 1.0, 1e-5)


def main():
    test_calc_MI(display=True)


if __name__ == '__main__':
    main()

Test for mutual information function
Test for mutual information function
