An entropy calculator for UVic SENG 474 - Data Mining

entropy.py

'''An entropy calculator for UVic SENG 474 - Data Mining'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from math import log

def main():
    '''Calculates the entropy H(X) for the data in a given attribute'''
    num_attr = int(input("How many classes in the attribute?\n"))
    num_el = int(input("How many elements (eg. total cardinality) in the attribute?\n"))
    entropy = 0
    attributes = []
    elements = ""
    total = 0

    for i in range(1, num_attr+1):
        prompt = "Enter number of elements (out of %d) in class %d:\n" % (num_el, i)
        numerator = int(input(prompt))
        fraction = numerator / num_el
        attributes.append(fraction)

        elements += "%d/%d" % (numerator, num_el)
        if i < num_attr:
            elements += ", "


    for attr in attributes:
        entropy -= attr * log(attr, 2)
        total += attr

    if total == 1:
        print("H(%s) = %.4f" % (elements, entropy))
    else:
        print("\n--- ERROR ---\n")
        print("The probabilities must add up to 1.  Please try again.")
        print("\n--- ERROR ---\n\n")
        reset()

def reset():
    '''Restarts the main function'''
    main()

if __name__ == "__main__":
    '''Runs file as a script if not imported as a module'''
    main()