Script to calculate entropy for any column in a file

entropy.py

# Script to calculate entropy for any column in a file.

from __future__ import print_function

import numpy as np

def entropy(file_path, sep, col_index, col_name):
    '''Method to calculate entropy for any col_index
    in a file where columns are seperated by sep'''
    distribution = np.asarray(list(read_column(file_path, sep, col_index)))
    probs = [np.mean(distribution == c) for c in set(distribution)]
    _entropy = np.sum(-p * np.log2(p) for p in probs)
    print(_entropy)
    print("Entropy for %s = %f" % (col_name, _entropy))


def read_column(file_path, sep, col_index):
    '''Method to read col_index column in a file where columns are seperated by
    sep. It returns an iterator over the list of entries in the column'''

    with open(file_path, "r") as f:
        for line in f:
            yield float(line.strip().split(sep)[col_index])

file_path = "file_path"
sep = "\t"
col_index = 0
col_name = "col_name"
entropy(file_path, sep, col_index, col_name)