refactor: improve acquisition and disagreement modules

Author: nikolay-bushkov

modAL/acquisition.py

@@ -1,13 +1,16 @@
 """
 Acquisition functions for Bayesian optimization.
 """
+from typing import Tuple
 
 import numpy as np
-
 from scipy.stats import norm
 from scipy.special import ndtr
 from sklearn.exceptions import NotFittedError
+
+from modAL.models import BayesianOptimizer
 from modAL.utils.selection import multi_argmax
+from modAL.utils.data import modALinput
 
 
 def PI(mean, std, max_val, tradeoff):
@@ -30,209 +33,129 @@ def UCB(mean, std, beta):
 """
 
 
-def optimizer_PI(optimizer, X, tradeoff=0):
+def optimizer_PI(optimizer: BayesianOptimizer, X: modALinput, tradeoff: float = 0) -> np.ndarray:
     """
     Probability of improvement acquisition function for Bayesian optimization.
 
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
-
-    :param X:
-        The samples for which the probability of improvement is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
-
-    :param tradeoff:
-        Value controlling the tradeoff parameter.
-    :type tradeoff:
-        float
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the probability of improvement is to be calculated.
+        tradeoff: Value controlling the tradeoff parameter.
 
-    :returns:
-      - **pi** *(numpy.ndarray of shape (n_samples, ))* --
+    Returns:
         Probability of improvement utility score.
     """
     try:
         mean, std = optimizer.predict(X, return_std=True)
         std = std.reshape(-1, 1)
     except NotFittedError:
-        mean, std = np.zeros(shape=(len(X), 1)), np.ones(shape=(len(X), 1))
+        mean, std = np.zeros(shape=(X.shape[0], 1)), np.ones(shape=(X.shape[0], 1))
 
     return PI(mean, std, optimizer.y_max, tradeoff)
 
 
-def optimizer_EI(optimizer, X, tradeoff=0):
+def optimizer_EI(optimizer: BayesianOptimizer, X: modALinput, tradeoff: float = 0) -> np.ndarray:
     """
     Expected improvement acquisition function for Bayesian optimization.
 
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the expected improvement is to be calculated.
+        tradeoff: Value controlling the tradeoff parameter.
 
-    :param X:
-        The samples for which the expected improvement is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
-
-    :param tradeoff:
-        Value controlling the tradeoff parameter.
-    :type tradeoff:
-        float
-
-    :returns:
-      - **ei** *(numpy.ndarray of shape (n_samples, ))* --
+    Returns:
         Expected improvement utility score.
     """
     try:
         mean, std = optimizer.predict(X, return_std=True)
         std = std.reshape(-1, 1)
     except NotFittedError:
-        mean, std = np.zeros(shape=(len(X), 1)), np.ones(shape=(len(X), 1))
+        mean, std = np.zeros(shape=(X.shape[0], 1)), np.ones(shape=(X.shape[0], 1))
 
     return EI(mean, std, optimizer.y_max, tradeoff)
 
 
-def optimizer_UCB(optimizer, X, beta=1):
+def optimizer_UCB(optimizer: BayesianOptimizer, X: modALinput, beta: float = 1) -> np.ndarray:
     """
     Upper confidence bound acquisition function for Bayesian optimization.
 
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
-
-    :param X:
-        The samples for which the upper confidence bound is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the upper confidence bound is to be calculated.
+        beta: Value controlling the beta parameter.
 
-    :param beta:
-        Value controlling the beta parameter.
-    :type beta:
-        float
-
-    :returns:
-      - **ucb** *(numpy.ndarray of shape (n_samples, ))* --
+    Returns:
         Upper confidence bound utility score.
     """
     try:
         mean, std = optimizer.predict(X, return_std=True)
         std = std.reshape(-1, 1)
     except NotFittedError:
-        mean, std = np.zeros(shape=(len(X), 1)), np.ones(shape=(len(X), 1))
+        mean, std = np.zeros(shape=(X.shape[0], 1)), np.ones(shape=(X.shape[0], 1))
 
     return UCB(mean, std, beta)
 
+
 """
 --------------------------------------------
 Query strategies using acquisition functions
 --------------------------------------------
 """
 
-def max_PI(optimizer, X, tradeoff=0, n_instances=1):
+
+def max_PI(optimizer: BayesianOptimizer, X: modALinput, tradeoff: float = 0,
+           n_instances: int = 1) -> Tuple[np.ndarray, modALinput]:
     """
     Maximum PI query strategy. Selects the instance with highest probability of improvement.
 
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
-
-    :param X:
-        The samples for which the probability of improvement is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
-
-    :param tradeoff:
-        Value controlling the tradeoff parameter.
-    :type tradeoff:
-        float
-
-    :param n_instances:
-        Number of samples to be queried.
-    :type n_instances:
-        int
-
-    :returns:
-      - **query_idx** *(numpy.ndarray of shape (n_instances, ))* --
-        The indices of the instances from X chosen to be labelled.
-      - **X[query_idx]** *(numpy.ndarray of shape (n_instances, n_features))* --
-        The instances from X chosen to be labelled.
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the probability of improvement is to be calculated.
+        tradeoff: Value controlling the tradeoff parameter.
+        n_instances: Number of samples to be queried.
+
+    Returns:
+        The indices of the instances from X chosen to be labelled; the instances from X chosen to be labelled.
     """
     pi = optimizer_PI(optimizer, X, tradeoff=tradeoff)
     query_idx = multi_argmax(pi, n_instances=n_instances)
 
     return query_idx, X[query_idx]
 
 
-def max_EI(optimizer, X, tradeoff=0, n_instances=1):
+def max_EI(optimizer: BayesianOptimizer, X: modALinput, tradeoff: float = 0,
+           n_instances: int = 1) -> Tuple[np.ndarray, modALinput]:
     """
     Maximum EI query strategy. Selects the instance with highest expected improvement.
 
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
-
-    :param X:
-        The samples for which the expected improvement is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
-
-    :param tradeoff:
-        Value controlling the tradeoff parameter.
-    :type tradeoff:
-        float
-
-    :param n_instances:
-        Number of samples to be queried.
-    :type n_instances:
-        int
-
-    :returns:
-      - **query_idx** *(numpy.ndarray of shape (n_instances, )*) --
-        The indices of the instances from X chosen to be labelled.
-      - **X[query_idx]** *(numpy.ndarray of shape (n_instances, n_features)*) --
-        The instances from X chosen to be labelled.
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the expected improvement is to be calculated.
+        tradeoff: Value controlling the tradeoff parameter.
+        n_instances: Number of samples to be queried.
+
+    Returns:
+        The indices of the instances from X chosen to be labelled; the instances from X chosen to be labelled.
     """
     ei = optimizer_EI(optimizer, X, tradeoff=tradeoff)
     query_idx = multi_argmax(ei, n_instances=n_instances)
 
     return query_idx, X[query_idx]
 
 
-def max_UCB(optimizer, X, beta=1, n_instances=1):
+def max_UCB(optimizer: BayesianOptimizer, X: modALinput, beta: float = 1,
+            n_instances: int = 1) -> Tuple[np.ndarray, modALinput]:
     """
-    Maximum UCB query strategy. Selects the instance with highest upper confidence
-    bound.
-
-    :param optimizer:
-        The BayesianEstimator object for which the utility is to be calculated.
-    :type optimizer:
-        modAL.models.BayesianEstimator object
-
-    :param X:
-        The samples for which the probability of improvement is to be calculated.
-    :type X:
-        numpy.ndarray of shape (n_samples, n_features)
-
-    :param beta:
-        Value controlling the beta parameter.
-    :type beta:
-        float
-
-    :param n_instances:
-        Number of samples to be queried.
-    :type n_instances:
-        int
-
-    :returns:
-      - **query_idx** *(numpy.ndarray of shape (n_instances, ))* --
-        The indices of the instances from X chosen to be labelled.
-      - **X[query_idx]** *(numpy.ndarray of shape (n_instances, n_features))* --
-        The instances from X chosen to be labelled.
+    Maximum UCB query strategy. Selects the instance with highest upper confidence bound.
+
+    Args:
+        optimizer: The :class:`~modAL.models.BayesianOptimizer` object for which the utility is to be calculated.
+        X: The samples for which the maximum upper confidence bound is to be calculated.
+        beta: Value controlling the beta parameter.
+        n_instances: Number of samples to be queried.
+
+    Returns:
+        The indices of the instances from X chosen to be labelled; the instances from X chosen to be labelled.
     """
     ucb = optimizer_UCB(optimizer, X, beta=beta)
     query_idx = multi_argmax(ucb, n_instances=n_instances)