feat: run trivial initialization in problem constructor

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -759,7 +759,7 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         kwargs1 = merge(kwargs1, (; tstops))
     end
 
-    return ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
+    return remake(ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...))
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]
 
@@ -963,8 +963,9 @@ function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem, du0map, u0map, tspan
         kwargs1 = merge(kwargs1, (; tstops))
     end
 
-    DAEProblem{iip}(f, du0, u0, tspan, p; differential_vars = differential_vars,
-        kwargs..., kwargs1...)
+    return remake(DAEProblem{iip}(
+        f, du0, u0, tspan, p; differential_vars = differential_vars,
+        kwargs..., kwargs1...))
 end
 
 function generate_history(sys::AbstractODESystem, u0; expression = Val{false}, kwargs...)
@@ -1008,7 +1009,7 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     if cbs !== nothing
         kwargs1 = merge(kwargs1, (callback = cbs,))
     end
-    DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...)
+    return remake(DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...))
 end
 
 function DiffEqBase.SDDEProblem(sys::AbstractODESystem, args...; kwargs...)
@@ -1057,9 +1058,9 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     else
         noise_rate_prototype = zeros(eltype(u0), size(noiseeqs))
     end
-    SDDEProblem{iip}(f, f.g, u0, h, tspan, p;
+    return remake(SDDEProblem{iip}(f, f.g, u0, h, tspan, p;
         noise_rate_prototype =
-        noise_rate_prototype, kwargs1..., kwargs...)
+        noise_rate_prototype, kwargs1..., kwargs...))
 end
 
 """

src/systems/diffeqs/sdesystem.jl

@@ -792,8 +792,8 @@ function DiffEqBase.SDEProblem{iip, specialize}(
 
     kwargs = filter_kwargs(kwargs)
 
-    SDEProblem{iip}(f, u0, tspan, p; callback = cbs, noise,
-        noise_rate_prototype = noise_rate_prototype, kwargs...)
+    return remake(SDEProblem{iip}(f, u0, tspan, p; callback = cbs, noise,
+        noise_rate_prototype = noise_rate_prototype, kwargs...))
 end
 
 function DiffEqBase.SDEProblem(sys::ODESystem, args...; kwargs...)

src/systems/nonlinear/nonlinearsystem.jl

@@ -519,7 +519,7 @@ function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem, u0map,
     f, u0, p = process_SciMLProblem(NonlinearFunction{iip}, sys, u0map, parammap;
         check_length, kwargs...)
     pt = something(get_metadata(sys), StandardNonlinearProblem())
-    NonlinearProblem{iip}(f, u0, p, pt; filter_kwargs(kwargs)...)
+    return remake(NonlinearProblem{iip}(f, u0, p, pt; filter_kwargs(kwargs)...))
 end
 
 """
@@ -548,7 +548,7 @@ function DiffEqBase.NonlinearLeastSquaresProblem{iip}(sys::NonlinearSystem, u0ma
     f, u0, p = process_SciMLProblem(NonlinearFunction{iip}, sys, u0map, parammap;
         check_length, kwargs...)
     pt = something(get_metadata(sys), StandardNonlinearProblem())
-    NonlinearLeastSquaresProblem{iip}(f, u0, p; filter_kwargs(kwargs)...)
+    return remake(NonlinearLeastSquaresProblem{iip}(f, u0, p; filter_kwargs(kwargs)...))
 end
 
 const TypeT = Union{DataType, UnionAll}

test/initializationsystem.jl

@@ -1388,3 +1388,29 @@ end
     integ1 = init(oprob1)
     @test integ1[X1] ≈ 1.0
 end
+
+@testset "Trivial initialization is run on problem construction" begin
+    @variables _x(..) y(t)
+    @brownian a
+    @parameters tot
+    x = _x(t)
+    @testset "$Problem" for (Problem, lhs, rhs) in [
+        (ODEProblem, D, 0.0),
+        (SDEProblem, D, a),
+        (DDEProblem, D, _x(t - 0.1)),
+        (SDDEProblem, D, _x(t - 0.1) + a)
+    ]
+        @mtkbuild sys = ModelingToolkit.System([lhs(x) ~ x + rhs, x + y ~ tot], t;
+            guesses = [tot => 1.0], defaults = [tot => missing])
+        prob = Problem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+        @test prob.ps[tot] ≈ 2.0
+    end
+    @testset "$Problem" for Problem in [NonlinearProblem, NonlinearLeastSquaresProblem]
+        @parameters p1 p2
+        @mtkbuild sys = NonlinearSystem([x^2 + y^2 ~ p1, (x - 1)^2 + (y - 1)^2 ~ p2];
+            parameter_dependencies = [p2 ~ 2p1],
+            guesses = [p1 => 0.0], defaults = [p1 => missing])
+        prob = Problem(sys, [x => 1.0, y => 1.0], [p2 => 6.0])
+        @test prob.ps[p1] ≈ 3.0
+    end
+end