ENH: support SciPy-style callback(xk) during optimization

[Pragati5-DEBUG]

Summary

• Support SciPy-style callback(xk) in minimize / maximize
• Pass it into InternalOptimizationProblem and call it on objective evaluations (works for all optimizers, not via SciPy's own callback)
• callback(intermediate_result) left for a follow-up

Test plan

• pytest tests/optimagic/optimization/test_scipy_aliases.py::test_callback_xk_is_called
• Related optimization tests (96 passed)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
...imagic/optimization/create_optimization_problem.py	88.13% <100.00%> (-0.10%)	⬇️
...agic/optimization/internal_optimization_problem.py	92.99% <100.00%> (+0.13%)	⬆️
src/optimagic/optimization/optimize.py	92.10% <ø> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[janosg]

Hi @Pragati5-DEBUG

I left a few comments. After reading up a bit more on SciPy callbacks, I am also realizing there are two conceptual problems with the implementation sketch I gave you.

1. SciPy callbacks are evaluated after each Iteration, not after each function evaluation. For gradient based optimizers this is typically the same but for gradient free ones it is not. Optimagic cannot know what an algorithm calls an iteration, therefore any algorithm agnostic implementation of callbacks can only work on function evaluations, which is useful but not exactly what SciPy does.

I therefore suggest that for all optimizers that support SciPy style callbacks (we'll need a flag in AlgoInfo for that), we use the algorithm's callback mechanism instead of bypassing it and just add parameter conversion to get from the internal parameter vector of the optimizer to the external one the callback expects. For all other optimizers we go with our own callback implementation that works on function evaluations. That way we get full SciPy compatibility and a useful extension for all non scipy optimizers.

2. The _pure functions are actually not the right level to implement callbacks. We would get problems with parallelizing optimizers. With n_cores > 1, the callback executes in worker processes. Common callback state such as xs.append(...), counters, or in-memory progress objects will not be updated in the parent process; the callback must also be serializable, and side effects can occur concurrently. This is the exact same reason why history collection is not implemented inside the _pure... functions. A rule of thumb is that callback evaluations should probably happen next to history append operations to avoid this problem.

I am realizing that this makes the feature much harder than anticipated. Let me and @timmens know if you want us to take over the PR.

[janosg]

This type hint could be much tighter; Currently we only accept callables that take a PyTree (the current parameter Vector) and doesn't return anything.

Note: In contrast to scipy xk is not necessarily a numpy array, it can be an arbitrary PyTree.

[janosg]

Currently you don't do any validation of the user passed callback. For example, if the user passes a callable with the wrong signature, this would lead to a hard to understand error in the first iteration. We need the same validation for callbacks as we have for all other user evaluated functions like fun, jac, etc.

We also need tests that invalid callbacks are handled well.

[janosg]

SciPy callbacks can raise StopIteration, which aborts an optimization gracefully. This is probably a bit out of scope for this PR, but we need to document that it is not handled.

[janosg]

We should not evaluate the callback at derivative evaluations. Only at objective evaluations.

[janosg]

You are calling the callback with the internal parameter vector (a flat numpy array with entries that a user might not be able to interpret). Instead the callback should be called with the external parameters (which could be dictionaries or arbitrary PyTrees. We have this pattern for any user supplied function.

[Pragati5-DEBUG]

Thanks @janosg for the detailed review — that all makes sense.

I’d like to keep working on this. Could I please have some time to read up on SciPy callbacks and understand the design better before implementing the changes? I’ll update the PR once I’m ready.

[janosg]

Sure, take as long as you need.