由样本集合生成概率图模型(chou-liu算法)

# encoding = utf-8
import numpy as np
import math
p_matrix = np.zeros((5, 5, 4), dtype=float)  # store the posibilities of five entity,0,0 0,1 1,0 1,1
i_matrix = np.zeros((5, 5), dtype=float)

def new_cal_pmatrix(arr):
    l = len(arr)
    for i in range(l):
        for j in range(i + 1, l):
            if arr[i] == 0:
                if arr[j] == 0:
                    p_matrix[i, j, 0] += 1
                if arr[j] == 1:
                    p_matrix[i, j, 1] += 1
            if arr[i] == 1:
                if arr[j] == 0:
                    p_matrix[i, j, 2] += 1
                if arr[j] == 1:
                    p_matrix[i, j, 3] += 1


def cal_imatrix(i_matrix, pxi):  ## ----------------------->calculate i_matrix
    i = 0
    dist = {0: (0, 0), 1: (0, 1), 2: (1, 0), 3: (1, 1)}
    while (i < 4):
        j = i + 1
        while (j <= 4):
            for k in range(4):
                i_matrix[i, j] += p_matrix[i, j, k] * math.log(
                        p_matrix[i, j, k] / (pxi[i, dist[k][0]] * pxi[j, dist[k][1]]),2)
            j += 1
        i += 1

def kruscal(i_matrix):
    result = []
    s = []
    temp = i_matrix
    k = 1
    while (k <= 4):
        i, j = np.unravel_index(temp.argmax(), temp.shape)  # calculate max value in temp
        temp[i, j] = 0
        if i in s and j in s:
            continue
        else:
            if i not in s:
                s.append(i)
            if j not in s:
                s.append(j)
        tup = (i, j)
        result.append(tup)
        k += 1
    return result


def main():
    global p_matrix
    trai = np.loadtxt('/home/zw/PycharmProjects/PGM/traindata.txt', dtype=int)
    train = trai[:, 1:]
    pxi = np.zeros((5, 2), dtype=float)  # store p(xi)
    len1 = train.shape[1]
    for sample in train:
        for i in range(5):
            if sample[i] == 1:
                pxi[i, 1] +=1
            else:
                pxi[i, 0] += 1
    # ------------------------->compute p_matrix
    for sample in train:
        new_cal_pmatrix(sample)
    # <------------------------------_compute p_matrix
    p_matrix = p_matrix / 100.0
    print p_matrix
    pxi = pxi / 100
    print pxi
    cal_imatrix(i_matrix, pxi)
    print i_matrix
    print kruscal(i_matrix)

main()