Merge pull request #3106 from AayushSabharwal/as/observed-discrete

fix: detect observed variables and dependent parameters dependent on discrete parameters

Author: ChrisRackauckas

.github/workflows/Downstream.yml

@@ -72,4 +72,4 @@ jobs:
         with:
           file: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: true
+          fail_ci_if_error: false

src/systems/abstractsystem.jl

@@ -731,7 +731,7 @@ end
 function has_observed_with_lhs(sys, sym)
     has_observed(sys) || return false
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
-        return any(isequal(sym), ic.observed_syms)
+        return haskey(ic.observed_syms_to_timeseries, sym)
     else
         return any(isequal(sym), [eq.lhs for eq in observed(sys)])
     end
@@ -740,7 +740,7 @@ end
 function has_parameter_dependency_with_lhs(sys, sym)
     has_parameter_dependencies(sys) || return false
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
-        return any(isequal(sym), ic.dependent_pars)
+        return haskey(ic.dependent_pars_to_timeseries, unwrap(sym))
     else
         return any(isequal(sym), [eq.lhs for eq in parameter_dependencies(sys)])
     end
@@ -762,11 +762,20 @@ for traitT in [
         allsyms = vars(sym; op = Symbolics.Operator)
         for s in allsyms
             s = unwrap(s)
-            if is_variable(sys, s) || is_independent_variable(sys, s) ||
-               has_observed_with_lhs(sys, s)
+            if is_variable(sys, s) || is_independent_variable(sys, s)
                 push!(ts_idxs, ContinuousTimeseries())
             elseif is_timeseries_parameter(sys, s)
                 push!(ts_idxs, timeseries_parameter_index(sys, s).timeseries_idx)
+            elseif is_time_dependent(sys) && iscall(s) && issym(operation(s)) &&
+                   is_variable(sys, operation(s)(get_iv(sys)))
+                # DDEs case, to detect x(t - k)
+                push!(ts_idxs, ContinuousTimeseries())
+            elseif has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+                if (ts = get(ic.observed_syms_to_timeseries, s, nothing)) !== nothing
+                    union!(ts_idxs, ts)
+                elseif (ts = get(ic.dependent_pars_to_timeseries, s, nothing)) !== nothing
+                    union!(ts_idxs, ts)
+                end
             end
         end
     end

src/systems/diffeqs/odesystem.jl

@@ -490,7 +490,16 @@ function build_explicit_observed_function(sys, ts;
     ivs = independent_variables(sys)
     dep_vars = scalarize(setdiff(vars, ivs))
 
-    obs = param_only ? Equation[] : observed(sys)
+    obs = observed(sys)
+    if param_only
+        if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+            obs = filter(obs) do eq
+                !(ContinuousTimeseries() in ic.observed_syms_to_timeseries[eq.lhs])
+            end
+        else
+            obs = Equation[]
+        end
+    end
 
     cs = collect_constants(obs)
     if !isempty(cs) > 0

src/systems/index_cache.jl

@@ -38,6 +38,7 @@ const UnknownIndexMap = Dict{
     BasicSymbolic, Union{Int, UnitRange{Int}, AbstractArray{Int}}}
 const TunableIndexMap = Dict{BasicSymbolic,
     Union{Int, UnitRange{Int}, Base.ReshapedArray{Int, N, UnitRange{Int}} where {N}}}
+const TimeseriesSetType = Set{Union{ContinuousTimeseries, Int}}
 
 struct IndexCache
     unknown_idx::UnknownIndexMap
@@ -48,8 +49,9 @@ struct IndexCache
     tunable_idx::TunableIndexMap
     constant_idx::ParamIndexMap
     nonnumeric_idx::NonnumericMap
-    observed_syms::Set{BasicSymbolic}
-    dependent_pars::Set{Union{BasicSymbolic, CallWithMetadata}}
+    observed_syms_to_timeseries::Dict{BasicSymbolic, TimeseriesSetType}
+    dependent_pars_to_timeseries::Dict{
+        Union{BasicSymbolic, CallWithMetadata}, TimeseriesSetType}
     discrete_buffer_sizes::Vector{Vector{BufferTemplate}}
     tunable_buffer_size::BufferTemplate
     constant_buffer_sizes::Vector{BufferTemplate}
@@ -91,20 +93,6 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
-    observed_syms = Set{BasicSymbolic}()
-    for eq in observed(sys)
-        if symbolic_type(eq.lhs) != NotSymbolic()
-            sym = eq.lhs
-            ttsym = default_toterm(sym)
-            rsym = renamespace(sys, sym)
-            rttsym = renamespace(sys, ttsym)
-            push!(observed_syms, sym)
-            push!(observed_syms, ttsym)
-            push!(observed_syms, rsym)
-            push!(observed_syms, rttsym)
-        end
-    end
-
     tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()
     nonnumeric_buffers = Dict{Any, Set{Union{BasicSymbolic, CallWithMetadata}}}()
@@ -267,38 +255,77 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
-    for sym in Iterators.flatten((keys(unk_idxs), keys(disc_idxs), keys(tunable_idxs),
-        keys(const_idxs), keys(nonnumeric_idxs),
-        observed_syms, independent_variable_symbols(sys)))
-        if hasname(sym) && (!iscall(sym) || operation(sym) !== getindex)
-            symbol_to_variable[getname(sym)] = sym
-        end
-    end
-
-    dependent_pars = Set{Union{BasicSymbolic, CallWithMetadata}}()
+    dependent_pars_to_timeseries = Dict{
+        Union{BasicSymbolic, CallWithMetadata}, TimeseriesSetType}()
 
     for eq in parameter_dependencies(sys)
         sym = eq.lhs
+        vs = vars(eq.rhs)
+        timeseries = TimeseriesSetType()
+        if is_time_dependent(sys)
+            for v in vs
+                if (idx = get(disc_idxs, v, nothing)) !== nothing
+                    push!(timeseries, idx.clock_idx)
+                end
+            end
+        end
         ttsym = default_toterm(sym)
         rsym = renamespace(sys, sym)
         rttsym = renamespace(sys, ttsym)
-        for s in [sym, ttsym, rsym, rttsym]
-            push!(dependent_pars, s)
+        for s in (sym, ttsym, rsym, rttsym)
+            dependent_pars_to_timeseries[s] = timeseries
             if hasname(s) && (!iscall(s) || operation(s) != getindex)
                 symbol_to_variable[getname(s)] = sym
             end
         end
     end
 
+    observed_syms_to_timeseries = Dict{BasicSymbolic, TimeseriesSetType}()
+    for eq in observed(sys)
+        if symbolic_type(eq.lhs) != NotSymbolic()
+            sym = eq.lhs
+            vs = vars(eq.rhs; op = Nothing)
+            timeseries = TimeseriesSetType()
+            if is_time_dependent(sys)
+                for v in vs
+                    if (idx = get(disc_idxs, v, nothing)) !== nothing
+                        push!(timeseries, idx.clock_idx)
+                    elseif haskey(observed_syms_to_timeseries, v)
+                        union!(timeseries, observed_syms_to_timeseries[v])
+                    elseif haskey(dependent_pars_to_timeseries, v)
+                        union!(timeseries, dependent_pars_to_timeseries[v])
+                    end
+                end
+                if isempty(timeseries)
+                    push!(timeseries, ContinuousTimeseries())
+                end
+            end
+            ttsym = default_toterm(sym)
+            rsym = renamespace(sys, sym)
+            rttsym = renamespace(sys, ttsym)
+            for s in (sym, ttsym, rsym, rttsym)
+                observed_syms_to_timeseries[s] = timeseries
+            end
+        end
+    end
+
+    for sym in Iterators.flatten((keys(unk_idxs), keys(disc_idxs), keys(tunable_idxs),
+        keys(const_idxs), keys(nonnumeric_idxs),
+        keys(observed_syms_to_timeseries), independent_variable_symbols(sys)))
+        if hasname(sym) && (!iscall(sym) || operation(sym) !== getindex)
+            symbol_to_variable[getname(sym)] = sym
+        end
+    end
+
     return IndexCache(
         unk_idxs,
         disc_idxs,
         callback_to_clocks,
         tunable_idxs,
         const_idxs,
         nonnumeric_idxs,
-        observed_syms,
-        dependent_pars,
+        observed_syms_to_timeseries,
+        dependent_pars_to_timeseries,
         disc_buffer_templates,
         BufferTemplate(Real, tunable_buffer_size),
         const_buffer_sizes,

test/odesystem.jl

@@ -1504,3 +1504,14 @@ end
     sys2 = complete(sys; split = false)
     @test ModelingToolkit.get_index_cache(sys2) === nothing
 end
+
+# https://github.com/SciML/SciMLBase.jl/issues/786
+@testset "Observed variables dependent on discrete parameters" begin
+    @variables x(t) obs(t)
+    @parameters c(t)
+    @mtkbuild sys = ODESystem(
+        [D(x) ~ c * cos(x), obs ~ c], t, [x], [c]; discrete_events = [1.0 => [c ~ c + 1]])
+    prob = ODEProblem(sys, [x => 0.0], (0.0, 2pi), [c => 1.0])
+    sol = solve(prob, Tsit5())
+    @test sol[obs] ≈ 1:7
+end