From 4ae195b9ccce121085914f36004a375b5223f815 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@juliahub.com>
Date: Mon, 1 Apr 2024 17:14:20 +0530
Subject: [PATCH] fix: fix namespacing of defaults and equations (#2594)

---
 src/systems/abstractsystem.jl | 22 ++++++++-----
 test/odesystem.jl             | 58 +++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
index c808aa4057..8753b30656 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -819,22 +819,24 @@ namespace_controls(sys::AbstractSystem) = controls(sys, controls(sys))
 
 function namespace_defaults(sys)
     defs = defaults(sys)
-    Dict((isparameter(k) ? parameters(sys, k) : unknowns(sys, k)) => namespace_expr(v, sys)
+    Dict((isparameter(k) ? parameters(sys, k) : unknowns(sys, k)) => namespace_expr(
+             v, sys; check = true)
     for (k, v) in pairs(defs))
 end
 
 function namespace_equations(sys::AbstractSystem, ivs = independent_variables(sys))
     eqs = equations(sys)
     isempty(eqs) && return Equation[]
-    map(eq -> namespace_equation(eq, sys; ivs), eqs)
+    map(eq -> namespace_equation(eq, sys; ivs, check = true), eqs)
 end
 
 function namespace_equation(eq::Equation,
         sys,
         n = nameof(sys);
-        ivs = independent_variables(sys))
-    _lhs = namespace_expr(eq.lhs, sys, n; ivs)
-    _rhs = namespace_expr(eq.rhs, sys, n; ivs)
+        ivs = independent_variables(sys),
+        check = false)
+    _lhs = namespace_expr(eq.lhs, sys, n; ivs, check)
+    _rhs = namespace_expr(eq.rhs, sys, n; ivs, check)
     _lhs ~ _rhs
 end
 
@@ -844,32 +846,36 @@ function namespace_assignment(eq::Assignment, sys)
     Assignment(_lhs, _rhs)
 end
 
-function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys))
+function namespace_expr(
+        O, sys, n = nameof(sys); ivs = independent_variables(sys), check = false)
     O = unwrap(O)
     if any(isequal(O), ivs)
         return O
     elseif istree(O)
         T = typeof(O)
         renamed = let sys = sys, n = n, T = T
-            map(a -> namespace_expr(a, sys, n; ivs)::Any, arguments(O))
+            map(a -> namespace_expr(a, sys, n; ivs, check)::Any, arguments(O))
         end
         if isvariable(O)
+            check && !is_variable(sys, O) && !is_parameter(sys, O) && return O
             # Use renamespace so the scope is correct, and make sure to use the
             # metadata from the rescoped variable
             rescoped = renamespace(n, O)
             similarterm(O, operation(rescoped), renamed,
                 metadata = metadata(rescoped))
         elseif Symbolics.isarraysymbolic(O)
+            check && !is_variable(sys, O) && !is_parameter(sys, O) && return O
             # promote_symtype doesn't work for array symbolics
             similarterm(O, operation(O), renamed, symtype(O), metadata = metadata(O))
         else
             similarterm(O, operation(O), renamed, metadata = metadata(O))
         end
     elseif isvariable(O)
+        check && !is_variable(sys, O) && !is_parameter(sys, O) && return O
         renamespace(n, O)
     elseif O isa Array
         let sys = sys, n = n
-            map(o -> namespace_expr(o, sys, n; ivs), O)
+            map(o -> namespace_expr(o, sys, n; ivs, check), O)
         end
     else
         O
diff --git a/test/odesystem.jl b/test/odesystem.jl
index f45987cdc5..ee4b22f5fc 100644
--- a/test/odesystem.jl
+++ b/test/odesystem.jl
@@ -995,3 +995,61 @@ let # Issue https://github.com/SciML/ModelingToolkit.jl/issues/2322
     sol = solve(prob, Rodas4())
     @test sol(1)[]≈0.6065307685451087 rtol=1e-4
 end
+
+# Issue#2344
+function FML2(; name)
+    @parameters begin
+        k2[1:1] = [1.0]
+    end
+    systems = @named begin
+        constant = Constant(k = k2[1])
+    end
+    @variables begin
+        x(t) = 0
+    end
+    eqs = [
+        D(x) ~ constant.output.u + k2[1]
+    ]
+    ODESystem(eqs, t; systems, name)
+end
+
+@mtkbuild model = FML2()
+
+@test isequal(ModelingToolkit.defaults(model)[model.constant.k], model.k2[1])
+@test_nowarn ODEProblem(model, [], (0.0, 10.0))
+
+# Issue#2477
+function RealExpression(; name, y)
+    vars = @variables begin
+        u(t)
+    end
+    eqns = [
+        u ~ y
+    ]
+    sys = ODESystem(eqns, t, vars, []; name)
+end
+
+function sys1(; name)
+    vars = @variables begin
+        x(t)
+        z(t)[1:1]
+    end # doing a collect on z doesn't work either. 
+    @named e1 = RealExpression(y = x) # This works perfectly. 
+    @named e2 = RealExpression(y = z[1]) # This bugs. However, `full_equations(e2)` works as expected. 
+    systems = [e1, e2]
+    ODESystem(Equation[], t, Iterators.flatten(vars), []; systems, name)
+end
+
+@named sys = sys1()
+sys = complete(sys)
+@test Set(equations(sys)) == Set([sys.e1.u ~ sys.x, sys.e2.u ~ sys.z[1]])
+
+# Issue#2522
+@parameters a[1:2]=[1, 2] b=4 c=1
+@variables x(t)=ParentScope(a[1]) y(t)=ParentScope(b)
+@named level0 = ODESystem([D(x) ~ ParentScope(a[2]),
+        D(y) ~ ParentScope(c)], t, [x, y], [])
+level1 = ODESystem(Equation[], t, [], [a..., b, c]; name = :level1) ∘ level0
+level1 = structural_simplify(level1)
+@test isequal(ModelingToolkit.defaults(level1)[level1.level0.x], level1.a[1])
+@test_nowarn ODEProblem(level1, [], (0, 1))