Merge pull request #2791 from AayushSabharwal/as/sym-obs

feat: store observed equation lhs in `symbol_to_variable` mapping

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -495,10 +495,40 @@ function SymbolicIndexingInterface.is_observed(sys::AbstractSystem, sym)
 end
 
 function SymbolicIndexingInterface.observed(sys::AbstractSystem, sym)
-    return let _fn = build_explicit_observed_function(sys, sym)
-        fn(u, p, t) = _fn(u, p, t)
-        fn(u, p::MTKParameters, t) = _fn(u, p..., t)
-        fn
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        if sym isa Symbol
+            _sym = get(ic.symbol_to_variable, sym, nothing)
+            if _sym === nothing
+                throw(ArgumentError("Symbol $sym does not exist in the system"))
+            end
+            sym = _sym
+        elseif sym isa AbstractArray && symbolic_type(sym) isa NotSymbolic &&
+               any(x -> x isa Symbol, sym)
+            sym = map(sym) do s
+                if s isa Symbol
+                    _s = get(ic.symbol_to_variable, s, nothing)
+                    if _s === nothing
+                        throw(ArgumentError("Symbol $s does not exist in the system"))
+                    end
+                    return _s
+                end
+                return unwrap(s)
+            end
+        end
+    end
+    _fn = build_explicit_observed_function(sys, sym)
+    if is_time_dependent(sys)
+        return let _fn = _fn
+            fn1(u, p, t) = _fn(u, p, t)
+            fn1(u, p::MTKParameters, t) = _fn(u, p..., t)
+            fn1
+        end
+    else
+        return let _fn = _fn
+            fn2(u, p) = _fn(u, p)
+            fn2(u, p::MTKParameters) = _fn(u, p...)
+            fn2
+        end
     end
 end
 

src/systems/index_cache.jl

@@ -80,6 +80,12 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
+    for eq in observed(sys)
+        if symbolic_type(eq.lhs) != NotSymbolic() && hasname(eq.lhs)
+            symbol_to_variable[getname(eq.lhs)] = eq.lhs
+        end
+    end
+
     disc_buffers = Dict{Any, Set{BasicSymbolic}}()
     tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()

test/symbolic_indexing_interface.jl

@@ -101,3 +101,16 @@ using SymbolicIndexingInterface
 prob = ODEProblem(complete(sys))
 get_dep = @test_nowarn getu(prob, 2p1)
 @test get_dep(prob) == [2.0, 4.0]
+
+@testset "Observed functions with variables as `Symbol`s" begin
+    @variables x(t) y(t) z(t)[1:2]
+    @parameters p1 p2[1:2, 1:2]
+    @mtkbuild sys = ODESystem([D(x) ~ x * t + p1, y ~ 2x, D(z) ~ p2 * z], t)
+    prob = ODEProblem(
+        sys, [x => 1.0, z => ones(2)], (0.0, 1.0), [p1 => 2.0, p2 => ones(2, 2)])
+    @test getu(prob, x)(prob) == getu(prob, :x)(prob)
+    @test getu(prob, [x, y])(prob) == getu(prob, [:x, :y])(prob)
+    @test getu(prob, z)(prob) == getu(prob, :z)(prob)
+    @test getu(prob, p1)(prob) == getu(prob, :p1)(prob)
+    @test getu(prob, p2)(prob) == getu(prob, :p2)(prob)
+end