Infinite Gradient Handling

src/optimagic/optimization/internal_optimization_problem.py

@@ -2,7 +2,7 @@
 import warnings
 from copy import copy
 from dataclasses import asdict, dataclass, replace
-from typing import Any, Callable, cast
+from typing import Any, Callable, Literal, cast
 
 import numpy as np
 from numpy.typing import NDArray
@@ -471,6 +471,7 @@
         out_jac = _process_jac_value(
             value=jac_value, direction=self._direction, converter=self._converter, x=x
         )
+        _assert_finite_jac(out_jac, jac_value, params, "jac")
 
         stop_time = time.perf_counter()
 
@@ -508,6 +509,7 @@
             p = self._converter.params_from_internal(x)
             return self._fun(p)
 
+        params = self._converter.params_from_internal(x)
         try:
             numdiff_res = first_derivative(
                 func,
@@ -543,6 +545,8 @@
                 warnings.warn(msg)
                 fun_value, jac_value = self._error_penalty_func(x)
 
+        _assert_finite_jac(jac_value, jac_value, params, "numerical")
+
         algo_fun_value, hist_fun_value = _process_fun_value(
             value=fun_value,  # type: ignore
             solver_type=self._solver_type,
@@ -682,6 +686,8 @@
         if self._direction == Direction.MAXIMIZE:
             out_jac = -out_jac
 
+        _assert_finite_jac(out_jac, jac_value, params, "fun_and_jac")
+
         stop_time = time.perf_counter()
 
         hist_entry = HistoryEntry(
@@ -705,6 +711,49 @@
         return (algo_fun_value, out_jac), hist_entry, log_entry
 
 
+def _assert_finite_jac(
+    out_jac: NDArray[np.float64],
+    jac_value: PyTree,
+    params: PyTree,
+    origin: Literal["numerical", "jac", "fun_and_jac"],
+) -> None:
+    """Check for infinite and NaN values in the Jacobian and raise an error if found.
+
+    Args:
+        out_jac: internal processed Jacobian to check for finiteness.
+        jac_value: original Jacobian value as returned by the user function,
+        params: user-facing parameter representation at evaluation point.
+
+    Raises:
+        UserFunctionRuntimeError:
+            If any infinite or NaN values are found in the Jacobian.
+
+    """
+    if not np.all(np.isfinite(out_jac)):
+        if origin == "jac":
+            msg = (
+                "The optimization failed because the derivative provided via "
+                "jac contains infinite or NaN values."
+                "\nPlease validate the derivative function."
+            )
+        elif origin == "fun_and_jac":
+            msg = (
+                "The optimization failed because the derivative provided via "
+                "fun_and_jac contains infinite or NaN values."
+                "\nPlease validate the derivative function."
+            )
+        elif origin == "numerical":
+            msg = (
+                "The optimization failed because the numerical derivative "
+                "(computed using fun) contains infinite or NaN values."
+                "\nPlease validate the criterion function or try a different optimizer."
+            )
+        msg += (
+            f"\nParameters at evaluation point: {params}\nJacobian values: {jac_value}"
+        )
+        raise UserFunctionRuntimeError(msg)
+
+
 def _process_fun_value(
     value: SpecificFunctionValue,
     solver_type: AggregationLevel,

src/optimagic/parameters/space_conversion.py

@@ -145,7 +145,6 @@ def get_space_converter(
         soft_lower_bounds=_soft_lower,
         soft_upper_bounds=_soft_upper,
     )
-
     return converter, params
 
 

tests/optimagic/optimization/test_invalid_jacobian_value.py

@@ -0,0 +1,105 @@
+import numpy as np
+import pytest
+
+from optimagic.exceptions import UserFunctionRuntimeError
+from optimagic.optimization.optimize import minimize
+
+
+def sphere(params):
+    return (
+        params["a"] ** 2
+        + (params["b"] ** 2).sum()
+        + params["c"]["x"] ** 2
+        + params["c"]["y"] ** 2
+    )
+
+
+def sphere_gradient(params):
+    return {
+        "a": 2 * params["a"],
+        "b": 2 * params["b"],
+        "c": {"x": 2 * params["c"]["x"], "y": 2 * params["c"]["y"]},
+    }
+
+
+def sphere_and_gradient(params):
+    return sphere(params), sphere_gradient(params)
+
+
+def params_norm(params):
+    squared_norm = (
+        params["a"] ** 2
+        + np.linalg.norm(params["b"]) ** 2
+        + params["c"]["x"] ** 2
+        + params["c"]["y"] ** 2
+    )
+    return np.sqrt(squared_norm)
+
+
+def get_invalid_jac(invalid_jac_value):
+    """Get function that returns invalid jac if the parameter norm < 1."""
+
+    def jac(params):
+        if params_norm(params) < 1:
+            return invalid_jac_value
+        else:
+            return sphere_gradient(params)
+
+    return jac
+
+
+def get_invalid_fun_and_jac(invalid_jac_value):
+    """Get function that returns invalid fun and jac if the parameter norm < 1."""
+
+    def fun_and_jac(params):
+        if params_norm(params) < 1:
+            return sphere(params), invalid_jac_value
+        else:
+            return sphere_and_gradient(params)
+
+    return fun_and_jac
+
+
+INVALID_JACOBIAN_VALUES = [
+    {"a": np.inf, "b": 2 * np.array([1, 2]), "c": {"x": 1, "y": 2}},
+    {"a": 1, "b": 2 * np.array([np.inf, 2]), "c": {"x": 1, "y": 2}},
+    {"a": np.nan, "b": 2 * np.array([1, 2]), "c": {"x": 1, "y": 2}},
+    {"a": 1, "b": 2 * np.array([np.nan, 2]), "c": {"x": 1, "y": 2}},
+    {"a": 1, "b": 2 * np.array([1, 2]), "c": {"x": np.inf, "y": 2}},
+    {"a": 1, "b": 2 * np.array([1, 2]), "c": {"x": 1, "y": np.nan}},
+]
+
+PARAMS = {"a": 1, "b": np.array([3, 4]), "c": {"x": 5, "y": 6}}
+
+
+@pytest.mark.parametrize("invalid_jac_value", INVALID_JACOBIAN_VALUES)
+def test_minimize_with_invalid_jac(invalid_jac_value):
+    with pytest.raises(
+        UserFunctionRuntimeError,
+        match=(
+            "The optimization failed because the derivative provided via jac "
+            "contains infinite or NaN values."
+        ),
+    ):
+        minimize(
+            fun=sphere,
+            params=PARAMS,
+            algorithm="scipy_lbfgsb",
+            jac=get_invalid_jac(invalid_jac_value),
+        )
+
+
+@pytest.mark.parametrize("invalid_jac_value", INVALID_JACOBIAN_VALUES)
+def test_minimize_with_invalid_fun_and_jac(invalid_jac_value):
+    with pytest.raises(
+        UserFunctionRuntimeError,
+        match=(
+            "The optimization failed because the derivative provided via fun_and_jac "
+            "contains infinite or NaN values."
+        ),
+    ):
+        minimize(
+            params=PARAMS,
+            algorithm="scipy_lbfgsb",
+            fun_and_jac=get_invalid_fun_and_jac(invalid_jac_value),
+        )