better-optimize 0.4.1__tar.gz → 0.4.2__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 (36) hide show
  1. {better_optimize-0.4.1 → better_optimize-0.4.2}/PKG-INFO +50 -2
  2. {better_optimize-0.4.1 → better_optimize-0.4.2}/README.md +49 -1
  3. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/__init__.py +2 -0
  4. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/_version.py +3 -3
  5. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/basinhopping/core.py +21 -8
  6. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/constants.py +3 -0
  7. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/differential_evolution.py +7 -3
  8. better_optimize-0.4.2/better_optimize/exceptions.py +10 -0
  9. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/minimize.py +12 -2
  10. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/multi_optimize.py +3 -1
  11. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/root.py +31 -2
  12. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/sequential_optimize.py +2 -1
  13. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/wrapper.py +83 -0
  14. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize.egg-info/PKG-INFO +50 -2
  15. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize.egg-info/SOURCES.txt +1 -0
  16. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_basinhopping.py +54 -0
  17. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_differential_evolution.py +29 -7
  18. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_minimize.py +96 -0
  19. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_root.py +60 -0
  20. better_optimize-0.4.2/tests/test_wrapper.py +196 -0
  21. better_optimize-0.4.1/tests/test_wrapper.py +0 -99
  22. {better_optimize-0.4.1 → better_optimize-0.4.2}/LICENSE +0 -0
  23. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/basinhopping/__init__.py +0 -0
  24. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/basinhopping/progress.py +0 -0
  25. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/basinhopping/runner.py +0 -0
  26. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize/utilities.py +0 -0
  27. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize.egg-info/dependency_links.txt +0 -0
  28. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize.egg-info/requires.txt +0 -0
  29. {better_optimize-0.4.1 → better_optimize-0.4.2}/better_optimize.egg-info/top_level.txt +0 -0
  30. {better_optimize-0.4.1 → better_optimize-0.4.2}/pyproject.toml +0 -0
  31. {better_optimize-0.4.1 → better_optimize-0.4.2}/setup.cfg +0 -0
  32. {better_optimize-0.4.1 → better_optimize-0.4.2}/setup.py +0 -0
  33. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_jax_integration.py +0 -0
  34. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_multistart.py +0 -0
  35. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_sequential_optimize.py +0 -0
  36. {better_optimize-0.4.1 → better_optimize-0.4.2}/tests/test_utilities.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: better_optimize
3
- Version: 0.4.1
3
+ Version: 0.4.2
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
@@ -89,6 +89,19 @@ All optimization routines in `better_optimize` can display a rich, informative p
89
89
  - Optionally accepts and stores failed minimizer results if they improve the global minimum.
90
90
  - Useful for noisy or non-smooth objective functions where local minimization may occasionally fail.
91
91
 
92
+ ### 6. Uniform Callback API
93
+
94
+ scipy passes callbacks a different argument for almost every method: `callback(xk)`, `callback(xk, state)`,
95
+ `callback(intermediate_result)`, or `callback(x, f)` for root finding. `better_optimize` normalizes them to a single
96
+ `callback(res)`, where `res` is an `OptimizeResult` with the current `res.x`, `res.fun`, and `res.nit` — plus `res.jac`
97
+ when a gradient is available, and solver-specific fields like `res.accept` (basinhopping) or `res.convergence`
98
+ (differential evolution).
99
+
100
+ - The return value is ignored, so a callback can report a value (an ELBO, a logged loss) without stopping the run.
101
+ - Raise `StopOptimization` to stop early; the best result so far is returned with `success=False`.
102
+ - The same callback works across `minimize`, `root`, `basinhopping`, and `differential_evolution`. `hybr` and `lm` are
103
+ the exception — scipy never calls a callback for them, so it is ignored with a warning.
104
+
92
105
  ---
93
106
 
94
107
  ## Example Usage
@@ -96,6 +109,7 @@ All optimization routines in `better_optimize` can display a rich, informative p
96
109
  ### Simple Example
97
110
 
98
111
  ```python
112
+ import numpy as np
99
113
  from better_optimize import minimize
100
114
 
101
115
  def rosenbrock(x):
@@ -103,7 +117,7 @@ def rosenbrock(x):
103
117
 
104
118
  result = minimize(
105
119
  rosenbrock,
106
- x0=[-1, 2],
120
+ x0=np.array([-1.0, 2.0]),
107
121
  method="L-BFGS-B",
108
122
  tol=1e-6,
109
123
  maxiter=1000,
@@ -119,6 +133,36 @@ result = minimize(
119
133
 
120
134
  The result object is a standard `OptimizeResult` from `scipy.optimize`, so there are no surprises there!
121
135
 
136
+ ### Callbacks and Early Stopping
137
+
138
+ Pass a `callback` and it runs after each iteration with an `OptimizeResult`. Its return value is ignored; raise
139
+ `StopOptimization` to stop and return the best result so far.
140
+
141
+ ```python
142
+ import numpy as np
143
+ from better_optimize import minimize, StopOptimization
144
+
145
+ def rosenbrock(x):
146
+ return sum(100.0*(x[1:] - x[:-1]**2.0)**2.0 + (1 - x[:-1])**2.0)
147
+
148
+ history = []
149
+
150
+ def callback(res):
151
+ history.append(res.fun) # res.x, res.fun, res.nit; res.jac when available
152
+ if res.fun < 1e-8:
153
+ raise StopOptimization
154
+
155
+ result = minimize(
156
+ rosenbrock,
157
+ x0=np.array([-1.0, 2.0]),
158
+ method="L-BFGS-B",
159
+ callback=callback,
160
+ )
161
+ ```
162
+
163
+ The same callback works for `root` (`res.fun` is the residual vector), `basinhopping` (`res.accept`), and
164
+ `differential_evolution` (`res.convergence`).
165
+
122
166
  ### Triple-Fused Function using Pytensor
123
167
 
124
168
  ```python
@@ -155,6 +199,7 @@ Real-world objectives often have multiple local minima. A common workaround is t
155
199
  the optimizer and keep the best result. `better_optimize` makes this painless with `multi_optimize`:
156
200
 
157
201
  ```python
202
+ import numpy as np
158
203
  from better_optimize import minimize, multi_optimize
159
204
 
160
205
  def rosenbrock(x):
@@ -179,6 +224,9 @@ print(result.fun_best) # Best objective value
179
224
  result.summary() # Rich table of all runs, ranked
180
225
  ```
181
226
 
227
+ > The `loky` (process) backend needs an `if __name__ == "__main__":` guard when run as a script on macOS or Windows.
228
+ > Notebooks and the `threading`/`sequential` backends don't.
229
+
182
230
  `multi_optimize` works with **any** solver that follows the `(x0, **kwargs) → OptimizeResult` signature — that includes
183
231
  `minimize`, `root`, `basinhopping`, or your own custom wrapper. It just calls `solver(x0=x0_i, **solver_kwargs)` for
184
232
  each starting point; it never inspects the solver internals.
@@ -57,6 +57,19 @@ All optimization routines in `better_optimize` can display a rich, informative p
57
57
  - Optionally accepts and stores failed minimizer results if they improve the global minimum.
58
58
  - Useful for noisy or non-smooth objective functions where local minimization may occasionally fail.
59
59
 
60
+ ### 6. Uniform Callback API
61
+
62
+ scipy passes callbacks a different argument for almost every method: `callback(xk)`, `callback(xk, state)`,
63
+ `callback(intermediate_result)`, or `callback(x, f)` for root finding. `better_optimize` normalizes them to a single
64
+ `callback(res)`, where `res` is an `OptimizeResult` with the current `res.x`, `res.fun`, and `res.nit` — plus `res.jac`
65
+ when a gradient is available, and solver-specific fields like `res.accept` (basinhopping) or `res.convergence`
66
+ (differential evolution).
67
+
68
+ - The return value is ignored, so a callback can report a value (an ELBO, a logged loss) without stopping the run.
69
+ - Raise `StopOptimization` to stop early; the best result so far is returned with `success=False`.
70
+ - The same callback works across `minimize`, `root`, `basinhopping`, and `differential_evolution`. `hybr` and `lm` are
71
+ the exception — scipy never calls a callback for them, so it is ignored with a warning.
72
+
60
73
  ---
61
74
 
62
75
  ## Example Usage
@@ -64,6 +77,7 @@ All optimization routines in `better_optimize` can display a rich, informative p
64
77
  ### Simple Example
65
78
 
66
79
  ```python
80
+ import numpy as np
67
81
  from better_optimize import minimize
68
82
 
69
83
  def rosenbrock(x):
@@ -71,7 +85,7 @@ def rosenbrock(x):
71
85
 
72
86
  result = minimize(
73
87
  rosenbrock,
74
- x0=[-1, 2],
88
+ x0=np.array([-1.0, 2.0]),
75
89
  method="L-BFGS-B",
76
90
  tol=1e-6,
77
91
  maxiter=1000,
@@ -87,6 +101,36 @@ result = minimize(
87
101
 
88
102
  The result object is a standard `OptimizeResult` from `scipy.optimize`, so there are no surprises there!
89
103
 
104
+ ### Callbacks and Early Stopping
105
+
106
+ Pass a `callback` and it runs after each iteration with an `OptimizeResult`. Its return value is ignored; raise
107
+ `StopOptimization` to stop and return the best result so far.
108
+
109
+ ```python
110
+ import numpy as np
111
+ from better_optimize import minimize, StopOptimization
112
+
113
+ def rosenbrock(x):
114
+ return sum(100.0*(x[1:] - x[:-1]**2.0)**2.0 + (1 - x[:-1])**2.0)
115
+
116
+ history = []
117
+
118
+ def callback(res):
119
+ history.append(res.fun) # res.x, res.fun, res.nit; res.jac when available
120
+ if res.fun < 1e-8:
121
+ raise StopOptimization
122
+
123
+ result = minimize(
124
+ rosenbrock,
125
+ x0=np.array([-1.0, 2.0]),
126
+ method="L-BFGS-B",
127
+ callback=callback,
128
+ )
129
+ ```
130
+
131
+ The same callback works for `root` (`res.fun` is the residual vector), `basinhopping` (`res.accept`), and
132
+ `differential_evolution` (`res.convergence`).
133
+
90
134
  ### Triple-Fused Function using Pytensor
91
135
 
92
136
  ```python
@@ -123,6 +167,7 @@ Real-world objectives often have multiple local minima. A common workaround is t
123
167
  the optimizer and keep the best result. `better_optimize` makes this painless with `multi_optimize`:
124
168
 
125
169
  ```python
170
+ import numpy as np
126
171
  from better_optimize import minimize, multi_optimize
127
172
 
128
173
  def rosenbrock(x):
@@ -147,6 +192,9 @@ print(result.fun_best) # Best objective value
147
192
  result.summary() # Rich table of all runs, ranked
148
193
  ```
149
194
 
195
+ > The `loky` (process) backend needs an `if __name__ == "__main__":` guard when run as a script on macOS or Windows.
196
+ > Notebooks and the `threading`/`sequential` backends don't.
197
+
150
198
  `multi_optimize` works with **any** solver that follows the `(x0, **kwargs) → OptimizeResult` signature — that includes
151
199
  `minimize`, `root`, `basinhopping`, or your own custom wrapper. It just calls `solver(x0=x0_i, **solver_kwargs)` for
152
200
  each starting point; it never inspects the solver internals.
@@ -3,6 +3,7 @@ import logging
3
3
  from better_optimize._version import get_versions
4
4
  from better_optimize.basinhopping import basinhopping
5
5
  from better_optimize.differential_evolution import differential_evolution
6
+ from better_optimize.exceptions import StopOptimization
6
7
  from better_optimize.minimize import minimize
7
8
  from better_optimize.multi_optimize import multi_optimize
8
9
  from better_optimize.root import root
@@ -19,6 +20,7 @@ if not logging.root.handlers:
19
20
 
20
21
  __version__ = get_versions()["version"]
21
22
  __all__ = [
23
+ "StopOptimization",
22
24
  "basinhopping",
23
25
  "differential_evolution",
24
26
  "minimize",
@@ -8,11 +8,11 @@ import json
8
8
 
9
9
  version_json = '''
10
10
  {
11
- "date": "2026-04-21T14:03:17-0500",
11
+ "date": "2026-05-23T05:38:17+0200",
12
12
  "dirty": false,
13
13
  "error": null,
14
- "full-revisionid": "ac00b49724e28bdbe53d12fcd222d003decfbd67",
15
- "version": "0.4.1"
14
+ "full-revisionid": "8613524a4578ba95c931ea6cf3091963c35f7a8f",
15
+ "version": "0.4.2"
16
16
  }
17
17
  ''' # END VERSION_JSON
18
18
 
@@ -67,7 +67,10 @@ def basinhopping(
67
67
  Accept the new point even if the minimizer fails. Will also update the global minimum if the point is
68
68
  better than the current global minimum, even if the minimizer fails.
69
69
  callback: Callable, optional
70
- User-supplied function to call after each iteration
70
+ Function called after each iteration as ``callback(res)``, where ``res`` is an
71
+ ``OptimizeResult`` carrying the trial minimum ``res.x``, its objective value ``res.fun``,
72
+ ``res.nit``, and ``res.accept`` (whether the step was accepted). The return value is
73
+ ignored; raise ``StopOptimization`` to stop early. Default None.
71
74
  interval: int
72
75
  Interval for how often to update the step size
73
76
  progressbar: bool
@@ -237,12 +240,20 @@ def basinhopping(
237
240
 
238
241
  # The wrapped minimizer is called once during construction of
239
242
  # BasinHoppingRunner, so run the callback
243
+ early_stop = False
240
244
  if callable(callback):
241
- callback(bh.storage.minres.x, bh.storage.minres.fun, True)
245
+ try:
246
+ callback(
247
+ OptimizeResult(x=bh.storage.minres.x, fun=bh.storage.minres.fun, nit=0, accept=True)
248
+ )
249
+ except StopIteration:
250
+ early_stop = True
242
251
 
243
252
  # start main iteration loop
244
253
  count, i = 0, 0
245
254
  message = ["requested number of basinhopping iterations completed" " successfully"]
255
+ if early_stop:
256
+ message = ["callback requested early stop"]
246
257
 
247
258
  grad_norm_at_min = 0.0
248
259
  if use_jac:
@@ -261,6 +272,8 @@ def basinhopping(
261
272
 
262
273
  with progress:
263
274
  for i in range(niter):
275
+ if early_stop:
276
+ break
264
277
  if progressbar:
265
278
  progress.reset(minimize_task, visible=True)
266
279
  new_global_min = bh.one_cycle()
@@ -277,12 +290,12 @@ def basinhopping(
277
290
  )
278
291
 
279
292
  if callable(callback):
280
- # should we pass a copy of x?
281
- val = callback(bh.xtrial, bh.energy_trial, bh.accept)
282
- if val is not None:
283
- if val:
284
- message = ["callback function requested stop early by" "returning True"]
285
- break
293
+ res = OptimizeResult(x=bh.xtrial, fun=bh.energy_trial, nit=i + 1, accept=bh.accept)
294
+ try:
295
+ callback(res) # return ignored; raise StopOptimization to stop early
296
+ except StopIteration:
297
+ message = ["callback requested early stop"]
298
+ break
286
299
 
287
300
  count += 1
288
301
  if new_global_min:
@@ -50,6 +50,9 @@ DE_STRATEGY_OPTIONS = (
50
50
 
51
51
  DE_INIT_OPTIONS = ("sobol", "halton", "latinhypercube", "random")
52
52
 
53
+ # SciPy's direct solvers ignore the callback argument; the iterative root finders honor it.
54
+ ROOT_METHODS_WITHOUT_CALLBACK = frozenset({"hybr", "lm"})
55
+
53
56
  MINIMIZE_MODE_KWARGS = {
54
57
  "nelder-mead": {
55
58
  "uses_grad": False,
@@ -86,6 +86,11 @@ def differential_evolution(
86
86
  When ``sequential_optimize`` pre-registers a task on the shared progress and passes
87
87
  the id here, that task is updated instead of creating a new one. Ignored for
88
88
  standalone calls.
89
+ callback : Callable, optional
90
+ Function called once per generation as ``callback(res)``, where ``res`` is an
91
+ ``OptimizeResult`` carrying the current best ``res.x`` and ``res.fun``, plus ``res.nit`` and
92
+ ``res.convergence``. The return value is ignored; raise ``StopOptimization`` to stop early.
93
+ Default None.
89
94
  rng : int, numpy.random.Generator, or None
90
95
  Seed / generator for reproducibility. Forwarded to scipy as the ``rng`` kwarg.
91
96
 
@@ -173,9 +178,8 @@ def differential_evolution(
173
178
  progress.update(de_task, advance=1, f_value=best_fun_seen[0], refresh=True)
174
179
 
175
180
  if callback is not None:
176
- user_stop = callback(xk, convergence=convergence)
177
- if user_stop:
178
- raise StopIteration("user callback returned True")
181
+ res = OptimizeResult(x=xk, fun=f_val, nit=iter_counter[0], convergence=convergence)
182
+ callback(res) # return ignored; raise StopOptimization to stop early
179
183
  return False
180
184
 
181
185
  scipy_kwargs: dict[str, Any] = {
@@ -0,0 +1,10 @@
1
+ class StopOptimization(StopIteration):
2
+ """Raise from inside a callback to request a graceful early stop.
3
+
4
+ Subclasses :class:`StopIteration` so that both SciPy's callback machinery and
5
+ ``better_optimize``'s own early-stopping wrapper halt on it. When raised, the solver returns
6
+ the best point seen so far with ``success=False``.
7
+ """
8
+
9
+
10
+ __all__ = ["StopOptimization"]
@@ -17,7 +17,11 @@ from better_optimize.utilities import (
17
17
  kwargs_to_options,
18
18
  validate_provided_functions_minimize,
19
19
  )
20
- from better_optimize.wrapper import ObjectiveWrapper, optimizer_early_stopping_wrapper
20
+ from better_optimize.wrapper import (
21
+ ObjectiveWrapper,
22
+ _compose_callback,
23
+ optimizer_early_stopping_wrapper,
24
+ )
21
25
 
22
26
 
23
27
  def minimize(
@@ -32,6 +36,7 @@ def minimize(
32
36
  progressbar_update_interval: int = 1,
33
37
  verbose: bool = False,
34
38
  args: tuple | None = None,
39
+ callback: Callable[..., bool | None] | None = None,
35
40
  **optimizer_kwargs,
36
41
  ) -> OptimizeResult:
37
42
  """
@@ -61,6 +66,11 @@ def minimize(
61
66
  verbose: bool
62
67
  If True, warnings about the provided configuration are displayed. These warnings are intended to help users
63
68
  understand potential configuration issues that may affect the optimization process, but can be safely ignored.
69
+ callback: Callable, optional
70
+ Function called after each iteration as ``callback(res)``, where ``res`` is an
71
+ ``OptimizeResult`` carrying the current ``res.x``, ``res.fun``, and ``res.nit`` (plus
72
+ ``res.jac`` when a gradient is available). The return value is ignored; raise
73
+ ``StopOptimization`` to stop early. Default None.
64
74
  optimizer_kwargs
65
75
  Additional keyword arguments to pass to the optimizer
66
76
 
@@ -115,7 +125,7 @@ def minimize(
115
125
  jac=True if has_fused_f_and_grad or jac is not None else None,
116
126
  hess=None if not use_hess else lambda x: hess(x, *args),
117
127
  hessp=None if not use_hessp else lambda x, p: hessp(x, p, *args),
118
- callback=objective.callback,
128
+ callback=_compose_callback(objective.callback, objective.callback_result, callback),
119
129
  **optimizer_kwargs,
120
130
  )
121
131
 
@@ -476,7 +476,9 @@ def multi_optimize(
476
476
  A single starting point used as the center for init_strategy, or an explicit list of
477
477
  starting points. When a list is provided, n_runs and init_strategy are ignored.
478
478
  solver_kwargs: dict, optional
479
- Keyword arguments forwarded verbatim to solver on every call
479
+ Keyword arguments forwarded verbatim to solver on every call. Pass a per-iteration callback here (e.g.
480
+ ``solver_kwargs={"callback": fn}``); it fires on every run, and under the ``loky`` backend that run is a
481
+ separate process, so the callback and anything it closes over must be picklable.
480
482
  sort_key: Callable, optional
481
483
  Function mapping an OptimizeResult to a float used for ranking. Defaults to
482
484
  lambda res: res.fun, which ranks results by their optimized objective value.
@@ -1,3 +1,5 @@
1
+ import logging
2
+
1
3
  from collections.abc import Callable
2
4
  from functools import partial
3
5
 
@@ -7,7 +9,7 @@ from rich.progress import Progress, TaskID
7
9
  from scipy.optimize import OptimizeResult
8
10
  from scipy.optimize import root as sp_root
9
11
 
10
- from better_optimize.constants import root_method
12
+ from better_optimize.constants import ROOT_METHODS_WITHOUT_CALLBACK, root_method
11
13
  from better_optimize.utilities import (
12
14
  LRUCache1,
13
15
  check_f_is_fused_root,
@@ -17,7 +19,13 @@ from better_optimize.utilities import (
17
19
  kwargs_to_options,
18
20
  validate_provided_functions_root,
19
21
  )
20
- from better_optimize.wrapper import ObjectiveWrapper, optimizer_early_stopping_wrapper
22
+ from better_optimize.wrapper import (
23
+ ObjectiveWrapper,
24
+ _compose_callback,
25
+ optimizer_early_stopping_wrapper,
26
+ )
27
+
28
+ _log = logging.getLogger(__name__)
21
29
 
22
30
 
23
31
  def root(
@@ -30,6 +38,7 @@ def root(
30
38
  progressbar_update_interval: int = 1,
31
39
  verbose: bool = False,
32
40
  args: tuple | None = None,
41
+ callback: Callable[..., bool | None] | None = None,
33
42
  **optimizer_kwargs,
34
43
  ) -> OptimizeResult:
35
44
  """
@@ -59,6 +68,11 @@ def root(
59
68
  verbose: bool
60
69
  If True, warnings about the provided configuration are displayed. These warnings are intended to help users
61
70
  understand potential configuration issues that may affect the optimization process, but can be safely ignored.
71
+ callback: Callable, optional
72
+ Function called after each iteration as ``callback(res)``, where ``res`` is an
73
+ ``OptimizeResult`` carrying the current ``res.x``, ``res.fun`` (the residual vector), and
74
+ ``res.nit``. The return value is ignored; raise ``StopOptimization`` to stop early. The
75
+ direct solvers ``hybr`` and ``lm`` ignore the callback and warn. Default None.
62
76
  optimizer_kwargs
63
77
  Additional keyword arguments to pass to the optimizer
64
78
 
@@ -93,12 +107,27 @@ def root(
93
107
  task=progress_task,
94
108
  )
95
109
 
110
+ if callback is not None and method in ROOT_METHODS_WITHOUT_CALLBACK:
111
+ _log.warning(
112
+ f"Method {method} does not support callbacks; the provided callback will be ignored."
113
+ )
114
+ callback = None
115
+
116
+ # Passing callback=None matches SciPy's default, so the direct methods never trigger SciPy's
117
+ # "does not accept callback" warning.
118
+ root_callback = (
119
+ _compose_callback(objective.callback, objective.callback_result, callback)
120
+ if callback is not None
121
+ else None
122
+ )
123
+
96
124
  f_optim = partial(
97
125
  sp_root,
98
126
  fun=objective,
99
127
  x0=x0,
100
128
  method=method,
101
129
  jac=True if has_fused_f_and_grad or jac is not None else None,
130
+ callback=root_callback,
102
131
  **optimizer_kwargs,
103
132
  )
104
133
 
@@ -265,7 +265,8 @@ def sequential_optimize(
265
265
  Initial guess for the first stage (unless the stage overrides via its own ``"x0"``).
266
266
  stages : list of dict
267
267
  Each dict must contain ``"solver": Callable`` and may include ``"name": str`` plus
268
- any kwargs the solver accepts (e.g. ``"method"``, ``"bounds"``).
268
+ any kwargs the solver accepts (e.g. ``"method"``, ``"bounds"``, or a per-stage
269
+ ``"callback"`` forwarded to that stage's solver when it accepts one).
269
270
  args : tuple
270
271
  Extra positional args forwarded identically to every stage's objective call.
271
272
  progressbar : bool
@@ -114,6 +114,9 @@ class ObjectiveWrapper:
114
114
 
115
115
  self.use_rayleigh = self.use_jac and (self.use_hess or self.use_hessp)
116
116
 
117
+ # Counts iterations for the uniform OptimizeResult handed to user callbacks.
118
+ self.n_iter = 0
119
+
117
120
  self.previous_x = None
118
121
  if isinstance(progressbar, bool):
119
122
  self.progressbar = progressbar
@@ -189,6 +192,54 @@ class ObjectiveWrapper:
189
192
  if self.interrupted:
190
193
  raise StopIteration
191
194
 
195
+ def callback_result(self, *args) -> OptimizeResult:
196
+ """Build the uniform ``OptimizeResult`` handed to a user callback.
197
+
198
+ SciPy passes callbacks several different argument shapes depending on the method (``xk``;
199
+ ``xk, state``; an ``intermediate_result``; or ``x, f`` for root finding); collapse them into
200
+ one result with ``x``, ``fun``, and ``nit``. ``fun`` is the scalar objective for
201
+ minimization and the residual vector for root finding. ``jac`` is present when a fused
202
+ value-and-gradient objective supplies it.
203
+ """
204
+ self.n_iter += 1
205
+
206
+ state = None
207
+ fun = None
208
+ if len(args) == 1 and isinstance(args[0], OptimizeResult):
209
+ state = args[0] # minimize new-protocol intermediate_result
210
+ x = np.asarray(state.x, dtype=float)
211
+ elif len(args) >= 2 and isinstance(args[1], OptimizeResult):
212
+ state = args[1] # trust-constr legacy (xk, state)
213
+ x = np.asarray(args[0], dtype=float)
214
+ elif self.root and len(args) >= 2:
215
+ x = np.asarray(args[0], dtype=float)
216
+ fun = np.asarray(args[1], dtype=float) # residual vector
217
+ else:
218
+ x = np.asarray(args[0], dtype=float)
219
+
220
+ nit = self.n_iter
221
+ jac = None
222
+ if state is not None:
223
+ # Prefer SciPy's own numbers when it gave us a populated result.
224
+ fun = getattr(state, "fun", fun)
225
+ nit = int(getattr(state, "nit", nit))
226
+ jac = getattr(state, "jac", None)
227
+
228
+ if fun is None:
229
+ out = self.f(x)
230
+ if self.has_fused_f_and_grad:
231
+ fun, grad = out
232
+ if jac is None and not self.root:
233
+ jac = grad
234
+ else:
235
+ fun = out
236
+ fun = fun if self.root else float(np.asarray(fun).reshape(()))
237
+
238
+ res = OptimizeResult(x=x, fun=fun, nit=nit)
239
+ if jac is not None:
240
+ res.jac = jac
241
+ return res
242
+
192
243
  def update_progressbar(
193
244
  self, value: float, grad: np.float64 = None, rayleigh: float | None = None, completed=False
194
245
  ) -> None:
@@ -247,6 +298,38 @@ class ObjectiveWrapper:
247
298
  )
248
299
 
249
300
 
301
+ def _compose_callback(
302
+ internal: Callable[..., None],
303
+ result_builder: Callable[..., OptimizeResult],
304
+ user: Callable | None,
305
+ ) -> Callable:
306
+ """Bundle the internal early-stop hook with an optional user-supplied callback.
307
+
308
+ The internal hook runs first, so ``maxeval`` and ``KeyboardInterrupt`` halting stays intact.
309
+ The user callback then receives a uniform ``OptimizeResult`` built by ``result_builder`` from
310
+ whatever SciPy passed. Its return value is ignored; it requests an early stop by raising
311
+ :class:`~better_optimize.exceptions.StopOptimization` (or any ``StopIteration``).
312
+
313
+ Parameters
314
+ ----------
315
+ internal : callable
316
+ The :meth:`ObjectiveWrapper.callback` early-stop hook.
317
+ result_builder : callable
318
+ Normalizes SciPy's per-method callback arguments into a uniform ``OptimizeResult``
319
+ (typically :meth:`ObjectiveWrapper.callback_result`).
320
+ user : callable or None
321
+ The user-supplied callback. When None, ``internal`` is returned unchanged.
322
+ """
323
+ if user is None:
324
+ return internal
325
+
326
+ def composite(*args):
327
+ internal(*args)
328
+ user(result_builder(*args))
329
+
330
+ return composite
331
+
332
+
250
333
  def optimizer_early_stopping_wrapper(f_optim: partial):
251
334
  objective = f_optim.keywords["fun"]
252
335
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: better_optimize
3
- Version: 0.4.1
3
+ Version: 0.4.2
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
@@ -89,6 +89,19 @@ All optimization routines in `better_optimize` can display a rich, informative p
89
89
  - Optionally accepts and stores failed minimizer results if they improve the global minimum.
90
90
  - Useful for noisy or non-smooth objective functions where local minimization may occasionally fail.
91
91
 
92
+ ### 6. Uniform Callback API
93
+
94
+ scipy passes callbacks a different argument for almost every method: `callback(xk)`, `callback(xk, state)`,
95
+ `callback(intermediate_result)`, or `callback(x, f)` for root finding. `better_optimize` normalizes them to a single
96
+ `callback(res)`, where `res` is an `OptimizeResult` with the current `res.x`, `res.fun`, and `res.nit` — plus `res.jac`
97
+ when a gradient is available, and solver-specific fields like `res.accept` (basinhopping) or `res.convergence`
98
+ (differential evolution).
99
+
100
+ - The return value is ignored, so a callback can report a value (an ELBO, a logged loss) without stopping the run.
101
+ - Raise `StopOptimization` to stop early; the best result so far is returned with `success=False`.
102
+ - The same callback works across `minimize`, `root`, `basinhopping`, and `differential_evolution`. `hybr` and `lm` are
103
+ the exception — scipy never calls a callback for them, so it is ignored with a warning.
104
+
92
105
  ---
93
106
 
94
107
  ## Example Usage
@@ -96,6 +109,7 @@ All optimization routines in `better_optimize` can display a rich, informative p
96
109
  ### Simple Example
97
110
 
98
111
  ```python
112
+ import numpy as np
99
113
  from better_optimize import minimize
100
114
 
101
115
  def rosenbrock(x):
@@ -103,7 +117,7 @@ def rosenbrock(x):
103
117
 
104
118
  result = minimize(
105
119
  rosenbrock,
106
- x0=[-1, 2],
120
+ x0=np.array([-1.0, 2.0]),
107
121
  method="L-BFGS-B",
108
122
  tol=1e-6,
109
123
  maxiter=1000,
@@ -119,6 +133,36 @@ result = minimize(
119
133
 
120
134
  The result object is a standard `OptimizeResult` from `scipy.optimize`, so there are no surprises there!
121
135
 
136
+ ### Callbacks and Early Stopping
137
+
138
+ Pass a `callback` and it runs after each iteration with an `OptimizeResult`. Its return value is ignored; raise
139
+ `StopOptimization` to stop and return the best result so far.
140
+
141
+ ```python
142
+ import numpy as np
143
+ from better_optimize import minimize, StopOptimization
144
+
145
+ def rosenbrock(x):
146
+ return sum(100.0*(x[1:] - x[:-1]**2.0)**2.0 + (1 - x[:-1])**2.0)
147
+
148
+ history = []
149
+
150
+ def callback(res):
151
+ history.append(res.fun) # res.x, res.fun, res.nit; res.jac when available
152
+ if res.fun < 1e-8:
153
+ raise StopOptimization
154
+
155
+ result = minimize(
156
+ rosenbrock,
157
+ x0=np.array([-1.0, 2.0]),
158
+ method="L-BFGS-B",
159
+ callback=callback,
160
+ )
161
+ ```
162
+
163
+ The same callback works for `root` (`res.fun` is the residual vector), `basinhopping` (`res.accept`), and
164
+ `differential_evolution` (`res.convergence`).
165
+
122
166
  ### Triple-Fused Function using Pytensor
123
167
 
124
168
  ```python
@@ -155,6 +199,7 @@ Real-world objectives often have multiple local minima. A common workaround is t
155
199
  the optimizer and keep the best result. `better_optimize` makes this painless with `multi_optimize`:
156
200
 
157
201
  ```python
202
+ import numpy as np
158
203
  from better_optimize import minimize, multi_optimize
159
204
 
160
205
  def rosenbrock(x):
@@ -179,6 +224,9 @@ print(result.fun_best) # Best objective value
179
224
  result.summary() # Rich table of all runs, ranked
180
225
  ```
181
226
 
227
+ > The `loky` (process) backend needs an `if __name__ == "__main__":` guard when run as a script on macOS or Windows.
228
+ > Notebooks and the `threading`/`sequential` backends don't.
229
+
182
230
  `multi_optimize` works with **any** solver that follows the `(x0, **kwargs) → OptimizeResult` signature — that includes
183
231
  `minimize`, `root`, `basinhopping`, or your own custom wrapper. It just calls `solver(x0=x0_i, **solver_kwargs)` for
184
232
  each starting point; it never inspects the solver internals.