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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. {better_optimize-0.2.1 → better_optimize-0.3}/PKG-INFO +4 -1
  2. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/__init__.py +1 -1
  3. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/_version.py +3 -3
  4. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/basinhopping.py +6 -1
  5. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/constants.py +1 -0
  6. better_optimize-0.3/better_optimize/multistart.py +452 -0
  7. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/root.py +4 -1
  8. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/wrapper.py +5 -1
  9. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize.egg-info/PKG-INFO +4 -1
  10. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize.egg-info/SOURCES.txt +2 -0
  11. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize.egg-info/requires.txt +3 -0
  12. {better_optimize-0.2.1 → better_optimize-0.3}/pyproject.toml +3 -0
  13. better_optimize-0.3/tests/test_multistart.py +527 -0
  14. {better_optimize-0.2.1 → better_optimize-0.3}/LICENSE +0 -0
  15. {better_optimize-0.2.1 → better_optimize-0.3}/README.md +0 -0
  16. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/minimize.py +0 -0
  17. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize/utilities.py +0 -0
  18. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize.egg-info/dependency_links.txt +0 -0
  19. {better_optimize-0.2.1 → better_optimize-0.3}/better_optimize.egg-info/top_level.txt +0 -0
  20. {better_optimize-0.2.1 → better_optimize-0.3}/setup.cfg +0 -0
  21. {better_optimize-0.2.1 → better_optimize-0.3}/setup.py +0 -0
  22. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_basinhopping.py +0 -0
  23. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_jax_integration.py +0 -0
  24. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_minimize.py +0 -0
  25. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_root.py +0 -0
  26. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_utilities.py +0 -0
  27. {better_optimize-0.2.1 → better_optimize-0.3}/tests/test_wrapper.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: better_optimize
3
- Version: 0.2.1
3
+ Version: 0.3
4
4
  Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
5
5
  License: MIT License
6
6
  Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -21,6 +21,9 @@ License-File: LICENSE
21
21
  Requires-Dist: numpy
22
22
  Requires-Dist: scipy>=1.15
23
23
  Requires-Dist: rich
24
+ Requires-Dist: threadpoolctl
25
+ Requires-Dist: joblib
26
+ Requires-Dist: pandas
24
27
  Provides-Extra: tests
25
28
  Requires-Dist: pytest; extra == "tests"
26
29
  Requires-Dist: pytest-cov; extra == "tests"
@@ -15,4 +15,4 @@ if not logging.root.handlers:
15
15
 
16
16
 
17
17
  __version__ = get_versions()["version"]
18
- __all__ = ["minimize", "root", "basinhopping"]
18
+ __all__ = ["minimize", "root", "basinhopping", "MultiStart", "MultiStartResult"]
@@ -8,11 +8,11 @@ import json
8
8
 
9
9
  version_json = '''
10
10
  {
11
- "date": "2026-03-02T01:36:39+0100",
11
+ "date": "2026-03-19T20:00:42+0100",
12
12
  "dirty": false,
13
13
  "error": null,
14
- "full-revisionid": "14b8b75021dd22fa3d38135eee74b2ab95c30a06",
15
- "version": "0.2.1"
14
+ "full-revisionid": "7c17be556adeb26794fee1e150f9b149cfa3d95c",
15
+ "version": "0.3"
16
16
  }
17
17
  ''' # END VERSION_JSON
18
18
 
@@ -21,6 +21,7 @@ from scipy.optimize._basinhopping import (
21
21
  check_random_state,
22
22
  )
23
23
 
24
+ from better_optimize.constants import CONSOLE_WIDTH
24
25
  from better_optimize.minimize import minimize
25
26
  from better_optimize.utilities import (
26
27
  LRUCache1,
@@ -60,7 +61,11 @@ def initialize_progress_bar(progressbar, use_jac=True, use_hess=False):
60
61
  ]
61
62
 
62
63
  return ToggleableProgress(
63
- *columns, expand=False, disable=not progressbar, console=Console(width=100), transient=True
64
+ *columns,
65
+ expand=False,
66
+ disable=not progressbar,
67
+ console=Console(width=CONSOLE_WIDTH),
68
+ transient=True,
64
69
  )
65
70
 
66
71
 
@@ -31,6 +31,7 @@ root_method = Literal[
31
31
  ]
32
32
 
33
33
  TOLERANCES = ["xtol", "ftol", "gtol", "fatol", "xatol"]
34
+ CONSOLE_WIDTH = 120
34
35
 
35
36
  MINIMIZE_MODE_KWARGS = {
36
37
  "nelder-mead": {
@@ -0,0 +1,452 @@
1
+ import contextlib
2
+ import logging
3
+ import multiprocessing as mp
4
+ import pickle
5
+ import threading
6
+
7
+ from collections.abc import Callable
8
+ from dataclasses import dataclass, field
9
+ from typing import ClassVar, Literal, NamedTuple
10
+
11
+ import numpy as np
12
+ import pandas as pd
13
+
14
+ from rich.console import Console
15
+ from rich.progress import BarColumn, MofNCompleteColumn, TaskID, TextColumn, TimeElapsedColumn
16
+ from rich.table import Column, Table
17
+ from scipy.optimize import OptimizeResult
18
+ from scipy.stats import qmc
19
+ from threadpoolctl import threadpool_limits
20
+
21
+ from better_optimize.constants import CONSOLE_WIDTH
22
+ from better_optimize.utilities import ToggleableProgress
23
+
24
+ _log = logging.getLogger(__name__)
25
+
26
+ _QUEUE_POLL_INTERVAL = 0.05 # seconds between progress-queue polls in the listener thread
27
+
28
+ InitStrategy = Literal["uniform", "normal", "sobol", "lhs"] | Callable
29
+
30
+
31
+ def generate_starts(
32
+ x0: np.ndarray,
33
+ n_runs: int,
34
+ strategy: InitStrategy,
35
+ bounds: tuple[np.ndarray, np.ndarray] | None,
36
+ init_scale: float,
37
+ rng: np.random.Generator,
38
+ ) -> list[np.ndarray]:
39
+ """Create ``n_runs`` starting points from ``x0`` according to ``strategy``."""
40
+ n_params = x0.shape[0]
41
+
42
+ # str is callable in Python, so we must exclude it explicitly
43
+ if callable(strategy) and not isinstance(strategy, str):
44
+ return [np.asarray(row) for row in strategy(x0, n_runs, rng)]
45
+
46
+ if strategy == "normal":
47
+ return [x0 + init_scale * rng.standard_normal(n_params) for _ in range(n_runs)]
48
+
49
+ if bounds is None:
50
+ raise ValueError(f"strategy={strategy!r} requires bounds")
51
+ low = np.broadcast_to(bounds[0], n_params)
52
+ high = np.broadcast_to(bounds[1], n_params)
53
+
54
+ if strategy == "uniform":
55
+ return [rng.uniform(low, high) for _ in range(n_runs)]
56
+
57
+ samplers = {"sobol": qmc.Sobol, "lhs": qmc.LatinHypercube}
58
+ if strategy not in samplers:
59
+ raise ValueError(f"Unknown init strategy: {strategy!r}")
60
+
61
+ raw = samplers[strategy](d=n_params, seed=rng).random(n_runs)
62
+ scaled = qmc.scale(raw, low, high)
63
+ return [scaled[i] for i in range(n_runs)]
64
+
65
+
66
+ class ProgressProxy:
67
+ """Mimics the ``ToggleableProgress`` interface but serializes calls over a ``Queue``."""
68
+
69
+ def __init__(self, queue: mp.Queue, task_id):
70
+ self._queue = queue
71
+ self._task_id = task_id
72
+
73
+ def update(self, task_id, **kwargs):
74
+ self._queue.put(("update", task_id, kwargs))
75
+
76
+ def add_task(self, *args, **kwargs):
77
+ return self._task_id
78
+
79
+ def reset(self, task_id, **kwargs):
80
+ self._queue.put(("reset", task_id, kwargs))
81
+
82
+ def __enter__(self):
83
+ return self
84
+
85
+ def __exit__(self, *args):
86
+ pass
87
+
88
+
89
+ def _drain_progress_queue(
90
+ queue: mp.Queue,
91
+ progress: ToggleableProgress,
92
+ stop: threading.Event,
93
+ ) -> None:
94
+ """Drain ``queue`` into ``progress`` until ``stop`` is set. Runs on a daemon thread."""
95
+ while not stop.is_set():
96
+ try:
97
+ action, task_id, kwargs = queue.get(timeout=_QUEUE_POLL_INTERVAL)
98
+ except Exception:
99
+ continue
100
+ getattr(progress, action)(task_id, **kwargs)
101
+
102
+
103
+ def _resolve_n_jobs(n_jobs: int) -> int:
104
+ """Translate joblib-style n_jobs into an actual core count."""
105
+ import os
106
+
107
+ if n_jobs < 0:
108
+ # joblib convention: -1 → cpu_count, -2 → cpu_count-1, …
109
+ return max(1, os.cpu_count() + 1 + n_jobs)
110
+ return max(1, n_jobs)
111
+
112
+
113
+ class BlasConfig(NamedTuple):
114
+ """Result of :func:`setup_blas_cores`."""
115
+
116
+ limiter: Callable[[], contextlib.AbstractContextManager]
117
+ per_worker: int | None
118
+
119
+
120
+ def setup_blas_cores(
121
+ blas_cores: int | None | str,
122
+ n_jobs: int,
123
+ mp_ctx: mp.context.BaseContext | None,
124
+ ) -> BlasConfig:
125
+ """Resolve the BLAS thread budget into a main-process limiter and a per-worker quota.
126
+
127
+ Returns
128
+ -------
129
+ BlasConfig
130
+ ``limiter`` – callable returning a context manager that caps total BLAS threads.
131
+ ``per_worker`` – per-worker thread count, or ``None`` (unlimited).
132
+ """
133
+ n_jobs = _resolve_n_jobs(n_jobs)
134
+
135
+ if isinstance(blas_cores, str):
136
+ if blas_cores != "auto":
137
+ raise ValueError(f"blas_cores must be int, 'auto', or None — got {blas_cores!r}")
138
+ blas_cores = n_jobs
139
+
140
+ if blas_cores is None:
141
+ return BlasConfig(contextlib.nullcontext, None)
142
+
143
+ per_worker = max(1, blas_cores // n_jobs)
144
+
145
+ # Forked children inherit the parent's thread-pool state, so a
146
+ # main-process limiter would be redundant (and potentially racy).
147
+ is_fork = mp_ctx is not None and mp_ctx.get_start_method() == "fork"
148
+ if is_fork:
149
+ return BlasConfig(contextlib.nullcontext, per_worker)
150
+
151
+ def joined_limiter():
152
+ return threadpool_limits(limits=blas_cores)
153
+
154
+ return BlasConfig(joined_limiter, per_worker)
155
+
156
+
157
+ def _is_pickle_error(exc: Exception) -> bool:
158
+ if isinstance(exc, pickle.PickleError):
159
+ return True
160
+ return isinstance(exc, AttributeError) and str(exc).startswith("Can't pickle")
161
+
162
+
163
+ def _run_single(
164
+ run_index: int,
165
+ x0: np.ndarray,
166
+ solver: Callable[..., OptimizeResult],
167
+ solver_kwargs: dict,
168
+ progress: ToggleableProgress | ProgressProxy,
169
+ task_id,
170
+ blas_cores_per_worker: int | None = None,
171
+ ) -> OptimizeResult:
172
+ """Execute one optimization run, optionally under a per-worker BLAS thread limit."""
173
+ limiter = (
174
+ threadpool_limits(limits=blas_cores_per_worker)
175
+ if blas_cores_per_worker is not None
176
+ else contextlib.nullcontext()
177
+ )
178
+ with limiter:
179
+ result = solver(x0=x0, progressbar=progress, progress_task=task_id, **solver_kwargs)
180
+
181
+ result.run_index = run_index
182
+ result.x0 = x0
183
+ return result
184
+
185
+
186
+ @dataclass
187
+ class MultiStartResult:
188
+ """Collected results from a multi-start optimization campaign."""
189
+
190
+ results: list[OptimizeResult]
191
+ sort_key: Callable[[OptimizeResult], float] = field(
192
+ default_factory=lambda: lambda res: float(res.fun)
193
+ )
194
+
195
+ def ranked(self, ascending: bool = True) -> list[OptimizeResult]:
196
+ return sorted(self.results, key=self.sort_key, reverse=not ascending)
197
+
198
+ @property
199
+ def best(self) -> OptimizeResult:
200
+ return min(self.results, key=self.sort_key)
201
+
202
+ @property
203
+ def x_best(self) -> np.ndarray:
204
+ return self.best.x
205
+
206
+ @property
207
+ def fun_best(self) -> float:
208
+ return self.sort_key(self.best)
209
+
210
+ def top_k(self, k: int = 5) -> list[OptimizeResult]:
211
+ return self.ranked()[:k]
212
+
213
+ @property
214
+ def success_rate(self) -> float:
215
+ return sum(r.success for r in self.results) / len(self.results)
216
+
217
+ def _format_result_row(self, rank: int, res: OptimizeResult) -> tuple[str, ...]:
218
+ """Format a single ``OptimizeResult`` into a table row tuple."""
219
+ grad_str = (
220
+ f"{np.linalg.norm(res.jac):.6e}" if hasattr(res, "jac") and res.jac is not None else ""
221
+ )
222
+ return (
223
+ str(rank),
224
+ str(getattr(res, "run_index", "?")),
225
+ f"{self.sort_key(res):.8e}",
226
+ grad_str,
227
+ "✓" if res.success else "✗",
228
+ str(getattr(res, "nit", "—")),
229
+ str(getattr(res, "message", "")),
230
+ )
231
+
232
+ _SUMMARY_COLUMNS: ClassVar[tuple[tuple[str, str], ...]] = (
233
+ ("Rank", "right"),
234
+ ("Run", "right"),
235
+ ("f(x*)", "right"),
236
+ ("||grad||", "right"),
237
+ ("Success", "center"),
238
+ ("Iterations", "right"),
239
+ ("Message", "left"),
240
+ )
241
+
242
+ def summary(self, top_k: int | None = None) -> None:
243
+ table = Table(title="Multi-Start Optimization Results", show_lines=True)
244
+ for col, justify in self._SUMMARY_COLUMNS:
245
+ table.add_column(col, justify=justify)
246
+
247
+ display = self.ranked()[:top_k]
248
+
249
+ for rank, res in enumerate(display, 1):
250
+ table.add_row(*self._format_result_row(rank, res))
251
+ Console().print(table)
252
+
253
+ def _result_to_dict(self, rank: int, res: OptimizeResult) -> dict:
254
+ """Convert a single ``OptimizeResult`` into a flat dict for DataFrame export."""
255
+ return {
256
+ "rank": rank,
257
+ "run_index": getattr(res, "run_index", None),
258
+ "fun": self.sort_key(res),
259
+ "success": res.success,
260
+ "nit": getattr(res, "nit", None),
261
+ "x": res.x,
262
+ "message": str(getattr(res, "message", "")),
263
+ }
264
+
265
+ def to_dataframe(self) -> pd.DataFrame:
266
+ return pd.DataFrame(
267
+ [self._result_to_dict(rank, res) for rank, res in enumerate(self.ranked(), 1)]
268
+ )
269
+
270
+
271
+ class MultiStart:
272
+ """Run a solver from many starting points and collect the results.
273
+
274
+ Parameters
275
+ ----------
276
+ solver
277
+ Any callable with signature ``(x0, **kwargs) → OptimizeResult``.
278
+ x0
279
+ A single starting point (used as center for ``init_strategy``) or an explicit
280
+ list of starting points (in which case ``n_runs`` and ``init_strategy`` are ignored).
281
+ solver_kwargs
282
+ Forwarded verbatim to ``solver`` on every call.
283
+ sort_key
284
+ Function mapping an ``OptimizeResult`` to a float used for ranking.
285
+ n_runs
286
+ Number of random restarts. Ignored when ``x0`` is a list.
287
+ init_strategy
288
+ How to generate starting points from ``x0``.
289
+ bounds
290
+ ``(low, high)`` arrays or scalars, used by bounded init strategies.
291
+ init_scale
292
+ Standard deviation for ``"normal"`` init strategy.
293
+ backend
294
+ ``"sequential"``, ``"loky"``, or ``"threading"``.
295
+ n_jobs
296
+ Number of parallel workers (``-1`` = all cores).
297
+ seed
298
+ Random seed for reproducible starting points.
299
+ progressbar
300
+ Whether to display per-run progress bars.
301
+ mp_ctx
302
+ Explicit ``multiprocessing`` context (``spawn`` / ``fork`` / ``forkserver``).
303
+ blas_cores
304
+ Total BLAS/OpenMP thread budget. ``"auto"`` = ``n_jobs``, ``None`` = unlimited.
305
+ """
306
+
307
+ def __init__(
308
+ self,
309
+ solver: Callable[..., OptimizeResult],
310
+ x0: np.ndarray | list[np.ndarray],
311
+ solver_kwargs: dict | None = None,
312
+ sort_key: Callable[[OptimizeResult], float] | None = None,
313
+ n_runs: int = 16,
314
+ init_strategy: InitStrategy = "normal",
315
+ bounds: tuple[np.ndarray, np.ndarray] | None = None,
316
+ init_scale: float = 1.0,
317
+ backend: Literal["sequential", "loky", "threading"] = "sequential",
318
+ n_jobs: int = -1,
319
+ seed: int | None = None,
320
+ progressbar: bool = True,
321
+ mp_ctx: mp.context.BaseContext | None = None,
322
+ blas_cores: int | None | Literal["auto"] = "auto",
323
+ ):
324
+ self._solver = solver
325
+ self._solver_kwargs = solver_kwargs or {}
326
+ self._sort_key = sort_key or (lambda res: float(res.fun))
327
+ self._backend = backend
328
+ self._n_jobs = n_jobs
329
+ self._progressbar = progressbar
330
+ self._mp_ctx = mp_ctx
331
+ self._blas_cores = blas_cores
332
+
333
+ if isinstance(x0, list):
334
+ self._starting_points = [np.asarray(xi) for xi in x0]
335
+ else:
336
+ self._starting_points = generate_starts(
337
+ np.asarray(x0),
338
+ n_runs,
339
+ init_strategy,
340
+ bounds,
341
+ init_scale,
342
+ np.random.default_rng(seed),
343
+ )
344
+
345
+ def run(self) -> MultiStartResult:
346
+ if self._backend == "sequential":
347
+ return self._run_sequential()
348
+
349
+ try:
350
+ return self._run_parallel()
351
+ except (pickle.PickleError, AttributeError) as exc:
352
+ if not _is_pickle_error(exc):
353
+ raise
354
+ _log.warning("Could not pickle objective — falling back to sequential execution.")
355
+ return self._run_sequential()
356
+
357
+ def _run_sequential(self) -> MultiStartResult:
358
+ blas_config = setup_blas_cores(self._blas_cores, n_jobs=1, mp_ctx=self._mp_ctx)
359
+ progress = self._build_progress()
360
+ task_ids = self._register_tasks(progress)
361
+
362
+ with blas_config.limiter(), progress:
363
+ results = [
364
+ _run_single(
365
+ run_index=i,
366
+ x0=x0_i,
367
+ solver=self._solver,
368
+ solver_kwargs=self._solver_kwargs,
369
+ progress=progress,
370
+ task_id=task_id,
371
+ )
372
+ for i, (x0_i, task_id) in enumerate(zip(self._starting_points, task_ids))
373
+ ]
374
+ return MultiStartResult(results=results, sort_key=self._sort_key)
375
+
376
+ def _run_parallel(self) -> MultiStartResult:
377
+ from joblib import Parallel, delayed
378
+
379
+ blas_config = setup_blas_cores(
380
+ self._blas_cores,
381
+ n_jobs=self._n_jobs,
382
+ mp_ctx=self._mp_ctx,
383
+ )
384
+ progress = self._build_progress()
385
+ task_ids = self._register_tasks(progress)
386
+
387
+ queue: mp.Queue = mp.Queue()
388
+ stop = threading.Event()
389
+ listener = threading.Thread(
390
+ target=_drain_progress_queue,
391
+ args=(queue, progress, stop),
392
+ daemon=True,
393
+ )
394
+
395
+ with blas_config.limiter(), progress:
396
+ listener.start()
397
+ results = Parallel(n_jobs=self._n_jobs, backend=self._backend)(
398
+ delayed(_run_single)(
399
+ run_index=i,
400
+ x0=x0_i,
401
+ solver=self._solver,
402
+ solver_kwargs=self._solver_kwargs,
403
+ progress=ProgressProxy(queue, task_id),
404
+ task_id=task_id,
405
+ blas_cores_per_worker=blas_config.per_worker,
406
+ )
407
+ for i, (x0_i, task_id) in enumerate(zip(self._starting_points, task_ids))
408
+ )
409
+ stop.set()
410
+ listener.join(timeout=1.0)
411
+
412
+ return MultiStartResult(results=results, sort_key=self._sort_key)
413
+
414
+ _SOLVER_LABELS = {
415
+ "minimize": "Minimizing",
416
+ "root": "Finding Roots",
417
+ "basinhopping": "Basinhopping",
418
+ }
419
+
420
+ @property
421
+ def _task_description(self) -> str:
422
+ name = getattr(self._solver, "__name__", "")
423
+ return self._SOLVER_LABELS.get(name, name.title() or "Optimizing")
424
+
425
+ def _build_progress(self) -> ToggleableProgress:
426
+ columns = [
427
+ BarColumn(bar_width=None, table_column=Column(self._task_description, ratio=2)),
428
+ TimeElapsedColumn(table_column=Column("Elapsed", ratio=1)),
429
+ MofNCompleteColumn(table_column=Column("Iteration")),
430
+ TextColumn("{task.fields[f_value]:0.8f}", table_column=Column("Objective", ratio=1)),
431
+ TextColumn("{task.fields[grad_norm]:0.8f}", table_column=Column("||grad||", ratio=1)),
432
+ ]
433
+ return ToggleableProgress(
434
+ *columns,
435
+ expand=False,
436
+ disable=not self._progressbar,
437
+ console=Console(width=CONSOLE_WIDTH),
438
+ )
439
+
440
+ def _register_tasks(self, progress: ToggleableProgress) -> list[TaskID | None]:
441
+ return [
442
+ progress.add_task(
443
+ description=self._task_description,
444
+ total=None,
445
+ f_value=0.0,
446
+ grad_norm=0.0,
447
+ )
448
+ for _ in self._starting_points
449
+ ]
450
+
451
+
452
+ __all__ = ["MultiStart", "MultiStartResult", "generate_starts"]
@@ -3,6 +3,7 @@ from functools import partial
3
3
 
4
4
  import numpy as np
5
5
 
6
+ from rich.progress import Progress, TaskID
6
7
  from scipy.optimize import OptimizeResult
7
8
  from scipy.optimize import root as sp_root
8
9
 
@@ -24,7 +25,8 @@ def root(
24
25
  x0: np.ndarray,
25
26
  method: root_method,
26
27
  jac: Callable[..., np.ndarray] | None = None,
27
- progressbar: bool = True,
28
+ progressbar: bool | Progress = True,
29
+ progress_task: TaskID | None = None,
28
30
  progressbar_update_interval: int = 1,
29
31
  verbose: bool = False,
30
32
  args: tuple | None = None,
@@ -88,6 +90,7 @@ def root(
88
90
  progressbar_update_interval=progressbar_update_interval,
89
91
  has_fused_f_and_grad=has_fused_f_and_grad,
90
92
  root=True,
93
+ task=progress_task,
91
94
  )
92
95
 
93
96
  f_optim = partial(
@@ -18,6 +18,7 @@ from rich.progress import (
18
18
  from rich.table import Column
19
19
  from scipy.optimize import OptimizeResult
20
20
 
21
+ from better_optimize.constants import CONSOLE_WIDTH
21
22
  from better_optimize.utilities import ToggleableProgress
22
23
 
23
24
  _log = logging.getLogger(__name__)
@@ -185,7 +186,10 @@ class ObjectiveWrapper:
185
186
  ]
186
187
 
187
188
  return ToggleableProgress(
188
- *columns, expand=False, disable=not self.progressbar, console=Console(width=100)
189
+ *columns,
190
+ expand=False,
191
+ disable=not self.progressbar,
192
+ console=Console(width=CONSOLE_WIDTH),
189
193
  )
190
194
 
191
195
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: better_optimize
3
- Version: 0.2.1
3
+ Version: 0.3
4
4
  Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
5
5
  License: MIT License
6
6
  Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -21,6 +21,9 @@ License-File: LICENSE
21
21
  Requires-Dist: numpy
22
22
  Requires-Dist: scipy>=1.15
23
23
  Requires-Dist: rich
24
+ Requires-Dist: threadpoolctl
25
+ Requires-Dist: joblib
26
+ Requires-Dist: pandas
24
27
  Provides-Extra: tests
25
28
  Requires-Dist: pytest; extra == "tests"
26
29
  Requires-Dist: pytest-cov; extra == "tests"
@@ -7,6 +7,7 @@ better_optimize/_version.py
7
7
  better_optimize/basinhopping.py
8
8
  better_optimize/constants.py
9
9
  better_optimize/minimize.py
10
+ better_optimize/multistart.py
10
11
  better_optimize/root.py
11
12
  better_optimize/utilities.py
12
13
  better_optimize/wrapper.py
@@ -18,6 +19,7 @@ better_optimize.egg-info/top_level.txt
18
19
  tests/test_basinhopping.py
19
20
  tests/test_jax_integration.py
20
21
  tests/test_minimize.py
22
+ tests/test_multistart.py
21
23
  tests/test_root.py
22
24
  tests/test_utilities.py
23
25
  tests/test_wrapper.py
@@ -1,6 +1,9 @@
1
1
  numpy
2
2
  scipy>=1.15
3
3
  rich
4
+ threadpoolctl
5
+ joblib
6
+ pandas
4
7
 
5
8
  [tests]
6
9
  pytest
@@ -29,6 +29,9 @@ dependencies = [
29
29
  'numpy',
30
30
  'scipy>=1.15',
31
31
  'rich',
32
+ 'threadpoolctl',
33
+ 'joblib',
34
+ 'pandas'
32
35
  ]
33
36
 
34
37
  [project.optional-dependencies]
@@ -0,0 +1,527 @@
1
+ import contextlib
2
+ import multiprocessing
3
+ import pickle
4
+ import threading
5
+ import time
6
+
7
+ from unittest.mock import MagicMock
8
+
9
+ import numpy as np
10
+ import pytest
11
+
12
+ from numpy.testing import assert_allclose
13
+ from scipy.optimize import OptimizeResult
14
+
15
+ from better_optimize.minimize import minimize
16
+ from better_optimize.multistart import (
17
+ MultiStart,
18
+ MultiStartResult,
19
+ ProgressProxy,
20
+ _drain_progress_queue,
21
+ _is_pickle_error,
22
+ _run_single,
23
+ generate_starts,
24
+ setup_blas_cores,
25
+ )
26
+ from better_optimize.root import root
27
+
28
+
29
+ def rosenbrock(x):
30
+ return float(np.sum(100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2))
31
+
32
+
33
+ def rosenbrock_fused(x):
34
+ f = float(np.sum(100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2))
35
+ g = np.zeros_like(x)
36
+ g[0] = -400 * x[0] * (x[1] - x[0] ** 2) - 2 * (1 - x[0])
37
+ g[-1] = 200 * (x[-1] - x[-2] ** 2)
38
+ g[1:-1] = (
39
+ 200 * (x[1:-1] - x[:-2] ** 2) - 400 * x[1:-1] * (x[2:] - x[1:-1] ** 2) - 2 * (1 - x[1:-1])
40
+ )
41
+ return f, g
42
+
43
+
44
+ def simple_system(x):
45
+ return np.array([x[0] ** 2 + x[1] - 1, x[0] - x[1] ** 2 + 1])
46
+
47
+
48
+ def trivial_solver(x0, **kwargs):
49
+ return OptimizeResult(x=x0, fun=float(np.sum(x0**2)), success=True)
50
+
51
+
52
+ @pytest.fixture
53
+ def rng():
54
+ return np.random.default_rng()
55
+
56
+
57
+ def make_results(*fun_values, success=True):
58
+ results = []
59
+ for i, f in enumerate(fun_values):
60
+ r = OptimizeResult(x=np.array([float(i)]), fun=f, success=success, nit=10)
61
+ r.run_index = i
62
+ results.append(r)
63
+ return results
64
+
65
+
66
+ def test_normal_starts_centered_on_x0(rng):
67
+ x0 = rng.standard_normal(3) * 10
68
+ starts = generate_starts(x0, 10_000, "normal", bounds=None, init_scale=1.0, rng=rng)
69
+
70
+ assert len(starts) == 10_000
71
+ assert all(s.shape == x0.shape for s in starts)
72
+ assert_allclose(np.mean(starts, axis=0), x0, atol=0.1)
73
+
74
+
75
+ @pytest.mark.parametrize("strategy", ["uniform", "sobol", "lhs"])
76
+ def test_bounded_strategies_respect_bounds(strategy, rng):
77
+ n_params, n_runs = 3, 16
78
+ low, high = -5.0, 5.0
79
+ starts = generate_starts(
80
+ np.zeros(n_params),
81
+ n_runs,
82
+ strategy,
83
+ bounds=(low, high),
84
+ init_scale=1.0,
85
+ rng=rng,
86
+ )
87
+
88
+ stacked = np.stack(starts)
89
+ assert stacked.shape == (n_runs, n_params)
90
+ assert np.all(stacked >= low)
91
+ assert np.all(stacked <= high)
92
+ assert len(np.unique(stacked, axis=0)) == n_runs
93
+
94
+
95
+ def test_lhs_stratification(rng):
96
+ """Each LHS dimension must have exactly one sample per stratum."""
97
+ n_params, n_runs = 2, 8
98
+ starts = generate_starts(
99
+ np.zeros(n_params),
100
+ n_runs,
101
+ "lhs",
102
+ bounds=(0.0, 1.0),
103
+ init_scale=1.0,
104
+ rng=rng,
105
+ )
106
+ stacked = np.stack(starts)
107
+
108
+ for dim in range(n_params):
109
+ bins = (stacked[:, dim] * n_runs).astype(int).clip(0, n_runs - 1)
110
+ assert len(set(bins)) == n_runs
111
+
112
+
113
+ def test_generate_starts_with_array_bounds(rng):
114
+ low = np.array([-10.0, 0.0, -1.0])
115
+ high = np.array([0.0, 5.0, 1.0])
116
+ starts = generate_starts(
117
+ np.zeros(3),
118
+ 64,
119
+ "uniform",
120
+ bounds=(low, high),
121
+ init_scale=1.0,
122
+ rng=rng,
123
+ )
124
+
125
+ stacked = np.stack(starts)
126
+ assert np.all(stacked >= low)
127
+ assert np.all(stacked <= high)
128
+
129
+
130
+ def test_callable_strategy_dispatches_to_user_function(rng):
131
+ x0 = np.array([1.0, 2.0])
132
+
133
+ def custom(x0, n, rng):
134
+ return [x0 * i for i in range(n)]
135
+
136
+ starts = generate_starts(x0, 5, custom, bounds=None, init_scale=1.0, rng=rng)
137
+
138
+ assert len(starts) == 5
139
+ assert_allclose(starts[0], [0.0, 0.0])
140
+ assert_allclose(starts[3], [3.0, 6.0])
141
+
142
+
143
+ @pytest.mark.parametrize("strategy", ["uniform", "sobol", "lhs"])
144
+ def test_bounded_strategy_without_bounds_raises(strategy):
145
+ with pytest.raises(ValueError, match="requires bounds"):
146
+ generate_starts(
147
+ np.zeros(2),
148
+ 4,
149
+ strategy,
150
+ bounds=None,
151
+ init_scale=1.0,
152
+ rng=np.random.default_rng(),
153
+ )
154
+
155
+
156
+ def test_unknown_strategy_raises():
157
+ with pytest.raises(ValueError, match="Unknown init strategy"):
158
+ generate_starts(
159
+ np.zeros(2),
160
+ 4,
161
+ "bogus",
162
+ bounds=(-1, 1),
163
+ init_scale=1.0,
164
+ rng=np.random.default_rng(),
165
+ )
166
+
167
+
168
+ @pytest.mark.parametrize("action", ["update", "reset"])
169
+ def test_progress_proxy_serializes_action_to_queue(action):
170
+ queue = multiprocessing.Queue()
171
+ proxy = ProgressProxy(queue, task_id=7)
172
+ payload = {"f_value": 1.23}
173
+
174
+ getattr(proxy, action)(7, **payload)
175
+
176
+ received_action, received_task_id, received_kwargs = queue.get(timeout=1)
177
+ assert received_action == action
178
+ assert received_task_id == 7
179
+ assert received_kwargs == payload
180
+
181
+
182
+ def test_progress_proxy_add_task_returns_stored_id():
183
+ queue = multiprocessing.Queue()
184
+ proxy = ProgressProxy(queue, task_id=99)
185
+ assert proxy.add_task("ignored", total=100) == 99
186
+
187
+
188
+ def test_drain_thread_exits_when_stop_is_set():
189
+ queue = multiprocessing.Queue()
190
+ stop = threading.Event()
191
+ stop.set()
192
+
193
+ thread = threading.Thread(
194
+ target=_drain_progress_queue,
195
+ args=(queue, MagicMock(), stop),
196
+ )
197
+ thread.start()
198
+ thread.join(timeout=0.5)
199
+ assert not thread.is_alive()
200
+
201
+
202
+ def test_drain_thread_dispatches_queued_messages():
203
+ queue = multiprocessing.Queue()
204
+ progress = MagicMock()
205
+ stop = threading.Event()
206
+
207
+ queue.put(("update", 0, {"f_value": 1.0}))
208
+ queue.put(("reset", 1, {"visible": False}))
209
+
210
+ thread = threading.Thread(target=_drain_progress_queue, args=(queue, progress, stop))
211
+ thread.start()
212
+
213
+ deadline = time.monotonic() + 2.0
214
+ while time.monotonic() < deadline:
215
+ if progress.update.called and progress.reset.called:
216
+ break
217
+ time.sleep(0.01)
218
+
219
+ stop.set()
220
+ thread.join(timeout=1.0)
221
+
222
+ progress.update.assert_called_once_with(0, f_value=1.0)
223
+ progress.reset.assert_called_once_with(1, visible=False)
224
+
225
+
226
+ def test_setup_blas_cores_auto_divides_evenly():
227
+ config = setup_blas_cores("auto", n_jobs=4, mp_ctx=None)
228
+ assert config.per_worker == 1
229
+
230
+
231
+ def test_setup_blas_cores_explicit_budget():
232
+ config = setup_blas_cores(8, n_jobs=4, mp_ctx=None)
233
+ assert config.per_worker == 2
234
+
235
+
236
+ def test_setup_blas_cores_none_disables_limiting():
237
+ config = setup_blas_cores(None, n_jobs=4, mp_ctx=None)
238
+ assert config.limiter is contextlib.nullcontext
239
+ assert config.per_worker is None
240
+
241
+
242
+ def test_setup_blas_cores_fork_skips_joined_limiter():
243
+ mock_ctx = MagicMock()
244
+ mock_ctx.get_start_method.return_value = "fork"
245
+
246
+ config = setup_blas_cores(8, n_jobs=4, mp_ctx=mock_ctx)
247
+ assert config.limiter is contextlib.nullcontext
248
+ assert config.per_worker == 2
249
+
250
+
251
+ def test_setup_blas_cores_invalid_string_raises():
252
+ with pytest.raises(ValueError, match="blas_cores must be"):
253
+ setup_blas_cores("bogus", n_jobs=4, mp_ctx=None)
254
+
255
+
256
+ def test_setup_blas_cores_per_worker_floors_to_one():
257
+ config = setup_blas_cores(1, n_jobs=8, mp_ctx=None)
258
+ assert config.per_worker == 1
259
+
260
+
261
+ def test_setup_blas_cores_negative_n_jobs():
262
+ config = setup_blas_cores("auto", n_jobs=-1, mp_ctx=None)
263
+ assert config.per_worker == 1
264
+ assert config.limiter is not contextlib.nullcontext
265
+
266
+
267
+ @pytest.mark.parametrize(
268
+ "exc, expected",
269
+ [
270
+ (pickle.PicklingError("can't pickle"), True),
271
+ (AttributeError("Can't pickle local object"), True),
272
+ (AttributeError("no attribute 'foo'"), False),
273
+ (ValueError("something"), False),
274
+ ],
275
+ ids=["PicklingError", "AttributeError-pickle", "AttributeError-unrelated", "ValueError"],
276
+ )
277
+ def test_is_pickle_error(exc, expected):
278
+ assert _is_pickle_error(exc) == expected
279
+
280
+
281
+ def test_run_single_attaches_metadata():
282
+ x0 = np.array([1.0, 2.0])
283
+ result = _run_single(
284
+ run_index=5,
285
+ x0=x0,
286
+ solver=trivial_solver,
287
+ solver_kwargs={},
288
+ progress=MagicMock(),
289
+ task_id=0,
290
+ )
291
+ assert result.run_index == 5
292
+ assert_allclose(result.x0, x0)
293
+
294
+
295
+ def test_run_single_forwards_solver_kwargs():
296
+ received = {}
297
+
298
+ def capturing_solver(x0, **kwargs):
299
+ received.update(kwargs)
300
+ return OptimizeResult(x=x0, fun=0.0, success=True)
301
+
302
+ _run_single(
303
+ run_index=0,
304
+ x0=np.array([1.0]),
305
+ solver=capturing_solver,
306
+ solver_kwargs={"method": "L-BFGS-B", "tol": 1e-8},
307
+ progress=MagicMock(),
308
+ task_id=0,
309
+ )
310
+ assert received["method"] == "L-BFGS-B"
311
+ assert received["tol"] == 1e-8
312
+
313
+
314
+ def test_run_single_applies_blas_thread_limit():
315
+ from threadpoolctl import threadpool_info
316
+
317
+ observed_limits = []
318
+
319
+ def inspecting_solver(x0, **kwargs):
320
+ for lib in threadpool_info():
321
+ observed_limits.append(lib["num_threads"])
322
+ return OptimizeResult(x=x0, fun=0.0, success=True)
323
+
324
+ _run_single(
325
+ run_index=0,
326
+ x0=np.array([1.0]),
327
+ solver=inspecting_solver,
328
+ solver_kwargs={},
329
+ progress=MagicMock(),
330
+ task_id=0,
331
+ blas_cores_per_worker=1,
332
+ )
333
+ if observed_limits:
334
+ assert all(n <= 1 for n in observed_limits)
335
+
336
+
337
+ def test_run_single_propagates_solver_exception():
338
+ def exploding_solver(x0, **kwargs):
339
+ raise RuntimeError("solver exploded")
340
+
341
+ with pytest.raises(RuntimeError, match="solver exploded"):
342
+ _run_single(
343
+ run_index=0,
344
+ x0=np.zeros(2),
345
+ solver=exploding_solver,
346
+ solver_kwargs={},
347
+ progress=MagicMock(),
348
+ task_id=0,
349
+ )
350
+
351
+
352
+ class TestMultiStartResult:
353
+ def test_best_selects_minimum_and_exposes_x_and_fun(self):
354
+ msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
355
+ assert msr.best.fun == 1.0
356
+ assert_allclose(msr.x_best, np.array([1.0]))
357
+ assert msr.fun_best == 1.0
358
+
359
+ @pytest.mark.parametrize("ascending", [True, False], ids=["ascending", "descending"])
360
+ def test_ranked_ordering(self, ascending):
361
+ msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
362
+ funs = [r.fun for r in msr.ranked(ascending=ascending)]
363
+ expected = [1.0, 2.0, 3.0] if ascending else [3.0, 2.0, 1.0]
364
+ assert funs == expected
365
+
366
+ def test_top_k_truncates(self):
367
+ msr = MultiStartResult(results=make_results(5.0, 1.0, 3.0, 2.0, 4.0))
368
+ top2 = msr.top_k(2)
369
+ assert len(top2) == 2
370
+ assert [r.fun for r in top2] == [1.0, 2.0]
371
+
372
+ def test_success_rate(self):
373
+ results = [OptimizeResult(x=np.zeros(1), fun=float(i), success=(i < 2)) for i in range(3)]
374
+ msr = MultiStartResult(results=results)
375
+ assert_allclose(msr.success_rate, 2.0 / 3.0)
376
+
377
+ def test_custom_sort_key_reverses_ranking(self):
378
+ msr = MultiStartResult(
379
+ results=make_results(3.0, 1.0, 2.0),
380
+ sort_key=lambda r: -float(r.fun),
381
+ )
382
+ assert msr.best.fun == 3.0
383
+
384
+ def test_summary_smoke(self):
385
+ """Output format is not part of the contract; just verify it doesn't crash."""
386
+ msr = MultiStartResult(results=make_results(3.0, 1.0, 2.0))
387
+ msr.summary()
388
+ msr.summary(top_k=1)
389
+
390
+ def test_to_dataframe(self):
391
+ pytest.importorskip("pandas")
392
+ msr = MultiStartResult(results=make_results(3.0, 1.0))
393
+ df = msr.to_dataframe()
394
+ expected_columns = {"rank", "run_index", "fun", "success", "nit", "x", "message"}
395
+ assert set(df.columns) == expected_columns
396
+ assert len(df) == 2
397
+ assert df.iloc[0]["fun"] == 1.0
398
+
399
+ def test_format_result_row_handles_missing_jac(self):
400
+ result = OptimizeResult(x=np.zeros(2), fun=1.5, success=True, nit=5)
401
+ result.run_index = 0
402
+ msr = MultiStartResult(results=[result])
403
+ row = msr._format_result_row(1, result)
404
+ grad_column = row[3]
405
+ assert grad_column == ""
406
+
407
+ def test_empty_results_raises_on_best(self):
408
+ msr = MultiStartResult(results=[])
409
+ with pytest.raises(ValueError):
410
+ _ = msr.best
411
+
412
+
413
+ @pytest.mark.parametrize(
414
+ "solver_name, expected_label",
415
+ [
416
+ ("minimize", "Minimizing"),
417
+ ("root", "Finding Roots"),
418
+ ("basinhopping", "Basinhopping"),
419
+ ],
420
+ )
421
+ def test_task_description_for_known_solvers(solver_name, expected_label):
422
+ stub = MagicMock(__name__=solver_name)
423
+ ms = MultiStart(solver=stub, x0=[np.zeros(2)], progressbar=False)
424
+ assert ms._task_description == expected_label
425
+
426
+
427
+ def test_task_description_falls_back_for_unknown_solver():
428
+ stub = MagicMock(__name__="my_custom_opt")
429
+ ms = MultiStart(solver=stub, x0=[np.zeros(2)], progressbar=False)
430
+ assert ms._task_description == "My_Custom_Opt"
431
+
432
+
433
+ def test_multistart_finds_rosenbrock_minimum():
434
+ result = MultiStart(
435
+ solver=minimize,
436
+ solver_kwargs=dict(f=rosenbrock, method="L-BFGS-B"),
437
+ x0=np.zeros(3),
438
+ n_runs=8,
439
+ init_strategy="uniform",
440
+ bounds=(-2, 2),
441
+ seed=7913,
442
+ progressbar=False,
443
+ ).run()
444
+
445
+ assert len(result.results) == 8
446
+ assert_allclose(result.x_best, np.ones(3), atol=1e-3)
447
+ assert result.fun_best < 1e-5
448
+
449
+
450
+ def test_multistart_with_explicit_starting_points():
451
+ x0_list = [np.array([1.3, 0.7, 0.8]), np.array([0.5, 0.5, 0.5])]
452
+ result = MultiStart(
453
+ solver=minimize,
454
+ solver_kwargs=dict(f=rosenbrock, method="L-BFGS-B"),
455
+ x0=x0_list,
456
+ progressbar=False,
457
+ ).run()
458
+
459
+ assert len(result.results) == 2
460
+ for result_i, expected_x0 in zip(result.results, x0_list):
461
+ assert_allclose(result_i.x0, expected_x0)
462
+
463
+
464
+ def test_multistart_with_root_solver():
465
+ result = MultiStart(
466
+ solver=root,
467
+ solver_kwargs=dict(f=simple_system, method="hybr"),
468
+ sort_key=lambda res: float(np.linalg.norm(res.fun)),
469
+ x0=np.zeros(2),
470
+ n_runs=6,
471
+ init_strategy="normal",
472
+ init_scale=2.0,
473
+ seed=9173,
474
+ progressbar=False,
475
+ ).run()
476
+
477
+ assert result.fun_best < 1e-6
478
+
479
+
480
+ def test_multistart_x0_list_ignores_n_runs():
481
+ result = MultiStart(
482
+ solver=minimize,
483
+ solver_kwargs=dict(f=rosenbrock, method="nelder-mead"),
484
+ x0=[np.zeros(2), np.ones(2)],
485
+ n_runs=999,
486
+ progressbar=False,
487
+ ).run()
488
+
489
+ assert len(result.results) == 2
490
+
491
+
492
+ def test_multistart_with_fused_objective():
493
+ result = MultiStart(
494
+ solver=minimize,
495
+ solver_kwargs=dict(f=rosenbrock_fused, method="L-BFGS-B"),
496
+ x0=np.zeros(3),
497
+ n_runs=4,
498
+ init_strategy="uniform",
499
+ bounds=(-2, 2),
500
+ seed=3917,
501
+ progressbar=False,
502
+ ).run()
503
+
504
+ assert_allclose(result.x_best, np.ones(3), atol=1e-3)
505
+
506
+
507
+ def test_multistart_pickle_fallback(caplog):
508
+ def unpicklable_closure():
509
+ return None
510
+
511
+ def solver_with_closure(x0, **kwargs):
512
+ unpicklable_closure()
513
+ return OptimizeResult(x=x0, fun=float(np.sum(x0**2)), success=True)
514
+
515
+ with caplog.at_level("WARNING", logger="better_optimize.multistart"):
516
+ result = MultiStart(
517
+ solver=solver_with_closure,
518
+ x0=[np.zeros(2), np.ones(2)],
519
+ backend="loky",
520
+ n_jobs=2,
521
+ progressbar=False,
522
+ blas_cores=None,
523
+ ).run()
524
+
525
+ # Fallback may or may not trigger depending on environment,
526
+ # but either way we must get valid results.
527
+ assert len(result.results) == 2
File without changes
File without changes
File without changes
File without changes