Merge pull request #4539 from SciML/as/faster-iprobpmap

fix: avoid generating massive functions in `get_mtkparameters_reconstructor`

Author: AayushSabharwal

lib/ModelingToolkitBase/Project.toml

@@ -73,6 +73,7 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
@@ -86,6 +87,7 @@ MTKLabelledArraysExt = "LabelledArrays"
 MTKLatexifyExt = "Latexify"
 MTKMooncakeExt = "Mooncake"
 MTKPyomoDynamicOptExt = "Pyomo"
+MTKTrackerExt = "Tracker"
 
 [compat]
 ADTypes = "1.14.0"

lib/ModelingToolkitBase/ext/MTKTrackerExt.jl

@@ -0,0 +1,14 @@
+module MTKTrackerExt
+
+import ModelingToolkitBase as MTKBase
+import Tracker
+
+function MTKBase.promote_type_with_nothing(::Type{Tracker.TrackedReal{T}}, x::Tracker.TrackedArray{T}) where {T}
+    return Tracker.TrackedReal{T}
+end
+
+function MTKBase.promote_with_nothing(::Type{Tracker.TrackedReal{T}}, x::Tracker.TrackedArray{T}) where {T}
+    return x 
+end
+
+end

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -626,12 +626,17 @@ end
 """
 Returns true if the parameter `p` is of the form `Initial(x)`.
 """
+function isinitial(p::SymbolicT)
+    p, _ = split_indexed_var(p)
+    Moshi.Match.@match p begin
+        BSImpl.Term(; f) => f isa Initial
+        _ => false
+    end
+end
 function isinitial(p)
-    p = unwrap(p)
-    return iscall(p) && (
-        operation(p) isa Initial ||
-            operation(p) === getindex && isinitial(arguments(p)[1])
-    )
+    up = unwrap(p)
+    up === p && return false
+    return isinitial(up)
 end
 
 """

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

@@ -812,14 +812,9 @@ function DiffEqBase.get_updated_symbolic_problem(
     end
 
     u0 = DiffEqBase.promote_u0(u0, buffer, t0)
+    u0 = ArrayInterface.restructure(u0, meta.get_updated_u0(prob, initdata.initializeprob))
 
-    if ArrayInterface.ismutable(u0)
-        T = typeof(u0)
-    else
-        T = StaticArrays.similar_type(u0)
-    end
-
-    return remake(prob; u0 = T(meta.get_updated_u0(prob, initdata.initializeprob)), p)
+    return remake(prob; u0, p)
 end
 
 """