Fix Float32/static-array mass-matrix OOP Rosenbrock JacReuseState W typing

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -1,5 +1,5 @@
 name = "OrdinaryDiffEqDifferentiation"
-version = "3.2.0"
+version = "3.2.1"
 uuid = "4302a76b-040a-498a-8c04-15b101fed76b"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>", "Yingbo Ma <mayingbo5@gmail.com>"]
 

lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

@@ -937,12 +937,17 @@ function calc_rosenbrock_differentiation(integrator, cache, dtgamma, repeat_step
     mass_matrix = integrator.f.mass_matrix
     update_coefficients!(mass_matrix, integrator.uprev, integrator.p, integrator.t)
 
-    # Rebuild W from cached J and current dtgamma.
+    # Rebuild W from cached J and current dtgamma, mirroring calc_W's branches
+    # so W (and hence the cached_W slot) keeps the same concrete type.
     # The Jacobian evaluation is the expensive part; W = J − M/(dt·γ) and
     # its factorization are comparatively cheap and keep step control accurate.
-    W = J - mass_matrix * inv(dtgamma)
-    if !isa(W, Number)
-        W = DiffEqBase.default_factorize(W)
+    if cache.W isa StaticWOperator
+        W = StaticWOperator(J - mass_matrix * inv(dtgamma))
+    else
+        W = J - mass_matrix * inv(dtgamma)
+        if !isa(W, Number)
+            W = DiffEqBase.default_factorize(W)
+        end
     end
     integrator.stats.nw += 1
     jac_reuse.cached_W = W
@@ -1067,13 +1072,38 @@ function update_W!(
     return nothing
 end
 
+"""
+    _dtgamma_prototype(t, dt, uEltypeNoUnits)
+
+A value carrying the type `dtgamma = dt * γ` has at solve time: `dt`'s full
+(possibly Dual) time type promoted with the tableau eltype
+(`constvalue(uEltypeNoUnits)`, which `constvalue` strips of Duals the same way
+tableau constructors do). Non-`Number` eltypes (e.g. array-of-array states)
+have no tableau-eltype contribution to promote with, so `dt`'s type is used
+as-is.
+"""
+@inline function _dtgamma_prototype(t, dt, ::Type{T}) where {T <: Number}
+    return promote(t, dt)[2] * one(constvalue(T))
+end
+@inline _dtgamma_prototype(t, dt, ::Type{T}) where {T} = promote(t, dt)[2]
+
 function build_J_W(
         alg, u, uprev, p, t, dt, f::F, jac_config, ::Type{uEltypeNoUnits},
         ::Val{IIP}
     ) where {IIP, uEltypeNoUnits, F}
     # TODO - make J, W AbstractSciMLOperators (lazily defined with scimlops functionality)
     # TODO - if jvp given, make it SciMLOperators.FunctionOperator
     # TODO - make mass matrix a SciMLOperator so it can be updated with time. Default to IdentityOperator
+    #
+    # W must already carry the type `calc_W` produces at solve time: there
+    # W = J - mass_matrix * inv(dtgamma) promotes the eltype past the state
+    # eltype (e.g. a Float32 state over a Float64 tspan gives a Float64 W),
+    # and WOperator's gamma slot holds dtgamma itself. OOP caches store W
+    # concretely (e.g. JacReuseState.cached_W, Newton caches), so a
+    # state-eltype W here would reject calc_W's result on the first
+    # assignment.
+    dtgamma_prototype = _dtgamma_prototype(t, dt, uEltypeNoUnits)
+    invdtgamma_prototype = inv(oneunit(dtgamma_prototype))
     islin, isode = islinearfunction(f, alg)
     if isdefined(f, :W_prototype) && (f.W_prototype isa AbstractSciMLOperator)
         # We use W_prototype when it is provided as a SciMLOperator, and in this case we require jac_prototype to be a SciMLOperator too.
@@ -1088,10 +1118,10 @@ function build_J_W(
             @assert SciMLBase.has_jac(f) "f needs to have an associated jacobian"
             J = MatrixOperator(J; update_func! = f.jac)
         end
-        W = WOperator{IIP}(f.mass_matrix, promote(t, dt)[2], J, _vec(u))
+        W = WOperator{IIP}(f.mass_matrix, dtgamma_prototype, J, _vec(u))
     elseif islin
         J = isode ? f.f : f.f1.f # unwrap the Jacobian accordingly
-        W = WOperator{IIP}(f.mass_matrix, dt, J, _vec(u))
+        W = WOperator{IIP}(f.mass_matrix, dtgamma_prototype, J, _vec(u))
     elseif IIP && f.jac_prototype !== nothing && concrete_jac(alg) === nothing &&
             (alg.linsolve === nothing || LinearSolve.needs_concrete_A(alg.linsolve))
 
@@ -1109,7 +1139,7 @@ function build_J_W(
         jacvec = JVPCache(f, copy(u), u, p, t, autodiff = alg_autodiff(alg))
 
         J = jacvec
-        W = WOperator{IIP}(f.mass_matrix, promote(t, dt)[2], J, _vec(u), jacvec)
+        W = WOperator{IIP}(f.mass_matrix, dtgamma_prototype, J, _vec(u), jacvec)
     elseif alg.linsolve !== nothing && !LinearSolve.needs_concrete_A(alg.linsolve) ||
             concrete_jac(alg) !== nothing && concrete_jac(alg)
         # The linear solver does not need a concrete Jacobian, but the user has
@@ -1136,11 +1166,14 @@ function build_J_W(
             deepcopy(f.jac_prototype)
         end
         W = if J isa StaticMatrix
-            StaticWOperator(J, false)
+            # callinv = false skips inverting the seed-valued matrix while
+            # producing the same concrete type as calc_W's
+            # `StaticWOperator(J - mass_matrix * inv(dtgamma))`.
+            StaticWOperator(J - f.mass_matrix * invdtgamma_prototype, false)
         else
             jacvec = JVPCache(f, copy(u), u, p, t, autodiff = alg_autodiff(alg))
 
-            WOperator{IIP}(f.mass_matrix, promote(t, dt)[2], J, _vec(u), jacvec)
+            WOperator{IIP}(f.mass_matrix, dtgamma_prototype, J, _vec(u), jacvec)
         end
     else
         J = if !IIP && SciMLBase.has_jac(f)
@@ -1166,6 +1199,8 @@ function build_J_W(
             deepcopy(f.jac_prototype)
         end
         W = if alg isa DAEAlgorithm
+            # DAE W is not built from J - M/(dt·γ) (calc_W returns J or a
+            # cj-scaled combination), so no dtgamma promotion applies.
             if IIP
                 similar(J)
             elseif J isa StaticMatrix
@@ -1176,9 +1211,9 @@ function build_J_W(
         elseif IIP
             similar(J)
         elseif J isa StaticMatrix
-            StaticWOperator(J, false)
+            StaticWOperator(J - f.mass_matrix * invdtgamma_prototype, false)
         else
-            ArrayInterface.lu_instance(J)
+            ArrayInterface.lu_instance(J - f.mass_matrix * invdtgamma_prototype)
         end
     end
     return J, W

lib/OrdinaryDiffEqRosenbrock/Project.toml

@@ -1,6 +1,6 @@
 name = "OrdinaryDiffEqRosenbrock"
 uuid = "43230ef6-c299-4910-a778-202eb28ce4ce"
-version = "2.3.0"
+version = "2.3.1"
 authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
 
 [deps]
@@ -46,7 +46,7 @@ MacroTools = "0.5.6"
 MuladdMacro = "0.2.1"
 ODEProblemLibrary = "1"
 OrdinaryDiffEqCore = "4"
-OrdinaryDiffEqDifferentiation = "3"
+OrdinaryDiffEqDifferentiation = "3.2.1"
 OrdinaryDiffEqNonlinearSolve = "2"
 OrdinaryDiffEqRosenbrockTableaus = "2.1"
 Pkg = "1"
@@ -56,6 +56,7 @@ Random = "<0.0.1, 1"
 RecursiveArrayTools = "4"
 Reexport = "1.2.2"
 SafeTestsets = "0.1"
+StaticArrays = "1.9.18"
 SciMLBase = "3.1.0"
 Test = "<0.0.1, 1"
 julia = "1.10"
@@ -68,7 +69,8 @@ OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DiffEqDevTools", "Random", "OrdinaryDiffEqNonlinearSolve", "SafeTestsets", "Test", "ODEProblemLibrary", "Enzyme", "Pkg"]
+test = ["DiffEqDevTools", "Random", "OrdinaryDiffEqNonlinearSolve", "SafeTestsets", "StaticArrays", "Test", "ODEProblemLibrary", "Enzyme", "Pkg"]

lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_caches.jl

@@ -11,7 +11,7 @@ Jacobian recomputations when conditions allow it.
 `T` is the element type of `dtgamma` at solve time — usually `Float64`, but
 `ForwardDiff.Dual` (including nested Duals under `ForwardDiff.hessian`) when
 the solve is being differentiated. Callers must construct this with a zero
-value of the right type (typically `zero(dt)`) so the same cache can hold
+value of the right type (typically `zero(dt * gamma)`) so the same cache can hold
 Dual-valued dtgammas without any `ForwardDiff.value` unwrapping — unconditional
 unwrapping is unsafe under nested AD because it collapses a still-active
 inner derivative into a primal.
@@ -342,13 +342,15 @@ function alg_cache(
     linprob = nothing #LinearProblem(W,copy(u); u0=copy(u))
     linsolve = nothing #init(linprob,alg.linsolve,alias_A=true,alias_b=true)
     tab = Rosenbrock23Tableau(constvalue(uBottomEltypeNoUnits))
-    # Seed JacReuseState with `zero(dt)` rather than a `constvalue`-stripped
-    # type: under nested ForwardDiff (e.g. `hessian`), dt is a Dual-of-Dual and
-    # dtgamma inherits that full type; a Float64 field would reject the assign.
+    # Seed JacReuseState with `zero(dt * tab.d)` rather than a
+    # `constvalue`-stripped type: this carries the exact type dtgamma has at
+    # solve time, including Dual-of-Dual under nested ForwardDiff
+    # (e.g. `hessian`), where a Float64 field would reject the assign.
+    dtgamma_seed = zero(dt * tab.d)
     return Rosenbrock23ConstantCache(
         tab.c₃₂, tab.d, tf, uf, J, W, linsolve, alg_autodiff(alg),
         _make_jac_reuse_state(
-            zero(dt), alg.max_jac_age, J, zero(rate_prototype), W
+            dtgamma_seed, alg.max_jac_age, J, zero(rate_prototype), W
         )
     )
 end
@@ -378,12 +380,13 @@ function alg_cache(
     linprob = nothing #LinearProblem(W,copy(u); u0=copy(u))
     linsolve = nothing #init(linprob,alg.linsolve,alias_A=true,alias_b=true)
     tab = Rosenbrock32Tableau(constvalue(uBottomEltypeNoUnits))
-    # See the Rosenbrock23 OOP alg_cache above for why we pass `zero(dt)` here
-    # rather than a `constvalue`-stripped type.
+    # See the Rosenbrock23 OOP alg_cache above for why we pass `zero(dt * tab.d)`
+    # here rather than a `constvalue`-stripped type.
+    dtgamma_seed = zero(dt * tab.d)
     return Rosenbrock32ConstantCache(
         tab.c₃₂, tab.d, tf, uf, J, W, linsolve, alg_autodiff(alg),
         _make_jac_reuse_state(
-            zero(dt), alg.max_jac_age, J, zero(rate_prototype), W
+            dtgamma_seed, alg.max_jac_age, J, zero(rate_prototype), W
         )
     )
 end
@@ -475,14 +478,15 @@ function alg_cache(
     else
         interp_order = H_rows
     end
-    # Seed JacReuseState with `zero(dt)` so its dtgamma fields carry the full
-    # (possibly ForwardDiff.Dual) type dtgamma will have at solve time.
+    # Seed JacReuseState with `zero(dt * tab.gamma)` so its dtgamma fields carry
+    # the full (possibly ForwardDiff.Dual) type dtgamma will have at solve time.
+    dtgamma_seed = zero(dt * tab.gamma)
     return RosenbrockCombinedConstantCache(
         tf, uf,
         tab, J, W, linsolve,
         alg_autodiff(alg), interp_order,
         _make_jac_reuse_state(
-            zero(dt), alg.max_jac_age, J, zero(rate_prototype), W
+            dtgamma_seed, alg.max_jac_age, J, zero(rate_prototype), W
         )
     )
 end

lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl

@@ -429,3 +429,73 @@ end
         @test err_interp < 100 * max(err_knot, 1.0e-4)
     end
 end
+
+# https://github.com/SciML/OrdinaryDiffEq.jl/issues/3721
+# OOP static-array mass-matrix solves were failing in two ways:
+# 1. With a Float32 state and a Float64 time type (e.g. `dt = 0.1` promoting
+#    tspan), calc_W produces a Float64-eltype StaticWOperator that cannot be
+#    assigned into JacReuseState's Float32-seeded `cached_W` slot.
+# 2. On rejected-step retries the OOP reuse branch rebuilt W via
+#    `default_factorize`, producing a `StaticArrays.LU` instead of the
+#    `StaticWOperator` type of the `cached_W` slot.
+using StaticArrays
+using OrdinaryDiffEqNonlinearSolve: BrownFullBasicInit
+
+@testset "Static-array mass-matrix OOP JacReuseState (#3721)" begin
+    function rober_static(u, p, t)
+        y₁, y₂, y₃ = u
+        k₁, k₂, k₃ = p
+        return @SVector [
+            -k₁ * y₁ + k₃ * y₂ * y₃,
+            k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃,
+            y₁ + y₂ + y₃ - 1,
+        ]
+    end
+
+    # Float64 reference
+    M64 = @SMatrix [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0]
+    prob64 = ODEProblem(
+        ODEFunction(rober_static, mass_matrix = M64),
+        @SVector([1.0, 0.0, 0.0]), (0.0, 1.0e5), (0.04, 3.0e7, 1.0e4)
+    )
+    ref = solve(
+        prob64, Rosenbrock23(), abstol = 1.0e-10, reltol = 1.0e-10,
+        initializealg = BrownFullBasicInit()
+    )
+
+    # Failure mode 1: Float32 state with Float64 dt (promotes the time type).
+    M32 = @SMatrix [1.0f0 0.0f0 0.0f0; 0.0f0 1.0f0 0.0f0; 0.0f0 0.0f0 0.0f0]
+    prob32 = ODEProblem(
+        ODEFunction(rober_static, mass_matrix = M32),
+        @SVector([1.0f0, 0.0f0, 0.0f0]), (0.0f0, 1.0f5), (0.04f0, 3.0f7, 1.0f4)
+    )
+    sol32 = solve(
+        prob32, Rosenbrock23(), dt = 0.1, abstol = 1.0f-5, reltol = 1.0f-5,
+        initializealg = BrownFullBasicInit()
+    )
+    @test SciMLBase.successful_retcode(sol32)
+    @test norm(sol32.u[end] - ref.u[end]) < 1.0e-4
+
+    # Failure mode 2: homogeneous Float64 at a tolerance that produces step
+    # rejections, whose retries take the cached-J reuse branch. nreject > 0 is
+    # asserted so this keeps covering that branch if step control changes.
+    sol64 = solve(
+        prob64, Rosenbrock23(), abstol = 1.0e-5, reltol = 1.0e-5,
+        initializealg = BrownFullBasicInit()
+    )
+    @test SciMLBase.successful_retcode(sol64)
+    @test sol64.stats.nreject > 0
+    @test norm(sol64.u[end] - ref.u[end]) < 1.0e-4
+
+    # The Rodas-family OOP constant cache goes through the same cached_W
+    # seeding path; Rodas23W and ROS34PW2 are W-methods so they additionally
+    # exercise the Jacobian-reuse decision logic on the static path.
+    for alg in (Rodas4(), Rodas23W(), ROS34PW2())
+        sol_rodas = solve(
+            prob32, alg, dt = 0.1, abstol = 1.0f-5, reltol = 1.0f-5,
+            initializealg = BrownFullBasicInit()
+        )
+        @test SciMLBase.successful_retcode(sol_rodas)
+        @test norm(sol_rodas.u[end] - ref.u[end]) < 1.0e-4
+    end
+end