Add observed equation to SDEfunction

Author: frankschae

src/systems/diffeqs/sdesystem.jl

@@ -221,6 +221,7 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
                                      u0 = nothing;
                                      version = nothing, tgrad = false, sparse = false,
                                      jac = false, Wfact = false, eval_expression = true,
+                                     checkbounds = false,
                                      kwargs...) where {iip}
     dvs = scalarize.(dvs)
     ps = scalarize.(ps)
@@ -279,14 +280,25 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
     M = calculate_massmatrix(sys)
     _M = (u0 === nothing || M == I) ? M : ArrayInterfaceCore.restructure(u0 .* u0', M)
 
+    obs = observed(sys)
+    observedfun = let sys = sys, dict = Dict()
+        function generated_observed(obsvar, u, p, t)
+            obs = get!(dict, value(obsvar)) do
+                build_explicit_observed_function(sys, obsvar; checkbounds = checkbounds)
+            end
+            obs(u, p, t)
+        end
+    end
+
     sts = states(sys)
     SDEFunction{iip}(f, g,
                      jac = _jac === nothing ? nothing : _jac,
                      tgrad = _tgrad === nothing ? nothing : _tgrad,
                      Wfact = _Wfact === nothing ? nothing : _Wfact,
                      Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,
                      mass_matrix = _M,
-                     syms = Symbol.(states(sys)))
+                     syms = Symbol.(states(sys)),
+                     observed = observedfun)
 end
 
 function DiffEqBase.SDEFunction(sys::SDESystem, args...; kwargs...)

test/sdesystem.jl

@@ -479,3 +479,37 @@ eqs = [D(x) ~ x]
 noiseeqs = [0.1 * x]
 @named de = SDESystem(eqs, noiseeqs, t, [x], [])
 @test nameof(rename(de, :newname)) == :newname
+
+@testset "observed functionality" begin
+    @parameters α β
+    @variables t x(t) y(t) z(t)
+    @variables weight(t)
+    D = Differential(t)
+
+    eqs = [D(x) ~ α * x]
+    noiseeqs = [β * x]
+    dt = 1 // 2^(7)
+    x0 = 0.1
+
+    u0map = [
+        x => x0,
+    ]
+
+    parammap = [
+        α => 1.5,
+        β => 1.0,
+    ]
+
+    @named de = SDESystem(eqs, noiseeqs, t, [x], [α, β], observed = [weight ~ x * 10])
+
+    prob = SDEProblem(de, u0map, (0.0, 1.0), parammap)
+    sol = solve(prob, EM(), dt = dt)
+    @test observed(de) == [weight ~ x * 10]
+    @test sol[weight] == 10 * sol[x]
+
+    @named ode = ODESystem(eqs, t, [x], [α, β], observed = [weight ~ x * 10])
+    odeprob = ODEProblem(ode, u0map, (0.0, 1.0), parammap)
+    solode = solve(odeprob, Tsit5())
+    @test observed(ode) == [weight ~ x * 10]
+    @test solode[weight] == 10 * solode[x]
+end