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,9 @@ def _pure_evaluate_jac(
         out_jac = _process_jac_value(
             value=jac_value, direction=self._direction, converter=self._converter, x=x
         )
+        _assert_finite_jac(
+            out_jac=out_jac, jac_value=jac_value, params=params, origin="jac"
+        )
 
         stop_time = time.perf_counter()
 
@@ -543,6 +546,13 @@ def func(x: NDArray[np.float64]) -> SpecificFunctionValue:
                 warnings.warn(msg)
                 fun_value, jac_value = self._error_penalty_func(x)
 
+        _assert_finite_jac(
+            out_jac=jac_value,
+            jac_value=jac_value,
+            params=self._converter.params_from_internal(x),
+            origin="numerical",
+        )
+
         algo_fun_value, hist_fun_value = _process_fun_value(
             value=fun_value,  # type: ignore
             solver_type=self._solver_type,
@@ -682,6 +692,10 @@ def _pure_evaluate_fun_and_jac(
         if self._direction == Direction.MAXIMIZE:
             out_jac = -out_jac
 
+        _assert_finite_jac(
+            out_jac=out_jac, jac_value=jac_value, params=params, origin="fun_and_jac"
+        )
+
         stop_time = time.perf_counter()
 
         hist_entry = HistoryEntry(
@@ -705,6 +719,44 @@ def _pure_evaluate_fun_and_jac(
         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.
+        origin: Source of Jacobian calculation, for the error message.
+
+    Raises:
+        UserFunctionRuntimeError:
+            If any infinite or NaN values are found in the Jacobian.
+
+    """
+    if not np.all(np.isfinite(out_jac)):
+        if origin == "jac" or "fun_and_jac":
+            msg = (
+                "The optimization failed because the derivative provided via "
+                f"{origin} 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,114 @@
+import numpy as np
+import pytest
+
+from optimagic.exceptions import UserFunctionRuntimeError
+from optimagic.optimization.optimize import minimize
+
+# ======================================================================================
+# Test setup:
+# --------------------------------------------------------------------------------------
+# We test that minimize raises an error if the user function returns a jacobian
+# containing invalid values (np.inf, np.nan). To test that this works not only at
+# the start parameters, we create jac functions that return invalid values if the
+# parameter norm becomes smaller than one.
+# ======================================================================================
+
+
+@pytest.fixture
+def params():
+    return {"a": 1, "b": np.array([3, 4])}
+
+
+def sphere(params):
+    return params["a"] ** 2 + (params["b"] ** 2).sum()
+
+
+def sphere_gradient(params):
+    return {
+        "a": 2 * params["a"],
+        "b": 2 * params["b"],
+    }
+
+
+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
+    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])},
+    {"a": 1, "b": 2 * np.array([np.inf, 2])},
+    {"a": np.nan, "b": 2 * np.array([1, 2])},
+    {"a": 1, "b": 2 * np.array([np.nan, 2])},
+]
+
+
+# ======================================================================================
+# Test Invalid Jacobian raises proper error with jac argument
+# ======================================================================================
+
+
+@pytest.mark.parametrize("invalid_jac_value", INVALID_JACOBIAN_VALUES)
+def test_minimize_with_invalid_jac(invalid_jac_value, params):
+    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),
+        )
+
+
+# ======================================================================================
+# Test Invalid Jacobian raises proper error with fun_and_jac argument
+# ======================================================================================
+
+
+@pytest.mark.parametrize("invalid_jac_value", INVALID_JACOBIAN_VALUES)
+def test_minimize_with_invalid_fun_and_jac(invalid_jac_value, params):
+    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),
+        )