From 50143e55f1d3ab0707cb05aa65d965a23aa790eb Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:03:58 -0700
Subject: [PATCH 01/21] Add type-annotations to Runner

---
 adaptive/runner.py | 285 ++++++++++++++++++++++++++-------------------
 1 file changed, 165 insertions(+), 120 deletions(-)

diff --git a/adaptive/runner.py b/adaptive/runner.py
index ecc7fe14b..2a345a952 100644
--- a/adaptive/runner.py
+++ b/adaptive/runner.py
@@ -13,19 +13,28 @@
 import traceback
 import warnings
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Set
 
 import loky
+from _asyncio import Future, Task
 
 from adaptive.notebook_integration import in_ipynb, live_info, live_plot
 
+_ThirdPartyClient = []
+_ThirdPartyExecutor = [loky.reusable_executor._ReusablePoolExecutor]
+_FutureTypes = [Set[Future], Set[Task]]
+
 if TYPE_CHECKING:
     from adaptive import BaseLearner
 
 try:
     import ipyparallel
+    from ipyparallel.client.asyncresult import AsyncResult
 
     with_ipyparallel = True
+    _ThirdPartyClient.append(ipyparallel.Client)
+    _ThirdPartyExecutor.append(ipyparallel.client.view.ViewExecutor)
+    _FutureTypes.append(AsyncResult)
 except ModuleNotFoundError:
     with_ipyparallel = False
 
@@ -33,6 +42,8 @@
     import distributed
 
     with_distributed = True
+    _ThirdPartyClient.append(distributed.Client)
+    _ThirdPartyExecutor.append(distributed.cfexecutor.ClientExecutor)
 except ModuleNotFoundError:
     with_distributed = False
 
@@ -40,16 +51,18 @@
     import mpi4py.futures
 
     with_mpi4py = True
+    _ThirdPartyExecutor.append(mpi4py.futures.MPIPoolExecutor)
 except ModuleNotFoundError:
     with_mpi4py = False
 
-
 with suppress(ModuleNotFoundError):
     import uvloop
 
     asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
+# -- Runner definitions
+
 if platform.system() == "Linux":
     _default_executor = concurrent.ProcessPoolExecutor
 else:
@@ -62,8 +75,79 @@
     _default_executor = loky.get_reusable_executor
 
 
+# -- Internal executor-related, things
+
+
+class SequentialExecutor(concurrent.Executor):
+    """A trivial executor that runs functions synchronously.
+
+    This executor is mainly for testing.
+    """
+
+    def submit(self, fn: Callable, *args, **kwargs) -> Future:
+        fut: concurrent.Future = concurrent.Future()
+        try:
+            fut.set_result(fn(*args, **kwargs))
+        except Exception as e:
+            fut.set_exception(e)
+        return fut
+
+    def map(self, fn, *iterable, timeout=None, chunksize=1):
+        return map(fn, iterable)
+
+    def shutdown(self, wait=True):
+        pass
+
+
+def _ensure_executor(
+    executor: _ThirdPartyClient | concurrent.Executor | None,
+) -> concurrent.Executor:
+    if executor is None:
+        executor = concurrent.ProcessPoolExecutor()
+
+    if isinstance(executor, concurrent.Executor):
+        return executor
+    elif with_ipyparallel and isinstance(executor, ipyparallel.Client):
+        return executor.executor()
+    elif with_distributed and isinstance(executor, distributed.Client):
+        return executor.get_executor()
+    else:
+        raise TypeError(
+            "Only a concurrent.futures.Executor, distributed.Client,"
+            " or ipyparallel.Client can be used."
+        )
+
+
+def _get_ncores(
+    ex: (
+        _ThirdPartyExecutor
+        | concurrent.ProcessPoolExecutor
+        | concurrent.ThreadPoolExecutor
+        | SequentialExecutor
+    ),
+):
+    """Return the maximum  number of cores that an executor can use."""
+    if with_ipyparallel and isinstance(ex, ipyparallel.client.view.ViewExecutor):
+        return len(ex.view)
+    elif isinstance(
+        ex, (concurrent.ProcessPoolExecutor, concurrent.ThreadPoolExecutor)
+    ):
+        return ex._max_workers  # not public API!
+    elif isinstance(ex, loky.reusable_executor._ReusablePoolExecutor):
+        return ex._max_workers  # not public API!
+    elif isinstance(ex, SequentialExecutor):
+        return 1
+    elif with_distributed and isinstance(ex, distributed.cfexecutor.ClientExecutor):
+        return sum(n for n in ex._client.ncores().values())
+    elif with_mpi4py and isinstance(ex, mpi4py.futures.MPIPoolExecutor):
+        ex.bootup()  # wait until all workers are up and running
+        return ex._pool.size  # not public API!
+    else:
+        raise TypeError(f"Cannot get number of cores for {ex.__class__}")
+
+
 class BaseRunner(metaclass=abc.ABCMeta):
-    r"""Base class for runners that use `concurrent.futures.Executors`.
+    r"""Base class for runners that use `concurrent.futures.Executor`\'s.
 
     Parameters
     ----------
@@ -120,53 +204,59 @@ class BaseRunner(metaclass=abc.ABCMeta):
 
     def __init__(
         self,
-        learner,
-        goal,
+        learner: BaseLearner,
+        goal: Callable,
         *,
-        executor=None,
-        ntasks=None,
-        log=False,
-        shutdown_executor=False,
-        retries=0,
-        raise_if_retries_exceeded=True,
-    ):
+        executor: (
+            _ThirdPartyClient
+            | _ThirdPartyExecutor
+            | concurrent.ProcessPoolExecutor
+            | concurrent.ThreadPoolExecutor
+            | SequentialExecutor
+        ) = None,
+        ntasks: int = None,
+        log: bool = False,
+        shutdown_executor: bool = False,
+        retries: int = 0,
+        raise_if_retries_exceeded: bool = True,
+    ) -> None:
 
         self.executor = _ensure_executor(executor)
         self.goal = goal
 
         self._max_tasks = ntasks
 
-        self._pending_tasks = {}  # mapping from concurrent.futures.Future → point id
+        self._pending_tasks: dict[concurrent.Future, int] = {}
 
         # if we instantiate our own executor, then we are also responsible
         # for calling 'shutdown'
         self.shutdown_executor = shutdown_executor or (executor is None)
 
         self.learner = learner
-        self.log = [] if log else None
+        self.log: list | None = [] if log else None
 
         # Timing
         self.start_time = time.time()
-        self.end_time = None
+        self.end_time: float | None = None
         self._elapsed_function_time = 0
 
         # Error handling attributes
         self.retries = retries
         self.raise_if_retries_exceeded = raise_if_retries_exceeded
-        self._to_retry = {}
-        self._tracebacks = {}
+        self._to_retry: dict[int, int] = {}
+        self._tracebacks: dict[int, str] = {}
 
-        self._id_to_point = {}
-        self._next_id = functools.partial(
+        self._id_to_point: dict[int, Any] = {}
+        self._next_id: Callable[[], int] = functools.partial(
             next, itertools.count()
         )  # some unique id to be associated with each point
 
-    def _get_max_tasks(self):
+    def _get_max_tasks(self) -> int:
         return self._max_tasks or _get_ncores(self.executor)
 
-    def _do_raise(self, e, i):
-        tb = self._tracebacks[i]
-        x = self._id_to_point[i]
+    def _do_raise(self, e: Exception, pid: int):
+        tb = self._tracebacks[pid]
+        x = self._id_to_point[pid]
         raise RuntimeError(
             "An error occured while evaluating "
             f'"learner.function({x})". '
@@ -174,10 +264,10 @@ def _do_raise(self, e, i):
         ) from e
 
     @property
-    def do_log(self):
+    def do_log(self) -> bool:
         return self.log is not None
 
-    def _ask(self, n):
+    def _ask(self, n: int) -> tuple[list[int], list[float]]:
         pending_ids = self._pending_tasks.values()
         # using generator here because we only need until `n`
         pids_gen = (pid for pid in self._to_retry.keys() if pid not in pending_ids)
@@ -217,7 +307,10 @@ def overhead(self):
         t_total = self.elapsed_time()
         return (1 - t_function / t_total) * 100
 
-    def _process_futures(self, done_futs):
+    def _process_futures(
+        self,
+        done_futs: _FutureTypes,
+    ) -> None:
         for fut in done_futs:
             pid = self._pending_tasks.pop(fut)
             try:
@@ -239,7 +332,9 @@ def _process_futures(self, done_futs):
                     self.log.append(("tell", x, y))
                 self.learner.tell(x, y)
 
-    def _get_futures(self):
+    def _get_futures(
+        self,
+    ) -> _FutureTypes:
         # Launch tasks to replace the ones that completed
         # on the last iteration, making sure to fill workers
         # that have started since the last iteration.
@@ -261,7 +356,7 @@ def _get_futures(self):
         futures = list(self._pending_tasks.keys())
         return futures
 
-    def _remove_unfinished(self):
+    def _remove_unfinished(self) -> list[Future]:
         # remove points with 'None' values from the learner
         self.learner.remove_unfinished()
         # cancel any outstanding tasks
@@ -270,7 +365,7 @@ def _remove_unfinished(self):
             fut.cancel()
         return remaining
 
-    def _cleanup(self):
+    def _cleanup(self) -> None:
         if self.shutdown_executor:
             # XXX: temporary set wait=True because of a bug with Python ≥3.7
             # and loky in any Python version.
@@ -282,7 +377,7 @@ def _cleanup(self):
         self.end_time = time.time()
 
     @property
-    def failed(self):
+    def failed(self) -> set[Any]:
         """Set of points that failed ``runner.retries`` times."""
         return set(self._tracebacks) - set(self._to_retry)
 
@@ -299,15 +394,15 @@ def _submit(self, x):
         """Is called in `_get_futures`."""
 
     @property
-    def tracebacks(self):
+    def tracebacks(self) -> list[tuple[int, str]]:
         return [(self._id_to_point[pid], tb) for pid, tb in self._tracebacks.items()]
 
     @property
-    def to_retry(self):
+    def to_retry(self) -> list[tuple[int, int]]:
         return [(self._id_to_point[pid], n) for pid, n in self._to_retry.items()]
 
     @property
-    def pending_points(self):
+    def pending_points(self) -> list[tuple[Future, Any]]:
         return [
             (fut, self._id_to_point[pid]) for fut, pid in self._pending_tasks.items()
         ]
@@ -375,16 +470,22 @@ class BlockingRunner(BaseRunner):
 
     def __init__(
         self,
-        learner,
-        goal,
+        learner: BaseLearner,
+        goal: Callable,
         *,
-        executor=None,
-        ntasks=None,
+        executor: (
+            _ThirdPartyClient
+            | _ThirdPartyExecutor
+            | concurrent.ProcessPoolExecutor
+            | concurrent.ThreadPoolExecutor
+            | SequentialExecutor
+        ) = None,
+        ntasks: int | None = None,
         log=False,
         shutdown_executor=False,
         retries=0,
         raise_if_retries_exceeded=True,
-    ):
+    ) -> None:
         if inspect.iscoroutinefunction(learner.function):
             raise ValueError("Coroutine functions can only be used with 'AsyncRunner'.")
         super().__init__(
@@ -399,10 +500,10 @@ def __init__(
         )
         self._run()
 
-    def _submit(self, x):
+    def _submit(self, x: tuple[float, ...] | float | int) -> Future:
         return self.executor.submit(self.learner.function, x)
 
-    def _run(self):
+    def _run(self) -> None:
         first_completed = concurrent.FIRST_COMPLETED
 
         if self._get_max_tasks() < 1:
@@ -506,17 +607,23 @@ class AsyncRunner(BaseRunner):
 
     def __init__(
         self,
-        learner,
-        goal=None,
+        learner: BaseLearner,
+        goal: Callable | None = None,
         *,
-        executor=None,
-        ntasks=None,
-        log=False,
-        shutdown_executor=False,
+        executor: (
+            _ThirdPartyClient
+            | _ThirdPartyExecutor
+            | concurrent.ProcessPoolExecutor
+            | concurrent.ThreadPoolExecutor
+            | SequentialExecutor
+        ) = None,
+        ntasks: int | None = None,
+        log: bool = False,
+        shutdown_executor: bool = False,
         ioloop=None,
-        retries=0,
-        raise_if_retries_exceeded=True,
-    ):
+        retries: int = 0,
+        raise_if_retries_exceeded: bool = True,
+    ) -> None:
 
         if goal is None:
 
@@ -564,7 +671,7 @@ def goal(_):
             self.executor.shutdown()  # Make sure we don't shoot ourselves later
 
         self.task = self.ioloop.create_task(self._run())
-        self.saving_task = None
+        self.saving_task: Task | None = None
         if in_ipynb() and not self.ioloop.is_running():
             warnings.warn(
                 "The runner has been scheduled, but the asyncio "
@@ -573,7 +680,9 @@ def goal(_):
                 "'adaptive.notebook_extension()'"
             )
 
-    def _submit(self, x):
+    def _submit(
+        self, x: tuple[int, int] | int | tuple[float, float] | float
+    ) -> Task | Future:
         ioloop = self.ioloop
         if inspect.iscoroutinefunction(self.learner.function):
             return ioloop.create_task(self.learner.function(x))
@@ -638,7 +747,7 @@ def live_info(self, *, update_interval=0.1):
         """
         return live_info(self, update_interval=update_interval)
 
-    async def _run(self):
+    async def _run(self) -> None:
         first_completed = asyncio.FIRST_COMPLETED
 
         if self._get_max_tasks() < 1:
@@ -656,7 +765,7 @@ async def _run(self):
                 await asyncio.wait(remaining)
             self._cleanup()
 
-    def elapsed_time(self):
+    def elapsed_time(self) -> float:
         """Return the total time elapsed since the runner
         was started."""
         if self.task.done():
@@ -720,7 +829,7 @@ async def _saver():
 Runner = AsyncRunner
 
 
-def simple(learner, goal):
+def simple(learner: BaseLearner, goal: Callable) -> None:
     """Run the learner until the goal is reached.
 
     Requests a single point from the learner, evaluates
@@ -746,7 +855,7 @@ def simple(learner, goal):
             learner.tell(x, y)
 
 
-def replay_log(learner, log):
+def replay_log(learner: BaseLearner, log) -> None:
     """Apply a sequence of method calls to a learner.
 
     This is useful for debugging runners.
@@ -765,7 +874,7 @@ def replay_log(learner, log):
 # --- Useful runner goals
 
 
-def stop_after(*, seconds=0, minutes=0, hours=0):
+def stop_after(*, seconds=0, minutes=0, hours=0) -> Callable:
     """Stop a runner after a specified time.
 
     For example, to specify a runner that should stop after
@@ -797,67 +906,3 @@ def stop_after(*, seconds=0, minutes=0, hours=0):
     """
     stop_time = time.time() + seconds + 60 * minutes + 3600 * hours
     return lambda _: time.time() > stop_time
-
-
-# -- Internal executor-related, things
-
-
-class SequentialExecutor(concurrent.Executor):
-    """A trivial executor that runs functions synchronously.
-
-    This executor is mainly for testing.
-    """
-
-    def submit(self, fn, *args, **kwargs):
-        fut = concurrent.Future()
-        try:
-            fut.set_result(fn(*args, **kwargs))
-        except Exception as e:
-            fut.set_exception(e)
-        return fut
-
-    def map(self, fn, *iterable, timeout=None, chunksize=1):
-        return map(fn, iterable)
-
-    def shutdown(self, wait=True):
-        pass
-
-
-def _ensure_executor(executor):
-    if executor is None:
-        executor = _default_executor()
-
-    if isinstance(executor, concurrent.Executor):
-        return executor
-    elif with_ipyparallel and isinstance(executor, ipyparallel.Client):
-        return executor.executor()
-    elif with_distributed and isinstance(
-        executor, (distributed.Client, distributed.client.Client)
-    ):
-        return executor.get_executor()
-    else:
-        raise TypeError(
-            "Only a concurrent.futures.Executor, distributed.Client,"
-            " or ipyparallel.Client can be used."
-        )
-
-
-def _get_ncores(ex):
-    """Return the maximum  number of cores that an executor can use."""
-    if with_ipyparallel and isinstance(ex, ipyparallel.client.view.ViewExecutor):
-        return len(ex.view)
-    elif isinstance(
-        ex, (concurrent.ProcessPoolExecutor, concurrent.ThreadPoolExecutor)
-    ):
-        return ex._max_workers  # not public API!
-    elif isinstance(ex, loky.reusable_executor._ReusablePoolExecutor):
-        return ex._max_workers  # not public API!
-    elif isinstance(ex, SequentialExecutor):
-        return 1
-    elif with_distributed and isinstance(ex, distributed.cfexecutor.ClientExecutor):
-        return sum(n for n in ex._client.ncores().values())
-    elif with_mpi4py and isinstance(ex, mpi4py.futures.MPIPoolExecutor):
-        ex.bootup()  # wait until all workers are up and running
-        return ex._pool.size  # not public API!
-    else:
-        raise TypeError(f"Cannot get number of cores for {ex.__class__}")

From 6cf01d8f07613eb58a86b7004ed822c1399c00d4 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:15 -0700
Subject: [PATCH 02/21] Add type-hints to adaptive/_version.py

---
 adaptive/_version.py | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/adaptive/_version.py b/adaptive/_version.py
index 21293fb39..bdb4875c7 100644
--- a/adaptive/_version.py
+++ b/adaptive/_version.py
@@ -1,5 +1,7 @@
 # This file is part of 'miniver': https://github.com/jbweston/miniver
 #
+from __future__ import annotations
+
 import os
 import subprocess
 from collections import namedtuple
@@ -10,7 +12,7 @@
 Version = namedtuple("Version", ("release", "dev", "labels"))
 
 # No public API
-__all__ = []
+__all__: list[str] = []
 
 package_root = os.path.dirname(os.path.realpath(__file__))
 package_name = os.path.basename(package_root)
@@ -26,10 +28,9 @@
 STATIC_VERSION_FILE = "_static_version.py"
 
 
-def get_version(version_file=STATIC_VERSION_FILE):
+def get_version(version_file: str = STATIC_VERSION_FILE) -> str:
     version_info = get_static_version_info(version_file)
-    version = version_info["version"]
-    if version == "__use_git__":
+    if version_info["version"] == "__use_git__":
         version = get_version_from_git()
         if not version:
             version = get_version_from_git_archive(version_info)
@@ -37,11 +38,11 @@ def get_version(version_file=STATIC_VERSION_FILE):
             version = Version("unknown", None, None)
         return pep440_format(version)
     else:
-        return version
+        return version_info["version"]
 
 
-def get_static_version_info(version_file=STATIC_VERSION_FILE):
-    version_info = {}
+def get_static_version_info(version_file: str = STATIC_VERSION_FILE) -> dict[str, str]:
+    version_info: dict[str, str] = {}
     with open(os.path.join(package_root, version_file), "rb") as f:
         exec(f.read(), {}, version_info)
     return version_info
@@ -51,7 +52,7 @@ def version_is_from_git(version_file=STATIC_VERSION_FILE):
     return get_static_version_info(version_file)["version"] == "__use_git__"
 
 
-def pep440_format(version_info):
+def pep440_format(version_info: Version) -> str:
     release, dev, labels = version_info
 
     version_parts = [release]
@@ -68,7 +69,7 @@ def pep440_format(version_info):
     return "".join(version_parts)
 
 
-def get_version_from_git():
+def get_version_from_git() -> Version:
     try:
         p = subprocess.Popen(
             ["git", "rev-parse", "--show-toplevel"],
@@ -77,14 +78,14 @@ def get_version_from_git():
             stderr=subprocess.PIPE,
         )
     except OSError:
-        return
+        return None
     if p.wait() != 0:
-        return
+        return None
     if not os.path.samefile(p.communicate()[0].decode().rstrip("\n"), distr_root):
         # The top-level directory of the current Git repository is not the same
         # as the root directory of the distribution: do not extract the
         # version from Git.
-        return
+        return None
 
     # git describe --first-parent does not take into account tags from branches
     # that were merged-in. The '--long' flag gets us the 'dev' version and
@@ -92,17 +93,17 @@ def get_version_from_git():
     for opts in [["--first-parent"], []]:
         try:
             p = subprocess.Popen(
-                ["git", "describe", "--long", "--always", "--tags"] + opts,
+                ["git", "describe", "--long", "--always", "--tags"] + opts,  # type: ignore
                 cwd=distr_root,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
             )
         except OSError:
-            return
+            return None
         if p.wait() == 0:
             break
     else:
-        return
+        return None
 
     description = (
         p.communicate()[0]
@@ -142,7 +143,7 @@ def get_version_from_git():
 #       Currently we can only tell the tag the current commit is
 #       pointing to, or its hash (with no version info)
 #       if it is not tagged.
-def get_version_from_git_archive(version_info):
+def get_version_from_git_archive(version_info) -> Version:
     try:
         refnames = version_info["refnames"]
         git_hash = version_info["git_hash"]
@@ -165,7 +166,7 @@ def get_version_from_git_archive(version_info):
         return Version("unknown", dev=None, labels=[f"g{git_hash}"])
 
 
-__version__ = get_version()
+__version__: str = get_version()
 
 
 # The following section defines a module global 'cmdclass',

From 1786f80d6f4970470a54e92cbc6fa1e3a1975c42 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:16 -0700
Subject: [PATCH 03/21] Add type-hints to adaptive/learner/__init__.py

---
 adaptive/learner/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/adaptive/learner/__init__.py b/adaptive/learner/__init__.py
index 74564773a..7a77504c7 100644
--- a/adaptive/learner/__init__.py
+++ b/adaptive/learner/__init__.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from contextlib import suppress
 
 from adaptive.learner.average_learner import AverageLearner

From 94900937956e867b717687861dd7974a0a085e3f Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:17 -0700
Subject: [PATCH 04/21] Add type-hints to adaptive/learner/balancing_learner.py

---
 adaptive/learner/balancing_learner.py | 105 +++++++++++++++++++-------
 1 file changed, 77 insertions(+), 28 deletions(-)

diff --git a/adaptive/learner/balancing_learner.py b/adaptive/learner/balancing_learner.py
index 0215b3af6..0debb3db0 100644
--- a/adaptive/learner/balancing_learner.py
+++ b/adaptive/learner/balancing_learner.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import itertools
+import numbers
 from collections import defaultdict
 from collections.abc import Iterable
 from contextlib import suppress
 from functools import partial
 from operator import itemgetter
+from typing import Any, Callable, Dict, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -13,20 +15,32 @@
 from adaptive.notebook_integration import ensure_holoviews
 from adaptive.utils import cache_latest, named_product, restore
 
+try:
+    from typing import Literal, TypeAlias
+except ImportError:
+    from typing_extensions import Literal, TypeAlias
+
 try:
     import pandas
 
     with_pandas = True
-
 except ModuleNotFoundError:
     with_pandas = False
 
 
-def dispatch(child_functions, arg):
+def dispatch(child_functions: list[Callable], arg: Any) -> Any:
     index, x = arg
     return child_functions[index](x)
 
 
+STRATEGY_TYPE: TypeAlias = Literal["loss_improvements", "loss", "npoints", "cycle"]
+
+CDIMS_TYPE: TypeAlias = Union[
+    Sequence[Dict[str, Any]],
+    Tuple[Sequence[str], Sequence[Tuple[Any, ...]]],
+]
+
+
 class BalancingLearner(BaseLearner):
     r"""Choose the optimal points from a set of learners.
 
@@ -78,13 +92,19 @@ class BalancingLearner(BaseLearner):
     behave in an undefined way. Change the `strategy` in that case.
     """
 
-    def __init__(self, learners, *, cdims=None, strategy="loss_improvements"):
+    def __init__(
+        self,
+        learners: list[BaseLearner],
+        *,
+        cdims: CDIMS_TYPE | None = None,
+        strategy: STRATEGY_TYPE = "loss_improvements",
+    ) -> None:
         self.learners = learners
 
         # Naively we would make 'function' a method, but this causes problems
         # when using executors from 'concurrent.futures' because we have to
         # pickle the whole learner.
-        self.function = partial(dispatch, [l.function for l in self.learners])
+        self.function = partial(dispatch, [l.function for l in self.learners])  # type: ignore
 
         self._ask_cache = {}
         self._loss = {}
@@ -96,7 +116,7 @@ def __init__(self, learners, *, cdims=None, strategy="loss_improvements"):
                 "A BalacingLearner can handle only one type" " of learners."
             )
 
-        self.strategy = strategy
+        self.strategy: STRATEGY_TYPE = strategy
 
     def new(self) -> BalancingLearner:
         """Create a new `BalancingLearner` with the same parameters."""
@@ -107,21 +127,21 @@ def new(self) -> BalancingLearner:
         )
 
     @property
-    def data(self):
+    def data(self) -> dict[tuple[int, Any], Any]:
         data = {}
         for i, l in enumerate(self.learners):
             data.update({(i, p): v for p, v in l.data.items()})
         return data
 
     @property
-    def pending_points(self):
+    def pending_points(self) -> set[tuple[int, Any]]:
         pending_points = set()
         for i, l in enumerate(self.learners):
             pending_points.update({(i, p) for p in l.pending_points})
         return pending_points
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         return sum(l.npoints for l in self.learners)
 
     @property
@@ -134,7 +154,7 @@ def nsamples(self):
             )
 
     @property
-    def strategy(self):
+    def strategy(self) -> STRATEGY_TYPE:
         """Can be either 'loss_improvements' (default), 'loss', 'npoints', or
         'cycle'. The points that the `BalancingLearner` choses can be either
         based on: the best 'loss_improvements', the smallest total 'loss' of
@@ -145,7 +165,7 @@ def strategy(self):
         return self._strategy
 
     @strategy.setter
-    def strategy(self, strategy):
+    def strategy(self, strategy: STRATEGY_TYPE) -> None:
         self._strategy = strategy
         if strategy == "loss_improvements":
             self._ask_and_tell = self._ask_and_tell_based_on_loss_improvements
@@ -162,7 +182,9 @@ def strategy(self, strategy):
                 ' strategy="npoints", or strategy="cycle" is implemented.'
             )
 
-    def _ask_and_tell_based_on_loss_improvements(self, n):
+    def _ask_and_tell_based_on_loss_improvements(
+        self, n: int
+    ) -> tuple[list[tuple[int, Any]], list[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -185,7 +207,9 @@ def _ask_and_tell_based_on_loss_improvements(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_loss(self, n):
+    def _ask_and_tell_based_on_loss(
+        self, n: int
+    ) -> tuple[list[tuple[int, Any]], list[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -206,7 +230,9 @@ def _ask_and_tell_based_on_loss(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_npoints(self, n):
+    def _ask_and_tell_based_on_npoints(
+        self, n: numbers.Integral
+    ) -> tuple[list[tuple[numbers.Integral, Any]], list[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -222,7 +248,9 @@ def _ask_and_tell_based_on_npoints(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_cycle(self, n):
+    def _ask_and_tell_based_on_cycle(
+        self, n: int
+    ) -> tuple[list[tuple[numbers.Integral, Any]], list[float]]:
         points, loss_improvements = [], []
         for _ in range(n):
             index = next(self._cycle)
@@ -233,7 +261,9 @@ def _ask_and_tell_based_on_cycle(self, n):
 
         return points, loss_improvements
 
-    def ask(self, n, tell_pending=True):
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> tuple[list[tuple[numbers.Integral, Any]], list[float]]:
         """Chose points for learners."""
         if n == 0:
             return [], []
@@ -244,20 +274,20 @@ def ask(self, n, tell_pending=True):
         else:
             return self._ask_and_tell(n)
 
-    def tell(self, x, y):
+    def tell(self, x: tuple[numbers.Integral, Any], y: Any) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self._pending_loss.pop(index, None)
         self.learners[index].tell(x, y)
 
-    def tell_pending(self, x):
+    def tell_pending(self, x: tuple[numbers.Integral, Any]) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self.learners[index].tell_pending(x)
 
-    def _losses(self, real=True):
+    def _losses(self, real: bool = True) -> list[float]:
         losses = []
         loss_dict = self._loss if real else self._pending_loss
 
@@ -269,11 +299,16 @@ def _losses(self, real=True):
         return losses
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         losses = self._losses(real)
         return max(losses)
 
-    def plot(self, cdims=None, plotter=None, dynamic=True):
+    def plot(
+        self,
+        cdims: CDIMS_TYPE | None = None,
+        plotter: Callable[[BaseLearner], Any] | None = None,
+        dynamic: bool = True,
+    ):
         """Returns a DynamicMap with sliders.
 
         Parameters
@@ -346,13 +381,19 @@ def plot_function(*args):
             vals = {d.name: d.values for d in dm.dimensions() if d.values}
             return hv.HoloMap(dm.select(**vals))
 
-    def remove_unfinished(self):
+    def remove_unfinished(self) -> None:
         """Remove uncomputed data from the learners."""
         for learner in self.learners:
             learner.remove_unfinished()
 
     @classmethod
-    def from_product(cls, f, learner_type, learner_kwargs, combos):
+    def from_product(
+        cls,
+        f,
+        learner_type: BaseLearner,
+        learner_kwargs: dict[str, Any],
+        combos: dict[str, Sequence[Any]],
+    ) -> BalancingLearner:
         """Create a `BalancingLearner` with learners of all combinations of
         named variables’ values. The `cdims` will be set correctly, so calling
         `learner.plot` will be a `holoviews.core.HoloMap` with the correct labels.
@@ -448,7 +489,11 @@ def load_dataframe(
         for i, gr in df.groupby(index_name):
             self.learners[i].load_dataframe(gr, **kwargs)
 
-    def save(self, fname, compress=True):
+    def save(
+        self,
+        fname: Callable[[BaseLearner], str] | Sequence[str],
+        compress: bool = True,
+    ) -> None:
         """Save the data of the child learners into pickle files
         in a directory.
 
@@ -486,7 +531,11 @@ def save(self, fname, compress=True):
             for l in self.learners:
                 l.save(fname(l), compress=compress)
 
-    def load(self, fname, compress=True):
+    def load(
+        self,
+        fname: Callable[[BaseLearner], str] | Sequence[str],
+        compress: bool = True,
+    ) -> None:
         """Load the data of the child learners from pickle files
         in a directory.
 
@@ -510,20 +559,20 @@ def load(self, fname, compress=True):
             for l in self.learners:
                 l.load(fname(l), compress=compress)
 
-    def _get_data(self):
+    def _get_data(self) -> list[Any]:
         return [l._get_data() for l in self.learners]
 
-    def _set_data(self, data):
+    def _set_data(self, data: list[Any]):
         for l, _data in zip(self.learners, data):
             l._set_data(_data)
 
-    def __getstate__(self):
+    def __getstate__(self) -> tuple[list[BaseLearner], CDIMS_TYPE, STRATEGY_TYPE]:
         return (
             self.learners,
             self._cdims_default,
             self.strategy,
         )
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: tuple[list[BaseLearner], CDIMS_TYPE, STRATEGY_TYPE]):
         learners, cdims, strategy = state
         self.__init__(learners, cdims=cdims, strategy=strategy)

From 6b3209f3f4bb643b38b6feef4a8cd896f6edb9a5 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:18 -0700
Subject: [PATCH 05/21] Add type-hints to adaptive/learner/base_learner.py

---
 adaptive/learner/base_learner.py | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/adaptive/learner/base_learner.py b/adaptive/learner/base_learner.py
index 2f3f09af1..71bb3aa77 100644
--- a/adaptive/learner/base_learner.py
+++ b/adaptive/learner/base_learner.py
@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 import abc
 from contextlib import suppress
+from typing import Any, Callable
 
 import cloudpickle
 
 from adaptive.utils import _RequireAttrsABCMeta, load, save
 
 
-def uses_nth_neighbors(n: int):
+def uses_nth_neighbors(n: int) -> Callable:
     """Decorator to specify how many neighboring intervals the loss function uses.
 
     Wraps loss functions to indicate that they expect intervals together
@@ -82,10 +85,15 @@ class BaseLearner(metaclass=_RequireAttrsABCMeta):
     """
 
     data: dict
-    npoints: int
     pending_points: set
+    function: Callable
+
+    @property
+    @abc.abstractmethod
+    def npoints(self) -> int:
+        """Number of learned points."""
 
-    def tell(self, x, y):
+    def tell(self, x: Any, y) -> None:
         """Tell the learner about a single value.
 
         Parameters
@@ -95,7 +103,7 @@ def tell(self, x, y):
         """
         self.tell_many([x], [y])
 
-    def tell_many(self, xs, ys):
+    def tell_many(self, xs: Any, ys: Any) -> None:
         """Tell the learner about some values.
 
         Parameters
@@ -116,7 +124,7 @@ def remove_unfinished(self):
         """Remove uncomputed data from the learner."""
 
     @abc.abstractmethod
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         """Return the loss for the current state of the learner.
 
         Parameters
@@ -128,7 +136,7 @@ def loss(self, real=True):
         """
 
     @abc.abstractmethod
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True):
         """Choose the next 'n' points to evaluate.
 
         Parameters
@@ -146,7 +154,7 @@ def _get_data(self):
         pass
 
     @abc.abstractmethod
-    def _set_data(self):
+    def _set_data(self, data: Any):
         pass
 
     @abc.abstractmethod
@@ -164,7 +172,7 @@ def copy_from(self, other):
         """
         self._set_data(other._get_data())
 
-    def save(self, fname, compress=True):
+    def save(self, fname: str, compress: bool = True) -> None:
         """Save the data of the learner into a pickle file.
 
         Parameters
@@ -178,7 +186,7 @@ def save(self, fname, compress=True):
         data = self._get_data()
         save(fname, data, compress)
 
-    def load(self, fname, compress=True):
+    def load(self, fname: str, compress: bool = True) -> None:
         """Load the data of a learner from a pickle file.
 
         Parameters
@@ -193,8 +201,8 @@ def load(self, fname, compress=True):
             data = load(fname, compress)
             self._set_data(data)
 
-    def __getstate__(self):
+    def __getstate__(self) -> dict[str, Any]:
         return cloudpickle.dumps(self.__dict__)
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__ = cloudpickle.loads(state)

From 8363aa6bc7af93bd38754b0fa54e868db3258943 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:19 -0700
Subject: [PATCH 06/21] Add type-hints to adaptive/learner/data_saver.py

---
 adaptive/learner/data_saver.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/adaptive/learner/data_saver.py b/adaptive/learner/data_saver.py
index 8da281d7f..609d0ba82 100644
--- a/adaptive/learner/data_saver.py
+++ b/adaptive/learner/data_saver.py
@@ -2,6 +2,8 @@
 
 import functools
 from collections import OrderedDict
+from operator import itemgetter
+from typing import Any
 
 from adaptive.learner.base_learner import BaseLearner
 from adaptive.utils import copy_docstring_from
@@ -39,7 +41,7 @@ class DataSaver:
     >>> learner = DataSaver(_learner, arg_picker=itemgetter('y'))
     """
 
-    def __init__(self, learner, arg_picker):
+    def __init__(self, learner: BaseLearner, arg_picker: itemgetter) -> None:
         self.learner = learner
         self.extra_data = OrderedDict()
         self.function = learner.function
@@ -49,7 +51,7 @@ def new(self) -> DataSaver:
         """Return a new `DataSaver` with the same `arg_picker` and `learner`."""
         return DataSaver(self.learner.new(), self.arg_picker)
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: str) -> Any:
         return getattr(self.learner, attr)
 
     @copy_docstring_from(BaseLearner.tell)
@@ -122,10 +124,17 @@ def load_dataframe(
             key = _to_key(x[:-1])
             self.extra_data[key] = x[-1]
 
-    def _get_data(self):
+    def _get_data(self) -> tuple[Any, OrderedDict]:
         return self.learner._get_data(), self.extra_data
 
-    def _set_data(self, data):
+    def _set_data(
+        self,
+        data: (
+            tuple[OrderedDict, OrderedDict]
+            | tuple[dict[int | float, float], OrderedDict]
+            | tuple[tuple[dict[int, float], int, float, float], OrderedDict]
+        ),
+    ) -> None:
         learner_data, self.extra_data = data
         self.learner._set_data(learner_data)
 

From 2b0497f46f5f8fb6cc6e5c2d526933ff6de81d95 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:20 -0700
Subject: [PATCH 07/21] Add type-hints to adaptive/learner/integrator_coeffs.py

---
 adaptive/learner/integrator_coeffs.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/adaptive/learner/integrator_coeffs.py b/adaptive/learner/integrator_coeffs.py
index 9ccc54be1..711f30b76 100644
--- a/adaptive/learner/integrator_coeffs.py
+++ b/adaptive/learner/integrator_coeffs.py
@@ -1,4 +1,5 @@
 # Based on an adaptive quadrature algorithm by Pedro Gonnet
+from __future__ import annotations
 
 from collections import defaultdict
 from fractions import Fraction
@@ -8,7 +9,7 @@
 import scipy.linalg
 
 
-def legendre(n):
+def legendre(n: int) -> list[list[Fraction]]:
     """Return the first n Legendre polynomials.
 
     The polynomials have *standard* normalization, i.e.
@@ -29,7 +30,7 @@ def legendre(n):
     return result
 
 
-def newton(n):
+def newton(n: int) -> np.ndarray:
     """Compute the monomial coefficients of the Newton polynomial over the
     nodes of the n-point Clenshaw-Curtis quadrature rule.
     """
@@ -86,7 +87,7 @@ def newton(n):
     return cf
 
 
-def scalar_product(a, b):
+def scalar_product(a: list[Fraction], b: list[Fraction]) -> Fraction:
     """Compute the polynomial scalar product int_-1^1 dx a(x) b(x).
 
     The args must be sequences of polynomial coefficients.  This
@@ -107,7 +108,7 @@ def scalar_product(a, b):
     return 2 * sum(c[i] / (i + 1) for i in range(0, lc, 2))
 
 
-def calc_bdef(ns):
+def calc_bdef(ns: tuple[int, int, int, int]) -> list[np.ndarray]:
     """Calculate the decompositions of Newton polynomials (over the nodes
     of the n-point Clenshaw-Curtis quadrature rule) in terms of
     Legandre polynomials.
@@ -133,7 +134,7 @@ def calc_bdef(ns):
     return result
 
 
-def calc_V(x, n):
+def calc_V(x: np.ndarray, n: int) -> np.ndarray:
     V = [np.ones(x.shape), x.copy()]
     for i in range(2, n):
         V.append((2 * i - 1) / i * x * V[-1] - (i - 1) / i * V[-2])

From c2b6cf35f35a5d1b6f962c134dc5cd196289a6ac Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:21 -0700
Subject: [PATCH 08/21] Add type-hints to
 adaptive/learner/integrator_learner.py

---
 adaptive/learner/integrator_learner.py | 86 ++++++++++++++------------
 1 file changed, 47 insertions(+), 39 deletions(-)

diff --git a/adaptive/learner/integrator_learner.py b/adaptive/learner/integrator_learner.py
index da9bd7ffc..f0712778d 100644
--- a/adaptive/learner/integrator_learner.py
+++ b/adaptive/learner/integrator_learner.py
@@ -5,6 +5,7 @@
 from collections import defaultdict
 from math import sqrt
 from operator import attrgetter
+from typing import TYPE_CHECKING, Callable
 
 import cloudpickle
 import numpy as np
@@ -25,7 +26,7 @@
     with_pandas = False
 
 
-def _downdate(c, nans, depth):
+def _downdate(c: np.ndarray, nans: list[int], depth: int) -> np.ndarray:
     # This is algorithm 5 from the thesis of Pedro Gonnet.
     b = coeff.b_def[depth].copy()
     m = coeff.ns[depth] - 1
@@ -45,7 +46,7 @@ def _downdate(c, nans, depth):
     return c
 
 
-def _zero_nans(fx):
+def _zero_nans(fx: np.ndarray) -> list[int]:
     """Caution: this function modifies fx."""
     nans = []
     for i in range(len(fx)):
@@ -55,7 +56,7 @@ def _zero_nans(fx):
     return nans
 
 
-def _calc_coeffs(fx, depth):
+def _calc_coeffs(fx: np.ndarray, depth: int) -> np.ndarray:
     """Caution: this function modifies fx."""
     nans = _zero_nans(fx)
     c_new = coeff.V_inv[depth] @ fx
@@ -135,19 +136,24 @@ class _Interval:
         "removed",
     ]
 
-    def __init__(self, a, b, depth, rdepth):
-        self.children = []
-        self.data = {}
+    def __init__(self, a: int | float, b: int | float, depth: int, rdepth: int) -> None:
+        self.children: list[_Interval] = []
+        self.data: dict[float, float] = {}
         self.a = a
         self.b = b
         self.depth = depth
         self.rdepth = rdepth
-        self.done_leaves = set()
-        self.depth_complete = None
+        self.done_leaves: set[_Interval] = set()
+        self.depth_complete: int | None = None
         self.removed = False
+        if TYPE_CHECKING:
+            self.ndiv: int
+            self.parent: _Interval | None
+            self.err: float
+            self.c: np.ndarray
 
     @classmethod
-    def make_first(cls, a, b, depth=2):
+    def make_first(cls, a: int, b: int, depth: int = 2) -> _Interval:
         ival = _Interval(a, b, depth, rdepth=1)
         ival.ndiv = 0
         ival.parent = None
@@ -155,7 +161,7 @@ def make_first(cls, a, b, depth=2):
         return ival
 
     @property
-    def T(self):
+    def T(self) -> np.ndarray:
         """Get the correct shift matrix.
 
         Should only be called on children of a split interval.
@@ -166,24 +172,24 @@ def T(self):
         assert left != right
         return coeff.T_left if left else coeff.T_right
 
-    def refinement_complete(self, depth):
+    def refinement_complete(self, depth: int) -> bool:
         """The interval has all the y-values to calculate the intergral."""
         if len(self.data) < coeff.ns[depth]:
             return False
         return all(p in self.data for p in self.points(depth))
 
-    def points(self, depth=None):
+    def points(self, depth: int | None = None) -> np.ndarray:
         if depth is None:
             depth = self.depth
         a = self.a
         b = self.b
         return (a + b) / 2 + (b - a) * coeff.xi[depth] / 2
 
-    def refine(self):
+    def refine(self) -> _Interval:
         self.depth += 1
         return self
 
-    def split(self):
+    def split(self) -> list[_Interval]:
         points = self.points()
         m = points[len(points) // 2]
         ivals = [
@@ -198,10 +204,10 @@ def split(self):
 
         return ivals
 
-    def calc_igral(self):
+    def calc_igral(self) -> None:
         self.igral = (self.b - self.a) * self.c[0] / sqrt(2)
 
-    def update_heuristic_err(self, value):
+    def update_heuristic_err(self, value: float) -> None:
         """Sets the error of an interval using a heuristic (half the error of
         the parent) when the actual error cannot be calculated due to its
         parents not being finished yet. This error is propagated down to its
@@ -214,7 +220,7 @@ def update_heuristic_err(self, value):
                 continue
             child.update_heuristic_err(value / 2)
 
-    def calc_err(self, c_old):
+    def calc_err(self, c_old: np.ndarray) -> float:
         c_new = self.c
         c_diff = np.zeros(max(len(c_old), len(c_new)))
         c_diff[: len(c_old)] = c_old
@@ -226,9 +232,9 @@ def calc_err(self, c_old):
                 child.update_heuristic_err(self.err / 2)
         return c_diff
 
-    def calc_ndiv(self):
+    def calc_ndiv(self) -> None:
         div = self.parent.c00 and self.c00 / self.parent.c00 > 2
-        self.ndiv += div
+        self.ndiv += int(div)
 
         if self.ndiv > coeff.ndiv_max and 2 * self.ndiv > self.rdepth:
             raise DivergentIntegralError
@@ -237,7 +243,7 @@ def calc_ndiv(self):
             for child in self.children:
                 child.update_ndiv_recursively()
 
-    def update_ndiv_recursively(self):
+    def update_ndiv_recursively(self) -> None:
         self.ndiv += 1
         if self.ndiv > coeff.ndiv_max and 2 * self.ndiv > self.rdepth:
             raise DivergentIntegralError
@@ -245,7 +251,7 @@ def update_ndiv_recursively(self):
         for child in self.children:
             child.update_ndiv_recursively()
 
-    def complete_process(self, depth):
+    def complete_process(self, depth: int) -> tuple[bool, bool] | tuple[bool, np.bool_]:
         """Calculate the integral contribution and error from this interval,
         and update the done leaves of all ancestor intervals."""
         assert self.depth_complete is None or self.depth_complete == depth - 1
@@ -322,7 +328,7 @@ def complete_process(self, depth):
 
         return force_split, remove
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         lst = [
             f"(a, b)=({self.a:.5f}, {self.b:.5f})",
             f"depth={self.depth}",
@@ -334,7 +340,7 @@ def __repr__(self):
 
 
 class IntegratorLearner(BaseLearner):
-    def __init__(self, function, bounds, tol):
+    def __init__(self, function: Callable, bounds: tuple[int, int], tol: float) -> None:
         """
         Parameters
         ----------
@@ -368,16 +374,18 @@ def __init__(self, function, bounds, tol):
         plot : hv.Scatter
             Plots all the points that are evaluated.
         """
-        self.function = function
+        self.function = function  # type: ignore
         self.bounds = bounds
         self.tol = tol
         self.max_ivals = 1000
-        self.priority_split = []
+        self.priority_split: list[_Interval] = []
         self.data = {}
         self.pending_points = set()
-        self._stack = []
-        self.x_mapping = defaultdict(lambda: SortedSet([], key=attrgetter("rdepth")))
-        self.ivals = set()
+        self._stack: list[float] = []
+        self.x_mapping: dict[float, SortedSet] = defaultdict(
+            lambda: SortedSet([], key=attrgetter("rdepth"))
+        )
+        self.ivals: set[_Interval] = set()
         ival = _Interval.make_first(*self.bounds)
         self.add_ival(ival)
         self.first_ival = ival
@@ -387,10 +395,10 @@ def new(self) -> IntegratorLearner:
         return IntegratorLearner(self.function, self.bounds, self.tol)
 
     @property
-    def approximating_intervals(self):
+    def approximating_intervals(self) -> set[_Interval]:
         return self.first_ival.done_leaves
 
-    def tell(self, point, value):
+    def tell(self, point: float, value: float) -> None:
         if point not in self.x_mapping:
             raise ValueError(f"Point {point} doesn't belong to any interval")
         self.data[point] = value
@@ -426,7 +434,7 @@ def tell(self, point, value):
     def tell_pending(self):
         pass
 
-    def propagate_removed(self, ival):
+    def propagate_removed(self, ival: _Interval) -> None:
         def _propagate_removed_down(ival):
             ival.removed = True
             self.ivals.discard(ival)
@@ -436,7 +444,7 @@ def _propagate_removed_down(ival):
 
         _propagate_removed_down(ival)
 
-    def add_ival(self, ival):
+    def add_ival(self, ival: _Interval) -> None:
         for x in ival.points():
             # Update the mappings
             self.x_mapping[x].add(ival)
@@ -447,7 +455,7 @@ def add_ival(self, ival):
                 self._stack.append(x)
         self.ivals.add(ival)
 
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True) -> tuple[list[float], list[float]]:
         """Choose points for learners."""
         if not tell_pending:
             with restore(self):
@@ -455,7 +463,7 @@ def ask(self, n, tell_pending=True):
         else:
             return self._ask_and_tell_pending(n)
 
-    def _ask_and_tell_pending(self, n):
+    def _ask_and_tell_pending(self, n: int) -> tuple[list[float], list[float]]:
         points, loss_improvements = self.pop_from_stack(n)
         n_left = n - len(points)
         while n_left > 0:
@@ -471,7 +479,7 @@ def _ask_and_tell_pending(self, n):
 
         return points, loss_improvements
 
-    def pop_from_stack(self, n):
+    def pop_from_stack(self, n: int) -> tuple[list[float], list[float]]:
         points = self._stack[:n]
         self._stack = self._stack[n:]
         loss_improvements = [
@@ -482,7 +490,7 @@ def pop_from_stack(self, n):
     def remove_unfinished(self):
         pass
 
-    def _fill_stack(self):
+    def _fill_stack(self) -> list[float]:
         # XXX: to-do if all the ivals have err=inf, take the interval
         # with the lowest rdepth and no children.
         force_split = bool(self.priority_split)
@@ -518,16 +526,16 @@ def _fill_stack(self):
         return self._stack
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @property
-    def igral(self):
+    def igral(self) -> float:
         return sum(i.igral for i in self.approximating_intervals)
 
     @property
-    def err(self):
+    def err(self) -> float:
         if self.approximating_intervals:
             err = sum(i.err for i in self.approximating_intervals)
             if err > sys.float_info.max:

From 3293f60c3616e697baedb8d51cbfa5592e4ba030 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:22 -0700
Subject: [PATCH 09/21] Add type-hints to adaptive/learner/learner2D.py

---
 adaptive/learner/learner2D.py | 78 +++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 31 deletions(-)

diff --git a/adaptive/learner/learner2D.py b/adaptive/learner/learner2D.py
index 6c5be845f..d167ed4d6 100644
--- a/adaptive/learner/learner2D.py
+++ b/adaptive/learner/learner2D.py
@@ -5,14 +5,17 @@
 from collections import OrderedDict
 from copy import copy
 from math import sqrt
+from typing import Any, Callable, Iterable
 
 import cloudpickle
 import numpy as np
 from scipy import interpolate
+from scipy.interpolate.interpnd import LinearNDInterpolator
 
 from adaptive.learner.base_learner import BaseLearner
 from adaptive.learner.triangulation import simplex_volume_in_embedding
 from adaptive.notebook_integration import ensure_holoviews
+from adaptive.types import Bool
 from adaptive.utils import (
     assign_defaults,
     cache_latest,
@@ -30,7 +33,7 @@
 # Learner2D and helper functions.
 
 
-def deviations(ip):
+def deviations(ip: LinearNDInterpolator) -> list[np.ndarray]:
     """Returns the deviation of the linear estimate.
 
     Is useful when defining custom loss functions.
@@ -68,7 +71,7 @@ def deviation(p, v, g):
     return devs
 
 
-def areas(ip):
+def areas(ip: LinearNDInterpolator) -> np.ndarray:
     """Returns the area per triangle of the triangulation inside
     a `LinearNDInterpolator` instance.
 
@@ -89,7 +92,7 @@ def areas(ip):
     return areas
 
 
-def uniform_loss(ip):
+def uniform_loss(ip: LinearNDInterpolator) -> np.ndarray:
     """Loss function that samples the domain uniformly.
 
     Works with `~adaptive.Learner2D` only.
@@ -120,7 +123,9 @@ def uniform_loss(ip):
     return np.sqrt(areas(ip))
 
 
-def resolution_loss_function(min_distance=0, max_distance=1):
+def resolution_loss_function(
+    min_distance: float = 0, max_distance: float = 1
+) -> Callable[[LinearNDInterpolator], np.ndarray]:
     """Loss function that is similar to the `default_loss` function, but you
     can set the maximimum and minimum size of a triangle.
 
@@ -159,7 +164,7 @@ def resolution_loss(ip):
     return resolution_loss
 
 
-def minimize_triangle_surface_loss(ip):
+def minimize_triangle_surface_loss(ip: LinearNDInterpolator) -> np.ndarray:
     """Loss function that is similar to the distance loss function in the
     `~adaptive.Learner1D`. The loss is the area spanned by the 3D
     vectors of the vertices.
@@ -205,7 +210,7 @@ def _get_vectors(points):
     return np.linalg.norm(np.cross(a, b) / 2, axis=1)
 
 
-def default_loss(ip):
+def default_loss(ip: LinearNDInterpolator) -> np.ndarray:
     """Loss function that combines `deviations` and `areas` of the triangles.
 
     Works with `~adaptive.Learner2D` only.
@@ -225,7 +230,7 @@ def default_loss(ip):
     return losses
 
 
-def choose_point_in_triangle(triangle, max_badness):
+def choose_point_in_triangle(triangle: np.ndarray, max_badness: int) -> np.ndarray:
     """Choose a new point in inside a triangle.
 
     If the ratio of the longest edge of the triangle squared
@@ -364,7 +369,12 @@ class Learner2D(BaseLearner):
     over each triangle.
     """
 
-    def __init__(self, function, bounds, loss_per_triangle=None):
+    def __init__(
+        self,
+        function: Callable,
+        bounds: tuple[tuple[int, int], tuple[int, int]],
+        loss_per_triangle: Callable | None = None,
+    ) -> None:
         self.ndim = len(bounds)
         self._vdim = None
         self.loss_per_triangle = loss_per_triangle or default_loss
@@ -379,7 +389,7 @@ def __init__(self, function, bounds, loss_per_triangle=None):
 
         self._bounds_points = list(itertools.product(*bounds))
         self._stack.update({p: np.inf for p in self._bounds_points})
-        self.function = function
+        self.function = function  # type: ignore
         self._ip = self._ip_combined = None
 
         self.stack_size = 10
@@ -388,7 +398,7 @@ def new(self) -> Learner2D:
         return Learner2D(self.function, self.bounds, self.loss_per_triangle)
 
     @property
-    def xy_scale(self):
+    def xy_scale(self) -> np.ndarray:
         xy_scale = self._xy_scale
         if self.aspect_ratio == 1:
             return xy_scale
@@ -486,21 +496,21 @@ def load_dataframe(
                 self.function, df, function_prefix
             )
 
-    def _scale(self, points):
+    def _scale(self, points: Any) -> np.ndarray:
         points = np.asarray(points, dtype=float)
         return (points - self.xy_mean) / self.xy_scale
 
-    def _unscale(self, points):
+    def _unscale(self, points: np.ndarray) -> np.ndarray:
         points = np.asarray(points, dtype=float)
         return points * self.xy_scale + self.xy_mean
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @property
-    def vdim(self):
+    def vdim(self) -> int:
         """Length of the output of ``learner.function``.
         If the output is unsized (when it's a scalar)
         then `vdim = 1`.
@@ -516,12 +526,14 @@ def vdim(self):
         return self._vdim or 1
 
     @property
-    def bounds_are_done(self):
+    def bounds_are_done(self) -> bool:
         return not any(
             (p in self.pending_points or p in self._stack) for p in self._bounds_points
         )
 
-    def interpolated_on_grid(self, n=None):
+    def interpolated_on_grid(
+        self, n: int = None
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Get the interpolated data on a grid.
 
         Parameters
@@ -553,7 +565,7 @@ def interpolated_on_grid(self, n=None):
         xs, ys = self._unscale(np.vstack([xs, ys]).T).T
         return xs, ys, zs
 
-    def _data_in_bounds(self):
+    def _data_in_bounds(self) -> tuple[np.ndarray, np.ndarray]:
         if self.data:
             points = np.array(list(self.data.keys()))
             values = np.array(list(self.data.values()), dtype=float)
@@ -562,7 +574,7 @@ def _data_in_bounds(self):
             return points[inds], values[inds].reshape(-1, self.vdim)
         return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
 
-    def _data_interp(self):
+    def _data_interp(self) -> tuple[np.ndarray, np.ndarray]:
         if self.pending_points:
             points = list(self.pending_points)
             if self.bounds_are_done:
@@ -575,7 +587,7 @@ def _data_interp(self):
             return points, values
         return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
 
-    def _data_combined(self):
+    def _data_combined(self) -> tuple[np.ndarray, np.ndarray]:
         points, values = self._data_in_bounds()
         if not self.pending_points:
             return points, values
@@ -584,7 +596,7 @@ def _data_combined(self):
         values_combined = np.vstack([values, values_interp])
         return points_combined, values_combined
 
-    def ip(self):
+    def ip(self) -> LinearNDInterpolator:
         """Deprecated, use `self.interpolator(scaled=True)`"""
         warnings.warn(
             "`learner.ip()` is deprecated, use `learner.interpolator(scaled=True)`."
@@ -593,7 +605,7 @@ def ip(self):
         )
         return self.interpolator(scaled=True)
 
-    def interpolator(self, *, scaled=False):
+    def interpolator(self, *, scaled: bool = False) -> LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data.
 
@@ -624,7 +636,7 @@ def interpolator(self, *, scaled=False):
             points, values = self._data_in_bounds()
             return interpolate.LinearNDInterpolator(points, values)
 
-    def _interpolator_combined(self):
+    def _interpolator_combined(self) -> LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data *and* interpolated data of
         the `pending_points`."""
@@ -634,12 +646,12 @@ def _interpolator_combined(self):
             self._ip_combined = interpolate.LinearNDInterpolator(points, values)
         return self._ip_combined
 
-    def inside_bounds(self, xy):
+    def inside_bounds(self, xy: tuple[float, float]) -> Bool:
         x, y = xy
         (xmin, xmax), (ymin, ymax) = self.bounds
         return xmin <= x <= xmax and ymin <= y <= ymax
 
-    def tell(self, point, value):
+    def tell(self, point: tuple[float, float], value: float | Iterable[float]) -> None:
         point = tuple(point)
         self.data[point] = value
         if not self.inside_bounds(point):
@@ -648,7 +660,7 @@ def tell(self, point, value):
         self._ip = None
         self._stack.pop(point, None)
 
-    def tell_pending(self, point):
+    def tell_pending(self, point: tuple[float, float]) -> None:
         point = tuple(point)
         if not self.inside_bounds(point):
             return
@@ -656,7 +668,9 @@ def tell_pending(self, point):
         self._ip_combined = None
         self._stack.pop(point, None)
 
-    def _fill_stack(self, stack_till=1):
+    def _fill_stack(
+        self, stack_till: int = 1
+    ) -> tuple[list[tuple[float, float]], list[float]]:
         if len(self.data) + len(self.pending_points) < self.ndim + 1:
             raise ValueError("too few points...")
 
@@ -695,7 +709,9 @@ def _fill_stack(self, stack_till=1):
 
         return points_new, losses_new
 
-    def ask(self, n, tell_pending=True):
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> tuple[list[tuple[float, float] | np.array], list[float]]:
         # Even if tell_pending is False we add the point such that _fill_stack
         # will return new points, later we remove these points if needed.
         points = list(self._stack.keys())
@@ -726,14 +742,14 @@ def ask(self, n, tell_pending=True):
         return points[:n], loss_improvements[:n]
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         if not self.bounds_are_done:
             return np.inf
         ip = self.interpolator(scaled=True) if real else self._interpolator_combined()
         losses = self.loss_per_triangle(ip)
         return losses.max()
 
-    def remove_unfinished(self):
+    def remove_unfinished(self) -> None:
         self.pending_points = set()
         for p in self._bounds_points:
             if p not in self.data:
@@ -807,10 +823,10 @@ def plot(self, n=None, tri_alpha=0):
 
         return im.opts(style=im_opts) * tris.opts(style=tri_opts, **no_hover)
 
-    def _get_data(self):
+    def _get_data(self) -> OrderedDict:
         return self.data
 
-    def _set_data(self, data):
+    def _set_data(self, data: OrderedDict) -> None:
         self.data = data
         # Remove points from stack if they already exist
         for point in copy(self._stack):

From 2b5e0d1dcf1d22382a99bdcafe138656ed8d1237 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:24 -0700
Subject: [PATCH 10/21] Add type-hints to adaptive/learner/learnerND.py

---
 adaptive/learner/learnerND.py | 141 +++++++++++++++++++++-------------
 1 file changed, 88 insertions(+), 53 deletions(-)

diff --git a/adaptive/learner/learnerND.py b/adaptive/learner/learnerND.py
index 014692f12..e0d843f44 100644
--- a/adaptive/learner/learnerND.py
+++ b/adaptive/learner/learnerND.py
@@ -6,14 +6,18 @@
 from collections import OrderedDict
 from collections.abc import Iterable
 from copy import deepcopy
+from typing import Any, Callable, Sequence
 
 import numpy as np
 import scipy.spatial
 from scipy import interpolate
+from scipy.spatial.qhull import ConvexHull
 from sortedcontainers import SortedKeyList
 
 from adaptive.learner.base_learner import BaseLearner, uses_nth_neighbors
 from adaptive.learner.triangulation import (
+    Point,
+    Simplex,
     Triangulation,
     circumsphere,
     fast_det,
@@ -21,6 +25,7 @@
     simplex_volume_in_embedding,
 )
 from adaptive.notebook_integration import ensure_holoviews, ensure_plotly
+from adaptive.types import Bool
 from adaptive.utils import (
     assign_defaults,
     cache_latest,
@@ -37,13 +42,13 @@
     with_pandas = False
 
 
-def to_list(inp):
+def to_list(inp: float) -> list[float]:
     if isinstance(inp, Iterable):
         return list(inp)
     return [inp]
 
 
-def volume(simplex, ys=None):
+def volume(simplex: Simplex, ys: None = None) -> float:
     # Notice the parameter ys is there so you can use this volume method as
     # as loss function
     matrix = np.subtract(simplex[:-1], simplex[-1], dtype=float)
@@ -54,14 +59,14 @@ def volume(simplex, ys=None):
     return vol
 
 
-def orientation(simplex):
+def orientation(simplex: np.ndarray):
     matrix = np.subtract(simplex[:-1], simplex[-1])
     # See https://www.jstor.org/stable/2315353
     sign, _logdet = np.linalg.slogdet(matrix)
     return sign
 
 
-def uniform_loss(simplex, values, value_scale):
+def uniform_loss(simplex: np.ndarray, values: np.ndarray, value_scale: float) -> float:
     """
     Uniform loss.
 
@@ -81,7 +86,7 @@ def uniform_loss(simplex, values, value_scale):
     return volume(simplex)
 
 
-def std_loss(simplex, values, value_scale):
+def std_loss(simplex: Simplex, values: np.ndarray, value_scale: float) -> np.ndarray:
     """
     Computes the loss of the simplex based on the standard deviation.
 
@@ -107,7 +112,7 @@ def std_loss(simplex, values, value_scale):
     return r.flat * np.power(vol, 1.0 / dim) + vol
 
 
-def default_loss(simplex, values, value_scale):
+def default_loss(simplex: np.ndarray, values: np.ndarray, value_scale: float) -> float:
     """
     Computes the average of the volumes of the simplex.
 
@@ -132,7 +137,13 @@ def default_loss(simplex, values, value_scale):
 
 
 @uses_nth_neighbors(1)
-def triangle_loss(simplex, values, value_scale, neighbors, neighbor_values):
+def triangle_loss(
+    simplex: np.ndarray,
+    values: np.ndarray,
+    value_scale: float,
+    neighbors: list[None | np.ndarray] | list[None] | list[np.ndarray],
+    neighbor_values: list[None | float] | list[None] | list[float],
+) -> int | float:
     """
     Computes the average of the volumes of the simplex combined with each
     neighbouring point.
@@ -169,7 +180,7 @@ def triangle_loss(simplex, values, value_scale, neighbors, neighbor_values):
     )
 
 
-def curvature_loss_function(exploration=0.05):
+def curvature_loss_function(exploration: float = 0.05) -> Callable:
     # XXX: add doc-string!
     @uses_nth_neighbors(1)
     def curvature_loss(simplex, values, value_scale, neighbors, neighbor_values):
@@ -206,7 +217,9 @@ def curvature_loss(simplex, values, value_scale, neighbors, neighbor_values):
     return curvature_loss
 
 
-def choose_point_in_simplex(simplex, transform=None):
+def choose_point_in_simplex(
+    simplex: Simplex, transform: np.ndarray | None = None
+) -> np.ndarray:
     """Choose a new point in inside a simplex.
 
     Pick the center of the simplex if the shape is nice (that is, the
@@ -247,7 +260,7 @@ def choose_point_in_simplex(simplex, transform=None):
     return point
 
 
-def _simplex_evaluation_priority(key):
+def _simplex_evaluation_priority(key: Any) -> Any:
     # We round the loss to 8 digits such that losses
     # are equal up to numerical precision will be considered
     # to be equal. This is needed because we want the learner
@@ -307,7 +320,12 @@ class LearnerND(BaseLearner):
     children based on volume.
     """
 
-    def __init__(self, func, bounds, loss_per_simplex=None):
+    def __init__(
+        self,
+        function: Callable,
+        bounds: Sequence[tuple[float, float]] | ConvexHull,
+        loss_per_simplex: Callable | None = None,
+    ) -> None:
         self._vdim = None
         self.loss_per_simplex = loss_per_simplex or default_loss
 
@@ -339,14 +357,16 @@ def __init__(self, func, bounds, loss_per_simplex=None):
 
         self.ndim = len(self._bbox)
 
-        self.function = func
-        self._tri = None
-        self._losses = dict()
+        self.function = function  # type: ignore
+        self._tri: Triangulation | None = None
+        self._losses: dict[Simplex, float] = dict()
 
-        self._pending_to_simplex = dict()  # vertex → simplex
+        self._pending_to_simplex: dict[Point, Simplex] = dict()  # vertex → simplex
 
         # triangulation of the pending points inside a specific simplex
-        self._subtriangulations = dict()  # simplex → triangulation
+        self._subtriangulations: dict[
+            Simplex, Triangulation
+        ] = dict()  # simplex → triangulation
 
         # scale to unit hypercube
         # for the input
@@ -381,12 +401,12 @@ def new(self) -> LearnerND:
         return LearnerND(self.function, self.bounds, self.loss_per_simplex)
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @property
-    def vdim(self):
+    def vdim(self) -> int:
         """Length of the output of ``learner.function``.
         If the output is unsized (when it's a scalar)
         then `vdim = 1`.
@@ -489,17 +509,17 @@ def load_dataframe(
             )
 
     @property
-    def bounds_are_done(self):
+    def bounds_are_done(self) -> bool:
         return all(p in self.data for p in self._bounds_points)
 
-    def _ip(self):
+    def _ip(self) -> interpolate.LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data."""
         # XXX: take our own triangulation into account when generating the _ip
         return interpolate.LinearNDInterpolator(self.points, self.values)
 
     @property
-    def tri(self):
+    def tri(self) -> Triangulation | None:
         """An `adaptive.learner.triangulation.Triangulation` instance
         with all the points of the learner."""
         if self._tri is not None:
@@ -517,16 +537,16 @@ def tri(self):
         return self._tri
 
     @property
-    def values(self):
+    def values(self) -> np.ndarray:
         """Get the values from `data` as a numpy array."""
         return np.array(list(self.data.values()), dtype=float)
 
     @property
-    def points(self):
+    def points(self) -> np.ndarray:
         """Get the points from `data` as a numpy array."""
         return np.array(list(self.data.keys()), dtype=float)
 
-    def tell(self, point, value):
+    def tell(self, point: tuple[float, ...], value: float | np.ndarray) -> None:
         point = tuple(point)
 
         if point in self.data:
@@ -545,16 +565,17 @@ def tell(self, point, value):
         self._update_range(value)
         if tri is not None:
             simplex = self._pending_to_simplex.get(point)
+            assert self.tri is not None
             if simplex is not None and not self._simplex_exists(simplex):
                 simplex = None
             to_delete, to_add = tri.add_point(point, simplex, transform=self._transform)
             self._update_losses(to_delete, to_add)
 
-    def _simplex_exists(self, simplex):
+    def _simplex_exists(self, simplex: Simplex) -> bool:
         simplex = tuple(sorted(simplex))
         return simplex in self.tri.simplices
 
-    def inside_bounds(self, point):
+    def inside_bounds(self, point: tuple[float, ...]) -> Bool:
         """Check whether a point is inside the bounds."""
         if self._interior is not None:
             return self._interior.find_simplex(point, tol=1e-8) >= 0
@@ -564,7 +585,7 @@ def inside_bounds(self, point):
                 (mn - eps) <= p <= (mx + eps) for p, (mn, mx) in zip(point, self._bbox)
             )
 
-    def tell_pending(self, point, *, simplex=None):
+    def tell_pending(self, point: tuple[float, ...], *, simplex=None) -> None:
         point = tuple(point)
         if not self.inside_bounds(point):
             return
@@ -591,7 +612,9 @@ def tell_pending(self, point, *, simplex=None):
                 continue
             self._update_subsimplex_losses(simpl, to_add)
 
-    def _try_adding_pending_point_to_simplex(self, point, simplex):
+    def _try_adding_pending_point_to_simplex(
+        self, point: Point, simplex: Simplex
+    ) -> Any:
         # try to insert it
         if not self.tri.point_in_simplex(point, simplex):
             return None, None
@@ -603,7 +626,9 @@ def _try_adding_pending_point_to_simplex(self, point, simplex):
         self._pending_to_simplex[point] = simplex
         return self._subtriangulations[simplex].add_point(point)
 
-    def _update_subsimplex_losses(self, simplex, new_subsimplices):
+    def _update_subsimplex_losses(
+        self, simplex: Simplex, new_subsimplices: set[Simplex]
+    ) -> None:
         loss = self._losses[simplex]
 
         loss_density = loss / self.tri.volume(simplex)
@@ -612,11 +637,11 @@ def _update_subsimplex_losses(self, simplex, new_subsimplices):
             subloss = subtriangulation.volume(subsimplex) * loss_density
             self._simplex_queue.add((subloss, simplex, subsimplex))
 
-    def _ask_and_tell_pending(self, n=1):
+    def _ask_and_tell_pending(self, n: int = 1) -> Any:
         xs, losses = zip(*(self._ask() for _ in range(n)))
         return list(xs), list(losses)
 
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True) -> Any:
         """Chose points for learners."""
         if not tell_pending:
             with restore(self):
@@ -624,7 +649,9 @@ def ask(self, n, tell_pending=True):
         else:
             return self._ask_and_tell_pending(n)
 
-    def _ask_bound_point(self):
+    def _ask_bound_point(
+        self,
+    ) -> tuple[Point, float]:
         # get the next bound point that is still available
         new_point = next(
             p
@@ -634,7 +661,9 @@ def _ask_bound_point(self):
         self.tell_pending(new_point)
         return new_point, np.inf
 
-    def _ask_point_without_known_simplices(self):
+    def _ask_point_without_known_simplices(
+        self,
+    ) -> tuple[Point, float]:
         assert not self._bounds_available
         # pick a random point inside the bounds
         # XXX: change this into picking a point based on volume loss
@@ -649,7 +678,7 @@ def _ask_point_without_known_simplices(self):
         self.tell_pending(p)
         return p, np.inf
 
-    def _pop_highest_existing_simplex(self):
+    def _pop_highest_existing_simplex(self) -> Any:
         # find the simplex with the highest loss, we do need to check that the
         # simplex hasn't been deleted yet
         while len(self._simplex_queue):
@@ -675,7 +704,9 @@ def _pop_highest_existing_simplex(self):
             "  be a simplex available if LearnerND.tri() is not None."
         )
 
-    def _ask_best_point(self):
+    def _ask_best_point(
+        self,
+    ) -> tuple[Point, float]:
         assert self.tri is not None
 
         loss, simplex, subsimplex = self._pop_highest_existing_simplex()
@@ -696,13 +727,15 @@ def _ask_best_point(self):
         return point_new, loss
 
     @property
-    def _bounds_available(self):
+    def _bounds_available(self) -> bool:
         return any(
             (p not in self.pending_points and p not in self.data)
             for p in self._bounds_points
         )
 
-    def _ask(self):
+    def _ask(
+        self,
+    ) -> tuple[Point, float]:
         if self._bounds_available:
             return self._ask_bound_point()  # O(1)
 
@@ -714,7 +747,7 @@ def _ask(self):
 
         return self._ask_best_point()  # O(log N)
 
-    def _compute_loss(self, simplex):
+    def _compute_loss(self, simplex: Simplex) -> float:
         # get the loss
         vertices = self.tri.get_vertices(simplex)
         values = [self.data[tuple(v)] for v in vertices]
@@ -753,7 +786,7 @@ def _compute_loss(self, simplex):
             )
         )
 
-    def _update_losses(self, to_delete: set, to_add: set):
+    def _update_losses(self, to_delete: set[Simplex], to_add: set[Simplex]) -> None:
         # XXX: add the points outside the triangulation to this as well
         pending_points_unbound = set()
 
@@ -799,7 +832,7 @@ def _update_losses(self, to_delete: set, to_add: set):
                     simplex, self._subtriangulations[simplex].simplices
                 )
 
-    def _recompute_all_losses(self):
+    def _recompute_all_losses(self) -> None:
         """Recompute all losses and pending losses."""
         # amortized O(N) complexity
         if self.tri is None:
@@ -823,11 +856,11 @@ def _recompute_all_losses(self):
             )
 
     @property
-    def _scale(self):
+    def _scale(self) -> float:
         # get the output scale
         return self._max_value - self._min_value
 
-    def _update_range(self, new_output):
+    def _update_range(self, new_output: list[int] | float | np.ndarray) -> bool:
         if self._min_value is None or self._max_value is None:
             # this is the first point, nothing to do, just set the range
             self._min_value = np.min(new_output)
@@ -863,12 +896,12 @@ def _update_range(self, new_output):
         return False
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         # XXX: compute pending loss if real == False
         losses = self._losses if self.tri is not None else dict()
         return max(losses.values()) if losses else float("inf")
 
-    def remove_unfinished(self):
+    def remove_unfinished(self) -> None:
         # XXX: implement this method
         self.pending_points = set()
         self._subtriangulations = dict()
@@ -878,7 +911,7 @@ def remove_unfinished(self):
     # Plotting related stuff #
     ##########################
 
-    def plot(self, n=None, tri_alpha=0):
+    def plot(self, n: int | None = None, tri_alpha: float = 0):
         """Plot the function we want to learn, only works in 2D.
 
         Parameters
@@ -939,7 +972,7 @@ def plot(self, n=None, tri_alpha=0):
 
         return im.opts(style=im_opts) * tris.opts(style=tri_opts, **no_hover)
 
-    def plot_slice(self, cut_mapping, n=None):
+    def plot_slice(self, cut_mapping: dict[int, float], n: int | None = None):
         """Plot a 1D or 2D interpolated slice of a N-dimensional function.
 
         Parameters
@@ -1009,7 +1042,7 @@ def plot_slice(self, cut_mapping, n=None):
         else:
             raise ValueError("Only 1 or 2-dimensional plots can be generated.")
 
-    def plot_3D(self, with_triangulation=False, return_fig=False):
+    def plot_3D(self, with_triangulation: bool = False, return_fig: bool = False):
         """Plot the learner's data in 3D using plotly.
 
         Does *not* work with the
@@ -1094,7 +1127,7 @@ def plot_3D(self, with_triangulation=False, return_fig=False):
 
         return fig if return_fig else plotly.offline.iplot(fig)
 
-    def _get_iso(self, level=0.0, which="surface"):
+    def _get_iso(self, level: float = 0.0, which: str = "surface"):
         if which == "surface":
             if self.ndim != 3 or self.vdim != 1:
                 raise Exception(
@@ -1165,7 +1198,9 @@ def _get_vertex_index(a, b):
 
         return vertices, faces_or_lines
 
-    def plot_isoline(self, level=0.0, n=None, tri_alpha=0):
+    def plot_isoline(
+        self, level: float = 0.0, n: int | None = None, tri_alpha: float = 0
+    ):
         """Plot the isoline at a specific level, only works in 2D.
 
         Parameters
@@ -1205,7 +1240,7 @@ def plot_isoline(self, level=0.0, n=None, tri_alpha=0):
         contour = contour.opts(style=contour_opts)
         return plot * contour
 
-    def plot_isosurface(self, level=0.0, hull_opacity=0.2):
+    def plot_isosurface(self, level: float = 0.0, hull_opacity: float = 0.2):
         """Plots a linearly interpolated isosurface.
 
         This is the 3D analog of an isoline. Does *not* work with the
@@ -1243,7 +1278,7 @@ def plot_isosurface(self, level=0.0, hull_opacity=0.2):
         hull_mesh = self._get_hull_mesh(opacity=hull_opacity)
         return plotly.offline.iplot([isosurface, hull_mesh])
 
-    def _get_hull_mesh(self, opacity=0.2):
+    def _get_hull_mesh(self, opacity: float = 0.2):
         plotly = ensure_plotly()
         hull = scipy.spatial.ConvexHull(self._bounds_points)
 
@@ -1282,9 +1317,9 @@ def _get_plane_color(simplex):
             lighting=lighting,
         )
 
-    def _get_data(self):
+    def _get_data(self) -> dict[str, Any]:
         return deepcopy(self.__dict__)
 
-    def _set_data(self, state):
+    def _set_data(self, state: dict[str, Any]) -> None:
         for k, v in state.items():
             setattr(self, k, v)

From 08ddfb20b0a5cee9793ac363455d1c2b6bcc9f7b Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:25 -0700
Subject: [PATCH 11/21] Add type-hints to adaptive/learner/sequence_learner.py

---
 adaptive/learner/sequence_learner.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/adaptive/learner/sequence_learner.py b/adaptive/learner/sequence_learner.py
index 959b233c9..a76b821e1 100644
--- a/adaptive/learner/sequence_learner.py
+++ b/adaptive/learner/sequence_learner.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from copy import copy
-from typing import Any, Tuple
+from typing import Any, Callable, Iterable, Tuple
 
 import cloudpickle
 from sortedcontainers import SortedDict, SortedSet
@@ -37,17 +37,17 @@ class _IgnoreFirstArgument:
     pickable.
     """
 
-    def __init__(self, function):
-        self.function = function
+    def __init__(self, function: Callable) -> None:
+        self.function = function  # type: ignore
 
     def __call__(self, index_point: PointType, *args, **kwargs):
         index, point = index_point
         return self.function(point, *args, **kwargs)
 
-    def __getstate__(self):
+    def __getstate__(self) -> Callable:
         return self.function
 
-    def __setstate__(self, function):
+    def __setstate__(self, function: Callable) -> None:
         self.__init__(function)
 
 
@@ -78,9 +78,9 @@ class SequenceLearner(BaseLearner):
     the added benefit of having results in the local kernel already.
     """
 
-    def __init__(self, function, sequence):
+    def __init__(self, function: Callable, sequence: Iterable) -> None:
         self._original_function = function
-        self.function = _IgnoreFirstArgument(function)
+        self.function = _IgnoreFirstArgument(function)  # type: ignore
         self._to_do_indices = SortedSet({i for i, _ in enumerate(sequence)})
         self._ntotal = len(sequence)
         self.sequence = copy(sequence)

From 5d8110f3d6e8c930672e583b2171d0dab6a7bcd0 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:26 -0700
Subject: [PATCH 12/21] Add type-hints to adaptive/learner/skopt_learner.py

---
 adaptive/learner/skopt_learner.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/adaptive/learner/skopt_learner.py b/adaptive/learner/skopt_learner.py
index e12f49daa..1c3c18fd7 100644
--- a/adaptive/learner/skopt_learner.py
+++ b/adaptive/learner/skopt_learner.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import collections
+from typing import Callable
 
 import numpy as np
 from skopt import Optimizer
@@ -25,8 +26,8 @@ class SKOptLearner(Optimizer, BaseLearner):
         Arguments to pass to ``skopt.Optimizer``.
     """
 
-    def __init__(self, function, **kwargs):
-        self.function = function
+    def __init__(self, function: Callable, **kwargs) -> None:
+        self.function = function  # type: ignore
         self.pending_points = set()
         self.data = collections.OrderedDict()
         self._kwargs = kwargs
@@ -36,7 +37,7 @@ def new(self) -> SKOptLearner:
         """Return a new `~adaptive.SKOptLearner` without the data."""
         return SKOptLearner(self.function, **self._kwargs)
 
-    def tell(self, x, y, fit=True):
+    def tell(self, x: float | list[float], y: float, fit: bool = True) -> None:
         if isinstance(x, collections.abc.Iterable):
             self.pending_points.discard(tuple(x))
             self.data[tuple(x)] = y
@@ -55,7 +56,7 @@ def remove_unfinished(self):
         pass
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         if not self.models:
             return np.inf
         else:
@@ -65,7 +66,12 @@ 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
+    ) -> (
+        tuple[list[float], list[float]]
+        | tuple[list[list[float]], list[float]]  # XXX: this indicates a bug!
+    ):
         if not tell_pending:
             raise NotImplementedError(
                 "Asking points is an irreversible "
@@ -79,7 +85,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)
 

From 7fd0588ecbe8446283aa74d82c8dac61abf79dd1 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:27 -0700
Subject: [PATCH 13/21] Add type-hints to adaptive/learner/triangulation.py

---
 adaptive/learner/triangulation.py | 144 ++++++++++++++++++++----------
 1 file changed, 96 insertions(+), 48 deletions(-)

diff --git a/adaptive/learner/triangulation.py b/adaptive/learner/triangulation.py
index 4eb5952d5..da24ec3af 100644
--- a/adaptive/learner/triangulation.py
+++ b/adaptive/learner/triangulation.py
@@ -1,7 +1,11 @@
+from __future__ import annotations
+
+import collections.abc
+import numbers
 from collections import Counter
-from collections.abc import Iterable, Sized
 from itertools import chain, combinations
 from math import factorial, sqrt
+from typing import Any, Iterable, Iterator, List, Sequence, Tuple, Union
 
 import scipy.spatial
 from numpy import abs as np_abs
@@ -13,6 +17,7 @@
     dot,
     eye,
     mean,
+    ndarray,
     ones,
     square,
     subtract,
@@ -22,8 +27,22 @@
 from numpy.linalg import det as ndet
 from numpy.linalg import matrix_rank, norm, slogdet, solve
 
+from adaptive.types import Bool
+
+try:
+    from typing import TypeAlias
+except ImportError:
+    # Remove this when we drop support for Python 3.9
+    from typing_extensions import TypeAlias
+
 
-def fast_norm(v):
+SimplexPoints: TypeAlias = Union[List[Tuple[float, ...]], ndarray]
+Simplex: TypeAlias = Union[Sequence[numbers.Integral], ndarray]
+Point: TypeAlias = Union[Tuple[float, ...], ndarray]
+Points: TypeAlias = Union[Sequence[Tuple[float, ...]], ndarray]
+
+
+def fast_norm(v: tuple[float, ...] | ndarray) -> float:
     """Take the vector norm for len 2, 3 vectors.
     Defaults to a square root of the dot product for larger vectors.
 
@@ -41,7 +60,9 @@ def fast_norm(v):
     return sqrt(dot(v, v))
 
 
-def fast_2d_point_in_simplex(point, simplex, eps=1e-8):
+def fast_2d_point_in_simplex(
+    point: Point, simplex: SimplexPoints, eps: float = 1e-8
+) -> Bool:
     (p0x, p0y), (p1x, p1y), (p2x, p2y) = simplex
     px, py = point
 
@@ -55,7 +76,7 @@ def fast_2d_point_in_simplex(point, simplex, eps=1e-8):
     return (t >= -eps) and (s + t <= 1 + eps)
 
 
-def point_in_simplex(point, simplex, eps=1e-8):
+def point_in_simplex(point: Point, simplex: SimplexPoints, eps: float = 1e-8) -> Bool:
     if len(point) == 2:
         return fast_2d_point_in_simplex(point, simplex, eps)
 
@@ -66,7 +87,7 @@ def point_in_simplex(point, simplex, eps=1e-8):
     return all(alpha > -eps) and sum(alpha) < 1 + eps
 
 
-def fast_2d_circumcircle(points):
+def fast_2d_circumcircle(points: Points) -> tuple[tuple[float, float], float]:
     """Compute the center and radius of the circumscribed circle of a triangle
 
     Parameters
@@ -79,7 +100,7 @@ def fast_2d_circumcircle(points):
     tuple
         (center point : tuple(float), radius: float)
     """
-    points = array(points)
+    points = array(points, dtype=float)
     # transform to relative coordinates
     pts = points[1:] - points[0]
 
@@ -102,7 +123,9 @@ def fast_2d_circumcircle(points):
     return (x + points[0][0], y + points[0][1]), radius
 
 
-def fast_3d_circumcircle(points):
+def fast_3d_circumcircle(
+    points: Points,
+) -> tuple[tuple[float, float, float], float]:
     """Compute the center and radius of the circumscribed sphere of a simplex.
 
     Parameters
@@ -142,7 +165,7 @@ def fast_3d_circumcircle(points):
     return center, radius
 
 
-def fast_det(matrix):
+def fast_det(matrix: ndarray) -> float:
     matrix = asarray(matrix, dtype=float)
     if matrix.shape == (2, 2):
         return matrix[0][0] * matrix[1][1] - matrix[1][0] * matrix[0][1]
@@ -153,7 +176,7 @@ def fast_det(matrix):
         return ndet(matrix)
 
 
-def circumsphere(pts):
+def circumsphere(pts: Simplex) -> tuple[tuple[float, ...], float]:
     """Compute the center and radius of a N dimension sphere which touches each point in pts.
 
     Parameters
@@ -201,7 +224,7 @@ def circumsphere(pts):
     return tuple(center), radius
 
 
-def orientation(face, origin):
+def orientation(face: tuple | ndarray, origin: tuple | ndarray) -> int:
     """Compute the orientation of the face with respect to a point, origin.
 
     Parameters
@@ -224,14 +247,14 @@ def orientation(face, origin):
     sign, logdet = slogdet(vectors - origin)
     if logdet < -50:  # assume it to be zero when it's close to zero
         return 0
-    return sign
+    return int(sign)
 
 
-def is_iterable_and_sized(obj):
-    return isinstance(obj, Iterable) and isinstance(obj, Sized)
+def is_iterable_and_sized(obj: Any) -> bool:
+    return isinstance(obj, collections.abc.Collection)
 
 
-def simplex_volume_in_embedding(vertices) -> float:
+def simplex_volume_in_embedding(vertices: Sequence[Point]) -> float:
     """Calculate the volume of a simplex in a higher dimensional embedding.
     That is: dim > len(vertices) - 1. For example if you would like to know the
     surface area of a triangle in a 3d space.
@@ -312,7 +335,7 @@ class Triangulation:
         or more simplices in the
     """
 
-    def __init__(self, coords):
+    def __init__(self, coords: Points) -> None:
         if not is_iterable_and_sized(coords):
             raise TypeError("Please provide a 2-dimensional list of points")
         coords = list(coords)
@@ -340,10 +363,10 @@ def __init__(self, coords):
                 "(the points are linearly dependent)"
             )
 
-        self.vertices = list(coords)
-        self.simplices = set()
+        self.vertices: list[Point] = list(coords)
+        self.simplices: set[Simplex] = set()
         # initialise empty set for each vertex
-        self.vertex_to_simplices = [set() for _ in coords]
+        self.vertex_to_simplices: list[set[Simplex]] = [set() for _ in coords]
 
         # find a Delaunay triangulation to start with, then we will throw it
         # away and continue with our own algorithm
@@ -351,27 +374,29 @@ def __init__(self, coords):
         for simplex in initial_tri.simplices:
             self.add_simplex(simplex)
 
-    def delete_simplex(self, simplex):
+    def delete_simplex(self, simplex: Simplex) -> None:
         simplex = tuple(sorted(simplex))
         self.simplices.remove(simplex)
         for vertex in simplex:
             self.vertex_to_simplices[vertex].remove(simplex)
 
-    def add_simplex(self, simplex):
+    def add_simplex(self, simplex: Simplex) -> None:
         simplex = tuple(sorted(simplex))
         self.simplices.add(simplex)
         for vertex in simplex:
             self.vertex_to_simplices[vertex].add(simplex)
 
-    def get_vertices(self, indices):
+    def get_vertices(self, indices: Iterable[numbers.Integral]) -> list[Point | None]:
         return [self.get_vertex(i) for i in indices]
 
-    def get_vertex(self, index):
+    def get_vertex(self, index: numbers.Integral | None) -> Point | None:
         if index is None:
             return None
         return self.vertices[index]
 
-    def get_reduced_simplex(self, point, simplex, eps=1e-8) -> list:
+    def get_reduced_simplex(
+        self, point: Point, simplex: Simplex, eps: float = 1e-8
+    ) -> list[numbers.Integral]:
         """Check whether vertex lies within a simplex.
 
         Returns
@@ -396,11 +421,13 @@ def get_reduced_simplex(self, point, simplex, eps=1e-8) -> list:
 
         return [simplex[i] for i in result]
 
-    def point_in_simplex(self, point, simplex, eps=1e-8):
+    def point_in_simplex(
+        self, point: Point, simplex: Simplex, eps: float = 1e-8
+    ) -> Bool:
         vertices = self.get_vertices(simplex)
         return point_in_simplex(point, vertices, eps)
 
-    def locate_point(self, point):
+    def locate_point(self, point: Point) -> Simplex:
         """Find to which simplex the point belongs.
 
         Return indices of the simplex containing the point.
@@ -412,10 +439,15 @@ def locate_point(self, point):
         return ()
 
     @property
-    def dim(self):
+    def dim(self) -> int:
         return len(self.vertices[0])
 
-    def faces(self, dim=None, simplices=None, vertices=None):
+    def faces(
+        self,
+        dim: int | None = None,
+        simplices: Iterable[Simplex] | None = None,
+        vertices: Iterable[int] | None = None,
+    ) -> Iterator[tuple[numbers.Integral, ...]]:
         """Iterator over faces of a simplex or vertex sequence."""
         if dim is None:
             dim = self.dim
@@ -436,11 +468,11 @@ def faces(self, dim=None, simplices=None, vertices=None):
         else:
             return faces
 
-    def containing(self, face):
+    def containing(self, face: tuple[int, ...]) -> set[Simplex]:
         """Simplices containing a face."""
         return set.intersection(*(self.vertex_to_simplices[i] for i in face))
 
-    def _extend_hull(self, new_vertex, eps=1e-8):
+    def _extend_hull(self, new_vertex: Point, eps: float = 1e-8) -> set[Simplex]:
         # count multiplicities in order to get all hull faces
         multiplicities = Counter(face for face in self.faces())
         hull_faces = [face for face, count in multiplicities.items() if count == 1]
@@ -480,7 +512,9 @@ def _extend_hull(self, new_vertex, eps=1e-8):
 
         return new_simplices
 
-    def circumscribed_circle(self, simplex, transform):
+    def circumscribed_circle(
+        self, simplex: Simplex, transform: ndarray
+    ) -> tuple[tuple[float, ...], float]:
         """Compute the center and radius of the circumscribed circle of a simplex.
 
         Parameters
@@ -496,7 +530,9 @@ def circumscribed_circle(self, simplex, transform):
         pts = dot(self.get_vertices(simplex), transform)
         return circumsphere(pts)
 
-    def point_in_cicumcircle(self, pt_index, simplex, transform):
+    def point_in_cicumcircle(
+        self, pt_index: int, simplex: Simplex, transform: ndarray
+    ) -> Bool:
         # return self.fast_point_in_circumcircle(pt_index, simplex, transform)
         eps = 1e-8
 
@@ -506,10 +542,15 @@ def point_in_cicumcircle(self, pt_index, simplex, transform):
         return norm(center - pt) < (radius * (1 + eps))
 
     @property
-    def default_transform(self):
+    def default_transform(self) -> ndarray:
         return eye(self.dim)
 
-    def bowyer_watson(self, pt_index, containing_simplex=None, transform=None):
+    def bowyer_watson(
+        self,
+        pt_index: int,
+        containing_simplex: Simplex | None = None,
+        transform: ndarray | None = None,
+    ) -> tuple[set[Simplex], set[Simplex]]:
         """Modified Bowyer-Watson point adding algorithm.
 
         Create a hole in the triangulation around the new point,
@@ -569,10 +610,10 @@ def bowyer_watson(self, pt_index, containing_simplex=None, transform=None):
         new_triangles = self.vertex_to_simplices[pt_index]
         return bad_triangles - new_triangles, new_triangles - bad_triangles
 
-    def _simplex_is_almost_flat(self, simplex):
+    def _simplex_is_almost_flat(self, simplex: Simplex) -> Bool:
         return self._relative_volume(simplex) < 1e-8
 
-    def _relative_volume(self, simplex):
+    def _relative_volume(self, simplex: Simplex) -> float:
         """Compute the volume of a simplex divided by the average (Manhattan)
         distance of its vertices. The advantage of this is that the relative
         volume is only dependent on the shape of the simplex and not on the
@@ -583,20 +624,25 @@ def _relative_volume(self, simplex):
         average_edge_length = mean(np_abs(vectors))
         return self.volume(simplex) / (average_edge_length**self.dim)
 
-    def add_point(self, point, simplex=None, transform=None):
+    def add_point(
+        self,
+        point: Point,
+        simplex: Simplex | None = None,
+        transform: ndarray | None = None,
+    ) -> tuple[set[Simplex], set[Simplex]]:
         """Add a new vertex and create simplices as appropriate.
 
         Parameters
         ----------
         point : float vector
             Coordinates of the point to be added.
-        transform : N*N matrix of floats
-            Multiplication matrix to apply to the point (and neighbouring
-            simplices) when running the Bowyer Watson method.
         simplex : tuple of ints, optional
             Simplex containing the point. Empty tuple indicates points outside
             the hull. If not provided, the algorithm costs O(N), so this should
             be used whenever possible.
+        transform : N*N matrix of floats
+            Multiplication matrix to apply to the point (and neighbouring
+            simplices) when running the Bowyer Watson method.
         """
         point = tuple(point)
         if simplex is None:
@@ -632,16 +678,16 @@ def add_point(self, point, simplex=None, transform=None):
             self.vertices.append(point)
             return self.bowyer_watson(pt_index, actual_simplex, transform)
 
-    def volume(self, simplex):
+    def volume(self, simplex: Simplex) -> float:
         prefactor = factorial(self.dim)
         vertices = array(self.get_vertices(simplex))
         vectors = vertices[1:] - vertices[0]
         return float(abs(fast_det(vectors)) / prefactor)
 
-    def volumes(self):
+    def volumes(self) -> list[float]:
         return [self.volume(sim) for sim in self.simplices]
 
-    def reference_invariant(self):
+    def reference_invariant(self) -> bool:
         """vertex_to_simplices and simplices are compatible."""
         for vertex in range(len(self.vertices)):
             if any(vertex not in tri for tri in self.vertex_to_simplices[vertex]):
@@ -655,26 +701,28 @@ def vertex_invariant(self, vertex):
         """Simplices originating from a vertex don't overlap."""
         raise NotImplementedError
 
-    def get_neighbors_from_vertices(self, simplex):
+    def get_neighbors_from_vertices(self, simplex: Simplex) -> set[Simplex]:
         return set.union(*[self.vertex_to_simplices[p] for p in simplex])
 
-    def get_face_sharing_neighbors(self, neighbors, simplex):
+    def get_face_sharing_neighbors(
+        self, neighbors: set[Simplex], simplex: Simplex
+    ) -> set[Simplex]:
         """Keep only the simplices sharing a whole face with simplex."""
         return {
             simpl for simpl in neighbors if len(set(simpl) & set(simplex)) == self.dim
         }  # they share a face
 
-    def get_simplices_attached_to_points(self, indices):
+    def get_simplices_attached_to_points(self, indices: Simplex) -> set[Simplex]:
         # Get all simplices that share at least a point with the simplex
         neighbors = self.get_neighbors_from_vertices(indices)
         return self.get_face_sharing_neighbors(neighbors, indices)
 
-    def get_opposing_vertices(self, simplex):
+    def get_opposing_vertices(self, simplex: Simplex) -> tuple[int, ...]:
         if simplex not in self.simplices:
             raise ValueError("Provided simplex is not part of the triangulation")
         neighbors = self.get_simplices_attached_to_points(simplex)
 
-        def find_opposing_vertex(vertex):
+        def find_opposing_vertex(vertex: int):
             # find the simplex:
             simp = next((x for x in neighbors if vertex not in x), None)
             if simp is None:
@@ -687,7 +735,7 @@ def find_opposing_vertex(vertex):
         return result
 
     @property
-    def hull(self):
+    def hull(self) -> set[numbers.Integral]:
         """Compute hull from triangulation.
 
         Parameters

From 2e0efc1b6abc3b01f57a998844ec93a78b2f814a Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:28 -0700
Subject: [PATCH 14/21] Add type-hints to adaptive/notebook_integration.py

---
 adaptive/notebook_integration.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/adaptive/notebook_integration.py b/adaptive/notebook_integration.py
index 426c04541..60329110e 100644
--- a/adaptive/notebook_integration.py
+++ b/adaptive/notebook_integration.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import datetime
 import importlib
@@ -76,7 +78,7 @@ def ensure_plotly():
         raise RuntimeError("plotly is not installed; plotting is disabled.")
 
 
-def in_ipynb():
+def in_ipynb() -> bool:
     try:
         # If we are running in IPython, then `get_ipython()` is always a global
         return get_ipython().__class__.__name__ == "ZMQInteractiveShell"

From 96e186f245def4b0d73e212f603e8bf35fa4ac19 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:30 -0700
Subject: [PATCH 15/21] Add type-hints to adaptive/tests/algorithm_4.py

---
 adaptive/tests/algorithm_4.py | 67 +++++++++++++++++++++--------------
 1 file changed, 40 insertions(+), 27 deletions(-)

diff --git a/adaptive/tests/algorithm_4.py b/adaptive/tests/algorithm_4.py
index fb401e866..4566c0fa1 100644
--- a/adaptive/tests/algorithm_4.py
+++ b/adaptive/tests/algorithm_4.py
@@ -2,7 +2,8 @@
 # Copyright 2017 Christoph Groth
 
 from collections import defaultdict
-from fractions import Fraction as Frac
+from fractions import Fraction
+from typing import Callable, List, Tuple, Union
 
 import numpy as np
 from numpy.testing import assert_allclose
@@ -11,7 +12,7 @@
 eps = np.spacing(1)
 
 
-def legendre(n):
+def legendre(n: int) -> List[List[Fraction]]:
     """Return the first n Legendre polynomials.
 
     The polynomials have *standard* normalization, i.e.
@@ -19,12 +20,12 @@ def legendre(n):
 
     The return value is a list of list of fraction.Fraction instances.
     """
-    result = [[Frac(1)], [Frac(0), Frac(1)]]
+    result = [[Fraction(1)], [Fraction(0), Fraction(1)]]
     if n <= 2:
         return result[:n]
     for i in range(2, n):
         # Use Bonnet's recursion formula.
-        new = (i + 1) * [Frac(0)]
+        new = (i + 1) * [Fraction(0)]
         new[1:] = (r * (2 * i - 1) for r in result[-1])
         new[:-2] = (n - r * (i - 1) for n, r in zip(new[:-2], result[-2]))
         new[:] = (n / i for n in new)
@@ -32,7 +33,7 @@ def legendre(n):
     return result
 
 
-def newton(n):
+def newton(n: int) -> np.ndarray:
     """Compute the monomial coefficients of the Newton polynomial over the
     nodes of the n-point Clenshaw-Curtis quadrature rule.
     """
@@ -89,7 +90,7 @@ def newton(n):
     return cf
 
 
-def scalar_product(a, b):
+def scalar_product(a: List[Fraction], b: List[Fraction]) -> Fraction:
     """Compute the polynomial scalar product int_-1^1 dx a(x) b(x).
 
     The args must be sequences of polynomial coefficients.  This
@@ -110,7 +111,7 @@ def scalar_product(a, b):
     return 2 * sum(c[i] / (i + 1) for i in range(0, lc, 2))
 
 
-def calc_bdef(ns):
+def calc_bdef(ns: Tuple[int, int, int, int]) -> List[np.ndarray]:
     """Calculate the decompositions of Newton polynomials (over the nodes
     of the n-point Clenshaw-Curtis quadrature rule) in terms of
     Legandre polynomials.
@@ -123,7 +124,7 @@ def calc_bdef(ns):
     result = []
     for n in ns:
         poly = []
-        a = list(map(Frac, newton(n)))
+        a = list(map(Fraction, newton(n)))
         for b in legs[: n + 1]:
             igral = scalar_product(a, b)
 
@@ -145,7 +146,7 @@ def calc_bdef(ns):
 b_def = calc_bdef(n)
 
 
-def calc_V(xi, n):
+def calc_V(xi: np.ndarray, n: int) -> np.ndarray:
     V = [np.ones(xi.shape), xi.copy()]
     for i in range(2, n):
         V.append((2 * i - 1) / i * xi * V[-1] - (i - 1) / i * V[-2])
@@ -183,7 +184,7 @@ def calc_V(xi, n):
 gamma = np.concatenate([[0, 0], np.sqrt(k[2:] ** 2 / (4 * k[2:] ** 2 - 1))])
 
 
-def _downdate(c, nans, depth):
+def _downdate(c: np.ndarray, nans: List[int], depth: int) -> None:
     # This is algorithm 5 from the thesis of Pedro Gonnet.
     b = b_def[depth].copy()
     m = n[depth] - 1
@@ -200,7 +201,7 @@ def _downdate(c, nans, depth):
         m -= 1
 
 
-def _zero_nans(fx):
+def _zero_nans(fx: np.ndarray) -> List[int]:
     nans = []
     for i in range(len(fx)):
         if not np.isfinite(fx[i]):
@@ -209,7 +210,7 @@ def _zero_nans(fx):
     return nans
 
 
-def _calc_coeffs(fx, depth):
+def _calc_coeffs(fx: np.ndarray, depth: int) -> np.ndarray:
     """Caution: this function modifies fx."""
     nans = _zero_nans(fx)
     c_new = V_inv[depth] @ fx
@@ -220,7 +221,7 @@ def _calc_coeffs(fx, depth):
 
 
 class DivergentIntegralError(ValueError):
-    def __init__(self, msg, igral, err, nr_points):
+    def __init__(self, msg: str, igral: float, err: None, nr_points: int) -> None:
         self.igral = igral
         self.err = err
         self.nr_points = nr_points
@@ -230,19 +231,23 @@ def __init__(self, msg, igral, err, nr_points):
 class _Interval:
     __slots__ = ["a", "b", "c", "fx", "igral", "err", "depth", "rdepth", "ndiv", "c00"]
 
-    def __init__(self, a, b, depth, rdepth):
+    def __init__(
+        self, a: Union[int, float], b: Union[int, float], depth: int, rdepth: int
+    ) -> None:
         self.a = a
         self.b = b
         self.depth = depth
         self.rdepth = rdepth
 
-    def points(self):
+    def points(self) -> np.ndarray:
         a = self.a
         b = self.b
         return (a + b) / 2 + (b - a) * xi[self.depth] / 2
 
     @classmethod
-    def make_first(cls, f, a, b, depth=2):
+    def make_first(
+        cls, f: Callable, a: int, b: int, depth: int = 2
+    ) -> Tuple["_Interval", int]:
         ival = _Interval(a, b, depth, 1)
         fx = f(ival.points())
         ival.c = _calc_coeffs(fx, depth)
@@ -251,7 +256,7 @@ def make_first(cls, f, a, b, depth=2):
         ival.ndiv = 0
         return ival, n[depth]
 
-    def calc_igral_and_err(self, c_old):
+    def calc_igral_and_err(self, c_old: np.ndarray) -> float:
         self.c = c_new = _calc_coeffs(self.fx, self.depth)
         c_diff = np.zeros(max(len(c_old), len(c_new)))
         c_diff[: len(c_old)] = c_old
@@ -262,7 +267,9 @@ def calc_igral_and_err(self, c_old):
         self.err = w * c_diff
         return c_diff
 
-    def split(self, f):
+    def split(
+        self, f: Callable
+    ) -> Union[Tuple[Tuple[float, float, float], int], Tuple[List["_Interval"], int]]:
         m = (self.a + self.b) / 2
         f_center = self.fx[(len(self.fx) - 1) // 2]
 
@@ -287,7 +294,7 @@ def split(self, f):
 
         return ivals, nr_points
 
-    def refine(self, f):
+    def refine(self, f: Callable) -> Tuple[np.ndarray, bool, int]:
         """Increase degree of interval."""
         self.depth = depth = self.depth + 1
         points = self.points()
@@ -299,7 +306,9 @@ def refine(self, f):
         return points, split, n[depth] - n[depth - 1]
 
 
-def algorithm_4(f, a, b, tol, N_loops=int(1e9)):
+def algorithm_4(
+    f: Callable, a: int, b: int, tol: float, N_loops: int = int(1e9)
+) -> Tuple[float, float, int, List["_Interval"]]:
     """ALGORITHM_4 evaluates an integral using adaptive quadrature. The
     algorithm uses Clenshaw-Curtis quadrature rules of increasing
     degree in each interval and bisects the interval if either the
@@ -403,10 +412,10 @@ def algorithm_4(f, a, b, tol, N_loops=int(1e9)):
     return igral, err, nr_points, ivals
 
 
-################ Tests ################
+# ############### Tests ################
 
 
-def f0(x):
+def f0(x: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
     return x * np.sin(1 / x) * np.sqrt(abs(1 - x))
 
 
@@ -414,18 +423,20 @@ def f7(x):
     return x**-0.5
 
 
-def f24(x):
+def f24(x: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
     return np.floor(np.exp(x))
 
 
-def f21(x):
+def f21(x: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
     y = 0
     for i in range(1, 4):
         y += 1 / np.cosh(20**i * (x - 2 * i / 10))
     return y
 
 
-def f63(x, alpha, beta):
+def f63(
+    x: Union[float, np.ndarray], alpha: float, beta: float
+) -> Union[float, np.ndarray]:
     return abs(x - beta) ** alpha
 
 
@@ -433,7 +444,7 @@ def F63(x, alpha, beta):
     return (x - beta) * abs(x - beta) ** alpha / (alpha + 1)
 
 
-def fdiv(x):
+def fdiv(x: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
     return abs(x - 0.987654321) ** -1.1
 
 
@@ -461,7 +472,9 @@ def test_scalar_product(n=33):
     selection = [0, 5, 7, n - 1]
     for i in selection:
         for j in selection:
-            assert scalar_product(legs[i], legs[j]) == ((i == j) and Frac(2, 2 * i + 1))
+            assert scalar_product(legs[i], legs[j]) == (
+                (i == j) and Fraction(2, 2 * i + 1)
+            )
 
 
 def simple_newton(n):

From ec3649aa839dabf1f5d576cbee8f2eaf9662a7bb Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:31 -0700
Subject: [PATCH 16/21] Add type-hints to
 adaptive/tests/test_average_learner1d.py

---
 adaptive/tests/test_average_learner1d.py | 22 ++++------------------
 1 file changed, 4 insertions(+), 18 deletions(-)

diff --git a/adaptive/tests/test_average_learner1d.py b/adaptive/tests/test_average_learner1d.py
index 8b2670d77..4286f55b9 100644
--- a/adaptive/tests/test_average_learner1d.py
+++ b/adaptive/tests/test_average_learner1d.py
@@ -1,6 +1,6 @@
-from itertools import chain
-
 import numpy as np
+import pandas as pd
+from pandas.testing import assert_series_equal
 
 from adaptive import AverageLearner1D
 from adaptive.tests.test_learners import (
@@ -11,27 +11,13 @@
 
 
 def almost_equal_dicts(a, b):
-    assert a.keys() == b.keys()
-    for k, v1 in a.items():
-        v2 = b[k]
-        if (
-            v1 is None
-            or v2 is None
-            or isinstance(v1, (tuple, list))
-            and any(x is None for x in chain(v1, v2))
-        ):
-            assert v1 == v2
-        else:
-            try:
-                np.testing.assert_almost_equal(v1, v2)
-            except TypeError:
-                raise AssertionError(f"{v1} != {v2}")
+    assert_series_equal(pd.Series(sorted(a.items())), pd.Series(sorted(b.items())))
 
 
 def test_tell_many_at_point():
     f = generate_random_parametrization(noisy_peak)
     learner = AverageLearner1D(f, bounds=(-2, 2))
-    control = learner.new()
+    control = AverageLearner1D(f, bounds=(-2, 2))
     learner._recompute_losses_factor = 1
     control._recompute_losses_factor = 1
     simple_run(learner, 100)

From 54265bb8ceb27e116979e992e2bf186487f9cb88 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:32 -0700
Subject: [PATCH 17/21] Add type-hints to adaptive/tests/test_learner1d.py

---
 adaptive/tests/test_learner1d.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/adaptive/tests/test_learner1d.py b/adaptive/tests/test_learner1d.py
index 59a8b81a8..7e990bd7b 100644
--- a/adaptive/tests/test_learner1d.py
+++ b/adaptive/tests/test_learner1d.py
@@ -277,8 +277,8 @@ def f_vec(x, offset=0.123214):
     def assert_equal_dicts(d1, d2):
         xs1, ys1 = zip(*sorted(d1.items()))
         xs2, ys2 = zip(*sorted(d2.items()))
-        ys1 = np.array(ys1, dtype=np.float)
-        ys2 = np.array(ys2, dtype=np.float)
+        ys1 = np.array(ys1, dtype=np.float64)
+        ys2 = np.array(ys2, dtype=np.float64)
         np.testing.assert_almost_equal(xs1, xs2)
         np.testing.assert_almost_equal(ys1, ys2)
 

From 9e52ecb04facbfeb0fd4fbee5dcdd4bef91c4e64 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:33 -0700
Subject: [PATCH 18/21] Add type-hints to adaptive/tests/test_learners.py

---
 adaptive/tests/test_learners.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/adaptive/tests/test_learners.py b/adaptive/tests/test_learners.py
index d393511fb..8c616e6bb 100644
--- a/adaptive/tests/test_learners.py
+++ b/adaptive/tests/test_learners.py
@@ -12,6 +12,7 @@
 
 import flaky
 import numpy as np
+import pandas
 import pytest
 import scipy.spatial
 
@@ -27,7 +28,6 @@
     LearnerND,
     SequenceLearner,
 )
-from adaptive.learner.learner1D import with_pandas
 from adaptive.runner import simple
 
 try:
@@ -708,7 +708,6 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-@pytest.mark.skipif(not with_pandas, reason="pandas is not installed")
 @run_with(
     Learner1D,
     Learner2D,
@@ -720,8 +719,6 @@ def wrapper(*args, **kwargs):
     with_all_loss_functions=False,
 )
 def test_to_dataframe(learner_type, f, learner_kwargs):
-    import pandas
-
     if learner_type is LearnerND:
         kw = {"point_names": tuple("xyz")[: len(learner_kwargs["bounds"])]}
     else:

From b4ba66ee4d18d78384e7fdbc604d3001257947ca Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:34 -0700
Subject: [PATCH 19/21] Add type-hints to adaptive/types.py

---
 adaptive/types.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/adaptive/types.py b/adaptive/types.py
index e2d57a44f..67268f822 100644
--- a/adaptive/types.py
+++ b/adaptive/types.py
@@ -11,3 +11,4 @@
 Float: TypeAlias = Union[float, np.float_]
 Int: TypeAlias = Union[int, np.int_]
 Real: TypeAlias = Union[Float, Int]
+Bool: TypeAlias = Union[bool, np.bool_]

From 3928c3d35a533a66ec6649a7f9459317cf1e95a6 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:35 -0700
Subject: [PATCH 20/21] Add type-hints to adaptive/utils.py

---
 adaptive/utils.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/adaptive/utils.py b/adaptive/utils.py
index a98af12a1..2c34f1778 100644
--- a/adaptive/utils.py
+++ b/adaptive/utils.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import abc
 import functools
 import gzip
@@ -5,20 +7,21 @@
 import os
 import pickle
 import warnings
-from contextlib import contextmanager
+from contextlib import _GeneratorContextManager, contextmanager
 from itertools import product
+from typing import Any, Callable, Mapping, Sequence
 
 import cloudpickle
 
 
-def named_product(**items):
+def named_product(**items: Mapping[str, Sequence[Any]]):
     names = items.keys()
     vals = items.values()
     return [dict(zip(names, res)) for res in product(*vals)]
 
 
 @contextmanager
-def restore(*learners):
+def restore(*learners) -> _GeneratorContextManager:
     states = [learner.__getstate__() for learner in learners]
     try:
         yield
@@ -27,7 +30,7 @@ def restore(*learners):
             learner.__setstate__(state)
 
 
-def cache_latest(f):
+def cache_latest(f: Callable) -> Callable:
     """Cache the latest return value of the function and add it
     as 'self._cache[f.__name__]'."""
 
@@ -42,7 +45,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def save(fname, data, compress=True):
+def save(fname: str, data: Any, compress: bool = True) -> None:
     fname = os.path.expanduser(fname)
     dirname = os.path.dirname(fname)
     if dirname:
@@ -71,14 +74,14 @@ def save(fname, data, compress=True):
     return True
 
 
-def load(fname, compress=True):
+def load(fname: str, compress: bool = True):
     fname = os.path.expanduser(fname)
     _open = gzip.open if compress else open
     with _open(fname, "rb") as f:
         return cloudpickle.load(f)
 
 
-def copy_docstring_from(other):
+def copy_docstring_from(other: Callable) -> Callable:
     def decorator(method):
         return functools.wraps(other)(method)
 

From 95cc29b518fe540971cc23cc079d384c50d6732a Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Tue, 11 Oct 2022 21:05:36 -0700
Subject: [PATCH 21/21] Add type-hints to setup.py

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 08a7948f5..21c3191d8 100644
--- a/setup.py
+++ b/setup.py
@@ -42,7 +42,6 @@ def get_version_and_cmdclass(package_name):
         "holoviews>=1.9.1",
         "ipywidgets",
         "bokeh",
-        "pandas",
         "matplotlib",
         "plotly",
     ],
@@ -53,6 +52,7 @@ def get_version_and_cmdclass(package_name):
         "pytest-randomly",
         "pytest-timeout",
         "pre_commit",
+        "pandas",
         "typeguard",
     ],
     "other": [