Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor probnum.typing #599

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions docs/source/development/adding_to_the_api_documentation.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
68 changes: 34 additions & 34 deletions docs/source/development/implementing_a_probnum_method.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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."
]
Expand All @@ -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)"
]
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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)"
]
},
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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."
Expand All @@ -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",
Expand Down Expand Up @@ -758,8 +758,8 @@
"QuadOptBeliefUpdateType = Callable[\n",
" [\n",
" randvars.RandomVariable,\n",
" FloatArgType,\n",
" FloatArgType,\n",
" FloatLike,\n",
" FloatLike,\n",
" ],\n",
" randvars.RandomVariable,\n",
"]\n",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand Down
10 changes: 5 additions & 5 deletions docs/source/development/quadopt_example/_probsolve_qp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down
6 changes: 3 additions & 3 deletions docs/source/development/quadopt_example/belief_updates.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
6 changes: 3 additions & 3 deletions docs/source/development/quadopt_example/policies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]],
]

Expand Down Expand Up @@ -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.

Expand All @@ -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.

Expand Down Expand Up @@ -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,
Expand Down
Loading