Update test functions to produce NaN/Inf gradients naturally

- Use sqrt and max to produce NaN when x goes negative
- Use 1/x pattern to produce Inf gradients
- Functions naturally produce problematic gradients during optimization

Author: ChrisRackauckas

lib/OptimizationOptimisers/test/runtests.jl

@@ -137,16 +137,15 @@ end
 
 @testset "NaN/Inf gradient handling" begin
     # Test that optimizer skips updates when gradients contain NaN or Inf
-    # Function that returns NaN when parameters are in certain regions
+    # Function that can produce NaN due to sqrt of negative number
     function weird_nan_function(x, p)
-        # Return NaN when x[1] is close to certain values to simulate numerical issues
-        if abs(x[1] - 0.3) < 0.05 || abs(x[1] + 0.3) < 0.05
-            return NaN
-        end
-        return (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+        val = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+        # sqrt of a value that can become negative produces NaN
+        val += sqrt(max(x[1], 0.0)) * 0.01
+        return val
     end
 
-    x0 = zeros(2)
+    x0 = [-0.5, 0.1]  # Start with negative x[1] to trigger sqrt of negative
     _p = [1.0, 100.0]
 
     optprob = OptimizationFunction(weird_nan_function, OptimizationBase.AutoZygote())
@@ -162,13 +161,12 @@ end
     @test all(!isnan, sol.u)
     @test all(isfinite, sol.u)
 
-    # Function that returns Inf when parameters are in certain regions
+    # Function with 1/x that can produce Inf gradient when x is very small
     function weird_inf_function(x, p)
-        # Return Inf when x[1] is close to certain values
-        if abs(x[1] - 0.2) < 0.05 || abs(x[1] + 0.2) < 0.05
-            return Inf
-        end
-        return (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+        val = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+        # 1/(x[1] + 0.01) can have very large gradient near x[1] = -0.01
+        val += 0.01 / (abs(x[1] - 0.1) + 1e-8)
+        return val
     end
 
     optprob_inf = OptimizationFunction(weird_inf_function, OptimizationBase.AutoZygote())