Refactor probnum.typing (#599)

* Rename `IntArgType` -> `IntLike`

* Rename `FloatArgType` -> `FloatLike`

* Rename `ShapeArgType` -> `ShapeLike`

* Documentation for `DTypeLike`

* Rename `ArrayLikeGetitemArgType` -> `ArrayIndicesLike`

* Rename `ScalarArgType` -> `ScalarLike`

* Documentation for `ArrayLike`

* Rename `LinearOperatorArgType` -> `LinearOperatorLike`

* Remove `ToleranceDiffusionType` and `DenseOutputLocationArgType`

* Restructure API types

* Add annotations future import

* `DTypeArgType` -> `DTypeLike`

* Fix refactoring bug

* isort fix

* isort fix

Author: marvinpfoertner

docs/source/development/adding_to_the_api_documentation.ipynb

@@ -45,27 +45,27 @@
     "import probnum  # pylint: disable=unused-import\n",
     "from probnum import linops, randvars, utils\n",
     "from probnum.linalg.solvers.matrixbased import SymmetricMatrixBasedSolver\n",
-    "from probnum.typing import LinearOperatorArgType\n",
+    "from probnum.typing import LinearOperatorLike\n",
     "\n",
     "# pylint: disable=too-many-branches\n",
     "\n",
     "\n",
     "def problinsolve(\n",
     "    A: Union[\n",
-    "        LinearOperatorArgType,\n",
-    "        \"randvars.RandomVariable[LinearOperatorArgType]\",\n",
+    "        LinearOperatorLike,\n",
+    "        \"randvars.RandomVariable[LinearOperatorLike]\",\n",
     "    ],\n",
     "    b: Union[np.ndarray, \"randvars.RandomVariable[np.ndarray]\"],\n",
     "    A0: Optional[\n",
     "        Union[\n",
-    "            LinearOperatorArgType,\n",
-    "            \"randvars.RandomVariable[LinearOperatorArgType]\",\n",
+    "            LinearOperatorLike,\n",
+    "            \"randvars.RandomVariable[LinearOperatorLike]\",\n",
     "        ]\n",
     "    ] = None,\n",
     "    Ainv0: Optional[\n",
     "        Union[\n",
-    "            LinearOperatorArgType,\n",
-    "            \"randvars.RandomVariable[LinearOperatorArgType]\",\n",
+    "            LinearOperatorLike,\n",
+    "            \"randvars.RandomVariable[LinearOperatorLike]\",\n",
     "        ]\n",
     "    ] = None,\n",
     "    x0: Optional[Union[np.ndarray, \"randvars.RandomVariable[np.ndarray]\"]] = None,\n",

docs/source/development/implementing_a_probnum_method.ipynb

@@ -51,7 +51,7 @@
    "source": [
     "### Method `probsolve_qp`\n",
     "\n",
-    "We will now take a closer look at the interface of our 1D noisy quadratic optimization method. At a basic level `probsolve_qp` takes a function of the type `Callable[[FloatArgType], FloatArgType]`. This hints that the optimization objective is a 1D function. Our prior knowledge about the parameters $(a,b,c)$ is encoded in the random variable `fun_params0`. However, we want to also give a user the option to not specify any prior knowledge or just a guess about the parameter values, hence this argument is optional or can be an `np.ndarray`. \n",
+    "We will now take a closer look at the interface of our 1D noisy quadratic optimization method. At a basic level `probsolve_qp` takes a function of the type `Callable[[FloatLike], FloatLike]`. This hints that the optimization objective is a 1D function. Our prior knowledge about the parameters $(a,b,c)$ is encoded in the random variable `fun_params0`. However, we want to also give a user the option to not specify any prior knowledge or just a guess about the parameter values, hence this argument is optional or can be an `np.ndarray`. \n",
     "\n",
     "The interface also has an `assume_fun` argument, which allows specification of the variant of the probabilistic numerical method to use based on the assumptions about the problem. For convenience, this can be inferred from the problem itself. The actual implementation of the PN method variant which is initialized in a modular fashion is separate from the interface and will be explained later. Finally, the actual optimization routine is called and the result is returned."
    ]
@@ -67,7 +67,7 @@
     "\n",
     "import probnum as pn\n",
     "from probnum import randvars, linops\n",
-    "from probnum.typing import FloatArgType, IntArgType\n",
+    "from probnum.typing import FloatLike, IntLike\n",
     "\n",
     "rng = np.random.default_rng(seed=123)"
    ]
@@ -83,14 +83,14 @@
     "# %load -s probsolve_qp quadopt_example/_probsolve_qp\n",
     "def probsolve_qp(\n",
     "    rng: np.random.Generator,\n",
-    "    fun: Callable[[FloatArgType], FloatArgType],\n",
+    "    fun: Callable[[FloatLike], FloatLike],\n",
     "    fun_params0: Optional[Union[np.ndarray, randvars.RandomVariable]] = None,\n",
     "    assume_fun: Optional[str] = None,\n",
-    "    tol: FloatArgType = 10 ** -5,\n",
-    "    maxiter: IntArgType = 10 ** 4,\n",
+    "    tol: FloatLike = 10 ** -5,\n",
+    "    maxiter: IntLike = 10 ** 4,\n",
     "    noise_cov: Optional[Union[np.ndarray, linops.LinearOperator]] = None,\n",
     "    callback: Optional[\n",
-    "        Callable[[FloatArgType, FloatArgType, randvars.RandomVariable], None]\n",
+    "        Callable[[FloatLike, FloatLike, randvars.RandomVariable], None]\n",
     "    ] = None,\n",
     ") -> Tuple[float, randvars.RandomVariable, randvars.RandomVariable, Dict]:\n",
     "    \"\"\"Probabilistic 1D Quadratic Optimization.\n",
@@ -316,24 +316,24 @@
     "# Type aliases for quadratic optimization\n",
     "QuadOptPolicyType = Callable[\n",
     "    [\n",
-    "        Callable[[FloatArgType], FloatArgType],\n",
+    "        Callable[[FloatLike], FloatLike],\n",
     "        randvars.RandomVariable,\n",
     "    ],\n",
-    "    FloatArgType,\n",
+    "    FloatLike,\n",
     "]\n",
     "QuadOptObservationOperatorType = Callable[\n",
-    "    [Callable[[FloatArgType], FloatArgType], FloatArgType], FloatArgType\n",
+    "    [Callable[[FloatLike], FloatLike], FloatLike], FloatLike\n",
     "]\n",
     "QuadOptBeliefUpdateType = Callable[\n",
     "    [\n",
     "        randvars.RandomVariable,\n",
-    "        FloatArgType,\n",
-    "        FloatArgType,\n",
+    "        FloatLike,\n",
+    "        FloatLike,\n",
     "    ],\n",
     "    randvars.RandomVariable,\n",
     "]\n",
     "QuadOptStoppingCriterionType = Callable[\n",
-    "    [Callable[[FloatArgType], FloatArgType], randvars.RandomVariable, IntArgType],\n",
+    "    [Callable[[FloatLike], FloatLike], randvars.RandomVariable, IntLike],\n",
     "    Tuple[bool, Union[str, None]],\n",
     "]\n",
     "\n",
@@ -430,7 +430,7 @@
     "            self.stopping_criteria = stopping_criteria\n",
     "\n",
     "    def has_converged(\n",
-    "        self, fun: Callable[[FloatArgType], FloatArgType], iteration: IntArgType\n",
+    "        self, fun: Callable[[FloatLike], FloatLike], iteration: IntLike\n",
     "    ) -> Tuple[bool, Union[str, None]]:\n",
     "        \"\"\"Check whether the optimizer has converged.\n",
     "\n",
@@ -451,7 +451,7 @@
     "\n",
     "    def optim_iterator(\n",
     "        self,\n",
-    "        fun: Callable[[FloatArgType], FloatArgType],\n",
+    "        fun: Callable[[FloatLike], FloatLike],\n",
     "    ) -> Tuple[float, float, randvars.RandomVariable]:\n",
     "        \"\"\"Generator implementing the optimization iteration.\n",
     "\n",
@@ -486,7 +486,7 @@
     "\n",
     "    def optimize(\n",
     "        self,\n",
-    "        fun: Callable[[FloatArgType], FloatArgType],\n",
+    "        fun: Callable[[FloatLike], FloatLike],\n",
     "        callback: Optional[\n",
     "            Callable[[float, float, randvars.RandomVariable], None]\n",
     "        ] = None,\n",
@@ -584,10 +584,10 @@
     "internal representation of those same objects. Canonical examples are different kinds of integer or float types, which might be passed by a user. These are all unified internally.\n",
     "\n",
     "```python\n",
-    "IntArgType = Union[int, numbers.Integral, np.integer]\n",
-    "FloatArgType = Union[float, numbers.Real, np.floating]\n",
+    "IntLike = Union[int, numbers.Integral, np.integer]\n",
+    "FloatLike = Union[float, numbers.Real, np.floating]\n",
     "\n",
-    "ShapeArgType = Union[IntArgType, Iterable[IntArgType]]\n",
+    "ShapeLike = Union[IntLike, Iterable[IntLike]]\n",
     "\"\"\"Type of a public API argument for supplying a shape. Values of this type should\n",
     "always be converted into :class:`ShapeType` using the function\n",
     ":func:`probnum.utils.as_shape` before further internal processing.\"\"\"\n",
@@ -602,11 +602,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from probnum.typing import ShapeType, IntArgType, ShapeArgType\n",
+    "from probnum.typing import ShapeType, IntLike, ShapeLike\n",
     "from probnum.utils import as_shape\n",
     "\n",
     "\n",
-    "def extend_shape(shape: ShapeArgType, extension: IntArgType) -> ShapeType:\n",
+    "def extend_shape(shape: ShapeLike, extension: IntLike) -> ShapeType:\n",
     "    return as_shape(shape) + as_shape(extension)"
    ]
   },
@@ -674,7 +674,7 @@
    "source": [
     "# %load -s explore_exploit_policy quadopt_example/policies\n",
     "def explore_exploit_policy(\n",
-    "    fun: Callable[[FloatArgType], FloatArgType],\n",
+    "    fun: Callable[[FloatLike], FloatLike],\n",
     "    fun_params0: randvars.RandomVariable,\n",
     "    rng: np.random.Generator,\n",
     ") -> float:\n",
@@ -704,16 +704,16 @@
     "```python\n",
     "QuadOptPolicyType = Callable[\n",
     "    [\n",
-    "        Callable[[FloatArgType], FloatArgType],\n",
+    "        Callable[[FloatLike], FloatLike],\n",
     "        randvars.RandomVariable\n",
     "    ],\n",
-    "    FloatArgType,\n",
+    "    FloatLike,\n",
     "]\n",
     "```\n",
     "The observation process for this problem is very simple. It just evaluates the objective function. \n",
     "```python\n",
     "QuadOptObservationOperatorType = Callable[\n",
-    "    [Callable[[FloatArgType], FloatArgType], FloatArgType], FloatArgType\n",
+    "    [Callable[[FloatLike], FloatLike], FloatLike], FloatLike\n",
     "]\n",
     "```\n",
     "One can imagine a different probabilistic optimization method which evaluates the gradient as well. In this case the different observation processes would all get the function, its gradient and an evaluation point / action as arguments."
@@ -727,7 +727,7 @@
    "source": [
     "# %load -s function_evaluation quadopt_example/observation_operators\n",
     "def function_evaluation(\n",
-    "    fun: Callable[[FloatArgType], FloatArgType], action: FloatArgType\n",
+    "    fun: Callable[[FloatLike], FloatLike], action: FloatLike\n",
     ") -> np.float_:\n",
     "    \"\"\"Observe a (noisy) function evaluation of the quadratic objective.\n",
     "\n",
@@ -758,8 +758,8 @@
     "QuadOptBeliefUpdateType = Callable[\n",
     "    [\n",
     "        randvars.RandomVariable,\n",
-    "        FloatArgType,\n",
-    "        FloatArgType,\n",
+    "        FloatLike,\n",
+    "        FloatLike,\n",
     "    ],\n",
     "    randvars.RandomVariable,\n",
     "]\n",
@@ -776,8 +776,8 @@
     "# %load -s gaussian_belief_update quadopt_example/belief_updates\n",
     "def gaussian_belief_update(\n",
     "    fun_params0: randvars.RandomVariable,\n",
-    "    action: FloatArgType,\n",
-    "    observation: FloatArgType,\n",
+    "    action: FloatLike,\n",
+    "    observation: FloatLike,\n",
     "    noise_cov: Union[np.ndarray, linops.LinearOperator],\n",
     ") -> randvars.RandomVariable:\n",
     "    \"\"\"Update the belief over the parameters with an observation.\n",
@@ -823,7 +823,7 @@
     "The stopping criteria are also implemented as simple methods, which return a `bool` determining convergence and a string giving the name of the criterion.\n",
     "```python\n",
     "QuadOptStoppingCriterionType = Callable[\n",
-    "    [Callable[[FloatArgType], FloatArgType], randvars.RandomVariable, IntArgType],\n",
+    "    [Callable[[FloatLike], FloatLike], randvars.RandomVariable, IntLike],\n",
     "    Tuple[bool, Union[str, None]],\n",
     "]\n",
     "```\n",
@@ -838,11 +838,11 @@
    "source": [
     "# %load -s parameter_uncertainty quadopt_example/stopping_criteria\n",
     "def parameter_uncertainty(\n",
-    "    fun: Callable[[FloatArgType], FloatArgType],\n",
+    "    fun: Callable[[FloatLike], FloatLike],\n",
     "    fun_params0: randvars.RandomVariable,\n",
-    "    current_iter: IntArgType,\n",
-    "    abstol: FloatArgType,\n",
-    "    reltol: FloatArgType,\n",
+    "    current_iter: IntLike,\n",
+    "    abstol: FloatLike,\n",
+    "    reltol: FloatLike,\n",
     ") -> Tuple[bool, Union[str, None]]:\n",
     "    \"\"\"Termination based on numerical uncertainty about the parameters.\n",
     "\n",

docs/source/development/quadopt_example/_probsolve_qp.py

@@ -6,7 +6,7 @@
 import probnum as pn
 import probnum.utils as _utils
 from probnum import linops, randvars
-from probnum.typing import FloatArgType, IntArgType
+from probnum.typing import FloatLike, IntLike
 
 from .belief_updates import gaussian_belief_update
 from .observation_operators import function_evaluation
@@ -17,14 +17,14 @@
 
 def probsolve_qp(
     rng: np.random.Generator,
-    fun: Callable[[FloatArgType], FloatArgType],
+    fun: Callable[[FloatLike], FloatLike],
     fun_params0: Optional[Union[np.ndarray, randvars.RandomVariable]] = None,
     assume_fun: Optional[str] = None,
-    tol: FloatArgType = 10 ** -5,
-    maxiter: IntArgType = 10 ** 4,
+    tol: FloatLike = 10 ** -5,
+    maxiter: IntLike = 10 ** 4,
     noise_cov: Optional[Union[np.ndarray, linops.LinearOperator]] = None,
     callback: Optional[
-        Callable[[FloatArgType, FloatArgType, randvars.RandomVariable], None]
+        Callable[[FloatLike, FloatLike, randvars.RandomVariable], None]
     ] = None,
 ) -> Tuple[float, randvars.RandomVariable, randvars.RandomVariable, Dict]:
     """Probabilistic 1D Quadratic Optimization.

docs/source/development/quadopt_example/belief_updates.py

@@ -7,13 +7,13 @@
 
 import probnum as pn
 from probnum import linops, randvars
-from probnum.typing import FloatArgType
+from probnum.typing import FloatLike
 
 
 def gaussian_belief_update(
     fun_params0: randvars.RandomVariable,
-    action: FloatArgType,
-    observation: FloatArgType,
+    action: FloatLike,
+    observation: FloatLike,
     noise_cov: Union[np.ndarray, linops.LinearOperator],
 ) -> randvars.RandomVariable:
     """Update the belief over the parameters with an observation.

docs/source/development/quadopt_example/observation_operators.py

@@ -5,11 +5,11 @@
 import numpy as np
 
 from probnum import utils
-from probnum.typing import FloatArgType
+from probnum.typing import FloatLike
 
 
 def function_evaluation(
-    fun: Callable[[FloatArgType], FloatArgType], action: FloatArgType
+    fun: Callable[[FloatLike], FloatLike], action: FloatLike
 ) -> np.float_:
     """Observe a (noisy) function evaluation of the quadratic objective.
 

docs/source/development/quadopt_example/policies.py

@@ -5,11 +5,11 @@
 import numpy as np
 
 from probnum import randvars
-from probnum.typing import FloatArgType
+from probnum.typing import FloatLike
 
 
 def explore_exploit_policy(
-    fun: Callable[[FloatArgType], FloatArgType],
+    fun: Callable[[FloatLike], FloatLike],
     fun_params0: randvars.RandomVariable,
     rng: np.random.Generator,
 ) -> float:
@@ -31,7 +31,7 @@ def explore_exploit_policy(
 
 
 def stochastic_policy(
-    fun: Callable[[FloatArgType], FloatArgType],
+    fun: Callable[[FloatLike], FloatLike],
     fun_params0: randvars.RandomVariable,
     rng: np.random.Generator,
 ) -> float:

docs/source/development/quadopt_example/probabilistic_quadratic_optimizer.py

@@ -7,7 +7,7 @@
 import probnum as pn
 import probnum.utils as _utils
 from probnum import linops, randvars
-from probnum.typing import FloatArgType, IntArgType
+from probnum.typing import FloatLike, IntLike
 
 from .belief_updates import gaussian_belief_update
 from .observation_operators import function_evaluation
@@ -17,24 +17,24 @@
 # Type aliases for quadratic optimization
 QuadOptPolicyType = Callable[
     [
-        Callable[[FloatArgType], FloatArgType],
+        Callable[[FloatLike], FloatLike],
         randvars.RandomVariable,
     ],
-    FloatArgType,
+    FloatLike,
 ]
 QuadOptObservationOperatorType = Callable[
-    [Callable[[FloatArgType], FloatArgType], FloatArgType], FloatArgType
+    [Callable[[FloatLike], FloatLike], FloatLike], FloatLike
 ]
 QuadOptBeliefUpdateType = Callable[
     [
         randvars.RandomVariable,
-        FloatArgType,
-        FloatArgType,
+        FloatLike,
+        FloatLike,
     ],
     randvars.RandomVariable,
 ]
 QuadOptStoppingCriterionType = Callable[
-    [Callable[[FloatArgType], FloatArgType], randvars.RandomVariable, IntArgType],
+    [Callable[[FloatLike], FloatLike], randvars.RandomVariable, IntLike],
     Tuple[bool, Union[str, None]],
 ]
 
@@ -131,7 +131,7 @@ def __init__(
             self.stopping_criteria = stopping_criteria
 
     def has_converged(
-        self, fun: Callable[[FloatArgType], FloatArgType], iteration: IntArgType
+        self, fun: Callable[[FloatLike], FloatLike], iteration: IntLike
     ) -> Tuple[bool, Union[str, None]]:
         """Check whether the optimizer has converged.
 
@@ -152,7 +152,7 @@ def has_converged(
 
     def optim_iterator(
         self,
-        fun: Callable[[FloatArgType], FloatArgType],
+        fun: Callable[[FloatLike], FloatLike],
     ) -> Tuple[float, float, randvars.RandomVariable]:
         """Generator implementing the optimization iteration.
 
@@ -187,7 +187,7 @@ def optim_iterator(
 
     def optimize(
         self,
-        fun: Callable[[FloatArgType], FloatArgType],
+        fun: Callable[[FloatLike], FloatLike],
         callback: Optional[
             Callable[[float, float, randvars.RandomVariable], None]
         ] = None,