add type annotations for adaptive/learner/skopt_learner.py

Author: basnijholt

adaptive/learner/skopt_learner.py

@@ -1,6 +1,8 @@
 import collections
+from typing import Callable, List, Tuple, Union
 
 import numpy as np
+from numpy import float64
 from skopt import Optimizer
 
 from adaptive.learner.base_learner import BaseLearner
@@ -23,13 +25,15 @@ class SKOptLearner(Optimizer, BaseLearner):
         Arguments to pass to ``skopt.Optimizer``.
     """
 
-    def __init__(self, function, **kwargs):
+    def __init__(self, function: Callable, **kwargs) -> None:
         self.function = function
         self.pending_points = set()
         self.data = collections.OrderedDict()
         super().__init__(**kwargs)
 
-    def tell(self, x, y, fit=True):
+    def tell(
+        self, x: Union[float64, List[float64]], y: float64, fit: bool = True
+    ) -> None:
         if isinstance(x, collections.abc.Iterable):
             self.pending_points.discard(tuple(x))
             self.data[tuple(x)] = y
@@ -48,7 +52,7 @@ def remove_unfinished(self):
         pass
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> Union[float64, float]:
         if not self.models:
             return np.inf
         else:
@@ -58,7 +62,14 @@ def loss(self, real=True):
             # estimator of loss, but it is the cheapest.
             return 1 - model.score(self.Xi, self.yi)
 
-    def ask(self, n, tell_pending=True):
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> Union[
+        Tuple[List[float64], List[float64]],
+        Tuple[List[List[float64]], List[float64]],
+        Tuple[List[List[float64]], List[float]],
+        Tuple[List[float64], List[float]],
+    ]:
         if not tell_pending:
             raise NotImplementedError(
                 "Asking points is an irreversible "
@@ -72,7 +83,7 @@ def ask(self, n, tell_pending=True):
             return [p[0] for p in points], [self.loss() / n] * n
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.Xi)