Fix bounds transform for problems with initialization data and quasi-Newton algorithms (#955)

* Fix bounds transform for problems with initialization data and quasi-Newton algs

Two bugs in the box-constraint bounds transform:

1. `transform_bounded_problem` unconditionally read `alg.autodiff`, which threw a
   `FieldError` for algorithms without that field (e.g. `QuasiNewtonAlgorithm`,
   reached as part of the default polyalgorithm). Guard the access.

2. The wrapped function carried over the original `initialization_data`, which is
   defined in the *bounded* coordinates. Running it on the transformed problem
   interpreted the unbounded iterate `t` as a bounded state and corrupted the
   solution (e.g. ModelingToolkit `NonlinearProblem`s, which always attach
   initialization data, converged to `from_unbounded(root)` instead of `root`).
   Initialization is now run once in the original bounded space before
   transforming, and stripped from the transformed function so it isn't re-run
   in `t`-space.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

* Test the bounds-transform fix in NonlinearSolveBase's own suite

The umbrella `test/bounds_tests.jl` integration test exercised the fix through a
full solve, but on Julia < 1.11 the umbrella resolves the *registered*
NonlinearSolveBase (the `[sources]` path redirect to the in-repo sublibrary is
ignored there), so the test ran against unpatched code and failed.

Add a `transform_bounded_problem` unit test to NonlinearSolveBase's own test
suite, which always runs against the in-repo code (incl. on lts), covering the
quasi-Newton `autodiff`-access guard with a dummy algorithm that has no
`autodiff` field. Gate the umbrella integration test on the loaded
NonlinearSolveBase version so it only runs where the fix is present.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

---------

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: ChrisRackauckas-Claude

lib/NonlinearSolveBase/Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolveBase"
 uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
-version = "2.30.2"
+version = "2.30.3"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

lib/NonlinearSolveBase/src/bounds_transform.jl

@@ -117,6 +117,13 @@ end
 # unbounded parameter `t` using the logistic function to map all values of `t`
 # into the interval (lb, ub).
 function transform_bounded_problem(prob, alg)
+    # `initialization_data` (e.g. from ModelingToolkit) is defined in the original,
+    # bounded coordinates. It must run *before* the transform so the resulting `u0`/`p`
+    # are consistent; running it on the transformed problem would interpret the unbounded
+    # iterate `t` as a bounded state and corrupt the solution. We run it here and strip it
+    # from the transformed function below so it isn't re-run in `t`-space.
+    prob = run_bounded_initialization(prob, alg)
+
     lb, ub = _normalize_bounds(prob.lb, prob.ub, prob.u0)
 
     # Clamp u0 into the interior of the bounds so that _to_unbounded doesn't hit log(0)
@@ -125,8 +132,10 @@ function transform_bounded_problem(prob, alg)
     u0_transformed = _to_unbounded.(u0_clamped, lb, ub)
 
     # PreallocationTools is only supported by ForwardDiff so we only use
-    # FixedSizeDiffCache if we're using ForwardDiff.
-    u_cache = if isnothing(alg) || isnothing(alg.autodiff) || alg.autodiff isa AutoForwardDiff
+    # FixedSizeDiffCache if we're using ForwardDiff. Not every algorithm has an
+    # `autodiff` field (e.g. `QuasiNewtonAlgorithm`), so guard the access.
+    alg_ad = alg !== nothing && hasproperty(alg, :autodiff) ? alg.autodiff : nothing
+    u_cache = if alg_ad === nothing || alg_ad isa AutoForwardDiff
         FixedSizeDiffCache(prob.u0)
     else
         similar(prob.u0)
@@ -148,7 +157,24 @@ function transform_bounded_problem(prob, alg)
         wrapped
     end
 
-    transformed_prob = remake(prob; f = new_f, u0 = u0_transformed, lb = nothing, ub = nothing)
+    transformed_prob = remake(
+        prob; f = new_f, u0 = u0_transformed, lb = nothing, ub = nothing,
+        build_initializeprob = Val{false}
+    )
 
     return transformed_prob
 end
+
+# Run problem initialization (if any) in the original bounded coordinates so that the
+# resulting `u0`/`p` are consistent before the bounds transform is applied. Returns the
+# problem unchanged when there is no `OverrideInitData` to run.
+function run_bounded_initialization(prob, alg)
+    SciMLBase.has_initialization_data(prob.f) || return prob
+    prob.f.initialization_data isa SciMLBase.OverrideInitData || return prob
+    iip = SciMLBase.isinplace(prob)
+    u0, p, _ = SciMLBase.get_initial_values(
+        prob, prob, prob.f, SciMLBase.OverrideInit(), Val(iip);
+        nlsolve_alg = alg, abstol = get_abstol(prob), reltol = get_reltol(prob)
+    )
+    return remake(prob; u0, p, build_initializeprob = Val{false})
+end

lib/NonlinearSolveBase/test/bounds_transform.jl

@@ -0,0 +1,31 @@
+module BoundsTransformTests
+
+using Test
+using NonlinearSolveBase
+using SciMLBase
+# Load ForwardDiff so PreallocationTools' `FixedSizeDiffCache` dual-cache extension is
+# active (the transform builds one for the default/ForwardDiff autodiff path).
+import ForwardDiff
+using NonlinearSolveBase: transform_bounded_problem, BoundedWrapper, _to_unbounded
+
+# An algorithm type without an `autodiff` field, mirroring `QuasiNewtonAlgorithm`
+# (which the default polyalgorithm reaches). `transform_bounded_problem` must not
+# assume every algorithm exposes `autodiff`.
+struct NoAutodiffAlg end
+
+@testset "bounds transform handles algorithms without an `autodiff` field" begin
+    f(u, p) = u .^ 2 .- p
+    prob = NonlinearProblem(f, [1.5, 1.5], [2.0, 2.0]; lb = [0.0, 0.0], ub = [3.0, 3.0])
+
+    for alg in (nothing, NoAutodiffAlg())
+        tprob = transform_bounded_problem(prob, alg)
+        @test tprob.f.f isa BoundedWrapper
+        # bounds are removed from the transformed (unconstrained) problem
+        @test tprob.lb === nothing
+        @test tprob.ub === nothing
+        # u0 is mapped into unbounded space
+        @test tprob.u0 ≈ _to_unbounded.(prob.u0, prob.lb, prob.ub)
+    end
+end
+
+end

lib/NonlinearSolveBase/test/runtests.jl

@@ -142,4 +142,8 @@ using InteractiveUtils, Test
     @testset "EnzymeExt algorithms are inactive_type" begin
         include("enzyme_inactive_algorithm.jl")
     end
+
+    @testset "Bounds transform (#955)" begin
+        include("bounds_transform.jl")
+    end
 end

test/bounds_tests.jl

@@ -157,3 +157,29 @@ end
         end
     end
 end
+
+@testitem "Bounds: polyalgorithm and quasi-Newton algorithms" tags = [:core, :bounds] begin
+    using SciMLBase
+    using NonlinearSolveBase
+
+    # The quasi-Newton `autodiff`-access fix lives in NonlinearSolveBase, so only
+    # exercise the full solve when a NonlinearSolveBase new enough to contain it is
+    # actually loaded. On Julia < 1.11 the umbrella resolves the *registered*
+    # NonlinearSolveBase (which predates the fix) because `[sources]` path redirects
+    # are ignored there; the fix itself is unit-tested in NonlinearSolveBase's own
+    # suite, which always runs against the in-repo code. See SciML/NonlinearSolve.jl#955.
+    if pkgversion(NonlinearSolveBase) >= v"2.30.3"
+        # `x^2 - 4x + 3` has roots at 1 and 3; bounds select which root is reachable.
+        f(u, p) = u .^ 2 .- 4 .* u .+ 3
+
+        # The default polyalgorithm tries quasi-Newton methods (which have no `autodiff`
+        # field) as part of its sequence; the bounds transform must not error on them.
+        for alg in (nothing, FastShortcutNonlinearPolyalg(), Broyden(), Klement(), NewtonRaphson())
+            prob = NonlinearProblem(f, [1.5], nothing; lb = [0.0], ub = [2.0])
+            sol = alg === nothing ? solve(prob) : solve(prob, alg)
+            @test SciMLBase.successful_retcode(sol)
+            @test sol.u[1] ≈ 1.0 atol = 1.0e-6
+            @test 0.0 <= sol.u[1] <= 2.0
+        end
+    end
+end