modAL.acquisition EI implemented

Author: cosmic-cortex

modAL/acquisition.py

@@ -7,17 +7,33 @@
 -----------------------------------------------
 """
 
+from scipy.stats import norm
 from scipy.special import ndtr
 from modAL.utils.selection import multi_argmax
 
 
 def PI(optimizer, X, tradeoff=0):
     mean, std = optimizer.predict(X, return_std=True)
+
     return ndtr((mean - optimizer.max_val - tradeoff)/std)
 
 
+def EI(optimizer, X, tradeoff=0):
+    mean, std = optimizer.predict(X, return_std=True)
+    z = (mean - optimizer.max_val - tradeoff)/std
+
+    return (mean - optimizer.max_val - tradeoff)*ndtr(z) + std*norm.pdf(z)
+
+
 def max_PI(optimizer, X, tradeoff=0, n_instances=1):
     pi = PI(optimizer, X, tradeoff=tradeoff)
     query_idx = multi_argmax(pi, n_instances=n_instances)
 
-    return query_idx, X[query_idx]
+    return query_idx, X[query_idx]
+
+
+def max_EI(optimizer, X, tradeoff=0, n_instances=1):
+    ei = EI(optimizer, X, tradeoff=tradeoff)
+    query_idx = multi_argmax(ei, n_instances=n_instances)
+
+    return query_idx, X[query_idx]

tests/core_tests.py

@@ -16,7 +16,7 @@
 from collections import namedtuple
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.metrics import confusion_matrix
-from scipy.stats import entropy
+from scipy.stats import entropy, norm
 from scipy.special import ndtr
 
 
@@ -145,6 +145,28 @@ def test_PI(self):
                 modAL.acquisition.PI(optimizer, np.random.rand(n_samples, 2), tradeoff)
             )
 
+    def test_EI(self):
+        for n_samples in range(1, 100):
+            mean = np.random.rand(n_samples, )
+            std = np.random.rand(n_samples, )
+            tradeoff = np.random.rand()
+            max_val = np.random.rand()
+
+            mock_estimator = mock.MockEstimator(
+                predict_return=(mean, std)
+            )
+
+            optimizer = modAL.models.BayesianOptimizer(estimator=mock_estimator)
+            optimizer._set_max([max_val])
+
+            true_EI = (mean - optimizer.max_val - tradeoff) * ndtr((mean - optimizer.max_val - tradeoff)/std)\
+                      + std * norm.pdf((mean - optimizer.max_val - tradeoff)/std)
+
+            np.testing.assert_almost_equal(
+                true_EI,
+                modAL.acquisition.EI(optimizer, np.random.rand(n_samples, 2), tradeoff)
+            )
+
 
 class TestUncertainties(unittest.TestCase):