From c4a5616971f1bf25891aa264b508d3fea7c3d6f9 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Tue, 19 Nov 2024 09:27:50 +0100
Subject: [PATCH 01/24] add capability to trace MTK dynamics with InfiniteOpt

---
 Project.toml                        |   3 +
 ext/MTKInfiniteOptExt.jl            |  19 ++++++
 test/extensions/Project.toml        |   3 +
 test/extensions/test_infiniteopt.jl | 102 ++++++++++++++++++++++++++++
 test/runtests.jl                    |   1 +
 5 files changed, 128 insertions(+)
 create mode 100644 ext/MTKInfiniteOptExt.jl
 create mode 100644 test/extensions/test_infiniteopt.jl

diff --git a/Project.toml b/Project.toml
index 18b4d55992..5848298418 100644
--- a/Project.toml
+++ b/Project.toml
@@ -64,6 +64,7 @@ BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
+InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [extensions]
@@ -72,6 +73,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKHomotopyContinuationExt = "HomotopyContinuation"
 MTKLabelledArraysExt = "LabelledArrays"
+MTKInfiniteOptExt = "InfiniteOpt"
 
 [compat]
 AbstractTrees = "0.3, 0.4"
@@ -104,6 +106,7 @@ FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
 HomotopyContinuation = "2.11"
+InfiniteOpt = "0.5"
 InteractiveUtils = "1"
 JuliaFormatter = "1.0.47"
 JumpProcesses = "9.13.1"
diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
new file mode 100644
index 0000000000..88bf724d14
--- /dev/null
+++ b/ext/MTKInfiniteOptExt.jl
@@ -0,0 +1,19 @@
+module MTKInfiniteOptExt
+import ModelingToolkit
+import SymbolicUtils
+import NaNMath
+import InfiniteOpt
+
+# This file contains method definitions to make it possible to trace through functions generated by MTK using JuMP variables
+
+for ff in [acos, log1p, acosh, log2, asin, lgamma, tan, atanh, cos, log, sin, log10, sqrt]
+    f = nameof(ff)
+    # These need to be defined so that JuMP can trace through functions built by Symbolics
+    @eval NaNMath.$f(x::GeneralVariableRef) = Base.$f(x)
+end
+
+# JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
+Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false
+Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false
+
+end
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 37614d9bfb..5f7afe222a 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -4,6 +4,9 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
+InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"
diff --git a/test/extensions/test_infiniteopt.jl b/test/extensions/test_infiniteopt.jl
new file mode 100644
index 0000000000..e45aa0f2fd
--- /dev/null
+++ b/test/extensions/test_infiniteopt.jl
@@ -0,0 +1,102 @@
+using ModelingToolkit, InfiniteOpt, JuMP, Ipopt
+using ModelingToolkit: D_nounits as D, t_nounits as t, varmap_to_vars
+
+@mtkmodel Pendulum begin
+    @parameters begin
+        g = 9.8
+        L = 0.4
+        K = 1.2
+        m = 0.3
+    end
+    @variables begin
+        θ(t) # state
+        ω(t) # state
+        τ(t) = 0 # input
+        y(t) # output
+    end
+    @equations begin
+        D(θ) ~ ω
+        D(ω) ~ -g / L * sin(θ) - K / m * ω + τ / m / L^2
+        y ~ θ * 180 / π
+    end
+end
+@named model = Pendulum()
+model = complete(model)
+
+inputs = [model.τ]
+(f_oop, f_ip), dvs, psym, io_sys = ModelingToolkit.generate_control_function(
+    model, inputs, split = false)
+
+outputs = [model.y]
+f_obs = ModelingToolkit.build_explicit_observed_function(io_sys, outputs; inputs = inputs)
+
+expected_state_order = [model.θ, model.ω]
+permutation = [findfirst(isequal(x), expected_state_order) for x in dvs] # This maps our expected state order to the actual state order
+
+##
+
+ub = varmap_to_vars([model.θ => 2pi, model.ω => 10], dvs)
+lb = varmap_to_vars([model.θ => -2pi, model.ω => -10], dvs)
+xf = varmap_to_vars([model.θ => pi, model.ω => 0], dvs)
+nx = length(dvs)
+nu = length(inputs)
+ny = length(outputs)
+
+##
+m = InfiniteModel(optimizer_with_attributes(Ipopt.Optimizer,
+    "print_level" => 0, "acceptable_tol" => 1e-3, "constr_viol_tol" => 1e-5, "max_iter" => 1000,
+    "tol" => 1e-5, "mu_strategy" => "monotone", "nlp_scaling_method" => "gradient-based",
+    "alpha_for_y" => "safer-min-dual-infeas", "bound_mult_init_method" => "mu-based", "print_user_options" => "yes"));
+
+@infinite_parameter(m, τ in [0, 1], num_supports=51,
+    derivative_method=OrthogonalCollocation(4)) # Time variable
+guess_xs = [t -> pi, t -> 0.1][permutation]
+guess_us = [t -> 0.1]
+InfiniteOpt.@variables(m,
+    begin
+        # state variables
+        (lb[i] <= x[i = 1:nx] <= ub[i], Infinite(τ), start = guess_xs[i]) # state variables
+        -10 <= u[i = 1:nu] <= 10, Infinite(τ), (start = guess_us[i]) # control variables
+        0 <= tf <= 10, (start = 5) # Final time
+        0.2 <= L <= 0.6, (start = 0.4) # Length parameter
+    end)
+
+# Trace the dynamics
+x0, p = ModelingToolkit.get_u0_p(io_sys, [model.θ => 0, model.ω => 0], [model.L => L])
+
+xp = f_oop(x, u, p, τ)
+cp = f_obs(x, u, p, τ) # Test that it's possible to trace through an observed function
+
+@objective(m, Min, tf)
+@constraint(m, [i = 1:nx], x[i](0)==x0[i]) # Initial condition
+@constraint(m, [i = 1:nx], x[i](1)==xf[i]) # Terminal state
+
+x_scale = varmap_to_vars([model.θ => 1
+                          model.ω => 1], dvs)
+
+# Add dynamics constraints
+@constraint(m, [i = 1:nx], (∂(x[i], τ) - tf * xp[i]) / x_scale[i]==0)
+
+optimize!(m)
+
+# Extract the optimal solution
+opt_tf = value(tf)
+opt_time = opt_tf * value(τ)
+opt_x = [value(x[i]) for i in permutation]
+opt_u = [value(u[i]) for i in 1:nu]
+opt_L = value(L)
+
+# Plot the results
+# using Plots
+# plot(opt_time, opt_x[1], label = "θ", xlabel = "Time [s]", layout=3)
+# plot!(opt_time, opt_x[2], label = "ω", sp=2)
+# plot!(opt_time, opt_u[1], label = "τ", sp=3)
+
+using Test
+@test opt_x[1][end]≈pi atol=1e-3
+@test opt_x[2][end]≈0 atol=1e-3
+
+@test opt_x[1][1]≈0 atol=1e-3
+@test opt_x[2][1]≈0 atol=1e-3
+
+@test opt_L≈0.2 atol=1e-3 # Smallest permissible length is optimal
diff --git a/test/runtests.jl b/test/runtests.jl
index 192ead935c..25b692f678 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -113,5 +113,6 @@ end
         @safetestset "Auto Differentiation Test" include("extensions/ad.jl")
         @safetestset "LabelledArrays Test" include("labelledarrays.jl")
         @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl")
+        @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
     end
 end

From 00bfd1d539977588416cfb6de8cdcbcbf8ce2c24 Mon Sep 17 00:00:00 2001
From: Christopher Rackauckas <accounts@chrisrackauckas.com>
Date: Tue, 19 Nov 2024 09:36:52 -0100
Subject: [PATCH 02/24] Update MTKInfiniteOptExt.jl

---
 ext/MTKInfiniteOptExt.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index 88bf724d14..79463f7e71 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -15,5 +15,6 @@ end
 # JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
 Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false
 Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false
-
+Base.isequal(::SymbolicUtils.Symbolic, ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) = false
+Base.isequal(::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, ::SymbolicUtils.Symbolic) = false
 end

From 72009993b81bc0e7908e35c20775b0c03589beb1 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Wed, 20 Nov 2024 18:39:08 +0100
Subject: [PATCH 03/24] rm lgamma

---
 ext/MTKInfiniteOptExt.jl | 16 +++++++++++-----
 test/odesystem.jl        |  2 +-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index 79463f7e71..3b482a1f03 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -3,18 +3,24 @@ import ModelingToolkit
 import SymbolicUtils
 import NaNMath
 import InfiniteOpt
+import InfiniteOpt: JuMP, GeneralVariableRef
 
 # This file contains method definitions to make it possible to trace through functions generated by MTK using JuMP variables
 
-for ff in [acos, log1p, acosh, log2, asin, lgamma, tan, atanh, cos, log, sin, log10, sqrt]
+for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt]
     f = nameof(ff)
     # These need to be defined so that JuMP can trace through functions built by Symbolics
     @eval NaNMath.$f(x::GeneralVariableRef) = Base.$f(x)
 end
 
 # JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
-Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false
-Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false
-Base.isequal(::SymbolicUtils.Symbolic, ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) = false
-Base.isequal(::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, ::SymbolicUtils.Symbolic) = false
+function Base.isequal(::SymbolicUtils.Symbolic,
+        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr})
+    false
+end
+function Base.isequal(
+        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr},
+        ::SymbolicUtils.Symbolic)
+    false
+end
 end
diff --git a/test/odesystem.jl b/test/odesystem.jl
index 90d1c4c578..02eb80ef85 100644
--- a/test/odesystem.jl
+++ b/test/odesystem.jl
@@ -1521,6 +1521,6 @@ end
     @parameters p[1:2] = [1.0, 2.0]
     @mtkbuild sys = ODESystem([D(x) ~ x, y^2 ~ x + sum(p)], t)
     prob = DAEProblem(sys, [D(x) => x, D(y) => D(x) / 2y], [], (0.0, 1.0))
-    sol = solve(prob, DFBDF(), abstol=1e-8, reltol=1e-8)
+    sol = solve(prob, DFBDF(), abstol = 1e-8, reltol = 1e-8)
     @test sol[x]≈sol[y^2 - sum(p)] atol=1e-5
 end

From d0748e6467285627531071e4377f43c9489cebfa Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 13:04:00 +0530
Subject: [PATCH 04/24] fix: fix nonnumeric parameters in `@mtkmodel`

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

diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
index 5b79eaa91b..5402659b1c 100644
--- a/src/systems/model_parsing.jl
+++ b/src/systems/model_parsing.jl
@@ -923,7 +923,7 @@ function parse_variable_arg(dict, mod, arg, varclass, kwargs, where_types)
         end
 
         push!(varexpr.args, metadata_expr)
-        return vv isa Num ? name : :($name...), varexpr
+        return symbolic_type(vv) == ScalarSymbolic() ? name : :($name...), varexpr
     else
         return vv
     end
diff --git a/test/model_parsing.jl b/test/model_parsing.jl
index fd223800e3..6ab7636b32 100644
--- a/test/model_parsing.jl
+++ b/test/model_parsing.jl
@@ -979,3 +979,14 @@ end
     @test MultipleExtend.structure[:extend][1] == [:inmodel, :b, :inmodel_b]
     @test tosymbol.(parameters(multiple_extend)) == [:b, :inmodel_b₊p, :inmodel₊p]
 end
+
+struct CustomStruct end
+@testset "Nonnumeric parameters" begin
+    @mtkmodel MyModel begin
+        @parameters begin
+            p::CustomStruct
+        end
+    end
+    @named sys = MyModel(p = CustomStruct())
+    @test ModelingToolkit.defaults(sys)[@nonamespace sys.p] == CustomStruct()
+end

From 472921dc215bd2ecf8785c4ca5a630e9095e42c5 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 14:31:12 +0530
Subject: [PATCH 05/24] fix: support callable parameters provided to discretes
 list of callback

---
 src/systems/index_cache.jl | 29 +++++++++++++++++++----------
 test/index_cache.jl        | 26 ++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/src/systems/index_cache.jl b/src/systems/index_cache.jl
index 1e3490ee69..ab0dd08764 100644
--- a/src/systems/index_cache.jl
+++ b/src/systems/index_cache.jl
@@ -40,10 +40,12 @@ const TunableIndexMap = Dict{BasicSymbolic,
     Union{Int, UnitRange{Int}, Base.ReshapedArray{Int, N, UnitRange{Int}} where {N}}}
 const TimeseriesSetType = Set{Union{ContinuousTimeseries, Int}}
 
+const SymbolicParam = Union{BasicSymbolic, CallWithMetadata}
+
 struct IndexCache
     unknown_idx::UnknownIndexMap
     # sym => (bufferidx, idx_in_buffer)
-    discrete_idx::Dict{BasicSymbolic, DiscreteIndex}
+    discrete_idx::Dict{SymbolicParam, DiscreteIndex}
     # sym => (clockidx, idx_in_clockbuffer)
     callback_to_clocks::Dict{Any, Vector{Int}}
     tunable_idx::TunableIndexMap
@@ -56,13 +58,13 @@ struct IndexCache
     tunable_buffer_size::BufferTemplate
     constant_buffer_sizes::Vector{BufferTemplate}
     nonnumeric_buffer_sizes::Vector{BufferTemplate}
-    symbol_to_variable::Dict{Symbol, Union{BasicSymbolic, CallWithMetadata}}
+    symbol_to_variable::Dict{Symbol, SymbolicParam}
 end
 
 function IndexCache(sys::AbstractSystem)
     unks = solved_unknowns(sys)
     unk_idxs = UnknownIndexMap()
-    symbol_to_variable = Dict{Symbol, Union{BasicSymbolic, CallWithMetadata}}()
+    symbol_to_variable = Dict{Symbol, SymbolicParam}()
 
     let idx = 1
         for sym in unks
@@ -95,7 +97,7 @@ function IndexCache(sys::AbstractSystem)
 
     tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()
-    nonnumeric_buffers = Dict{Any, Set{Union{BasicSymbolic, CallWithMetadata}}}()
+    nonnumeric_buffers = Dict{Any, Set{SymbolicParam}}()
 
     function insert_by_type!(buffers::Dict{Any, S}, sym, ctype) where {S}
         sym = unwrap(sym)
@@ -103,10 +105,10 @@ function IndexCache(sys::AbstractSystem)
         push!(buf, sym)
     end
 
-    disc_param_callbacks = Dict{BasicSymbolic, Set{Int}}()
+    disc_param_callbacks = Dict{SymbolicParam, Set{Int}}()
     events = vcat(continuous_events(sys), discrete_events(sys))
     for (i, event) in enumerate(events)
-        discs = Set{BasicSymbolic}()
+        discs = Set{SymbolicParam}()
         affs = affects(event)
         if !(affs isa AbstractArray)
             affs = [affs]
@@ -130,26 +132,32 @@ function IndexCache(sys::AbstractSystem)
                isequal(only(arguments(sym)), get_iv(sys))
                 clocks = get!(() -> Set{Int}(), disc_param_callbacks, sym)
                 push!(clocks, i)
-            else
+            elseif is_variable_floatingpoint(sym)
                 insert_by_type!(constant_buffers, sym, symtype(sym))
+            else
+                stype = symtype(sym)
+                if stype <: FnType
+                    stype = fntype_to_function_type(stype)
+                end
+                insert_by_type!(nonnumeric_buffers, sym, stype)
             end
         end
     end
     clock_partitions = unique(collect(values(disc_param_callbacks)))
     disc_symtypes = unique(symtype.(keys(disc_param_callbacks)))
     disc_symtype_idx = Dict(disc_symtypes .=> eachindex(disc_symtypes))
-    disc_syms_by_symtype = [BasicSymbolic[] for _ in disc_symtypes]
+    disc_syms_by_symtype = [SymbolicParam[] for _ in disc_symtypes]
     for sym in keys(disc_param_callbacks)
         push!(disc_syms_by_symtype[disc_symtype_idx[symtype(sym)]], sym)
     end
-    disc_syms_by_symtype_by_partition = [Vector{BasicSymbolic}[] for _ in disc_symtypes]
+    disc_syms_by_symtype_by_partition = [Vector{SymbolicParam}[] for _ in disc_symtypes]
     for (i, buffer) in enumerate(disc_syms_by_symtype)
         for partition in clock_partitions
             push!(disc_syms_by_symtype_by_partition[i],
                 [sym for sym in buffer if disc_param_callbacks[sym] == partition])
         end
     end
-    disc_idxs = Dict{BasicSymbolic, DiscreteIndex}()
+    disc_idxs = Dict{SymbolicParam, DiscreteIndex}()
     callback_to_clocks = Dict{
         Union{SymbolicContinuousCallback, SymbolicDiscreteCallback}, Set{Int}}()
     for (typei, disc_syms_by_partition) in enumerate(disc_syms_by_symtype_by_partition)
@@ -191,6 +199,7 @@ function IndexCache(sys::AbstractSystem)
         end
         haskey(disc_idxs, p) && continue
         haskey(constant_buffers, ctype) && p in constant_buffers[ctype] && continue
+        haskey(nonnumeric_buffers, ctype) && p in nonnumeric_buffers[ctype] && continue
         insert_by_type!(
             if ctype <: Real || ctype <: AbstractArray{<:Real}
                 if istunable(p, true) && Symbolics.shape(p) != Symbolics.Unknown() &&
diff --git a/test/index_cache.jl b/test/index_cache.jl
index c479075797..455203d759 100644
--- a/test/index_cache.jl
+++ b/test/index_cache.jl
@@ -92,3 +92,29 @@ end
     reorder_dimension_by_tunables!(dst, sys, src, [r, q, p]; dim = 2)
     @test dst ≈ stack([vcat(4ones(4), 3ones(3), 1.0) for i in 1:5]; dims = 1)
 end
+
+mutable struct ParamTest
+    y::Any
+end
+(pt::ParamTest)(x) = pt.y - x
+@testset "Issue#3215: Callable discrete parameter" begin
+    function update_affect!(integ, u, p, ctx)
+        integ.p[p.p_1].y = integ.t
+    end
+
+    tp1 = typeof(ParamTest(1))
+    @parameters (p_1::tp1)(..) = ParamTest(1)
+    @variables x(ModelingToolkit.t_nounits) = 0
+
+    event1 = [1.0, 2, 3] => (update_affect!, [], [p_1], [p_1], nothing)
+
+    @named sys = ODESystem([
+            ModelingToolkit.D_nounits(x) ~ p_1(x)
+        ],
+        ModelingToolkit.t_nounits;
+        discrete_events = [event1]
+    )
+    ss = @test_nowarn complete(sys)
+    @test length(parameters(ss)) == 1
+    @test !is_timeseries_parameter(ss, p_1)
+end

From 585de359502db53e50609d656e97d3696e2d6300 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 15:05:31 +0530
Subject: [PATCH 06/24] fix: fix SDEs with noise dependent on observed
 variables

---
 .../symbolics_tearing.jl                      |  8 ++++++++
 src/systems/systems.jl                        |  2 ++
 test/sdesystem.jl                             | 19 +++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index a854acb9b1..8161a9572b 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -88,6 +88,14 @@ function tearing_sub(expr, dict, s)
     s ? simplify(expr) : expr
 end
 
+function tearing_substitute_expr(sys::AbstractSystem, expr; simplify = false)
+    empty_substitutions(sys) && return expr
+    substitutions = get_substitutions(sys)
+    @unpack subs = substitutions
+    solved = Dict(eq.lhs => eq.rhs for eq in subs)
+    return tearing_sub(expr, solved, simplify)
+end
+
 """
 $(TYPEDSIGNATURES)
 
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 848980605f..4205a4d207 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -152,6 +152,8 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
             noise_eqs = sorted_g_rows
             is_scalar_noise = false
         end
+
+        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)
diff --git a/test/sdesystem.jl b/test/sdesystem.jl
index c258a4142b..78d7a0418b 100644
--- a/test/sdesystem.jl
+++ b/test/sdesystem.jl
@@ -780,3 +780,22 @@ end
     prob = @test_nowarn SDEProblem(sys, nothing, (0.0, 1.0))
     @test_nowarn solve(prob, ImplicitEM())
 end
+
+@testset "Issue#3212: Noise dependent on observed" begin
+    sts = @variables begin
+        x(t) = 1.0
+        input(t)
+        [input = true]
+    end
+    ps = @parameters a = 2
+    @brownian η
+
+    eqs = [D(x) ~ -a * x + (input + 1) * η
+           input ~ 0.0]
+
+    sys = System(eqs, t, sts, ps; name = :name)
+    sys = structural_simplify(sys)
+    @test ModelingToolkit.get_noiseeqs(sys) ≈ [1.0]
+    prob = SDEProblem(sys, [], (0.0, 1.0), [])
+    @test_nowarn solve(prob, RKMil())
+end

From 398849613e4fd4bc7f9bcdc0960d54c2bb2c270e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Mon, 25 Nov 2024 13:06:07 +0200
Subject: [PATCH 07/24] fix: fix InitializationProblems that have vectors in
 u0map

---
 src/systems/diffeqs/abstractodesystem.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 1dc2198cc2..18aac952f4 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1356,6 +1356,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         u0T = promote_type(u0T, typeof(fullmap[eq.lhs]))
     end
     if u0T != Union{}
+        u0T = eltype(u0T)
         u0map = Dict(k => if symbolic_type(v) == NotSymbolic() && !is_array_of_symbolics(v)
                          v isa AbstractArray ? u0T.(v) : u0T(v)
                      else

From d4e0d0f1634d6b322ef0cfa4ee110e0f9d7cd693 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 12:37:39 +0530
Subject: [PATCH 08/24] feat: enable `structural_simplify` to automatically
 handle non-fully-determined systems

Uses the result of `check_consistency`
---
 src/structural_transformation/utils.jl | 58 +++++++++++++++++++-------
 src/systems/systemstructure.jl         |  9 +++-
 2 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 8b2ac1f69c..405a3ebed3 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -58,7 +58,41 @@ end
 ###
 ### Structural check
 ###
-function check_consistency(state::TransformationState, orig_inputs)
+
+"""
+    $(TYPEDSIGNATURES)
+
+Check if the `state` represents a singular system, and return the unmatched variables.
+"""
+function singular_check(state::TransformationState)
+    @unpack graph, var_to_diff = state.structure
+    fullvars = get_fullvars(state)
+    # This is defined to check if Pantelides algorithm terminates. For more
+    # details, check the equation (15) of the original paper.
+    extended_graph = (@set graph.fadjlist = Vector{Int}[graph.fadjlist;
+                                                        map(collect, edges(var_to_diff))])
+    extended_var_eq_matching = maximal_matching(extended_graph)
+
+    nvars = ndsts(graph)
+    unassigned_var = []
+    for (vj, eq) in enumerate(extended_var_eq_matching)
+        vj > nvars && break
+        if eq === unassigned && !isempty(𝑑neighbors(graph, vj))
+            push!(unassigned_var, fullvars[vj])
+        end
+    end
+    return unassigned_var
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Check the consistency of `state`, given the inputs `orig_inputs`. If `nothrow == false`,
+throws an error if the system is under-/over-determined or singular. In this case, if the
+function returns it will return `true`. If `nothrow == true`, it will return `false`
+instead of throwing an error. The singular case will print a warning.
+"""
+function check_consistency(state::TransformationState, orig_inputs; nothrow = false)
     fullvars = get_fullvars(state)
     neqs = n_concrete_eqs(state)
     @unpack graph, var_to_diff = state.structure
@@ -72,6 +106,7 @@ function check_consistency(state::TransformationState, orig_inputs)
     is_balanced = n_highest_vars == neqs
 
     if neqs > 0 && !is_balanced
+        nothrow && return false
         varwhitelist = var_to_diff .== nothing
         var_eq_matching = maximal_matching(graph, eq -> true, v -> varwhitelist[v]) # not assigned
         # Just use `error_reporting` to do conditional
@@ -85,20 +120,7 @@ function check_consistency(state::TransformationState, orig_inputs)
         error_reporting(state, bad_idxs, n_highest_vars, iseqs, orig_inputs)
     end
 
-    # This is defined to check if Pantelides algorithm terminates. For more
-    # details, check the equation (15) of the original paper.
-    extended_graph = (@set graph.fadjlist = Vector{Int}[graph.fadjlist;
-                                                        map(collect, edges(var_to_diff))])
-    extended_var_eq_matching = maximal_matching(extended_graph)
-
-    nvars = ndsts(graph)
-    unassigned_var = []
-    for (vj, eq) in enumerate(extended_var_eq_matching)
-        vj > nvars && break
-        if eq === unassigned && !isempty(𝑑neighbors(graph, vj))
-            push!(unassigned_var, fullvars[vj])
-        end
-    end
+    unassigned_var = singular_check(state)
 
     if !isempty(unassigned_var) || !is_balanced
         io = IOBuffer()
@@ -107,10 +129,14 @@ function check_consistency(state::TransformationState, orig_inputs)
         errmsg = "The system is structurally singular! " *
                  "Here are the problematic variables: \n" *
                  unassigned_var_str
+        if nothrow
+            @warn errmsg
+            return false
+        end
         throw(InvalidSystemException(errmsg))
     end
 
-    return nothing
+    return true
 end
 
 ###
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 71bdf6fe30..1c0ad8b8ae 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -677,7 +677,11 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
         check_consistency = true, fully_determined = true, warn_initialize_determined = false,
         dummy_derivative = true,
         kwargs...)
-    check_consistency &= fully_determined
+    if fully_determined isa Bool
+        check_consistency &= fully_determined
+    else
+        check_consistency = true
+    end
     has_io = io !== nothing
     orig_inputs = Set()
     if has_io
@@ -690,7 +694,8 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
     end
     sys, mm = ModelingToolkit.alias_elimination!(state; kwargs...)
     if check_consistency
-        ModelingToolkit.check_consistency(state, orig_inputs)
+        fully_determined = ModelingToolkit.check_consistency(
+            state, orig_inputs; nothrow = fully_determined === nothing)
     end
     if fully_determined && dummy_derivative
         sys = ModelingToolkit.dummy_derivative(

From beb16ec624e1f68a834a0aff7e7731c489e6fabc Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 12:38:02 +0530
Subject: [PATCH 09/24] feat: simplify initialization system with
 `fully_determined=nothing` by default

---
 src/systems/diffeqs/abstractodesystem.jl | 6 +++++-
 src/systems/problem_utils.jl             | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 1dc2198cc2..4019a868d8 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1295,7 +1295,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         check_length = true,
         warn_initialize_determined = true,
         initialization_eqs = [],
-        fully_determined = false,
+        fully_determined = nothing,
         check_units = true,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
@@ -1313,6 +1313,10 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
                 sys; u0map, initialization_eqs, check_units, pmap = parammap); fully_determined)
     end
 
+    if !isempty(StructuralTransformations.singular_check(get_tearing_state(isys)))
+        @warn "Since the initialization system is singular, the guess values may significantly affect the initial values of the ODE"
+    end
+
     uninit = setdiff(unknowns(sys), [unknowns(isys); getfield.(observed(isys), :lhs)])
 
     # TODO: throw on uninitialized arrays
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
index ff1d69e2bc..1cf110df13 100644
--- a/src/systems/problem_utils.jl
+++ b/src/systems/problem_utils.jl
@@ -498,7 +498,7 @@ function process_SciMLProblem(
         constructor, sys::AbstractSystem, u0map, pmap; build_initializeprob = true,
         implicit_dae = false, t = nothing, guesses = AnyDict(),
         warn_initialize_determined = true, initialization_eqs = [],
-        eval_expression = false, eval_module = @__MODULE__, fully_determined = false,
+        eval_expression = false, eval_module = @__MODULE__, fully_determined = nothing,
         check_initialization_units = false, tofloat = true, use_union = false,
         u0_constructor = identity, du0map = nothing, check_length = true,
         symbolic_u0 = false, warn_cyclic_dependency = false,

From 6632a2995ad99e129622da4c4910e0a36d2c93da Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 12:38:13 +0530
Subject: [PATCH 10/24] test: test that singular initialization systems throw a
 warning

---
 test/initializationsystem.jl | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index f39ec3c2f2..6d3e677c1e 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -947,3 +947,14 @@ end
 
     @test_nowarn remake(prob, p = prob.p)
 end
+
+@testset "Singular initialization prints a warning" begin
+    @parameters g
+    @variables x(t) y(t) [state_priority = 10] λ(t)
+    eqs = [D(D(x)) ~ λ * x
+           D(D(y)) ~ λ * y - g
+           x^2 + y^2 ~ 1]
+    @mtkbuild pend = ODESystem(eqs, t)
+    @test_warn ["structurally singular", "initialization", "guess"] ODEProblem(
+        pend, [x => 1, y => 0], (0.0, 1.5), [g => 1], guesses = [λ => 1])
+end

From 1892f0a10900e124587ab034e5842a1f4ec565da Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 12:43:28 +0530
Subject: [PATCH 11/24] refactor: format

---
 src/systems/abstractsystem.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
index 497b5e6476..3a9bbd33e1 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -3335,7 +3335,8 @@ function parse_variable(sys::AbstractSystem, str::AbstractString)
     # I'd write a regex to validate `str`, but https://xkcd.com/1171/
     str = strip(str)
     derivative_level = 0
-    while ((cond1 = startswith(str, "D(")) || startswith(str, "Differential(")) && endswith(str, ")")
+    while ((cond1 = startswith(str, "D(")) || startswith(str, "Differential(")) &&
+        endswith(str, ")")
         if cond1
             derivative_level += 1
             str = _string_view_inner(str, 2, 1)

From 2b7a6b695a33a09634a6039f27eabf414f2c7530 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 25 Nov 2024 18:03:16 +0530
Subject: [PATCH 12/24] refactor: only display singular warning if
 `warn_initialize_determined`

---
 src/structural_transformation/utils.jl   |  7 +++----
 src/systems/diffeqs/abstractodesystem.jl | 13 +++++++++++--
 test/initializationsystem.jl             |  2 +-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 405a3ebed3..bd24a1d017 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -123,16 +123,15 @@ function check_consistency(state::TransformationState, orig_inputs; nothrow = fa
     unassigned_var = singular_check(state)
 
     if !isempty(unassigned_var) || !is_balanced
+        if nothrow
+            return false
+        end
         io = IOBuffer()
         Base.print_array(io, unassigned_var)
         unassigned_var_str = String(take!(io))
         errmsg = "The system is structurally singular! " *
                  "Here are the problematic variables: \n" *
                  unassigned_var_str
-        if nothrow
-            @warn errmsg
-            return false
-        end
         throw(InvalidSystemException(errmsg))
     end
 
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 4019a868d8..5c13572aef 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1313,8 +1313,17 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
                 sys; u0map, initialization_eqs, check_units, pmap = parammap); fully_determined)
     end
 
-    if !isempty(StructuralTransformations.singular_check(get_tearing_state(isys)))
-        @warn "Since the initialization system is singular, the guess values may significantly affect the initial values of the ODE"
+    ts = get_tearing_state(isys)
+    if warn_initialize_determined &&
+       (unassigned_vars = StructuralTransformations.singular_check(ts); !isempty(unassigned_vars))
+        errmsg = """
+        The initialization system is structurally singular. Guess values may \
+        significantly affect the initial values of the ODE. The problematic variables \
+        are $unassigned_vars.
+
+        Note that the identification of problematic variables is a best-effort heuristic.
+        """
+        @warn errmsg
     end
 
     uninit = setdiff(unknowns(sys), [unknowns(isys); getfield.(observed(isys), :lhs)])
diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 6d3e677c1e..920a73a723 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -955,6 +955,6 @@ end
            D(D(y)) ~ λ * y - g
            x^2 + y^2 ~ 1]
     @mtkbuild pend = ODESystem(eqs, t)
-    @test_warn ["structurally singular", "initialization", "guess"] ODEProblem(
+    @test_warn ["structurally singular", "initialization", "Guess", "heuristic"] ODEProblem(
         pend, [x => 1, y => 0], (0.0, 1.5), [g => 1], guesses = [λ => 1])
 end

From 9c7fbc7fda33b687553f51549af2e66f536cc407 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 22 Nov 2024 14:23:43 +0530
Subject: [PATCH 13/24] test: make initializaton tests runnable from
 OrdinaryDiffEq

---
 test/runtests.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 192ead935c..0ba23fc4f7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -33,7 +33,6 @@ end
             @safetestset "Dynamic Quantities Test" include("dq_units.jl")
             @safetestset "Unitful Quantities Test" include("units.jl")
             @safetestset "Mass Matrix Test" include("mass_matrix.jl")
-            @safetestset "InitializationSystem Test" include("initializationsystem.jl")
             @safetestset "Guess Propagation" include("guess_propagation.jl")
             @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
             @safetestset "Reduction Test" include("reduction.jl")
@@ -64,6 +63,10 @@ end
         end
     end
 
+    if GROUP == "All" || GROUP == "InterfaceI" || GROUP == "Initialization"
+            @safetestset "InitializationSystem Test" include("initializationsystem.jl")
+    end
+
     if GROUP == "All" || GROUP == "InterfaceII"
         @testset "InterfaceII" begin
             @safetestset "IndexCache Test" include("index_cache.jl")

From 7eba3438d96d3a0a1c275edd3f61c47e49b2788f Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 26 Nov 2024 17:56:51 +0530
Subject: [PATCH 14/24] test: also move guess propagation, hierarchical
 initialization equations and initial values testsets

---
 test/runtests.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 0ba23fc4f7..4f6c20e04b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -33,8 +33,6 @@ end
             @safetestset "Dynamic Quantities Test" include("dq_units.jl")
             @safetestset "Unitful Quantities Test" include("units.jl")
             @safetestset "Mass Matrix Test" include("mass_matrix.jl")
-            @safetestset "Guess Propagation" include("guess_propagation.jl")
-            @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
             @safetestset "Reduction Test" include("reduction.jl")
             @safetestset "Split Parameters Test" include("split_parameters.jl")
             @safetestset "StaticArrays Test" include("static_arrays.jl")
@@ -57,14 +55,16 @@ end
             @safetestset "Constants Test" include("constants.jl")
             @safetestset "Parameter Dependency Test" include("parameter_dependencies.jl")
             @safetestset "Generate Custom Function Test" include("generate_custom_function.jl")
-            @safetestset "Initial Values Test" include("initial_values.jl")
             @safetestset "Equation Type Accessors Test" include("equation_type_accessors.jl")
             @safetestset "Equations with complex values" include("complex.jl")
         end
     end
 
     if GROUP == "All" || GROUP == "InterfaceI" || GROUP == "Initialization"
+            @safetestset "Guess Propagation" include("guess_propagation.jl")
+            @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
             @safetestset "InitializationSystem Test" include("initializationsystem.jl")
+            @safetestset "Initial Values Test" include("initial_values.jl")
     end
 
     if GROUP == "All" || GROUP == "InterfaceII"

From d73342d458fe6ca0b85fbfb6c6f8cb274dfd5ba0 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 27 Nov 2024 11:27:37 +0530
Subject: [PATCH 15/24] test: run `Initialization` as its own test group

---
 .github/workflows/Tests.yml | 1 +
 test/runtests.jl            | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
index 609ab2fb3d..a95f19a085 100644
--- a/.github/workflows/Tests.yml
+++ b/.github/workflows/Tests.yml
@@ -32,6 +32,7 @@ jobs:
         group:
           - InterfaceI
           - InterfaceII
+          - Initialization
           - SymbolicIndexingInterface
           - Extended
           - Extensions
diff --git a/test/runtests.jl b/test/runtests.jl
index 4f6c20e04b..06a167bdd9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -60,7 +60,7 @@ end
         end
     end
 
-    if GROUP == "All" || GROUP == "InterfaceI" || GROUP == "Initialization"
+    if GROUP == "All" || GROUP == "Initialization"
             @safetestset "Guess Propagation" include("guess_propagation.jl")
             @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
             @safetestset "InitializationSystem Test" include("initializationsystem.jl")

From 237cb82ef42252c6619179befa918ad580647ee9 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 27 Nov 2024 13:10:31 +0530
Subject: [PATCH 16/24] refactor: format

---
 test/runtests.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 06a167bdd9..244393d351 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -61,10 +61,10 @@ end
     end
 
     if GROUP == "All" || GROUP == "Initialization"
-            @safetestset "Guess Propagation" include("guess_propagation.jl")
-            @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
-            @safetestset "InitializationSystem Test" include("initializationsystem.jl")
-            @safetestset "Initial Values Test" include("initial_values.jl")
+        @safetestset "Guess Propagation" include("guess_propagation.jl")
+        @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
+        @safetestset "InitializationSystem Test" include("initializationsystem.jl")
+        @safetestset "Initial Values Test" include("initial_values.jl")
     end
 
     if GROUP == "All" || GROUP == "InterfaceII"

From 5b702cab0620e79b92637929bb773ef2a2805a21 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 21 Nov 2024 17:37:24 +0530
Subject: [PATCH 17/24] fix: propagate `initializeprobpmap` and
 `update_initializeprob!` to `DAEFunction`

---
 src/systems/diffeqs/abstractodesystem.jl |  6 +++++-
 test/initializationsystem.jl             | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index b6d750c621..cbf835f48d 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -498,6 +498,8 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         checkbounds = false,
         initializeprob = nothing,
         initializeprobmap = nothing,
+        initializeprobpmap = nothing,
+        update_initializeprob! = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
@@ -551,7 +553,9 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         jac_prototype = jac_prototype,
         observed = observedfun,
         initializeprob = initializeprob,
-        initializeprobmap = initializeprobmap)
+        initializeprobmap = initializeprobmap,
+        initializeprobpmap = initializeprobpmap,
+        update_initializeprob! = update_initializeprob!)
 end
 
 function DiffEqBase.DDEFunction(sys::AbstractODESystem, args...; kwargs...)
diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 920a73a723..770f9f12bd 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -958,3 +958,20 @@ end
     @test_warn ["structurally singular", "initialization", "Guess", "heuristic"] ODEProblem(
         pend, [x => 1, y => 0], (0.0, 1.5), [g => 1], guesses = [λ => 1])
 end
+
+@testset "DAEProblem initialization" begin
+    @variables x(t) [guess = 1.0] y(t) [guess = 1.0]
+    @parameters p=missing [guess = 1.0] q=missing [guess = 1.0]
+    @mtkbuild sys = ODESystem(
+        [D(x) ~ p * y + q * t, x^3 + y^3 ~ 5], t; initialization_eqs = [p^2 + q^3 ~ 3])
+
+    # FIXME: solve for du0
+    prob = DAEProblem(
+        sys, [D(x) => cbrt(4), D(y) => -1 / cbrt(4)], [x => 1.0], (0.0, 1.0), [p => 1.0])
+
+    integ = init(prob, DImplicitEuler())
+    @test integ[x] ≈ 1.0
+    @test integ[y]≈cbrt(4) rtol=1e-6
+    @test integ.ps[p] ≈ 1.0
+    @test integ.ps[q]≈cbrt(2) rtol=1e-6
+end

From 57ea5fc3ffae436d516624a285b928b0d226be51 Mon Sep 17 00:00:00 2001
From: Andrew Leonard <amleonard211@gmail.com>
Date: Wed, 27 Nov 2024 08:47:15 -0500
Subject: [PATCH 18/24] add checks kwarg to OptimizationProblem, use in
 ConstraintSystem

---
 src/systems/optimization/optimizationsystem.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
index 39c5c1e7a7..176a2bdfed 100644
--- a/src/systems/optimization/optimizationsystem.jl
+++ b/src/systems/optimization/optimizationsystem.jl
@@ -284,6 +284,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         linenumbers = true, parallel = SerialForm(),
         eval_expression = false, eval_module = @__MODULE__,
         use_union = false,
+        checks = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `OptimizationSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `OptimizationProblem`")
@@ -393,7 +394,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
     observedfun = ObservedFunctionCache(sys; eval_expression, eval_module)
 
     if length(cstr) > 0
-        @named cons_sys = ConstraintsSystem(cstr, dvs, ps)
+        @named cons_sys = ConstraintsSystem(cstr, dvs, ps; checks)
         cons_sys = complete(cons_sys)
         cons, lcons_, ucons_ = generate_function(cons_sys, checkbounds = checkbounds,
             linenumbers = linenumbers,

From d9943261b68954456de1de36865dc05d39722cb7 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 27 Nov 2024 20:05:45 +0530
Subject: [PATCH 19/24] fix: fix constraint function not wrapped in
 `OptimizationProblem`

---
 src/systems/optimization/optimizationsystem.jl | 9 +++++++--
 test/optimizationsystem.jl                     | 8 ++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
index 176a2bdfed..801e7b05f3 100644
--- a/src/systems/optimization/optimizationsystem.jl
+++ b/src/systems/optimization/optimizationsystem.jl
@@ -399,7 +399,12 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         cons, lcons_, ucons_ = generate_function(cons_sys, checkbounds = checkbounds,
             linenumbers = linenumbers,
             expression = Val{true})
-        cons = eval_or_rgf.(cons; eval_expression, eval_module)
+        cons = let (cons_oop, cons_iip) = eval_or_rgf.(cons; eval_expression, eval_module)
+            _cons(u, p) = cons_oop(u, p)
+            _cons(resid, u, p) = cons_iip(resid, u, p)
+            _cons(u, p::MTKParameters) = cons_oop(u, p...)
+            _cons(resid, u, p::MTKParameters) = cons_iip(resid, u, p...)
+        end
         if cons_j
             _cons_j = let (cons_jac_oop, cons_jac_iip) = eval_or_rgf.(
                     generate_jacobian(cons_sys;
@@ -465,7 +470,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
             grad = _grad,
             hess = _hess,
             hess_prototype = hess_prototype,
-            cons = cons[2],
+            cons = cons,
             cons_j = _cons_j,
             cons_h = _cons_h,
             cons_jac_prototype = cons_jac_prototype,
diff --git a/test/optimizationsystem.jl b/test/optimizationsystem.jl
index a8e2be936d..23d07dcb85 100644
--- a/test/optimizationsystem.jl
+++ b/test/optimizationsystem.jl
@@ -368,3 +368,11 @@ end
     @test is_variable(sys, x[2])
     @test is_variable(sys, x[3])
 end
+
+@testset "Constraints work with nonnumeric parameters" begin
+    @variables x
+    @parameters p f(::Real)
+    @mtkbuild sys = OptimizationSystem(x^2 + f(x) * p, [x], [f, p]; constraints = [2.0 ≲ f(x) + p])
+    prob = OptimizationProblem(sys, [x => 1.0], [p => 1.0, f => (x -> 2x)])
+    @test abs(prob.f.cons(prob.u0, prob.p)[1]) ≈ 1.0
+end

From d6add0eefe9f4d1232762e975cd023464c49cc01 Mon Sep 17 00:00:00 2001
From: Andrew Leonard <amleonard211@gmail.com>
Date: Wed, 27 Nov 2024 10:45:30 -0500
Subject: [PATCH 20/24] formatting

---
 test/optimizationsystem.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/optimizationsystem.jl b/test/optimizationsystem.jl
index 23d07dcb85..bb59fb09d9 100644
--- a/test/optimizationsystem.jl
+++ b/test/optimizationsystem.jl
@@ -372,7 +372,8 @@ end
 @testset "Constraints work with nonnumeric parameters" begin
     @variables x
     @parameters p f(::Real)
-    @mtkbuild sys = OptimizationSystem(x^2 + f(x) * p, [x], [f, p]; constraints = [2.0 ≲ f(x) + p])
+    @mtkbuild sys = OptimizationSystem(
+        x^2 + f(x) * p, [x], [f, p]; constraints = [2.0 ≲ f(x) + p])
     prob = OptimizationProblem(sys, [x => 1.0], [p => 1.0, f => (x -> 2x)])
     @test abs(prob.f.cons(prob.u0, prob.p)[1]) ≈ 1.0
 end

From f20c7941f18b88cd4b7d50754a83e414bfba6eff Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 28 Nov 2024 12:41:19 +0530
Subject: [PATCH 21/24] build: bump minor version

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index dcad7c5e4d..1cb204a41b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Yingbo Ma <mayingbo5@gmail.com>", "Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "9.53.0"
+version = "9.54.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

From 28858155db1b668ea82547fc505da68131854f67 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 28 Nov 2024 13:20:27 +0530
Subject: [PATCH 22/24] feat: add automatic variable discovery for
 `OptimizationSystem`

---
 .../optimization/optimizationsystem.jl        | 27 +++++++++++++++++++
 src/utils.jl                                  | 13 ++++++++-
 test/optimizationsystem.jl                    | 11 ++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
index 801e7b05f3..3425216b5f 100644
--- a/src/systems/optimization/optimizationsystem.jl
+++ b/src/systems/optimization/optimizationsystem.jl
@@ -144,6 +144,33 @@ function OptimizationSystem(op, unknowns, ps;
         checks = checks)
 end
 
+function OptimizationSystem(objective; constraints = [], kwargs...)
+    allunknowns = OrderedSet()
+    ps = OrderedSet()
+    collect_vars!(allunknowns, ps, objective, nothing)
+    for cons in constraints
+        collect_vars!(allunknowns, ps, cons, nothing)
+    end
+    for ssys in get(kwargs, :systems, OptimizationSystem[])
+        collect_scoped_vars!(allunknowns, ps, ssys, nothing)
+    end
+    new_ps = OrderedSet()
+    for p in ps
+        if iscall(p) && operation(p) === getindex
+            par = arguments(p)[begin]
+            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
+               all(par[i] in ps for i in eachindex(par))
+                push!(new_ps, par)
+            else
+                push!(new_ps, p)
+            end
+        else
+            push!(new_ps, p)
+        end
+    end
+    return OptimizationSystem(objective, collect(allunknowns), collect(new_ps); constraints, kwargs...)
+end
+
 function flatten(sys::OptimizationSystem)
     systems = get_systems(sys)
     isempty(systems) && return sys
diff --git a/src/utils.jl b/src/utils.jl
index 830ec98e44..1555cd624e 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -516,6 +516,15 @@ function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Dif
             end
         end
     end
+    if has_constraints(sys)
+        for eq in get_constraints(sys)
+            eqtype_supports_collect_vars(eq) || continue
+            collect_vars!(unknowns, parameters, eq, iv; depth, op)
+        end
+    end
+    if has_op(sys)
+        collect_vars!(unknowns, parameters, get_op(sys), iv; depth, op)
+    end
     newdepth = depth == -1 ? depth : depth + 1
     for ssys in get_systems(sys)
         collect_scoped_vars!(unknowns, parameters, ssys, iv; depth = newdepth, op)
@@ -544,9 +553,10 @@ Can be dispatched by higher-level libraries to indicate support.
 """
 eqtype_supports_collect_vars(eq) = false
 eqtype_supports_collect_vars(eq::Equation) = true
+eqtype_supports_collect_vars(eq::Inequality) = true
 eqtype_supports_collect_vars(eq::Pair) = true
 
-function collect_vars!(unknowns, parameters, eq::Equation, iv;
+function collect_vars!(unknowns, parameters, eq::Union{Equation, Inequality}, iv;
         depth = 0, op = Differential)
     collect_vars!(unknowns, parameters, eq.lhs, iv; depth, op)
     collect_vars!(unknowns, parameters, eq.rhs, iv; depth, op)
@@ -559,6 +569,7 @@ function collect_vars!(unknowns, parameters, p::Pair, iv; depth = 0, op = Differ
     return nothing
 end
 
+
 function collect_var!(unknowns, parameters, var, iv; depth = 0)
     isequal(var, iv) && return nothing
     check_scope_depth(getmetadata(var, SymScope, LocalScope()), depth) || return nothing
diff --git a/test/optimizationsystem.jl b/test/optimizationsystem.jl
index bb59fb09d9..c20613441a 100644
--- a/test/optimizationsystem.jl
+++ b/test/optimizationsystem.jl
@@ -377,3 +377,14 @@ end
     prob = OptimizationProblem(sys, [x => 1.0], [p => 1.0, f => (x -> 2x)])
     @test abs(prob.f.cons(prob.u0, prob.p)[1]) ≈ 1.0
 end
+
+@testset "Variable discovery" begin
+    @variables x1 x2
+    @parameters p1 p2
+    @named sys1 = OptimizationSystem(x1^2; constraints = [p1 * x1 ≲ 2.0])
+    @named sys2 = OptimizationSystem(x2^2; constraints = [p2 * x2 ≲ 2.0], systems = [sys1])
+    @test isequal(only(unknowns(sys1)), x1)
+    @test isequal(only(parameters(sys1)), p1)
+    @test all(y -> any(x -> isequal(x, y), unknowns(sys2)), [x2, sys1.x1])
+    @test all(y -> any(x -> isequal(x, y), parameters(sys2)), [p2, sys1.p1])
+end

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 23/24] 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

From 12b02c0bc5881873c5bfc27c141895c1a006b7e2 Mon Sep 17 00:00:00 2001
From: Christopher Rackauckas <accounts@chrisrackauckas.com>
Date: Sun, 1 Dec 2024 12:28:48 -0800
Subject: [PATCH 24/24] Fix initialization_eqs example

It was trivially wrong, fixes https://github.com/SciML/ModelingToolkit.jl/issues/3245
---
 docs/src/tutorials/initialization.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/tutorials/initialization.md b/docs/src/tutorials/initialization.md
index c3dbf8bbb0..c01310eb06 100644
--- a/docs/src/tutorials/initialization.md
+++ b/docs/src/tutorials/initialization.md
@@ -113,7 +113,7 @@ the `initialization_eqs` keyword argument, for example:
 
 ```@example init
 prob = ODEProblem(pend, [x => 1], (0.0, 1.5), [g => 1], guesses = [λ => 0, y => 1],
-    initialization_eqs = [y ~ 1])
+    initialization_eqs = [y ~ 0])
 sol = solve(prob, Rodas5P())
 plot(sol, idxs = (x, y))
 ```