✨ complete scipy.optimize.differential_evolution

scipy-stubs/optimize/_differentialevolution.pyi

@@ -1,49 +1,70 @@
-from collections.abc import Sequence
-from typing import Literal, type_check_only
+from collections.abc import Callable, Iterable
+from typing import Concatenate, Literal, TypeAlias, TypeVar
 
 import numpy as np
-import numpy.typing as npt
 import optype.numpy as onp
-from scipy._typing import Seed, Untyped, UntypedCallable
-from scipy.optimize import OptimizeResult
+from scipy._typing import Seed
+from scipy.optimize import OptimizeResult as _OptimizeResult
+from ._constraints import Bounds, LinearConstraint, NonlinearConstraint
 
 __all__ = ["differential_evolution"]
 
-@type_check_only
-class _OptimizeResult(OptimizeResult):
-    message: str
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+_Float2D: TypeAlias = onp.Array2D[np.float64]
+
+_StrategyName: TypeAlias = Literal[
+    "best1bin",
+    "best1exp",
+    "best2exp",
+    "best2bin",
+    "rand1bin",
+    "rand1exp",
+    "rand2bin",
+    "rand2exp",
+    "randtobest1bin",
+    "randtobest1exp",
+    "currenttobest1bin",
+    "currenttobest1exp",
+]
+
+_S = TypeVar("_S")
+_T = TypeVar("_T")
+
+###
+
+class OptimizeResult(_OptimizeResult):
+    x: _Float1D
+    fun: float | np.float64
+    population: _Float2D
+    population_energies: _Float1D
+    jac: _Float2D  # only if `polish=True`
+
     success: bool
-    fun: float
-    x: onp.ArrayND[np.float64]  # 1d
+    message: str
     nit: int
     nfev: int
-    population: onp.ArrayND[np.float64]  # 2d
-    population_energies: onp.ArrayND[np.float64]  # 1d
-    jac: onp.ArrayND[np.float64]  # 1d
-
-###
 
 def differential_evolution(
-    func: UntypedCallable,
-    bounds: Untyped,
+    func: Callable[Concatenate[_Float1D, ...], onp.ToFloat],
+    bounds: tuple[onp.ToFloat | onp.ToFloat1D, onp.ToFloat | onp.ToFloat1D] | Bounds,
     args: tuple[object, ...] = (),
-    strategy: str | UntypedCallable = "best1bin",
-    maxiter: int = 1000,
-    popsize: int = 15,
+    strategy: _StrategyName | Callable[[int, _Float2D, np.random.Generator], onp.ToFloat1D] = "best1bin",
+    maxiter: onp.ToJustInt = 1000,
+    popsize: onp.ToJustInt = 15,
     tol: onp.ToFloat = 0.01,
     mutation: onp.ToFloat | tuple[onp.ToFloat, onp.ToFloat] = (0.5, 1),
     recombination: onp.ToFloat = 0.7,
     seed: Seed | None = None,
-    callback: UntypedCallable | None = None,
-    disp: bool = False,
-    polish: bool = True,
-    init: str | npt.ArrayLike = "latinhypercube",
+    callback: Callable[[OptimizeResult], None] | Callable[[_Float1D, onp.ToFloat], None] | None = None,
+    disp: onp.ToBool = False,
+    polish: onp.ToBool = True,
+    init: onp.ToFloat2D | Literal["sobol", "halton", "random", "latinhypercube"] = "latinhypercube",
     atol: onp.ToFloat = 0,
     updating: Literal["immediate", "deferred"] = "immediate",
-    workers: int | UntypedCallable = 1,
-    constraints: Untyped = (),
-    x0: npt.ArrayLike | None = None,
+    workers: Callable[[Callable[[_S], _T], Iterable[_S]], Iterable[_T]] | onp.ToJustInt = 1,
+    constraints: NonlinearConstraint | LinearConstraint | Bounds | tuple[()] = (),
+    x0: onp.ToArray1D | None = None,
     *,
-    integrality: Sequence[bool] | onp.ArrayND[np.bool_] | None = None,
-    vectorized: bool = False,
-) -> _OptimizeResult: ...
+    integrality: onp.ToBool1D | None = None,
+    vectorized: onp.ToBool = False,
+) -> OptimizeResult: ...