From 338ad7dcd4ace73aab767695e24bc42ccc6f1b48 Mon Sep 17 00:00:00 2001 From: jorenham Date: Thu, 28 Nov 2024 16:37:21 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20complete=20`scipy.optimize.=5Fdiffe?= =?UTF-8?q?rentiable=5Ffunctions`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../optimize/_differentiable_functions.pyi | 288 +++++++++++++----- 1 file changed, 213 insertions(+), 75 deletions(-) diff --git a/scipy-stubs/optimize/_differentiable_functions.pyi b/scipy-stubs/optimize/_differentiable_functions.pyi index 70309bca..e8200680 100644 --- a/scipy-stubs/optimize/_differentiable_functions.pyi +++ b/scipy-stubs/optimize/_differentiable_functions.pyi @@ -1,91 +1,229 @@ -from typing import Final +from collections.abc import Callable, Sequence +from types import ModuleType +from typing import Any, Concatenate, Final, Generic, Literal, TypeAlias, overload +from typing_extensions import TypeVar -from scipy._typing import Untyped +import numpy as np +import optype.numpy as onp +from scipy.sparse import csr_matrix, sparray, spmatrix +from scipy.sparse.linalg import LinearOperator +from ._hessian_update_strategy import HessianUpdateStrategy -FD_METHODS: Final = ("2-point", "3-point", "cs") +_XT = TypeVar("_XT", bound=np.floating[Any], default=np.floating[Any]) +_XT_contra = TypeVar("_XT_contra", bound=np.floating[Any], default=np.floating[Any], contravariant=True) -class ScalarFunction: - xp: Untyped - x: Untyped - x_dtype: Untyped - n: Untyped +_ToFloat64Vec: TypeAlias = Sequence[float | np.float64 | np.integer[Any] | np.bool_] | onp.CanArrayND[np.float64] +_ToJac: TypeAlias = onp.ToFloat2D | spmatrix | sparray +_ToHess: TypeAlias = _ToJac | LinearOperator + +_Vec: TypeAlias = onp.Array1D[_XT] +_Jac: TypeAlias = onp.Array2D[_XT] | csr_matrix +_Hess: TypeAlias = _Jac[_XT] | LinearOperator + +_ScalarFun: TypeAlias = Callable[Concatenate[onp.Array1D[_XT], ...], onp.ToFloat] +_VectorFun: TypeAlias = Callable[Concatenate[onp.Array1D[_XT], ...], onp.ToFloat1D] +_JacFun: TypeAlias = Callable[Concatenate[onp.Array1D[_XT], ...], _ToJac] +_HessFun: TypeAlias = Callable[Concatenate[onp.Array1D[_XT], ...], _ToHess] + +_FDMethod: TypeAlias = Literal["2-point", "3-point", "cs"] +_FDBounds: TypeAlias = onp.ToFloat1D | onp.ToFloat2D # len-2 array-like of scalar- or vector-likes + +_ToGradFun: TypeAlias = _VectorFun[_XT_contra] | _FDMethod +_ToJacFun: TypeAlias = _JacFun[_XT_contra] | _FDMethod +_ToHessFun: TypeAlias = _HessFun[_XT_contra] | _FDMethod | HessianUpdateStrategy + +### + +FD_METHODS: Final = "2-point", "3-point", "cs" + +# TODO(jorenham): Array API compatibility +# https://github.com/jorenham/scipy-stubs/issues/140 + +class ScalarFunction(Generic[_XT_contra]): + xp: Final[ModuleType] + + _args: Final[tuple[object, ...]] + + n: Final[int] + x_dtype: _XT_contra # readonly + x: _Vec[_XT_contra] + x_prev: _Vec[_XT_contra] | None + _lowest_x: _Vec[_XT_contra] | None + + _orig_fun: _ScalarFun[_XT_contra] # readonly + _wrapped_fun: Callable[[onp.Array1D[_XT_contra]], onp.ToFloat] # readonly + _nfev: Final[list[int]] # size 1 + _lowest_f: onp.ToFloat f_updated: bool + + _orig_grad: _ToGradFun[_XT_contra] # readonly + _wrapped_grad: Callable[[onp.Array1D[_XT_contra]], _Vec] # readonly + _ngev: Final[list[int]] # size 1 + g_prev: _Vec | None g_updated: bool + + _orig_hess: _ToHessFun[_XT_contra] # readonly + _wrapped_hess: Callable[[onp.Array1D[_XT_contra]], _Hess] # readonly + _nhev: Final[list[int]] # size 1 + H: _Hess H_updated: bool - H: Untyped - x_prev: Untyped - g_prev: Untyped - def __init__( - self, - fun: Untyped, - x0: Untyped, - args: Untyped, - grad: Untyped, - hess: Untyped, - finite_diff_rel_step: Untyped, - finite_diff_bounds: Untyped, - epsilon: Untyped | None = None, - ) -> None: ... + + # @property - def nfev(self) -> Untyped: ... + def nfev(self, /) -> int: ... @property - def ngev(self) -> Untyped: ... + def ngev(self, /) -> int: ... @property - def nhev(self) -> Untyped: ... - def fun(self, x: Untyped) -> Untyped: ... - def grad(self, x: Untyped) -> Untyped: ... - def hess(self, x: Untyped) -> Untyped: ... - def fun_and_grad(self, x: Untyped) -> Untyped: ... - -class VectorFunction: - xp: Untyped - x: Untyped - x_dtype: Untyped - n: Untyped - nfev: int - njev: int - nhev: int + def nhev(self, /) -> int: ... + + # + @overload + def __init__( + self: ScalarFunction[np.float64], + /, + fun: _ScalarFun[np.float64], + x0: _ToFloat64Vec, + args: tuple[object, ...], + grad: _ToGradFun[np.float64], + hess: _ToHessFun[np.float64], + finite_diff_rel_step: onp.ToFloat | onp.ToFloat1D | None, + finite_diff_bounds: _FDBounds, + epsilon: onp.ToFloat | onp.ToFloat1D | None = None, + ) -> None: ... + @overload + def __init__( + self, + /, + fun: _ScalarFun[_XT_contra], + x0: _Vec[_XT_contra] | onp.ToFloat1D, + args: tuple[object, ...], + grad: _ToGradFun[_XT_contra], + hess: _ToHessFun[_XT_contra], + finite_diff_rel_step: onp.ToFloat | onp.ToFloat1D | None, + finite_diff_bounds: _FDBounds, + epsilon: onp.ToFloat | onp.ToFloat1D | None = None, + ) -> None: ... + + # + def _update_x(self, /, x: onp.ToFloat1D) -> None: ... + def _update_fun(self, /) -> None: ... + def _update_grad(self, /) -> None: ... + def _update_hess(self, /) -> None: ... + + # + def fun(self, /, x: onp.ToFloat1D) -> float | np.floating[Any]: ... + def grad(self, /, x: onp.ToFloat1D) -> _Vec: ... + def hess(self, /, x: onp.ToFloat1D) -> _Hess: ... + def fun_and_grad(self, /, x: onp.ToFloat1D) -> tuple[float | np.floating[Any], _Vec]: ... + +class VectorFunction(Generic[_XT_contra]): + xp: Final[ModuleType] + + n: Final[int] + x_dtype: _XT_contra # readonly + x: _Vec[_XT_contra] + x_diff: _Vec[_XT_contra] + x_prev: _Vec[_XT_contra] | None + + m: Final[int] + f: _Vec + v: _Vec f_updated: bool + nfev: int + + sparse_jacobian: Final[bool] + J: _Jac + J_prev: _Jac | None J_updated: bool + njev: int + + H: _Hess H_updated: bool - x_diff: Untyped - f: Untyped - v: Untyped - m: Untyped - J: Untyped - sparse_jacobian: bool - H: Untyped - x_prev: Untyped - J_prev: Untyped + nhev: int + + @overload + def __init__( + self: VectorFunction[np.float64], + /, + fun: _VectorFun[np.float64], + x0: _ToFloat64Vec, + jac: _ToJacFun[np.float64], + hess: _ToHessFun[np.float64], + finite_diff_rel_step: onp.ToFloat | onp.ToFloat1D | None, + finite_diff_jac_sparsity: _ToJac | None, + finite_diff_bounds: _FDBounds, + sparse_jacobian: onp.ToBool | None, + ) -> None: ... + @overload def __init__( self, - fun: Untyped, - x0: Untyped, - jac: Untyped, - hess: Untyped, - finite_diff_rel_step: Untyped, - finite_diff_jac_sparsity: Untyped, - finite_diff_bounds: Untyped, - sparse_jacobian: Untyped, + /, + fun: _VectorFun[_XT_contra], + x0: _Vec[_XT_contra] | onp.ToFloat1D, + jac: _ToJacFun[_XT_contra], + hess: _ToHessFun[_XT_contra], + finite_diff_rel_step: onp.ToFloat | onp.ToFloat1D | None, + finite_diff_jac_sparsity: _ToJac | None, + finite_diff_bounds: _FDBounds, + sparse_jacobian: onp.ToBool | None, ) -> None: ... - def fun(self, x: Untyped) -> Untyped: ... - def jac(self, x: Untyped) -> Untyped: ... - def hess(self, x: Untyped, v: Untyped) -> Untyped: ... - -class LinearVectorFunction: - J: Untyped - sparse_jacobian: bool - xp: Untyped - x: Untyped - x_dtype: Untyped - f: Untyped + + # + def _update_v(self, /, v: onp.ToFloat1D) -> None: ... + def _update_x(self, /, x: onp.ToFloat1D) -> None: ... + def _update_fun(self, /) -> None: ... + def _update_jac(self, /) -> None: ... + def _update_hess(self, /) -> None: ... + + # + def fun(self, /, x: onp.ToFloat1D) -> _Vec: ... + def jac(self, /, x: onp.ToFloat1D) -> _Jac: ... + def hess(self, /, x: onp.ToFloat1D, v: onp.ToFloat1D) -> _Hess: ... + +class LinearVectorFunction(Generic[_XT_contra]): + xp: Final[ModuleType] + + n: Final[int] + x_dtype: _XT_contra # readonly + x: _Vec[_XT_contra] + + m: Final[int] + f: _Vec f_updated: bool - v: Untyped - H: Untyped - def __init__(self, A: Untyped, x0: Untyped, sparse_jacobian: Untyped) -> None: ... - def fun(self, x: Untyped) -> Untyped: ... - def jac(self, x: Untyped) -> Untyped: ... - def hess(self, x: Untyped, v: Untyped) -> Untyped: ... - -class IdentityVectorFunction(LinearVectorFunction): - def __init__(self, x0: Untyped, sparse_jacobian: Untyped) -> None: ... + + sparse_jacobian: Final[bool] + J: Final[_Jac] + + H: Final[csr_matrix] + v: _Vec[np.float64] + + @overload + def __init__( + self: LinearVectorFunction[np.float64], + /, + A: onp.ToFloat2D | spmatrix | sparray, + x0: _ToFloat64Vec, + sparse_jacobian: onp.ToBool | None, + ) -> None: ... + @overload + def __init__( + self, + /, + A: onp.ToFloat2D | spmatrix | sparray, + x0: _Vec[_XT_contra] | onp.ToFloat1D, + sparse_jacobian: onp.ToBool | None, + ) -> None: ... + + # + def _update_x(self, /, x: onp.ToFloat1D) -> None: ... + + # + def fun(self, /, x: onp.ToFloat1D) -> _Vec: ... + def jac(self, /, x: onp.ToFloat1D) -> _Jac: ... + def hess(self, /, x: onp.ToFloat1D, v: _Vec[np.float64]) -> csr_matrix: ... + +class IdentityVectorFunction(LinearVectorFunction[_XT_contra]): + @overload + def __init__(self: IdentityVectorFunction[np.float64], /, x0: _ToFloat64Vec, sparse_jacobian: onp.ToBool | None) -> None: ... + @overload + def __init__(self, /, x0: onp.CanArrayND[_XT_contra] | onp.ToFloat1D, sparse_jacobian: onp.ToBool | None) -> None: ...