Fixed a bug in DecisionTreeLearner in count function

agarwl · agarwl · commit 6368474a2f3e · 2017-03-04T00:52:28.000+05:30
diff --git a/learning.py b/learning.py
@@ -370,7 +370,7 @@ def plurality_value(examples):
 
     def count(attr, val, examples):
         "Count the number of examples that have attr = val."
-        return len(e[attr] == val for e in examples) #count(e[attr] == val for e in examples)
+        return sum(e[attr] == val for e in examples)
 
     def all_same_class(examples):
         "Are all these examples in the same target class?"
diff --git a/tests/test_learning.py b/tests/test_learning.py
@@ -2,8 +2,10 @@
 import math
 from learning import (
     parse_csv, weighted_mode, weighted_replicate,
-    rms_error, manhattan_distance, mean_boolean_error, mean_error
+    rms_error, manhattan_distance, mean_boolean_error, mean_error,
+    information_content
 )
+from utils import rounder
 
 
 def test_parse_csv():
@@ -44,3 +46,14 @@ def test_mean_error():
     assert mean_error([1,0], [0,1]) ==  1
     assert mean_error([0,0], [0,-1]) ==  0.5
     assert mean_error([0,0.5], [0,-0.5]) == 0.5
+
+def test_information_content():
+    # Entropy of a 4 sided fair-die
+    assert information_content([1,1,1,1]) == 2
+    # Entropy of a fair coin
+    assert information_content([1,1]) == 1
+    # Entropy of a biased coin with heads probability 0.99
+    assert rounder(information_content([0.99,0.01])) == 0.0808
+
+if __name__ == '__main__':
+    pytest.main()
