Merge pull request #219 from jorenham/optimize.differential_evolution

`optimize`: complete `differential_evolution`

scipy-stubs/optimize/_differentialevolution.pyi

@@ -1,128 +1,70 @@
-from collections.abc import Sequence
-from typing import Literal, type_check_only
-from typing_extensions import Self
+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 EnterSelfMixin, 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: ...
-
-# undocumented
-class DifferentialEvolutionSolver(EnterSelfMixin):
-    mutation_func: Untyped
-    strategy: Untyped
-    callback: Untyped
-    polish: Untyped
-    vectorized: Untyped
-    scale: Untyped
-    dither: Untyped
-    cross_over_probability: Untyped
-    func: Untyped
-    args: Untyped
-    limits: Untyped
-    maxiter: Untyped
-    maxfun: Untyped
-    parameter_count: Untyped
-    random_number_generator: Untyped
-    integrality: Untyped
-    num_population_members: Untyped
-    population_shape: Untyped
-    constraints: Untyped
-    total_constraints: Untyped
-    constraint_violation: Untyped
-    feasible: Untyped
-    disp: Untyped
-    def __init__(
-        self,
-        /,
-        func: Untyped,
-        bounds: Untyped,
-        args: Untyped = (),
-        strategy: str = "best1bin",
-        maxiter: int = 1000,
-        popsize: int = 15,
-        tol: float = 0.01,
-        mutation: Untyped = (0.5, 1),
-        recombination: float = 0.7,
-        seed: Untyped | None = None,
-        maxfun: Untyped = ...,
-        callback: Untyped | None = None,
-        disp: bool = False,
-        polish: bool = True,
-        init: str = "latinhypercube",
-        atol: int = 0,
-        updating: str = "immediate",
-        workers: int = 1,
-        constraints: Untyped = (),
-        x0: Untyped | None = None,
-        *,
-        integrality: Untyped | None = None,
-        vectorized: bool = False,
-    ) -> None: ...
-    population: Untyped
-    population_energies: Untyped
-    def init_population_lhs(self, /) -> None: ...
-    def init_population_qmc(self, /, qmc_engine: Untyped) -> None: ...
-    def init_population_random(self, /) -> None: ...
-    def init_population_array(self, /, init: Untyped) -> None: ...
-    @property
-    def x(self, /) -> Untyped: ...
-    @property
-    def convergence(self, /) -> Untyped: ...
-    def converged(self, /) -> Untyped: ...
-    def solve(self, /) -> Untyped: ...
-    def __iter__(self, /) -> Self: ...
-    def __next__(self, /) -> Untyped: ...
-
-# undocumented
-class _ConstraintWrapper:
-    constraint: Untyped
-    fun: Untyped
-    num_constr: Untyped
-    parameter_count: Untyped
-    bounds: Untyped
-    def __init__(self, /, constraint: Untyped, x0: Untyped) -> None: ...
-    def __call__(self, /, x: Untyped) -> Untyped: ...
-    def violation(self, /, x: Untyped) -> Untyped: ...
+    integrality: onp.ToBool1D | None = None,
+    vectorized: onp.ToBool = False,
+) -> OptimizeResult: ...