From 0f878b34528bee0452589dc625df283a9f4f333f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20P=2E=20D=C3=BCrholt?= Date: Tue, 2 Jun 2026 22:30:02 +0200 Subject: [PATCH 1/5] increase scipy compatibility --- python/pounce/_minimize.py | 71 ++++++++++++++++++++++++++++++++------ 1 file changed, 61 insertions(+), 10 deletions(-) diff --git a/python/pounce/_minimize.py b/python/pounce/_minimize.py index e6c2f8f4..aa6808a0 100644 --- a/python/pounce/_minimize.py +++ b/python/pounce/_minimize.py @@ -6,8 +6,12 @@ Notes ----- -* When ``jac`` is omitted we fall back to forward finite differences - (step ``sqrt(eps)``). Production callers should provide a Jacobian. +* When ``jac`` is omitted (or ``False``) we fall back to forward finite + differences (step ``sqrt(eps)``). Production callers should provide + a Jacobian. +* When ``jac=True`` (BoTorch convention) ``fun(x, *args)`` must return + ``(f, grad)``; the pair is cached so each Ipopt iterate triggers only + one forward pass. * When ``hess`` is omitted, or when constraints are present, the solver is driven with ``hessian_approximation = limited-memory``. * Equality / inequality dicts are concatenated into a single ``g(x)`` @@ -60,6 +64,40 @@ def _finite_diff_grad(fun: Callable, x: np.ndarray) -> np.ndarray: return g +class _FunAndGradCache: + """Memoize ``(f, g) = fun(x, *args)`` on the most recent ``x``. + + Ipopt evaluates ``objective`` and ``gradient`` as separate calls + (often at the same point). When ``jac=True``, the user-supplied + ``fun`` returns both in one forward pass — caching here preserves + that single-pass guarantee across the two Ipopt callbacks. + """ + + def __init__(self, fun: Callable, args: tuple): + self._fun = fun + self._args = args + self._x: np.ndarray | None = None + self._f: float | None = None + self._g: np.ndarray | None = None + + def _ensure(self, x: np.ndarray) -> None: + if self._x is None or self._x.shape != x.shape or not np.array_equal(self._x, x): + f, g = self._fun(x, *self._args) + self._x = x.copy() + self._f = float(f) + self._g = _to_array(g).ravel() + + def f(self, x: np.ndarray) -> float: + self._ensure(x) + assert self._f is not None + return self._f + + def g(self, x: np.ndarray) -> np.ndarray: + self._ensure(x) + assert self._g is not None + return self._g + + def _finite_diff_jac(g_fun: Callable, x: np.ndarray, m: int) -> np.ndarray: g0 = _to_array(g_fun(x)) J = np.empty((m, x.size)) @@ -139,7 +177,8 @@ def _build_problem_obj( fun: Callable, n: int, m: int, - jac: Callable | None, + args: tuple, + jac: Callable | bool | None, hess: Callable | None, g: Callable | None, jac_g: Callable | None, @@ -150,13 +189,23 @@ def _build_problem_obj( members: dict[str, Any] = {} - def objective(self, x): - return float(fun(x)) + if jac is True: + cache = _FunAndGradCache(fun, args) + + def objective(self, x, _c=cache): + return _c.f(x) + + def gradient(self, x, _c=cache): + return _c.g(x) + else: + + def objective(self, x): + return float(fun(x, *args)) - def gradient(self, x): - if jac is None: - return _finite_diff_grad(fun, x) - return _to_array(jac(x)).ravel() + def gradient(self, x): + if jac is None or jac is False: + return _finite_diff_grad(lambda x: fun(x, *args), x) + return _to_array(jac(x, *args)).ravel() members["objective"] = objective members["gradient"] = gradient @@ -199,7 +248,8 @@ def hessian(self, x, lam, obj_factor): def minimize( fun: Callable[[np.ndarray], float], x0: np.ndarray, - jac: Callable | None = None, + args: tuple = (), + jac: Callable | bool | None = None, hess: Callable | None = None, bounds: Sequence | None = None, constraints: Sequence | dict | None = None, @@ -215,6 +265,7 @@ def minimize( fun=fun, n=n, m=m, + args=args, jac=jac, hess=hess, g=g_combined, From a3164e8836b47f43f77cbceaa427e192bff5714f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20P=2E=20D=C3=BCrholt?= Date: Tue, 2 Jun 2026 22:37:52 +0200 Subject: [PATCH 2/5] tidy up note --- python/pounce/_minimize.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/pounce/_minimize.py b/python/pounce/_minimize.py index aa6808a0..9379ef0e 100644 --- a/python/pounce/_minimize.py +++ b/python/pounce/_minimize.py @@ -9,7 +9,7 @@ * When ``jac`` is omitted (or ``False``) we fall back to forward finite differences (step ``sqrt(eps)``). Production callers should provide a Jacobian. -* When ``jac=True`` (BoTorch convention) ``fun(x, *args)`` must return +* When ``jac=True``, ``fun(x, *args)`` must return ``(f, grad)``; the pair is cached so each Ipopt iterate triggers only one forward pass. * When ``hess`` is omitted, or when constraints are present, the solver @@ -81,7 +81,11 @@ def __init__(self, fun: Callable, args: tuple): self._g: np.ndarray | None = None def _ensure(self, x: np.ndarray) -> None: - if self._x is None or self._x.shape != x.shape or not np.array_equal(self._x, x): + if ( + self._x is None + or self._x.shape != x.shape + or not np.array_equal(self._x, x) + ): f, g = self._fun(x, *self._args) self._x = x.copy() self._f = float(f) From 0a423710168aab8c9bbcc609a25129085281fc20 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20P=2E=20D=C3=BCrholt?= Date: Wed, 3 Jun 2026 11:09:44 +0200 Subject: [PATCH 3/5] add args for hessian --- python/pounce/_minimize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pounce/_minimize.py b/python/pounce/_minimize.py index 9379ef0e..4d5de988 100644 --- a/python/pounce/_minimize.py +++ b/python/pounce/_minimize.py @@ -238,7 +238,7 @@ def hessianstructure(self): return (r, c) def hessian(self, x, lam, obj_factor): - H = obj_factor * _to_array(hess(x)) + H = obj_factor * _to_array(hess(x, *args)) r, c = np.tril_indices(n) return H[r, c] From 2c9236779acd83b502c818cec8d649ef5aac4d02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20P=2E=20D=C3=BCrholt?= Date: Wed, 3 Jun 2026 11:17:35 +0200 Subject: [PATCH 4/5] add args for constraints --- python/pounce/_minimize.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/pounce/_minimize.py b/python/pounce/_minimize.py index 4d5de988..db6c3ed2 100644 --- a/python/pounce/_minimize.py +++ b/python/pounce/_minimize.py @@ -136,29 +136,32 @@ def _wrap_constraints(constraints, n: int): if isinstance(constraints, dict): constraints = [constraints] - funs, jacs = [], [] + funs, jacs, cargs = [], [], [] for c in constraints: kind = c["type"] if kind not in ("eq", "ineq"): raise ValueError(f"unknown constraint type {kind!r}") funs.append(c["fun"]) jacs.append(c.get("jac")) + cargs.append(tuple(c.get("args", ()))) probe = np.zeros(n) - sizes = [int(_to_array(fn(probe)).size) for fn in funs] + sizes = [int(_to_array(fn(probe, *ca)).size) for fn, ca in zip(funs, cargs)] m_total = int(sum(sizes)) def g_combined(x): - return np.concatenate([_to_array(fn(x)).ravel() for fn in funs]) + return np.concatenate( + [_to_array(fn(x, *ca)).ravel() for fn, ca in zip(funs, cargs)] + ) def jac_combined(x): rows = [] - for fn, jc in zip(funs, jacs): + for fn, jc, ca in zip(funs, jacs, cargs): if jc is not None: - rows.append(np.atleast_2d(_to_array(jc(x)))) + rows.append(np.atleast_2d(_to_array(jc(x, *ca)))) else: - m_i = _to_array(fn(x)).size - rows.append(_finite_diff_jac(fn, x, m_i)) + m_i = _to_array(fn(x, *ca)).size + rows.append(_finite_diff_jac(lambda xx, fn=fn, ca=ca: fn(xx, *ca), x, m_i)) return np.vstack(rows) cl = np.empty(m_total) From 9eb953cf5ca61701054393d71a4d01079a0658a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20P=2E=20D=C3=BCrholt?= Date: Wed, 3 Jun 2026 15:42:14 +0200 Subject: [PATCH 5/5] add callback --- python/pounce/_minimize.py | 116 +++++++++++++++++++++++++++++++++++-- 1 file changed, 112 insertions(+), 4 deletions(-) diff --git a/python/pounce/_minimize.py b/python/pounce/_minimize.py index db6c3ed2..7d1cb48a 100644 --- a/python/pounce/_minimize.py +++ b/python/pounce/_minimize.py @@ -17,10 +17,16 @@ * Equality / inequality dicts are concatenated into a single ``g(x)`` with bound vectors ``cl`` / ``cu``. Constraint Jacobian is dense by design — sparse Jacobians belong on the :class:`Problem` API. +* ``callback`` accepts both scipy signatures (chosen by parameter-name + introspection): ``callback(intermediate_result=OptimizeResult)`` or + ``callback(xk)``. Raise ``StopIteration`` to terminate early. + ``intermediate_result.x`` is read from a cache populated by the + objective evaluation that precedes Ipopt's intermediate hook. """ from __future__ import annotations +import inspect from dataclasses import dataclass, field from typing import Any, Callable, Mapping, Sequence @@ -102,6 +108,52 @@ def g(self, x: np.ndarray) -> np.ndarray: return self._g +class _LastXCache: + """Stash the most recent ``x`` seen by ``objective`` for the callback shim. + + Pounce's Rust ``intermediate`` hook doesn't pass the primal iterate to + python. Ipopt evaluates the objective at the accepted iterate just + before firing ``intermediate``, so the latest cached ``x`` is the right + point to surface as ``OptimizeResult.x`` to a scipy-style callback. + """ + + def __init__(self) -> None: + self.x: np.ndarray | None = None + + def remember(self, x: np.ndarray) -> None: + self.x = np.asarray(x, dtype=np.float64).copy() + + +def _wrap_callback(callback: Callable | None) -> Callable | None: + """Normalize a scipy-style callback to ``(OptimizeResult) -> bool``. + + Returned shim returns ``True`` to continue, ``False`` to stop. Raising + ``StopIteration`` inside the user callback also stops the solve. + Mirrors ``scipy.optimize._minimize._wrap_callback`` introspection: a + single parameter literally named ``intermediate_result`` selects the + new-style signature; anything else gets the old-style ``xk``. + """ + if callback is None: + return None + try: + sig = inspect.signature(callback) + use_new = set(sig.parameters) == {"intermediate_result"} + except (TypeError, ValueError): + use_new = False + + def wrapped(res: OptimizeResult) -> bool: + try: + if use_new: + callback(intermediate_result=res) + else: + callback(np.copy(res.x)) + except StopIteration: + return False + return True + + return wrapped + + def _finite_diff_jac(g_fun: Callable, x: np.ndarray, m: int) -> np.ndarray: g0 = _to_array(g_fun(x)) J = np.empty((m, x.size)) @@ -161,7 +213,9 @@ def jac_combined(x): rows.append(np.atleast_2d(_to_array(jc(x, *ca)))) else: m_i = _to_array(fn(x, *ca)).size - rows.append(_finite_diff_jac(lambda xx, fn=fn, ca=ca: fn(xx, *ca), x, m_i)) + rows.append( + _finite_diff_jac(lambda xx, fn=fn, ca=ca: fn(xx, *ca), x, m_i) + ) return np.vstack(rows) cl = np.empty(m_total) @@ -189,24 +243,30 @@ def _build_problem_obj( hess: Callable | None, g: Callable | None, jac_g: Callable | None, + callback: Callable | None, ): """Build a problem-object-with-methods on the fly. Only attaches ``hessian`` / ``hessianstructure`` when ``hess`` is provided so - Problem's ``hasattr`` probe correctly falls back to L-BFGS.""" + Problem's ``hasattr`` probe correctly falls back to L-BFGS. Likewise, + ``intermediate`` is only attached when ``callback`` is provided so the + no-callback case has zero per-iter Python overhead.""" members: dict[str, Any] = {} + xcache = _LastXCache() if jac is True: cache = _FunAndGradCache(fun, args) - def objective(self, x, _c=cache): + def objective(self, x, _c=cache, _xc=xcache): + _xc.remember(x) return _c.f(x) def gradient(self, x, _c=cache): return _c.g(x) else: - def objective(self, x): + def objective(self, x, _xc=xcache): + _xc.remember(x) return float(fun(x, *args)) def gradient(self, x): @@ -217,6 +277,52 @@ def gradient(self, x): members["objective"] = objective members["gradient"] = gradient + if callback is not None: + wrapped_cb = _wrap_callback(callback) + assert wrapped_cb is not None + + def intermediate( + self, + *, + alg_mod, + iter_count, + obj_value, + inf_pr, + inf_du, + mu, + d_norm, + regularization_size, + alpha_du, + alpha_pr, + ls_trials, + _cb=wrapped_cb, + _xc=xcache, + _n=n, + ): + x = _xc.x if _xc.x is not None else np.full(_n, np.nan) + res = OptimizeResult( + x=x, + fun=float(obj_value), + success=False, + status=0, + message="intermediate", + nit=int(iter_count), + info={ + "alg_mod": int(alg_mod), + "inf_pr": float(inf_pr), + "inf_du": float(inf_du), + "mu": float(mu), + "d_norm": float(d_norm), + "regularization_size": float(regularization_size), + "alpha_du": float(alpha_du), + "alpha_pr": float(alpha_pr), + "ls_trials": int(ls_trials), + }, + ) + return _cb(res) + + members["intermediate"] = intermediate + if m > 0: def constraints(self, x): @@ -260,6 +366,7 @@ def minimize( hess: Callable | None = None, bounds: Sequence | None = None, constraints: Sequence | dict | None = None, + callback: Callable | None = None, options: Mapping[str, Any] | None = None, ) -> OptimizeResult: """scipy.optimize.minimize-style facade over pounce.""" @@ -277,6 +384,7 @@ def minimize( hess=hess, g=g_combined, jac_g=jac_combined, + callback=callback, ) problem = Problem(