better-optimize 0.2.0__tar.gz → 0.3__tar.gz

{better_optimize-0.2.0 → better_optimize-0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: better_optimize
-Version: 0.2.0
+Version: 0.3
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -21,6 +21,9 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy>=1.15
 Requires-Dist: rich
+Requires-Dist: threadpoolctl
+Requires-Dist: joblib
+Requires-Dist: pandas
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/__init__.py

@@ -15,4 +15,4 @@ if not logging.root.handlers:
 
 
 __version__ = get_versions()["version"]
-__all__ = ["minimize", "root", "basinhopping"]
+__all__ = ["minimize", "root", "basinhopping", "MultiStart", "MultiStartResult"]

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-12-13T05:36:24+0100",
+ "date": "2026-03-19T20:00:42+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "487a1c99b33650b7ba7264b19373ec2d54ca2e1f",
- "version": "0.2.0"
+ "full-revisionid": "7c17be556adeb26794fee1e150f9b149cfa3d95c",
+ "version": "0.3"
 }
 '''  # END VERSION_JSON
 

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/basinhopping.py

@@ -21,6 +21,7 @@ from scipy.optimize._basinhopping import (
     check_random_state,
 )
 
+from better_optimize.constants import CONSOLE_WIDTH
 from better_optimize.minimize import minimize
 from better_optimize.utilities import (
     LRUCache1,
@@ -60,7 +61,11 @@ def initialize_progress_bar(progressbar, use_jac=True, use_hess=False):
         ]
 
     return ToggleableProgress(
-        *columns, expand=False, disable=not progressbar, console=Console(width=100), transient=True
+        *columns,
+        expand=False,
+        disable=not progressbar,
+        console=Console(width=CONSOLE_WIDTH),
+        transient=True,
     )
 
 

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/constants.py

@@ -31,6 +31,7 @@ root_method = Literal[
 ]
 
 TOLERANCES = ["xtol", "ftol", "gtol", "fatol", "xatol"]
+CONSOLE_WIDTH = 120
 
 MINIMIZE_MODE_KWARGS = {
     "nelder-mead": {

better_optimize-0.3/better_optimize/multistart.py

@@ -0,0 +1,452 @@
+import contextlib
+import logging
+import multiprocessing as mp
+import pickle
+import threading
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import ClassVar, Literal, NamedTuple
+
+import numpy as np
+import pandas as pd
+
+from rich.console import Console
+from rich.progress import BarColumn, MofNCompleteColumn, TaskID, TextColumn, TimeElapsedColumn
+from rich.table import Column, Table
+from scipy.optimize import OptimizeResult
+from scipy.stats import qmc
+from threadpoolctl import threadpool_limits
+
+from better_optimize.constants import CONSOLE_WIDTH
+from better_optimize.utilities import ToggleableProgress
+
+_log = logging.getLogger(__name__)
+
+_QUEUE_POLL_INTERVAL = 0.05  # seconds between progress-queue polls in the listener thread
+
+InitStrategy = Literal["uniform", "normal", "sobol", "lhs"] | Callable
+
+
+def generate_starts(
+    x0: np.ndarray,
+    n_runs: int,
+    strategy: InitStrategy,
+    bounds: tuple[np.ndarray, np.ndarray] | None,
+    init_scale: float,
+    rng: np.random.Generator,
+) -> list[np.ndarray]:
+    """Create ``n_runs`` starting points from ``x0`` according to ``strategy``."""
+    n_params = x0.shape[0]
+
+    # str is callable in Python, so we must exclude it explicitly
+    if callable(strategy) and not isinstance(strategy, str):
+        return [np.asarray(row) for row in strategy(x0, n_runs, rng)]
+
+    if strategy == "normal":
+        return [x0 + init_scale * rng.standard_normal(n_params) for _ in range(n_runs)]
+
+    if bounds is None:
+        raise ValueError(f"strategy={strategy!r} requires bounds")
+    low = np.broadcast_to(bounds[0], n_params)
+    high = np.broadcast_to(bounds[1], n_params)
+
+    if strategy == "uniform":
+        return [rng.uniform(low, high) for _ in range(n_runs)]
+
+    samplers = {"sobol": qmc.Sobol, "lhs": qmc.LatinHypercube}
+    if strategy not in samplers:
+        raise ValueError(f"Unknown init strategy: {strategy!r}")
+
+    raw = samplers[strategy](d=n_params, seed=rng).random(n_runs)
+    scaled = qmc.scale(raw, low, high)
+    return [scaled[i] for i in range(n_runs)]
+
+
+class ProgressProxy:
+    """Mimics the ``ToggleableProgress`` interface but serializes calls over a ``Queue``."""
+
+    def __init__(self, queue: mp.Queue, task_id):
+        self._queue = queue
+        self._task_id = task_id
+
+    def update(self, task_id, **kwargs):
+        self._queue.put(("update", task_id, kwargs))
+
+    def add_task(self, *args, **kwargs):
+        return self._task_id
+
+    def reset(self, task_id, **kwargs):
+        self._queue.put(("reset", task_id, kwargs))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+def _drain_progress_queue(
+    queue: mp.Queue,
+    progress: ToggleableProgress,
+    stop: threading.Event,
+) -> None:
+    """Drain ``queue`` into ``progress`` until ``stop`` is set.  Runs on a daemon thread."""
+    while not stop.is_set():
+        try:
+            action, task_id, kwargs = queue.get(timeout=_QUEUE_POLL_INTERVAL)
+        except Exception:
+            continue
+        getattr(progress, action)(task_id, **kwargs)
+
+
+def _resolve_n_jobs(n_jobs: int) -> int:
+    """Translate joblib-style n_jobs into an actual core count."""
+    import os
+
+    if n_jobs < 0:
+        # joblib convention: -1 → cpu_count, -2 → cpu_count-1, …
+        return max(1, os.cpu_count() + 1 + n_jobs)
+    return max(1, n_jobs)
+
+
+class BlasConfig(NamedTuple):
+    """Result of :func:`setup_blas_cores`."""
+
+    limiter: Callable[[], contextlib.AbstractContextManager]
+    per_worker: int | None
+
+
+def setup_blas_cores(
+    blas_cores: int | None | str,
+    n_jobs: int,
+    mp_ctx: mp.context.BaseContext | None,
+) -> BlasConfig:
+    """Resolve the BLAS thread budget into a main-process limiter and a per-worker quota.
+
+    Returns
+    -------
+    BlasConfig
+        ``limiter`` – callable returning a context manager that caps total BLAS threads.
+        ``per_worker`` – per-worker thread count, or ``None`` (unlimited).
+    """
+    n_jobs = _resolve_n_jobs(n_jobs)
+
+    if isinstance(blas_cores, str):
+        if blas_cores != "auto":
+            raise ValueError(f"blas_cores must be int, 'auto', or None — got {blas_cores!r}")
+        blas_cores = n_jobs
+
+    if blas_cores is None:
+        return BlasConfig(contextlib.nullcontext, None)
+
+    per_worker = max(1, blas_cores // n_jobs)
+
+    # Forked children inherit the parent's thread-pool state, so a
+    # main-process limiter would be redundant (and potentially racy).
+    is_fork = mp_ctx is not None and mp_ctx.get_start_method() == "fork"
+    if is_fork:
+        return BlasConfig(contextlib.nullcontext, per_worker)
+
+    def joined_limiter():
+        return threadpool_limits(limits=blas_cores)
+
+    return BlasConfig(joined_limiter, per_worker)
+
+
+def _is_pickle_error(exc: Exception) -> bool:
+    if isinstance(exc, pickle.PickleError):
+        return True
+    return isinstance(exc, AttributeError) and str(exc).startswith("Can't pickle")
+
+
+def _run_single(
+    run_index: int,
+    x0: np.ndarray,
+    solver: Callable[..., OptimizeResult],
+    solver_kwargs: dict,
+    progress: ToggleableProgress | ProgressProxy,
+    task_id,
+    blas_cores_per_worker: int | None = None,
+) -> OptimizeResult:
+    """Execute one optimization run, optionally under a per-worker BLAS thread limit."""
+    limiter = (
+        threadpool_limits(limits=blas_cores_per_worker)
+        if blas_cores_per_worker is not None
+        else contextlib.nullcontext()
+    )
+    with limiter:
+        result = solver(x0=x0, progressbar=progress, progress_task=task_id, **solver_kwargs)
+
+    result.run_index = run_index
+    result.x0 = x0
+    return result
+
+
+@dataclass
+class MultiStartResult:
+    """Collected results from a multi-start optimization campaign."""
+
+    results: list[OptimizeResult]
+    sort_key: Callable[[OptimizeResult], float] = field(
+        default_factory=lambda: lambda res: float(res.fun)
+    )
+
+    def ranked(self, ascending: bool = True) -> list[OptimizeResult]:
+        return sorted(self.results, key=self.sort_key, reverse=not ascending)
+
+    @property
+    def best(self) -> OptimizeResult:
+        return min(self.results, key=self.sort_key)
+
+    @property
+    def x_best(self) -> np.ndarray:
+        return self.best.x
+
+    @property
+    def fun_best(self) -> float:
+        return self.sort_key(self.best)
+
+    def top_k(self, k: int = 5) -> list[OptimizeResult]:
+        return self.ranked()[:k]
+
+    @property
+    def success_rate(self) -> float:
+        return sum(r.success for r in self.results) / len(self.results)
+
+    def _format_result_row(self, rank: int, res: OptimizeResult) -> tuple[str, ...]:
+        """Format a single ``OptimizeResult`` into a table row tuple."""
+        grad_str = (
+            f"{np.linalg.norm(res.jac):.6e}" if hasattr(res, "jac") and res.jac is not None else ""
+        )
+        return (
+            str(rank),
+            str(getattr(res, "run_index", "?")),
+            f"{self.sort_key(res):.8e}",
+            grad_str,
+            "✓" if res.success else "✗",
+            str(getattr(res, "nit", "—")),
+            str(getattr(res, "message", "")),
+        )
+
+    _SUMMARY_COLUMNS: ClassVar[tuple[tuple[str, str], ...]] = (
+        ("Rank", "right"),
+        ("Run", "right"),
+        ("f(x*)", "right"),
+        ("||grad||", "right"),
+        ("Success", "center"),
+        ("Iterations", "right"),
+        ("Message", "left"),
+    )
+
+    def summary(self, top_k: int | None = None) -> None:
+        table = Table(title="Multi-Start Optimization Results", show_lines=True)
+        for col, justify in self._SUMMARY_COLUMNS:
+            table.add_column(col, justify=justify)
+
+        display = self.ranked()[:top_k]
+
+        for rank, res in enumerate(display, 1):
+            table.add_row(*self._format_result_row(rank, res))
+        Console().print(table)
+
+    def _result_to_dict(self, rank: int, res: OptimizeResult) -> dict:
+        """Convert a single ``OptimizeResult`` into a flat dict for DataFrame export."""
+        return {
+            "rank": rank,
+            "run_index": getattr(res, "run_index", None),
+            "fun": self.sort_key(res),
+            "success": res.success,
+            "nit": getattr(res, "nit", None),
+            "x": res.x,
+            "message": str(getattr(res, "message", "")),
+        }
+
+    def to_dataframe(self) -> pd.DataFrame:
+        return pd.DataFrame(
+            [self._result_to_dict(rank, res) for rank, res in enumerate(self.ranked(), 1)]
+        )
+
+
+class MultiStart:
+    """Run a solver from many starting points and collect the results.
+
+    Parameters
+    ----------
+    solver
+        Any callable with signature ``(x0, **kwargs) → OptimizeResult``.
+    x0
+        A single starting point (used as center for ``init_strategy``) or an explicit
+        list of starting points (in which case ``n_runs`` and ``init_strategy`` are ignored).
+    solver_kwargs
+        Forwarded verbatim to ``solver`` on every call.
+    sort_key
+        Function mapping an ``OptimizeResult`` to a float used for ranking.
+    n_runs
+        Number of random restarts.  Ignored when ``x0`` is a list.
+    init_strategy
+        How to generate starting points from ``x0``.
+    bounds
+        ``(low, high)`` arrays or scalars, used by bounded init strategies.
+    init_scale
+        Standard deviation for ``"normal"`` init strategy.
+    backend
+        ``"sequential"``, ``"loky"``, or ``"threading"``.
+    n_jobs
+        Number of parallel workers (``-1`` = all cores).
+    seed
+        Random seed for reproducible starting points.
+    progressbar
+        Whether to display per-run progress bars.
+    mp_ctx
+        Explicit ``multiprocessing`` context (``spawn`` / ``fork`` / ``forkserver``).
+    blas_cores
+        Total BLAS/OpenMP thread budget.  ``"auto"`` = ``n_jobs``, ``None`` = unlimited.
+    """
+
+    def __init__(
+        self,
+        solver: Callable[..., OptimizeResult],
+        x0: np.ndarray | list[np.ndarray],
+        solver_kwargs: dict | None = None,
+        sort_key: Callable[[OptimizeResult], float] | None = None,
+        n_runs: int = 16,
+        init_strategy: InitStrategy = "normal",
+        bounds: tuple[np.ndarray, np.ndarray] | None = None,
+        init_scale: float = 1.0,
+        backend: Literal["sequential", "loky", "threading"] = "sequential",
+        n_jobs: int = -1,
+        seed: int | None = None,
+        progressbar: bool = True,
+        mp_ctx: mp.context.BaseContext | None = None,
+        blas_cores: int | None | Literal["auto"] = "auto",
+    ):
+        self._solver = solver
+        self._solver_kwargs = solver_kwargs or {}
+        self._sort_key = sort_key or (lambda res: float(res.fun))
+        self._backend = backend
+        self._n_jobs = n_jobs
+        self._progressbar = progressbar
+        self._mp_ctx = mp_ctx
+        self._blas_cores = blas_cores
+
+        if isinstance(x0, list):
+            self._starting_points = [np.asarray(xi) for xi in x0]
+        else:
+            self._starting_points = generate_starts(
+                np.asarray(x0),
+                n_runs,
+                init_strategy,
+                bounds,
+                init_scale,
+                np.random.default_rng(seed),
+            )
+
+    def run(self) -> MultiStartResult:
+        if self._backend == "sequential":
+            return self._run_sequential()
+
+        try:
+            return self._run_parallel()
+        except (pickle.PickleError, AttributeError) as exc:
+            if not _is_pickle_error(exc):
+                raise
+            _log.warning("Could not pickle objective — falling back to sequential execution.")
+            return self._run_sequential()
+
+    def _run_sequential(self) -> MultiStartResult:
+        blas_config = setup_blas_cores(self._blas_cores, n_jobs=1, mp_ctx=self._mp_ctx)
+        progress = self._build_progress()
+        task_ids = self._register_tasks(progress)
+
+        with blas_config.limiter(), progress:
+            results = [
+                _run_single(
+                    run_index=i,
+                    x0=x0_i,
+                    solver=self._solver,
+                    solver_kwargs=self._solver_kwargs,
+                    progress=progress,
+                    task_id=task_id,
+                )
+                for i, (x0_i, task_id) in enumerate(zip(self._starting_points, task_ids))
+            ]
+        return MultiStartResult(results=results, sort_key=self._sort_key)
+
+    def _run_parallel(self) -> MultiStartResult:
+        from joblib import Parallel, delayed
+
+        blas_config = setup_blas_cores(
+            self._blas_cores,
+            n_jobs=self._n_jobs,
+            mp_ctx=self._mp_ctx,
+        )
+        progress = self._build_progress()
+        task_ids = self._register_tasks(progress)
+
+        queue: mp.Queue = mp.Queue()
+        stop = threading.Event()
+        listener = threading.Thread(
+            target=_drain_progress_queue,
+            args=(queue, progress, stop),
+            daemon=True,
+        )
+
+        with blas_config.limiter(), progress:
+            listener.start()
+            results = Parallel(n_jobs=self._n_jobs, backend=self._backend)(
+                delayed(_run_single)(
+                    run_index=i,
+                    x0=x0_i,
+                    solver=self._solver,
+                    solver_kwargs=self._solver_kwargs,
+                    progress=ProgressProxy(queue, task_id),
+                    task_id=task_id,
+                    blas_cores_per_worker=blas_config.per_worker,
+                )
+                for i, (x0_i, task_id) in enumerate(zip(self._starting_points, task_ids))
+            )
+            stop.set()
+            listener.join(timeout=1.0)
+
+        return MultiStartResult(results=results, sort_key=self._sort_key)
+
+    _SOLVER_LABELS = {
+        "minimize": "Minimizing",
+        "root": "Finding Roots",
+        "basinhopping": "Basinhopping",
+    }
+
+    @property
+    def _task_description(self) -> str:
+        name = getattr(self._solver, "__name__", "")
+        return self._SOLVER_LABELS.get(name, name.title() or "Optimizing")
+
+    def _build_progress(self) -> ToggleableProgress:
+        columns = [
+            BarColumn(bar_width=None, table_column=Column(self._task_description, ratio=2)),
+            TimeElapsedColumn(table_column=Column("Elapsed", ratio=1)),
+            MofNCompleteColumn(table_column=Column("Iteration")),
+            TextColumn("{task.fields[f_value]:0.8f}", table_column=Column("Objective", ratio=1)),
+            TextColumn("{task.fields[grad_norm]:0.8f}", table_column=Column("||grad||", ratio=1)),
+        ]
+        return ToggleableProgress(
+            *columns,
+            expand=False,
+            disable=not self._progressbar,
+            console=Console(width=CONSOLE_WIDTH),
+        )
+
+    def _register_tasks(self, progress: ToggleableProgress) -> list[TaskID | None]:
+        return [
+            progress.add_task(
+                description=self._task_description,
+                total=None,
+                f_value=0.0,
+                grad_norm=0.0,
+            )
+            for _ in self._starting_points
+        ]
+
+
+__all__ = ["MultiStart", "MultiStartResult", "generate_starts"]

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/root.py

@@ -3,6 +3,7 @@ from functools import partial
 
 import numpy as np
 
+from rich.progress import Progress, TaskID
 from scipy.optimize import OptimizeResult
 from scipy.optimize import root as sp_root
 
@@ -24,7 +25,8 @@ def root(
     x0: np.ndarray,
     method: root_method,
     jac: Callable[..., np.ndarray] | None = None,
-    progressbar: bool = True,
+    progressbar: bool | Progress = True,
+    progress_task: TaskID | None = None,
     progressbar_update_interval: int = 1,
     verbose: bool = False,
     args: tuple | None = None,
@@ -88,6 +90,7 @@ def root(
         progressbar_update_interval=progressbar_update_interval,
         has_fused_f_and_grad=has_fused_f_and_grad,
         root=True,
+        task=progress_task,
     )
 
     f_optim = partial(

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize/wrapper.py

@@ -4,6 +4,7 @@ from collections.abc import Callable
 from functools import partial
 
 import numpy as np
+import scipy.sparse as sp
 
 from rich.console import Console
 from rich.progress import (
@@ -17,6 +18,7 @@ from rich.progress import (
 from rich.table import Column
 from scipy.optimize import OptimizeResult
 
+from better_optimize.constants import CONSOLE_WIDTH
 from better_optimize.utilities import ToggleableProgress
 
 _log = logging.getLogger(__name__)
@@ -134,10 +136,16 @@ class ObjectiveWrapper:
 
         value_dict = {"f_value": value}
         if grad is not None:
-            grad_norm = np.linalg.norm(grad)
+            if sp.issparse(grad):
+                grad_norm = sp.linalg.norm(grad)
+            else:
+                grad_norm = np.linalg.norm(grad)
             value_dict["grad_norm"] = grad_norm
         if hess is not None:
-            hess_norm = np.linalg.norm(hess)
+            if sp.issparse(hess):
+                hess_norm = sp.linalg.norm(hess)
+            else:
+                hess_norm = np.linalg.norm(hess)
             value_dict["hess_norm"] = hess_norm
 
         if self.n_eval == 0 and self.task is None:
@@ -178,7 +186,10 @@ class ObjectiveWrapper:
             ]
 
         return ToggleableProgress(
-            *columns, expand=False, disable=not self.progressbar, console=Console(width=100)
+            *columns,
+            expand=False,
+            disable=not self.progressbar,
+            console=Console(width=CONSOLE_WIDTH),
         )
 
 

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: better_optimize
-Version: 0.2.0
+Version: 0.3
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -21,6 +21,9 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy>=1.15
 Requires-Dist: rich
+Requires-Dist: threadpoolctl
+Requires-Dist: joblib
+Requires-Dist: pandas
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ better_optimize/_version.py
 better_optimize/basinhopping.py
 better_optimize/constants.py
 better_optimize/minimize.py
+better_optimize/multistart.py
 better_optimize/root.py
 better_optimize/utilities.py
 better_optimize/wrapper.py
@@ -18,6 +19,7 @@ better_optimize.egg-info/top_level.txt
 tests/test_basinhopping.py
 tests/test_jax_integration.py
 tests/test_minimize.py
+tests/test_multistart.py
 tests/test_root.py
 tests/test_utilities.py
 tests/test_wrapper.py

{better_optimize-0.2.0 → better_optimize-0.3}/better_optimize.egg-info/requires.txt

@@ -1,6 +1,9 @@
 numpy
 scipy>=1.15
 rich
+threadpoolctl
+joblib
+pandas
 
 [tests]
 pytest

{better_optimize-0.2.0 → better_optimize-0.3}/pyproject.toml

@@ -29,6 +29,9 @@ dependencies = [
   'numpy',
   'scipy>=1.15',
   'rich',
+  'threadpoolctl',
+  'joblib',
+  'pandas'
 ]
 
 [project.optional-dependencies]

better_optimize-0.3/tests/test_multistart.py

@@ -0,0 +1,527 @@
+import contextlib
+import multiprocessing
+import pickle
+import threading
+import time
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from numpy.testing import assert_allclose
+from scipy.optimize import OptimizeResult
+
+from better_optimize.minimize import minimize
+from better_optimize.multistart import (
+    MultiStart,
+    MultiStartResult,
+    ProgressProxy,
+    _drain_progress_queue,
+    _is_pickle_error,
+    _run_single,
+    generate_starts,
+    setup_blas_cores,
+)
+from better_optimize.root import root
+
+
+def rosenbrock(x):
+    return float(np.sum(100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2))
+
+
+def rosenbrock_fused(x):
+    f = float(np.sum(100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2))
+    g = np.zeros_like(x)
+    g[0] = -400 * x[0] * (x[1] - x[0] ** 2) - 2 * (1 - x[0])
+    g[-1] = 200 * (x[-1] - x[-2] ** 2)
+    g[1:-1] = (
+        200 * (x[1:-1] - x[:-2] ** 2) - 400 * x[1:-1] * (x[2:] - x[1:-1] ** 2) - 2 * (1 - x[1:-1])
+    )
+    return f, g
+
+
+def simple_system(x):
+    return np.array([x[0] ** 2 + x[1] - 1, x[0] - x[1] ** 2 + 1])
+
+
+def trivial_solver(x0, **kwargs):
+    return OptimizeResult(x=x0, fun=float(np.sum(x0**2)), success=True)
+
+
+@pytest.fixture
+def rng():
+    return np.random.default_rng()
+
+
+def make_results(*fun_values, success=True):
+    results = []
+    for i, f in enumerate(fun_values):
+        r = OptimizeResult(x=np.array([float(i)]), fun=f, success=success, nit=10)
+        r.run_index = i
+        results.append(r)
+    return results
+
+
+def test_normal_starts_centered_on_x0(rng):
+    x0 = rng.standard_normal(3) * 10
+    starts = generate_starts(x0, 10_000, "normal", bounds=None, init_scale=1.0, rng=rng)
+
+    assert len(starts) == 10_000
+    assert all(s.shape == x0.shape for s in starts)
+    assert_allclose(np.mean(starts, axis=0), x0, atol=0.1)
+
+
+@pytest.mark.parametrize("strategy", ["uniform", "sobol", "lhs"])
+def test_bounded_strategies_respect_bounds(strategy, rng):
+    n_params, n_runs = 3, 16
+    low, high = -5.0, 5.0
+    starts = generate_starts(
+        np.zeros(n_params),
+        n_runs,
+        strategy,
+        bounds=(low, high),
+        init_scale=1.0,
+        rng=rng,
+    )
+
+    stacked = np.stack(starts)
+    assert stacked.shape == (n_runs, n_params)
+    assert np.all(stacked >= low)
+    assert np.all(stacked <= high)
+    assert len(np.unique(stacked, axis=0)) == n_runs
+
+
+def test_lhs_stratification(rng):
+    """Each LHS dimension must have exactly one sample per stratum."""
+    n_params, n_runs = 2, 8
+    starts = generate_starts(
+        np.zeros(n_params),
+        n_runs,
+        "lhs",
+        bounds=(0.0, 1.0),
+        init_scale=1.0,
+        rng=rng,
+    )
+    stacked = np.stack(starts)
+
+    for dim in range(n_params):
+        bins = (stacked[:, dim] * n_runs).astype(int).clip(0, n_runs - 1)
+        assert len(set(bins)) == n_runs
+
+
+def test_generate_starts_with_array_bounds(rng):
+    low = np.array([-10.0, 0.0, -1.0])
+    high = np.array([0.0, 5.0, 1.0])
+    starts = generate_starts(
+        np.zeros(3),
+        64,
+        "uniform",
+        bounds=(low, high),
+        init_scale=1.0,
+        rng=rng,
+    )
+
+    stacked = np.stack(starts)
+    assert np.all(stacked >= low)
+    assert np.all(stacked <= high)
+
+
+def test_callable_strategy_dispatches_to_user_function(rng):
+    x0 = np.array([1.0, 2.0])
+
+    def custom(x0, n, rng):
+        return [x0 * i for i in range(n)]
+
+    starts = generate_starts(x0, 5, custom, bounds=None, init_scale=1.0, rng=rng)
+
+    assert len(starts) == 5
+    assert_allclose(starts[0], [0.0, 0.0])
+    assert_allclose(starts[3], [3.0, 6.0])
+
+
+@pytest.mark.parametrize("strategy", ["uniform", "sobol", "lhs"])
+def test_bounded_strategy_without_bounds_raises(strategy):
+    with pytest.raises(ValueError, match="requires bounds"):
+        generate_starts(
+            np.zeros(2),
+            4,
+            strategy,
+            bounds=None,
+            init_scale=1.0,
+            rng=np.random.default_rng(),
+        )
+
+
+def test_unknown_strategy_raises():
+    with pytest.raises(ValueError, match="Unknown init strategy"):
+        generate_starts(
+            np.zeros(2),
+            4,
+            "bogus",
+            bounds=(-1, 1),
+            init_scale=1.0,
+            rng=np.random.default_rng(),
+        )
+
+
+@pytest.mark.parametrize("action", ["update", "reset"])
+def test_progress_proxy_serializes_action_to_queue(action):
+    queue = multiprocessing.Queue()
+    proxy = ProgressProxy(queue, task_id=7)
+    payload = {"f_value": 1.23}
+
+    getattr(proxy, action)(7, **payload)
+
+    received_action, received_task_id, received_kwargs = queue.get(timeout=1)
+    assert received_action == action
+    assert received_task_id == 7
+    assert received_kwargs == payload
+
+
+def test_progress_proxy_add_task_returns_stored_id():
+    queue = multiprocessing.Queue()
+    proxy = ProgressProxy(queue, task_id=99)
+    assert proxy.add_task("ignored", total=100) == 99
+
+
+def test_drain_thread_exits_when_stop_is_set():
+    queue = multiprocessing.Queue()
+    stop = threading.Event()
+    stop.set()
+
+    thread = threading.Thread(
+        target=_drain_progress_queue,
+        args=(queue, MagicMock(), stop),
+    )
+    thread.start()
+    thread.join(timeout=0.5)
+    assert not thread.is_alive()
+
+
+def test_drain_thread_dispatches_queued_messages():
+    queue = multiprocessing.Queue()
+    progress = MagicMock()
+    stop = threading.Event()
+
+    queue.put(("update", 0, {"f_value": 1.0}))
+    queue.put(("reset", 1, {"visible": False}))
+
+    thread = threading.Thread(target=_drain_progress_queue, args=(queue, progress, stop))
+    thread.start()
+
+    deadline = time.monotonic() + 2.0
+    while time.monotonic() < deadline:
+        if progress.update.called and progress.reset.called:
+            break
+        time.sleep(0.01)
+
+    stop.set()
+    thread.join(timeout=1.0)
+
+    progress.update.assert_called_once_with(0, f_value=1.0)
+    progress.reset.assert_called_once_with(1, visible=False)
+
+
+def test_setup_blas_cores_auto_divides_evenly():
+    config = setup_blas_cores("auto", n_jobs=4, mp_ctx=None)
+    assert config.per_worker == 1
+
+
+def test_setup_blas_cores_explicit_budget():
+    config = setup_blas_cores(8, n_jobs=4, mp_ctx=None)
+    assert config.per_worker == 2
+
+
+def test_setup_blas_cores_none_disables_limiting():
+    config = setup_blas_cores(None, n_jobs=4, mp_ctx=None)
+    assert config.limiter is contextlib.nullcontext
+    assert config.per_worker is None
+
+
+def test_setup_blas_cores_fork_skips_joined_limiter():
+    mock_ctx = MagicMock()
+    mock_ctx.get_start_method.return_value = "fork"
+
+    config = setup_blas_cores(8, n_jobs=4, mp_ctx=mock_ctx)
+    assert config.limiter is contextlib.nullcontext
+    assert config.per_worker == 2
+
+
+def test_setup_blas_cores_invalid_string_raises():
+    with pytest.raises(ValueError, match="blas_cores must be"):
+        setup_blas_cores("bogus", n_jobs=4, mp_ctx=None)
+
+
+def test_setup_blas_cores_per_worker_floors_to_one():
+    config = setup_blas_cores(1, n_jobs=8, mp_ctx=None)
+    assert config.per_worker == 1
+
+
+def test_setup_blas_cores_negative_n_jobs():
+    config = setup_blas_cores("auto", n_jobs=-1, mp_ctx=None)
+    assert config.per_worker == 1
+    assert config.limiter is not contextlib.nullcontext
+
+
+@pytest.mark.parametrize(
+    "exc, expected",
+    [
+        (pickle.PicklingError("can't pickle"), True),
+        (AttributeError("Can't pickle local object"), True),
+        (AttributeError("no attribute 'foo'"), False),
+        (ValueError("something"), False),
+    ],
+    ids=["PicklingError", "AttributeError-pickle", "AttributeError-unrelated", "ValueError"],
+)
+def test_is_pickle_error(exc, expected):
+    assert _is_pickle_error(exc) == expected
+
+
+def test_run_single_attaches_metadata():
+    x0 = np.array([1.0, 2.0])
+    result = _run_single(
+        run_index=5,
+        x0=x0,
+        solver=trivial_solver,
+        solver_kwargs={},
+        progress=MagicMock(),
+        task_id=0,
+    )
+    assert result.run_index == 5
+    assert_allclose(result.x0, x0)
+
+
+def test_run_single_forwards_solver_kwargs():
+    received = {}
+
+    def capturing_solver(x0, **kwargs):
+        received.update(kwargs)
+        return OptimizeResult(x=x0, fun=0.0, success=True)
+
+    _run_single(
+        run_index=0,
+        x0=np.array([1.0]),
+        solver=capturing_solver,
+        solver_kwargs={"method": "L-BFGS-B", "tol": 1e-8},
+        progress=MagicMock(),
+        task_id=0,
+    )
+    assert received["method"] == "L-BFGS-B"
+    assert received["tol"] == 1e-8
+
+
+def test_run_single_applies_blas_thread_limit():
+    from threadpoolctl import threadpool_info
+
+    observed_limits = []
+
+    def inspecting_solver(x0, **kwargs):
+        for lib in threadpool_info():
+            observed_limits.append(lib["num_threads"])
+        return OptimizeResult(x=x0, fun=0.0, success=True)
+
+    _run_single(
+        run_index=0,
+        x0=np.array([1.0]),
+        solver=inspecting_solver,
+        solver_kwargs={},
+        progress=MagicMock(),
+        task_id=0,
+        blas_cores_per_worker=1,
+    )
+    if observed_limits:
+        assert all(n <= 1 for n in observed_limits)
+
+
+def test_run_single_propagates_solver_exception():
+    def exploding_solver(x0, **kwargs):
+        raise RuntimeError("solver exploded")
+
+    with pytest.raises(RuntimeError, match="solver exploded"):
+        _run_single(
+            run_index=0,
+            x0=np.zeros(2),
+            solver=exploding_solver,
+            solver_kwargs={},
+            progress=MagicMock(),
+            task_id=0,
+        )
+
+
+class TestMultiStartResult:
+    def test_best_selects_minimum_and_exposes_x_and_fun(self):
+        msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
+        assert msr.best.fun == 1.0
+        assert_allclose(msr.x_best, np.array([1.0]))
+        assert msr.fun_best == 1.0
+
+    @pytest.mark.parametrize("ascending", [True, False], ids=["ascending", "descending"])
+    def test_ranked_ordering(self, ascending):
+        msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
+        funs = [r.fun for r in msr.ranked(ascending=ascending)]
+        expected = [1.0, 2.0, 3.0] if ascending else [3.0, 2.0, 1.0]
+        assert funs == expected
+
+    def test_top_k_truncates(self):
+        msr = MultiStartResult(results=make_results(5.0, 1.0, 3.0, 2.0, 4.0))
+        top2 = msr.top_k(2)
+        assert len(top2) == 2
+        assert [r.fun for r in top2] == [1.0, 2.0]
+
+    def test_success_rate(self):
+        results = [OptimizeResult(x=np.zeros(1), fun=float(i), success=(i < 2)) for i in range(3)]
+        msr = MultiStartResult(results=results)
+        assert_allclose(msr.success_rate, 2.0 / 3.0)
+
+    def test_custom_sort_key_reverses_ranking(self):
+        msr = MultiStartResult(
+            results=make_results(3.0, 1.0, 2.0),
+            sort_key=lambda r: -float(r.fun),
+        )
+        assert msr.best.fun == 3.0
+
+    def test_summary_smoke(self):
+        """Output format is not part of the contract; just verify it doesn't crash."""
+        msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
+        msr.summary()
+        msr.summary(top_k=1)
+
+    def test_to_dataframe(self):
+        pytest.importorskip("pandas")
+        msr = MultiStartResult(results=make_results(3.0, 1.0))
+        df = msr.to_dataframe()
+        expected_columns = {"rank", "run_index", "fun", "success", "nit", "x", "message"}
+        assert set(df.columns) == expected_columns
+        assert len(df) == 2
+        assert df.iloc[0]["fun"] == 1.0
+
+    def test_format_result_row_handles_missing_jac(self):
+        result = OptimizeResult(x=np.zeros(2), fun=1.5, success=True, nit=5)
+        result.run_index = 0
+        msr = MultiStartResult(results=[result])
+        row = msr._format_result_row(1, result)
+        grad_column = row[3]
+        assert grad_column == ""
+
+    def test_empty_results_raises_on_best(self):
+        msr = MultiStartResult(results=[])
+        with pytest.raises(ValueError):
+            _ = msr.best
+
+
+@pytest.mark.parametrize(
+    "solver_name, expected_label",
+    [
+        ("minimize", "Minimizing"),
+        ("root", "Finding Roots"),
+        ("basinhopping", "Basinhopping"),
+    ],
+)
+def test_task_description_for_known_solvers(solver_name, expected_label):
+    stub = MagicMock(__name__=solver_name)
+    ms = MultiStart(solver=stub, x0=[np.zeros(2)], progressbar=False)
+    assert ms._task_description == expected_label
+
+
+def test_task_description_falls_back_for_unknown_solver():
+    stub = MagicMock(__name__="my_custom_opt")
+    ms = MultiStart(solver=stub, x0=[np.zeros(2)], progressbar=False)
+    assert ms._task_description == "My_Custom_Opt"
+
+
+def test_multistart_finds_rosenbrock_minimum():
+    result = MultiStart(
+        solver=minimize,
+        solver_kwargs=dict(f=rosenbrock, method="L-BFGS-B"),
+        x0=np.zeros(3),
+        n_runs=8,
+        init_strategy="uniform",
+        bounds=(-2, 2),
+        seed=7913,
+        progressbar=False,
+    ).run()
+
+    assert len(result.results) == 8
+    assert_allclose(result.x_best, np.ones(3), atol=1e-3)
+    assert result.fun_best < 1e-5
+
+
+def test_multistart_with_explicit_starting_points():
+    x0_list = [np.array([1.3, 0.7, 0.8]), np.array([0.5, 0.5, 0.5])]
+    result = MultiStart(
+        solver=minimize,
+        solver_kwargs=dict(f=rosenbrock, method="L-BFGS-B"),
+        x0=x0_list,
+        progressbar=False,
+    ).run()
+
+    assert len(result.results) == 2
+    for result_i, expected_x0 in zip(result.results, x0_list):
+        assert_allclose(result_i.x0, expected_x0)
+
+
+def test_multistart_with_root_solver():
+    result = MultiStart(
+        solver=root,
+        solver_kwargs=dict(f=simple_system, method="hybr"),
+        sort_key=lambda res: float(np.linalg.norm(res.fun)),
+        x0=np.zeros(2),
+        n_runs=6,
+        init_strategy="normal",
+        init_scale=2.0,
+        seed=9173,
+        progressbar=False,
+    ).run()
+
+    assert result.fun_best < 1e-6
+
+
+def test_multistart_x0_list_ignores_n_runs():
+    result = MultiStart(
+        solver=minimize,
+        solver_kwargs=dict(f=rosenbrock, method="nelder-mead"),
+        x0=[np.zeros(2), np.ones(2)],
+        n_runs=999,
+        progressbar=False,
+    ).run()
+
+    assert len(result.results) == 2
+
+
+def test_multistart_with_fused_objective():
+    result = MultiStart(
+        solver=minimize,
+        solver_kwargs=dict(f=rosenbrock_fused, method="L-BFGS-B"),
+        x0=np.zeros(3),
+        n_runs=4,
+        init_strategy="uniform",
+        bounds=(-2, 2),
+        seed=3917,
+        progressbar=False,
+    ).run()
+
+    assert_allclose(result.x_best, np.ones(3), atol=1e-3)
+
+
+def test_multistart_pickle_fallback(caplog):
+    def unpicklable_closure():
+        return None
+
+    def solver_with_closure(x0, **kwargs):
+        unpicklable_closure()
+        return OptimizeResult(x=x0, fun=float(np.sum(x0**2)), success=True)
+
+    with caplog.at_level("WARNING", logger="better_optimize.multistart"):
+        result = MultiStart(
+            solver=solver_with_closure,
+            x0=[np.zeros(2), np.ones(2)],
+            backend="loky",
+            n_jobs=2,
+            progressbar=False,
+            blas_cores=None,
+        ).run()
+
+    # Fallback may or may not trigger depending on environment,
+    # but either way we must get valid results.
+    assert len(result.results) == 2

{better_optimize-0.2.0 → better_optimize-0.3}/tests/test_wrapper.py

@@ -7,6 +7,7 @@ from functools import partial
 
 import numpy as np
 import pytest
+import scipy.sparse as sp
 
 from scipy.optimize import minimize as scipy_minimize
 from scipy.optimize import root as scipy_root
@@ -66,3 +67,32 @@ def test_execption_breaks_optimization(root, monkeypatch):
     # Non-KeyboardInterrupt exceptions should break execution.
     with pytest.raises(Exception, match="Simulated error"):
         optimizer_early_stopping_wrapper(f_optim)
+
+
+def func2(x, a, b):
+    f = [a * x[0] * np.cos(x[1]) - 4, x[1] * x[0] - b * x[1] - 5]
+
+    return np.array(f)
+
+
+def func2_sparse_jac(x, a, b):
+    df = np.array([[a * np.cos(x[1]), -a * x[0] * np.sin(x[1])], [x[1], x[0] - b]])
+    return sp.csr_matrix(df)
+
+
+def test_wrapper_compatible_with_sparse_outputs():
+    x0 = np.array([0.8, 0.8])
+
+    objective = ObjectiveWrapper(
+        f=partial(func2, a=1, b=1),
+        jac=partial(func2_sparse_jac, a=1, b=1),
+        maxeval=100,
+        progressbar=True,
+        root=True,
+    )
+
+    with objective.progress:
+        result = objective(x0)
+
+    value, grad = result
+    assert sp.issparse(grad)