Merge branch 'SciML:master' into dg/paramgrad

Author: DhairyaLGandhi

.github/workflows/Downgrade.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        version: ['1']
+        version: ['min']
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <[email protected]> and contributors"]
-version = "2.79.0"
+version = "2.82.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -60,11 +60,10 @@ ChainRules = "1.58.0"
 ChainRulesCore = "1.18"
 CommonSolve = "0.2.4"
 ConstructionBase = "1.5"
-DataFrames = "1.6"
 Distributed = "1.10"
 DocStringExtensions = "0.9"
 EnumX = "1"
-ForwardDiff = "0.10.36"
+ForwardDiff = "0.10.36, 1"
 FunctionWrappersWrappers = "0.1.3"
 IteratorInterfaceExtensions = "^1"
 LinearAlgebra = "1.10"
@@ -73,7 +72,6 @@ MLStyle = "0.4.17"
 Makie = "0.20, 0.21, 0.22"
 Markdown = "1.10"
 Moshi = "0.3"
-NonlinearSolve = "3, 4"
 PartialFunctions = "1.1"
 PrecompileTools = "1.2"
 Preferences = "1.3"
@@ -99,12 +97,8 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
-NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -113,11 +107,10 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Pkg", "Plots", "UnicodePlots", "SafeTestsets", "Serialization", "Test", "StableRNGs", "StaticArrays", "StochasticDiffEq", "Aqua", "Zygote", "PartialFunctions", "DataFrames", "NonlinearSolve", "OrdinaryDiffEq", "ForwardDiff", "Tables", "MLStyle"]
+test = ["Aqua", "ForwardDiff", "MLStyle", "PartialFunctions", "Pkg", "Plots", "SafeTestsets", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"]

ext/SciMLBaseZygoteExt.jl

@@ -116,6 +116,14 @@ function obs_grads(VA, sym, obs_idx, Δ)
     back(Δobs)
 end
 
+function obs_grads2(VA::SciMLBase.NonlinearSolution, sym, obs_idx, Δ)
+    y, back = Zygote.pullback(VA) do sol
+        getindex.(Ref(sol), sym[obs_idx])
+    end
+    Δobs = Δ[obs_idx, :]
+    back(Δobs)
+end
+
 function obs_grads(VA, sym, ::Nothing, Δ)
     Zygote.nt_nothing(VA)
 end
@@ -154,6 +162,31 @@ end
     VA[sym], ODESolution_getindex_pullback
 end
 
+@adjoint function Base.getindex(VA::SciMLBase.NonlinearSolution, sym)
+    function NonlinearSolution_getindex_pullback(Δ)
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
+        if is_observed(VA, sym)
+            f = observed(VA, sym)
+            p = parameter_values(VA)
+            u = state_values(VA)
+            _, back = Zygote.pullback(u, p) do u, p
+                f.f_oop(u, p)
+            end
+            gs = back(Δ)
+            ((u = gs[1], prob = (p = gs[2],),), nothing)
+        elseif 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
+            VA = recursivecopy(VA)
+            recursivefill!(VA, zero(eltype(VA)))
+            v = view(VA, i, ntuple(_ -> :, ndims(VA) - 1)...)
+            copyto!(v, Δ)
+            (VA, nothing)
+        end
+    end
+    VA[sym], NonlinearSolution_getindex_pullback
+end
+
 @adjoint function ODESolution{
         T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15}(u,
         args...) where {T1, T2, T3, T4, T5, T6, T7, T8,

src/problems/rode_problems.jl

@@ -110,7 +110,6 @@ When a keyword argument is `nothing`, the default behaviour of the solver is use
 * `alias::Union{Bool, Nothing}`: sets all fields of the `RODEAliasSpecifier` to `alias`
 
 """
-
 struct RODEAliasSpecifier <: AbstractAliasSpecifier
     alias_p::Union{Bool, Nothing}
     alias_f::Union{Bool, Nothing}

src/remake.jl

@@ -191,7 +191,7 @@ function remake(
                 props = @delete props._func_cache
                 props = @insert props._func_cache = forig._func_cache
             end
-            
+
             args = (args..., f2)
         end
     end

src/scimlfunctions.jl

@@ -3067,7 +3067,7 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
                    f.observed :
                    DEFAULT_OBSERVED,
         resid_prototype = __has_resid_prototype(f) ?
-                          f.resid_prototype : 
+                          f.resid_prototype :
                           nothing,
         sys = __has_sys(f) ? f.sys : nothing,
         initialization_data = __has_initialization_data(f) ? f.initialization_data :
@@ -3083,11 +3083,12 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
             analytic,
             observed,
             sys,
-            resid_prototype, 
+            resid_prototype,
             initialization_data)
     else
         ImplicitDiscreteFunction{
-                                 iip, specialize, typeof(_f), typeof(analytic), typeof(observed), typeof(sys), typeof(resid_prototype),
+            iip, specialize, typeof(_f), typeof(analytic),
+            typeof(observed), typeof(sys), typeof(resid_prototype),
             typeof(initialization_data)}(
             _f, analytic, observed, sys, resid_prototype, initialization_data)
     end
@@ -3097,7 +3098,9 @@ function ImplicitDiscreteFunction{iip}(f; kwargs...) where {iip}
     ImplicitDiscreteFunction{iip, FullSpecialize}(f; kwargs...)
 end
 ImplicitDiscreteFunction{iip}(f::ImplicitDiscreteFunction; kwargs...) where {iip} = f
-function ImplicitDiscreteFunction(f; resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing, kwargs...)
+function ImplicitDiscreteFunction(
+        f; resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing,
+        kwargs...)
     ImplicitDiscreteFunction{isinplace(f, 5), FullSpecialize}(f; resid_prototype, kwargs...)
 end
 ImplicitDiscreteFunction(f::ImplicitDiscreteFunction; kwargs...) = f
@@ -3107,11 +3110,13 @@ function unwrapped_f(f::ImplicitDiscreteFunction, newf = unwrapped_f(f.f))
 
     if specialize === NoSpecialize
         ImplicitDiscreteFunction{isinplace(f, 5), specialize, Any, Any, Any,
-            Any, Any, Any}(newf, f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
+            Any, Any, Any}(
+            newf, f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     else
         ImplicitDiscreteFunction{isinplace(f, 5), specialize, typeof(newf),
             typeof(f.analytic),
-            typeof(f.observed), typeof(f.sys), typeof(resid_prototype), typeof(f.initialization_data)}(newf,
+            typeof(f.observed), typeof(f.sys), typeof(resid_prototype), typeof(f.initialization_data)}(
+            newf,
             f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     end
 end

src/solutions/rode_solutions.jl

@@ -32,14 +32,15 @@ https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/
   exited due to an error. For more details, see
   [the return code documentation](https://docs.sciml.ai/SciMLBase/stable/interfaces/Solutions/#retcodes).
 """
-struct RODESolution{T, N, uType, uType2, DType, tType, randType, P, A, IType, S,
+struct RODESolution{T, N, uType, uType2, DType, tType, randType, discType, P, A, IType, S,
     AC <: Union{Nothing, Vector{Int}}, V} <:
        AbstractRODESolution{T, N, uType}
     u::uType
     u_analytic::uType2
     errors::DType
     t::tType
     W::randType
+    discretes::discType
     prob::P
     alg::A
     interp::IType
@@ -63,9 +64,10 @@ function ConstructionBase.setproperties(sol::RODESolution, patch::NamedTuple)
     patch = merge(getproperties(sol), patch)
     return RODESolution{
         T, N, typeof(patch.u), typeof(patch.u_analytic), typeof(patch.errors),
-        typeof(patch.t), typeof(patch.W), typeof(patch.prob), typeof(patch.alg), typeof(patch.interp),
+        typeof(patch.t), typeof(patch.W), typeof(patch.discretes),
+        typeof(patch.prob), typeof(patch.alg), typeof(patch.interp),
         typeof(patch.stats), typeof(patch.alg_choice), typeof(patch.saved_subsystem)}(
-        patch.u, patch.u_analytic, patch.errors, patch.t, patch.W,
+        patch.u, patch.u_analytic, patch.errors, patch.t, patch.W, patch.discretes,
         patch.prob, patch.alg, patch.interp, patch.dense, patch.tslocation, patch.stats,
         patch.alg_choice, patch.retcode, patch.seed, patch.saved_subsystem)
 end
@@ -120,16 +122,28 @@ function build_solution(prob::Union{AbstractRODEProblem, AbstractSDDEProblem},
         Base.depwarn(msg, :build_solution)
     end
 
+    ps = parameter_values(prob)
+    if has_sys(prob.f)
+        sswf = if saved_subsystem === nothing
+            prob.f.sys
+        else
+            SavedSubsystemWithFallback(saved_subsystem, prob.f.sys)
+        end
+        discretes = create_parameter_timeseries_collection(sswf, ps, prob.tspan)
+    else
+        discretes = nothing
+    end
     if has_analytic(f)
         u_analytic = Vector{typeof(prob.u0)}()
         errors = Dict{Symbol, real(eltype(prob.u0))}()
         sol = RODESolution{T, N, typeof(u), typeof(u_analytic), typeof(errors), typeof(t),
-            typeof(W),
+            typeof(W), typeof(discretes),
             typeof(prob), typeof(alg), typeof(interp), typeof(stats),
             typeof(alg_choice), typeof(saved_subsystem)}(u,
             u_analytic,
             errors,
             t, W,
+            discretes,
             prob,
             alg,
             interp,
@@ -149,15 +163,37 @@ function build_solution(prob::Union{AbstractRODEProblem, AbstractSDDEProblem},
         return sol
     else
         return RODESolution{T, N, typeof(u), Nothing, Nothing, typeof(t),
-            typeof(W), typeof(prob), typeof(alg), typeof(interp),
+            typeof(W), typeof(discretes), typeof(prob), typeof(alg), typeof(interp),
             typeof(stats), typeof(alg_choice), typeof(saved_subsystem)}(
-            u, nothing, nothing, t, W,
+            u, nothing, nothing, t, W, discretes,
             prob, alg, interp,
             dense, 0, stats,
             alg_choice, retcode, seed, saved_subsystem)
     end
 end
 
+function save_discretes!(sol::AbstractRODESolution, t, vals, timeseries_idx)
+    RecursiveArrayTools.has_discretes(sol) || return
+    disc = RecursiveArrayTools.get_discretes(sol)
+    _save_discretes_internal!(disc[timeseries_idx], t, vals)
+end
+
+function get_interpolated_discretes(sol::AbstractRODESolution, t, deriv, continuity)
+    is_parameter_timeseries(sol) == Timeseries() || return nothing
+
+    discs::ParameterTimeseriesCollection = RecursiveArrayTools.get_discretes(sol)
+    interp_discs = map(discs) do partition
+        hold_discrete(partition.u, partition.t, t)
+    end
+    return ParameterTimeseriesCollection(interp_discs, parameter_values(discs))
+end
+
+function SymbolicIndexingInterface.is_parameter_timeseries(::Type{S}) where {
+        T1, T2, T3, T4, T5, T6, T7,
+        S <: RODESolution{T1, T2, T3, T4, T5, T6, T7, <:ParameterTimeseriesCollection}}
+    Timeseries()
+end
+
 function calculate_solution_errors!(sol::AbstractRODESolution; fill_uanalytic = true,
         timeseries_errors = true, dense_errors = true)
     if sol.prob.f isa Tuple

test/aqua.jl

@@ -1,6 +1,14 @@
 using Test
 using SciMLBase
 using Aqua
+using Pkg
+
+# yes this is horrible, we'll fix it when Pkg or Base provides a decent API
+manifest = Pkg.Types.EnvCache().manifest
+# these are good sentinels to test whether someone has added a heavy SciML package to the test deps
+if haskey(manifest.deps, "NonlinearSolveBase") || haskey(manifest.deps, "DiffEqBase")
+    error("Don't put Downstream Packages in non Downstream CI")
+end
 
 # https://github.com/JuliaArrays/FillArrays.jl/pull/163
 @test isempty(detect_ambiguities(SciMLBase))

test/downstream/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 BoundaryValueDiffEq = "764a87c0-6b3e-53db-9096-fe964310641d"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -30,6 +31,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 BoundaryValueDiffEq = "5"
+DataFrames = "1.6"
 DelayDiffEq = "5"
 DiffEqCallbacks = "3, 4"
 ForwardDiff = "0.10"

test/downstream/comprehensive_indexing.jl

@@ -977,3 +977,19 @@ end
     sol = solve(prob, ImplicitEM())
     @test sol[sym] ≈ sol(sol.t .- sol.ps[delay]; idxs = original)
 end
+
+@testset "RODESolutions save discretes" begin
+    @parameters k(t)
+    @variables A(t)
+    function affect2!(integ, u, p, ctx)
+        integ.ps[p.k] += 1.0
+    end
+    db = 1.0 => (affect2!, [], [k], [k], nothing)
+
+    @named ssys = SDESystem(D(A) ~ k * A, [0.0], t, [A], [k], discrete_events = db)
+    ssys = complete(ssys)
+    prob = SDEProblem(ssys, [A => 1.0], (0.0, 4.0), [k => 1.0])
+    sol = solve(prob, RI5())
+    @test sol[k] isa AbstractVector
+    @test sol[k] == [1.0, 2.0, 3.0, 4.0]
+end