From 6729cdaa202e272e922465f2d44f6dbca18870c3 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 28 Nov 2024 15:16:17 +0530
Subject: [PATCH] fix: retain observed equations after `structural_simplify` of
 `SDESystem`

---
 src/systems/systems.jl |  2 +-
 test/sdesystem.jl      | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 4205a4d207..a54206d1dd 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -156,6 +156,6 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
         noise_eqs = StructuralTransformations.tearing_substitute_expr(ode_sys, noise_eqs)
         return SDESystem(full_equations(ode_sys), noise_eqs,
             get_iv(ode_sys), unknowns(ode_sys), parameters(ode_sys);
-            name = nameof(ode_sys), is_scalar_noise)
+            name = nameof(ode_sys), is_scalar_noise, observed = observed(ode_sys))
     end
 end
diff --git a/test/sdesystem.jl b/test/sdesystem.jl
index 78d7a0418b..749aca86a7 100644
--- a/test/sdesystem.jl
+++ b/test/sdesystem.jl
@@ -799,3 +799,13 @@ end
     prob = SDEProblem(sys, [], (0.0, 1.0), [])
     @test_nowarn solve(prob, RKMil())
 end
+
+@testset "Observed variables retained after `structural_simplify`" begin
+    @variables x(t) y(t) z(t)
+    @brownian a
+    @mtkbuild sys = System([D(x) ~ x + a, D(y) ~ y + a, z ~ x + y], t)
+    @test sys isa SDESystem
+    @test length(observed(sys)) == 1
+    prob = SDEProblem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+    @test prob[z] ≈ 2.0
+end