Merge branch 'master' into dg/cr

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "2.35.0"
+version = "2.36.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -65,7 +65,7 @@ FunctionWrappersWrappers = "0.1.3"
 IteratorInterfaceExtensions = "^1"
 LinearAlgebra = "1.10"
 Logging = "1.10"
-Makie = "0.20"
+Makie = "0.20, 0.21"
 Markdown = "1.10"
 ModelingToolkit = "8.75, 9"
 PartialFunctions = "1.1"
@@ -76,15 +76,15 @@ PyCall = "1.96"
 PythonCall = "0.9.15"
 RCall = "0.14.0"
 RecipesBase = "1.3.4"
-RecursiveArrayTools = "3.8.0"
+RecursiveArrayTools = "3.14.0"
 Reexport = "1"
 RuntimeGeneratedFunctions = "0.5.12"
 SciMLOperators = "0.3.7"
 SciMLStructures = "1.1"
 StaticArrays = "1.7"
 StaticArraysCore = "1.4"
 Statistics = "1.10"
-SymbolicIndexingInterface = "0.3.15"
+SymbolicIndexingInterface = "0.3.20"
 Tables = "1.11"
 Zygote = "0.6.67"
 julia = "1.10"

ext/SciMLBaseChainRulesCoreExt.jl

@@ -3,6 +3,7 @@ module SciMLBaseChainRulesCoreExt
 using SciMLBase
 import ChainRulesCore
 import ChainRulesCore: NoTangent, @non_differentiable
+using SymbolicIndexingInterface
 
 function ChainRulesCore.rrule(
         config::ChainRulesCore.RuleConfig{
@@ -13,7 +14,7 @@ function ChainRulesCore.rrule(
         sym,
         j::Integer)
     function ODESolution_getindex_pullback(Δ)
-        i = symbolic_type(sym) != NotSymbolic() ? sym_to_index(sym, VA) : sym
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
             getter = getobserved(VA)
             grz = rrule_via_ad(config, getter, sym, VA.u[j], VA.prob.p, VA.t[j])[2](Δ)
@@ -66,7 +67,7 @@ end
 
 function ChainRulesCore.rrule(::typeof(getindex), VA::ODESolution, sym)
     function ODESolution_getindex_pullback(Δ)
-        i = symbolic_type(sym) != NotSymbolic() ? sym_to_index(sym, VA) : sym
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
             throw(error("AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))
         else
@@ -109,18 +110,19 @@ function ChainRulesCore.rrule(
     ODESolutionAdjoint
 end
 
-# function ChainRulesCore.rrule(
-#         ::Type{
-#             <:SDESolution{uType, tType, isinplace, P, NP, F, G, K,
-#             ND
-#         }}, u,
-#         args...) where {uType, tType, isinplace, P, NP, F, G, K, ND}
-#     function SDESolutionAdjoint(ȳ)
-#         (NoTangent(), ȳ, ntuple(_ -> NoTangent(), length(args))...)
-#     end
-# 
-#     SDESolution{uType, tType, isinplace, P, NP, F, G, K, ND}(u, args...), SDESolutionAdjoint
-# end
+function ChainRulesCore.rrule(
+        ::Type{
+            <:RODESolution{uType, tType, isinplace, P, NP, F, G, K,
+            ND
+        }}, u,
+        args...) where {uType, tType, isinplace, P, NP, F, G, K, ND}
+    function RODESolutionAdjoint(ȳ)
+        (NoTangent(), ȳ, ntuple(_ -> NoTangent(), length(args))...)
+    end
+
+    RODESolution{uType, tType, isinplace, P, NP, F, G, K, ND}(u, args...),
+    RODESolutionAdjoint
+end
 
 function ChainRulesCore.rrule(::SciMLBase.EnsembleSolution, sim, time, converged)
     out = EnsembleSolution(sim, time, converged)

ext/SciMLBaseZygoteExt.jl

@@ -13,7 +13,7 @@ using RecursiveArrayTools
 # This method resolves the ambiguity with the pullback defined in
 # RecursiveArrayToolsZygoteExt
 # https://github.com/SciML/RecursiveArrayTools.jl/blob/d06ecb856f43bc5e37cbaf50e5f63c578bf3f1bd/ext/RecursiveArrayToolsZygoteExt.jl#L67
-@adjoint function getindex(VA::ODESolution, i::Int, j::Int)
+@adjoint function Base.getindex(VA::ODESolution, i::Int, j::Int)
     function ODESolution_getindex_pullback(Δ)
         du = [m == j ? [i == k ? Δ : zero(VA.u[1][1]) for k in 1:length(VA.u[1])] :
               zero(VA.u[1]) for m in 1:length(VA.u)]
@@ -38,7 +38,7 @@ using RecursiveArrayTools
     VA[i, j], ODESolution_getindex_pullback
 end
 
-@adjoint function getindex(VA::ODESolution, sym, j::Int)
+@adjoint function Base.getindex(VA::ODESolution, sym, j::Int)
     function ODESolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         du, dprob = if i === nothing
@@ -92,7 +92,7 @@ end
     out, EnsembleSolution_adjoint
 end
 
-@adjoint function getindex(VA::ODESolution, i::Int)
+@adjoint function Base.getindex(VA::ODESolution, i::Int)
     function ODESolution_getindex_pullback(Δ)
         Δ′ = [(i == j ? Δ : Zygote.FillArrays.Fill(zero(eltype(x)), size(x)))
               for (x, j) in zip(VA.u, 1:length(VA))]
@@ -106,7 +106,7 @@ end
     sim.u, p̄ -> (EnsembleSolution(p̄, 0.0, true, sim.stats),)
 end
 
-@adjoint function getindex(VA::ODESolution, sym)
+@adjoint function Base.getindex(VA::ODESolution, sym)
     function ODESolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
@@ -120,6 +120,30 @@ end
     VA[sym], ODESolution_getindex_pullback
 end
 
+@adjoint function Base.getindex(
+        VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}) where {T}
+    function ODESolution_getindex_pullback(Δ)
+        sym = sym isa Tuple ? collect(sym) : sym
+        i = map(x -> symbolic_type(x) != NotSymbolic() ? variable_index(VA, x) : x, sym)
+        if i === nothing
+            throw(error("Zygote AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))
+        else
+            Δ′ = map(enumerate(VA.u)) do (t_idx, us)
+                map(enumerate(us)) do (u_idx, u)
+                    if u_idx in i
+                        idx = findfirst(isequal(u_idx), i)
+                        Δ[t_idx][idx]
+                    else
+                        zero(T)
+                    end
+                end
+            end
+            (Δ′, nothing)
+        end
+    end
+    VA[sym], ODESolution_getindex_pullback
+end
+
 @adjoint function ODESolution{T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12
 }(u,
         args...) where {T1, T2, T3, T4, T5, T6, T7, T8,
@@ -135,11 +159,11 @@ end
 @adjoint function SDEProblem{uType, tType, isinplace, P, NP, F, G, K, ND}(u,
         args...) where
         {uType, tType, isinplace, P, NP, F, G, K, ND}
-    function SDESolutionAdjoint(ȳ)
+    function SDEProblemAdjoint(ȳ)
         (ȳ, ntuple(_ -> nothing, length(args))...)
     end
 
-    SDESolution{uType, tType, isinplace, P, NP, F, G, K, ND}(u, args...), SDESolutionAdjoint
+    SDEProblem{uType, tType, isinplace, P, NP, F, G, K, ND}(u, args...), SDEProblemAdjoint
 end
 
 @adjoint function NonlinearSolution{T, N, uType, R, P, A, O, uType2}(u,

src/ensemble/ensemble_solutions.jl

@@ -211,8 +211,7 @@ end
     end
 end
 
-Base.@propagate_inbounds function RecursiveArrayTools._getindex(
-        x::AbstractEnsembleSolution, ::Union{ScalarSymbolic, ArraySymbolic}, s, ::Colon)
+Base.@propagate_inbounds function Base.getindex(x::AbstractEnsembleSolution, s, ::Colon)
     return [xi[s] for xi in x.u]
 end
 

src/integrator_interface.jl

@@ -643,8 +643,8 @@ Same as `check_error` but also set solution's return code
 """
 function check_error!(integrator::DEIntegrator)
     code = check_error(integrator)
+    integrator.sol = solution_new_retcode(integrator.sol, code)
     if code != ReturnCode.Success
-        integrator.sol = solution_new_retcode(integrator.sol, code)
         postamble!(integrator)
     end
     return code

src/solutions/ode_solutions.jl

@@ -146,6 +146,18 @@ function ODESolution{T, N}(u, u_analytic, errors, t, k, prob, alg, interp, dense
         dense, tslocation, stats, alg_choice, retcode, resid, original)
 end
 
+error_if_observed_derivative(_, _, ::Type{Val{0}}) = nothing
+function error_if_observed_derivative(sys, idx, ::Type)
+    if symbolic_type(idx) != NotSymbolic() && is_observed(sys, idx) ||
+       symbolic_type(idx) == NotSymbolic() && any(x -> is_observed(sys, x), idx)
+        error("""
+        Cannot interpolate derivatives of observed variables. A possible solution could be
+        interpolating the symbolic expression that evaluates to the derivative of the
+        observed variable or using DataInterpolations.jl.
+        """)
+    end
+end
+
 function (sol::AbstractODESolution)(t, ::Type{deriv} = Val{0}; idxs = nothing,
         continuity = :left) where {deriv}
     sol(t, deriv, idxs, continuity)
@@ -172,6 +184,9 @@ end
 function (sol::AbstractODESolution)(t::Number, ::Type{deriv},
         idxs::AbstractVector{<:Integer},
         continuity) where {deriv}
+    if eltype(sol.u) <: Number
+        idxs = only(idxs)
+    end
     sol.interp(t, idxs, deriv, sol.prob.p, continuity)
 end
 function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
@@ -183,6 +198,9 @@ end
 function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
         idxs::AbstractVector{<:Integer},
         continuity) where {deriv}
+    if eltype(sol.u) <: Number
+        idxs = only(idxs)
+    end
     A = sol.interp(t, idxs, deriv, sol.prob.p, continuity)
     p = hasproperty(sol.prob, :p) ? sol.prob.p : nothing
     return DiffEqArray(A.u, A.t, p, sol)
@@ -191,6 +209,7 @@ end
 function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs,
         continuity) where {deriv}
     symbolic_type(idxs) == NotSymbolic() && error("Incorrect specification of `idxs`")
+    error_if_observed_derivative(sol, idxs, deriv)
     if is_parameter(sol, idxs)
         return getp(sol, idxs)(sol)
     else
@@ -200,15 +219,19 @@ end
 
 function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs::AbstractVector,
         continuity) where {deriv}
-    all(!isequal(NotSymbolic()), symbolic_type.(idxs)) ||
+    if symbolic_type(idxs) == NotSymbolic() &&
+       any(isequal(NotSymbolic()), symbolic_type.(idxs))
         error("Incorrect specification of `idxs`")
+    end
+    error_if_observed_derivative(sol, idxs, deriv)
     interp_sol = augment(sol.interp([t], nothing, deriv, sol.prob.p, continuity), sol)
-    [is_parameter(sol, idx) ? getp(sol, idx)(sol) : first(interp_sol[idx]) for idx in idxs]
+    first(interp_sol[idxs])
 end
 
 function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv}, idxs,
         continuity) where {deriv}
     symbolic_type(idxs) == NotSymbolic() && error("Incorrect specification of `idxs`")
+    error_if_observed_derivative(sol, idxs, deriv)
     if is_parameter(sol, idxs)
         return getp(sol, idxs)(sol)
     else
@@ -222,10 +245,12 @@ function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
         idxs::AbstractVector, continuity) where {deriv}
     all(!isequal(NotSymbolic()), symbolic_type.(idxs)) ||
         error("Incorrect specification of `idxs`")
+    error_if_observed_derivative(sol, idxs, deriv)
     interp_sol = augment(sol.interp(t, nothing, deriv, sol.prob.p, continuity), sol)
     p = hasproperty(sol.prob, :p) ? sol.prob.p : nothing
+    indexed_sol = interp_sol[idxs]
     return DiffEqArray(
-        [[interp_sol[idx][i] for idx in idxs] for i in 1:length(t)], t, p, sol)
+        [indexed_sol[i] for i in 1:length(t)], t, p, sol)
 end
 
 function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem},

src/solutions/solution_interface.jl

@@ -436,6 +436,18 @@ function solplot_vecs_and_labels(dims, vars, plott, sol, plot_analytic,
     plot_vecs = []
     labels = String[]
     varsyms = variable_symbols(sol)
+    batch_symbolic_vars = []
+    for x in vars
+        for j in 2:length(x)
+            if (x[j] isa Integer && x[j] == 0) || isequal(x[j], getindepsym_defaultt(sol))
+            else
+                push!(batch_symbolic_vars, x[j])
+            end
+        end
+    end
+    batch_symbolic_vars = identity.(batch_symbolic_vars)
+    indexed_solution = sol(plott; idxs = batch_symbolic_vars)
+    idxx = 0
     for x in vars
         tmp = []
         strs = String[]
@@ -444,7 +456,8 @@ function solplot_vecs_and_labels(dims, vars, plott, sol, plot_analytic,
                 push!(tmp, plott)
                 push!(strs, "t")
             else
-                push!(tmp, sol(plott; idxs = x[j]))
+                idxx += 1
+                push!(tmp, indexed_solution[idxx, :])
                 if !isempty(varsyms) && x[j] isa Integer
                     push!(strs, String(getname(varsyms[x[j]])))
                 elseif hasname(x[j])

test/downstream/adjoints.jl

@@ -0,0 +1,92 @@
+using ModelingToolkit, OrdinaryDiffEq, SymbolicIndexingInterface, Zygote, Test
+using ModelingToolkit: t_nounits as t, D_nounits as D
+
+@parameters σ ρ β
+@variables x(t) y(t) z(t)
+
+eqs = [D(x) ~ σ * (y - x),
+    D(y) ~ x * (ρ - z) - y,
+    D(z) ~ x * y - β * z]
+
+@named lorenz1 = ODESystem(eqs, t)
+@named lorenz2 = ODESystem(eqs, t)
+
+@parameters γ
+@variables a(t) α(t)
+connections = [0 ~ lorenz1.x + lorenz2.y + a * γ,
+    α ~ 2lorenz1.x + a * γ]
+@mtkbuild sys = ODESystem(connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
+
+u0 = [lorenz1.x => 1.0,
+    lorenz1.y => 0.0,
+    lorenz1.z => 0.0,
+    lorenz2.x => 0.0,
+    lorenz2.y => 1.0,
+    lorenz2.z => 0.0,
+    a => 2.0]
+
+p = [lorenz1.σ => 10.0,
+    lorenz1.ρ => 28.0,
+    lorenz1.β => 8 / 3,
+    lorenz2.σ => 10.0,
+    lorenz2.ρ => 28.0,
+    lorenz2.β => 8 / 3,
+    γ => 2.0]
+
+tspan = (0.0, 100.0)
+prob = ODEProblem(sys, u0, tspan, p)
+sol = solve(prob, Rodas4())
+
+gs_sym, = Zygote.gradient(sol) do sol
+    sum(sol[lorenz1.x])
+end
+idx_sym = SymbolicIndexingInterface.variable_index(sys, lorenz1.x)
+true_grad_sym = zeros(length(ModelingToolkit.unknowns(sys)))
+true_grad_sym[idx_sym] = 1.0
+
+@test all(map(x -> x == true_grad_sym, gs_sym))
+
+gs_vec, = Zygote.gradient(sol) do sol
+    sum(sum.(sol[[lorenz1.x, lorenz2.x]]))
+end
+idx_vecsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lorenz2.x])
+true_grad_vecsym = zeros(length(ModelingToolkit.unknowns(sys)))
+true_grad_vecsym[idx_vecsym] .= 1.0
+
+@test all(map(x -> x == true_grad_vecsym, gs_vec))
+
+gs_tup, = Zygote.gradient(sol) do sol
+    sum(sum.(collect.(sol[(lorenz1.x, lorenz2.x)])))
+end
+idx_tupsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lorenz2.x])
+true_grad_tupsym = zeros(length(ModelingToolkit.unknowns(sys)))
+true_grad_tupsym[idx_tupsym] .= 1.0
+
+@test all(map(x -> x == true_grad_tupsym, gs_tup))
+
+gs_ts, = Zygote.gradient(sol) do sol
+    sum(sum.(sol[[lorenz1.x, lorenz2.x], :]))
+end
+
+@test all(map(x -> x == true_grad_vecsym, gs_ts))
+
+# BatchedInterface AD
+@variables x(t)=1.0 y(t)=1.0 z(t)=1.0 w(t)=1.0
+@named sys1 = ODESystem([D(x) ~ x + y, D(y) ~ y * z, D(z) ~ z * t * x], t)
+sys1 = complete(sys1)
+prob1 = ODEProblem(sys1, [], (0.0, 10.0))
+@named sys2 = ODESystem([D(x) ~ x + w, D(y) ~ w * t, D(w) ~ x + y + w], t)
+sys2 = complete(sys2)
+prob2 = ODEProblem(sys2, [], (0.0, 10.0))
+
+bi = BatchedInterface((sys1, [x, y, z]), (sys2, [x, y, w]))
+getter = getu(bi)
+
+p1grad, p2grad = Zygote.gradient(prob1, prob2) do prob1, prob2
+    sum(getter(prob1, prob2))
+end
+
+@test p1grad.u0 ≈ ones(3)
+testp2grad = zeros(3)
+testp2grad[variable_index(prob2, w)] = 1.0
+@test p2grad.u0 ≈ testp2grad