Handle infinite values in gradients

[janosg]

Current situation

Optimagic does not make a check for infinite values in gradients. Some optimizers do and raise an error if any infinite values are encountered. The error message raised by optimizers is necessarily based on internal (flat) gradients.

Possible improvements

1. Make the check in optimagic (fun and fun_and_jac of the InternalOptimizationProblem) and raise an error message based on external (non-flat) gradients
2. Integrate this with our error handling. If error_handling is set to continue we could treat infinite values in gradients as error and replace the function value and gradient with penalty versions.

Example

This topic came up in a skillmodels application by @hmgaudecker using the fides optimizer

│ /home/skanehealth/skane-struct-bw/.pixi/envs/default/lib/python3.13/site-packages/fides/minimize │
│ .py:793 in check_finite                                                                          │
│                                                                                                  │
│   790 │   │   if not np.isfinite(grad).all():                                                    │
│   791 │   │   │   self.exitflag = ExitFlag.NOT_FINITE                                            │
│   792 │   │   │   ix = np.where(np.logical_not(np.isfinite(grad)))                               │
│ ❱ 793 │   │   │   raise RuntimeError('Encountered non-finite gradient entries'                   │
│   794 │   │   │   │   │   │   │      f' {grad[ix]} for indices {ix} '                            │
│   795 │   │   │   │   │   │   │      f'{pointstr}')                                              │
│   796                                                                                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Encountered non-finite gradient entries [
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan
 nan nan nan nan nan nan nan nan nan nan nan nan nan nan
] for indices (
    array([
        0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
        13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,
        26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,
        39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,
        52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,
        65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,
        78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,
        91,  92,  93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103,
       104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,
       117, 118, 119, 120, 121
    ]),
) at initial point.

[Aziz-Shameem]

Hi @janosg
I can take this up.
If you could provide an example of the sample code that fails, that'd be helpful.
Thanks

[janosg]

Hi @Aziz-Shameem, thanks for your offer. Here are a few pointers to address point 1 of the issue. I think step two requires too much knowledge about optimagic's internals.

Reproduction

The code example above can only run with private data, so we cannot use it as a test case. However, you could create a simple test case based on the following:
- Use the sphere function we use everywhere in the documentation as objective function
- Use start parameters quite far away from zero (e.g. np.arange(10) + 400)
- Implement a custom derivative that returns np.inf or -np.inf for all parameters with an absolute value below 1 and is otherwise just the derivative of the sphere function.

This should allow you reproduce the error above with a manimal example. If it does not, switch to a slightly more complex objective function from here

Changes

Adjust all relevant methods of InternalOptimizationProblem to raise a custom Error instead of the algorithm specific you got from fides. Your changes should be in the private methods that start with _pure and do some kind of jac evaluation. Focus on providing really helpful error messages based on params the result of the user provided jac or fun_and_jac function. Don't use any of the flattened internal information like x.

Testing

This needs to be tested with test cases similar to the one you created for the Reproduction but we need tests for all possible ways of providing custom derivatives (e.g. only jac provided; only fun_and_jac` provided, both of them provided).

For now the goal is just to get a better error message with more useful debugging information. We probably want to use a UserFunctionRuntimeError.

[Aziz-Shameem]

Hi.
Thanks for the response.
I am yet to test my changes exhaustively, but is this the expected behavior ?

I have returned the parameter values and the faulty(?) gradients in the error message :

msg = ( "Infinite values found in user-provided gradient at the following "
               f"parameter positions:\nParameters: {params}\nGradient: {jac_value}" )

[janosg]

Yes, this looks good. It would be even easier to see if you use a dictionary intsead of a list or array as params because then the difference between x and params is easier to see. You could take that example from here

[Aziz-Shameem]

It seems to work correctly on all input types (Array, Dict, Scalar) and using both - jac and fun_and_juc.
All the pre-commit hooks also seem to pass.
However, while running pytest, I get the following error and I am not able to track its source. Could you direct me on what may the issue be ?

==================================================== short test summary info ===================================================== 
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-False-jac1-sos_ls] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-False-jac1-typed_sos_ls] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-False-jac1-sos_ls_with_info] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-True-jac1-sos_ls] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-True-jac1-typed_sos_ls] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_function_formats_ls.py::test_least_squares_minimize_dict[scipy_ls_lm-True-jac1-sos_ls_with_info] - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_internal_optimization_problem.py::test_pytree_problem_jac - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating jac during optimization.
FAILED tests/optimagic/optimization/test_internal_optimization_problem.py::test_pytree_problem_fun_and_jac - optimagic.exceptions.UserFunctionRuntimeError: An error occurred when evaluating fun_and_jac during optimization.
========================= 8 failed, 2539 passed, 64 skipped, 1 xfailed, 28 warnings in 94.24s (0:01:34) ========================== 

Also, related, should I include checking for nan values in the gradient along with checking for infinite values, for the sake of completeness ? I don't know if that is already checked in the current flow.

Thanks

[Aziz-Shameem]

Hi @janosg
I would like to try and tackle the second part of this as well.
Could you give some direction on how I can start ?

Thanks

[janosg]

So far we have achieved that the user gets a better error message if the derivative is undefined or infinite. However, the optimization will stop in that case.

The goal for the next step is to continue with the optimization, even though the derivative is invalid. We already do this for errors in the objective function. Our solution is described here. In short: If the objective function raises an error and the error_handling argument is set to continue, we catch the error and replace the objective value and the derivative with carefully designed values that guide the optimizer back towards the start values.

We would now extend this exact same approach to invalid derivatives, i.e. we use the same error_handling argument to decide what we want to do and the same penalty functions as in the existing case.

The steps are:

• Understand how over error handlnig works and why it works by reading the [documentation]
• Write tests for the new behavior. By default error_handling is set to raise, so the new behavior is switched off. Thus all the tests you wrote previously need to stay in place. Now you can add copies of those tests with error_handling set to continue and assert that they find the optimum (i.e. the zero vector in the case of the sphere function)
• Add the new error handling to the relevant methods in InternalOptimizationProblem. For this you probably have to refactor your previous changes a bit.

[Aziz-Shameem]

hey
Thanks for the brief.
I think I'll work on issue 593 first, I can take this up after that.